leocrm/alembic/versions/0035_comm_search_index.py

"""Add search_tsv and embedding columns to comm_messages for full-text search indexing.

Revision ID: 0035
Revises: 0034_automation_config
Create Date: 2026-07-25
"""

from __future__ import annotations

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import TSVECTOR

revision: str = "0035_comm_search_index"
down_revision: Union[str, None] = "0034_automation_config"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # Add search_tsv column for full-text search
    op.execute("ALTER TABLE IF EXISTS comm_messages ADD COLUMN IF NOT EXISTS search_tsv tsvector")
    # Add embedding column for vector search (768 dimensions matching pgvector)
    op.execute(
        "ALTER TABLE IF EXISTS comm_messages ADD COLUMN IF NOT EXISTS embedding vector(768)"
    )
    # Create GIN index on search_tsv for fast FTS queries (only if table exists)
    op.execute("""
        DO $$ BEGIN
          IF EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'comm_messages') THEN
            CREATE INDEX IF NOT EXISTS ix_comm_messages_search_tsv ON comm_messages (search_tsv);
          END IF;
        END $$
    """)
    # Create IVFFlat index on embedding for fast vector search (only if table exists)
    op.execute("""
        DO $$ BEGIN
          IF EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'comm_messages') THEN
            CREATE INDEX IF NOT EXISTS ix_comm_messages_embedding
            ON comm_messages USING ivfflat (embedding vector_cosine_ops)
            WITH (lists = 100);
          END IF;
        END $$
    """)


def downgrade() -> None:
    op.drop_index("ix_comm_messages_embedding", table_name="comm_messages")
    op.drop_index("ix_comm_messages_search_tsv", table_name="comm_messages")
    op.drop_column("comm_messages", "embedding")
    op.drop_column("comm_messages", "search_tsv")