74 lines
2.2 KiB
Python
74 lines
2.2 KiB
Python
"""add_bm25_fulltext_search
|
|
|
|
Revision ID: 1a35a4fa1950
|
|
Revises: 01f989db9079
|
|
Create Date: 2025-11-06 11:19:48.627698
|
|
|
|
"""
|
|
from typing import Sequence, Union
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = '1a35a4fa1950'
|
|
down_revision: Union[str, Sequence[str], None] = '01f989db9079'
|
|
branch_labels: Union[str, Sequence[str], None] = None
|
|
depends_on: Union[str, Sequence[str], None] = None
|
|
|
|
|
|
def upgrade() -> None:
|
|
"""Upgrade schema."""
|
|
# Add tsvector column for full-text search
|
|
op.execute("""
|
|
ALTER TABLE memory_units
|
|
ADD COLUMN search_vector tsvector
|
|
""")
|
|
|
|
# Populate tsvector with existing data (text + context combined)
|
|
op.execute("""
|
|
UPDATE memory_units
|
|
SET search_vector =
|
|
setweight(to_tsvector('english', COALESCE(text, '')), 'A') ||
|
|
setweight(to_tsvector('english', COALESCE(context, '')), 'B')
|
|
""")
|
|
|
|
# Create GIN index for fast full-text search
|
|
op.execute("""
|
|
CREATE INDEX idx_memory_units_search_vector
|
|
ON memory_units
|
|
USING GIN(search_vector)
|
|
""")
|
|
|
|
# Create trigger to auto-update tsvector on INSERT/UPDATE
|
|
op.execute("""
|
|
CREATE OR REPLACE FUNCTION memory_units_search_vector_trigger() RETURNS trigger AS $$
|
|
BEGIN
|
|
NEW.search_vector :=
|
|
setweight(to_tsvector('english', COALESCE(NEW.text, '')), 'A') ||
|
|
setweight(to_tsvector('english', COALESCE(NEW.context, '')), 'B');
|
|
RETURN NEW;
|
|
END
|
|
$$ LANGUAGE plpgsql;
|
|
""")
|
|
|
|
op.execute("""
|
|
CREATE TRIGGER update_memory_units_search_vector
|
|
BEFORE INSERT OR UPDATE ON memory_units
|
|
FOR EACH ROW
|
|
EXECUTE FUNCTION memory_units_search_vector_trigger();
|
|
""")
|
|
|
|
|
|
def downgrade() -> None:
|
|
"""Downgrade schema."""
|
|
# Drop trigger
|
|
op.execute("DROP TRIGGER IF EXISTS update_memory_units_search_vector ON memory_units")
|
|
op.execute("DROP FUNCTION IF EXISTS memory_units_search_vector_trigger()")
|
|
|
|
# Drop index
|
|
op.execute("DROP INDEX IF EXISTS idx_memory_units_search_vector")
|
|
|
|
# Drop column
|
|
op.execute("ALTER TABLE memory_units DROP COLUMN IF EXISTS search_vector")
|