Agent Zero
54 lines
1.9 KiB
Python
54 lines
1.9 KiB
Python
"""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")
|