149 lines
5.2 KiB
Python
149 lines
5.2 KiB
Python
"""message search index outside the hypertable, versions searchable
|
|
|
|
Revision ID: d8f1a6c3e5b7
|
|
Revises: c5e7a9b1d3f6
|
|
Create Date: 2026-10-19 03:00:00.000000
|
|
|
|
"""
|
|
|
|
from collections.abc import Sequence
|
|
|
|
from alembic import op
|
|
|
|
revision: str = "d8f1a6c3e5b7"
|
|
down_revision: str | None = "c5e7a9b1d3f6"
|
|
branch_labels: str | Sequence[str] | None = None
|
|
depends_on: str | Sequence[str] | None = None
|
|
|
|
_VECTOR = """
|
|
CREATE FUNCTION search_vector(body text) RETURNS tsvector
|
|
LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $$
|
|
SELECT to_tsvector('russian', translate(coalesce(body, ''), 'Ёё', 'Ее'))
|
|
|| to_tsvector('simple', translate(coalesce(body, ''), 'Ёё', 'Ее'))
|
|
$$
|
|
"""
|
|
|
|
_TABLE = """
|
|
CREATE TABLE message_search (
|
|
account_id integer NOT NULL,
|
|
chat_id bigint NOT NULL,
|
|
message_id bigint NOT NULL,
|
|
date timestamptz NOT NULL,
|
|
sender_id bigint,
|
|
text text,
|
|
has_media boolean NOT NULL DEFAULT false,
|
|
edited_at timestamptz,
|
|
deleted_at timestamptz,
|
|
tsv tsvector GENERATED ALWAYS AS (search_vector(text)) STORED,
|
|
PRIMARY KEY (account_id, chat_id, message_id, date)
|
|
)
|
|
"""
|
|
|
|
_BACKFILL = """
|
|
INSERT INTO message_search (account_id, chat_id, message_id, date, sender_id,
|
|
text, has_media, edited_at, deleted_at)
|
|
SELECT account_id, chat_id, message_id, date, sender_id,
|
|
text, has_media, edited_at, deleted_at
|
|
FROM messages
|
|
"""
|
|
|
|
_SYNC = """
|
|
CREATE FUNCTION message_search_sync() RETURNS trigger LANGUAGE plpgsql AS $$
|
|
BEGIN
|
|
IF TG_OP = 'DELETE' THEN
|
|
DELETE FROM message_search
|
|
WHERE account_id = OLD.account_id AND chat_id = OLD.chat_id
|
|
AND message_id = OLD.message_id AND date = OLD.date;
|
|
RETURN NULL;
|
|
END IF;
|
|
INSERT INTO message_search (account_id, chat_id, message_id, date,
|
|
sender_id, text, has_media, edited_at,
|
|
deleted_at)
|
|
VALUES (NEW.account_id, NEW.chat_id, NEW.message_id, NEW.date,
|
|
NEW.sender_id, NEW.text, NEW.has_media, NEW.edited_at,
|
|
NEW.deleted_at)
|
|
ON CONFLICT (account_id, chat_id, message_id, date) DO UPDATE SET
|
|
sender_id = EXCLUDED.sender_id,
|
|
text = EXCLUDED.text,
|
|
has_media = EXCLUDED.has_media,
|
|
edited_at = EXCLUDED.edited_at,
|
|
deleted_at = EXCLUDED.deleted_at;
|
|
RETURN NULL;
|
|
END;
|
|
$$
|
|
"""
|
|
|
|
_CHANGED = (
|
|
"OLD.text IS DISTINCT FROM NEW.text "
|
|
"OR OLD.sender_id IS DISTINCT FROM NEW.sender_id "
|
|
"OR OLD.has_media IS DISTINCT FROM NEW.has_media "
|
|
"OR OLD.edited_at IS DISTINCT FROM NEW.edited_at "
|
|
"OR OLD.deleted_at IS DISTINCT FROM NEW.deleted_at"
|
|
)
|
|
|
|
|
|
def upgrade() -> None:
|
|
op.execute(_VECTOR)
|
|
op.execute(_TABLE)
|
|
op.execute(_BACKFILL)
|
|
op.execute("CREATE INDEX ix_message_search_tsv ON message_search USING gin (tsv)")
|
|
op.execute(
|
|
"CREATE INDEX ix_message_search_recent ON message_search "
|
|
"(account_id, date DESC, chat_id DESC, message_id DESC)"
|
|
)
|
|
op.execute(
|
|
"CREATE INDEX ix_message_search_sender ON message_search "
|
|
"(account_id, sender_id, date DESC)"
|
|
)
|
|
for flag in ("deleted", "edited"):
|
|
op.execute(
|
|
f"CREATE INDEX ix_message_search_{flag} ON message_search "
|
|
"(account_id, date DESC, chat_id DESC, message_id DESC) "
|
|
f"WHERE {flag}_at IS NOT NULL"
|
|
)
|
|
op.execute("ALTER TABLE message_search ALTER COLUMN tsv SET STATISTICS 1000")
|
|
op.execute(_SYNC)
|
|
op.execute(
|
|
"CREATE TRIGGER messages_search_insert AFTER INSERT ON messages "
|
|
"FOR EACH ROW EXECUTE FUNCTION message_search_sync()"
|
|
)
|
|
op.execute(
|
|
"CREATE TRIGGER messages_search_update AFTER UPDATE ON messages "
|
|
f"FOR EACH ROW WHEN ({_CHANGED}) EXECUTE FUNCTION message_search_sync()"
|
|
)
|
|
op.execute(
|
|
"CREATE TRIGGER messages_search_delete AFTER DELETE ON messages "
|
|
"FOR EACH ROW EXECUTE FUNCTION message_search_sync()"
|
|
)
|
|
op.execute(
|
|
"CREATE INDEX ix_message_versions_tsv ON message_versions USING gin "
|
|
"(search_vector(text))"
|
|
)
|
|
op.execute("DROP INDEX IF EXISTS ix_media_tsv")
|
|
op.execute(
|
|
"CREATE INDEX ix_media_tsv ON media USING gin (search_vector(extracted_text))"
|
|
)
|
|
op.execute("DROP INDEX IF EXISTS ix_messages_tsv")
|
|
for index in ("ix_message_versions_tsv", "ix_media_tsv"):
|
|
op.execute(f"ALTER INDEX {index} ALTER COLUMN 1 SET STATISTICS 1000")
|
|
op.execute("ANALYZE message_search, message_versions, media")
|
|
|
|
|
|
def downgrade() -> None:
|
|
op.execute(
|
|
"CREATE INDEX ix_messages_tsv ON messages USING gin "
|
|
"(to_tsvector('russian', coalesce(text, '')))"
|
|
)
|
|
op.execute("DROP INDEX IF EXISTS ix_media_tsv")
|
|
op.execute(
|
|
"CREATE INDEX ix_media_tsv ON media USING gin "
|
|
"(to_tsvector('russian', coalesce(extracted_text, '')))"
|
|
)
|
|
op.execute("DROP INDEX IF EXISTS ix_message_versions_tsv")
|
|
op.execute("DROP TRIGGER IF EXISTS messages_search_delete ON messages")
|
|
op.execute("DROP TRIGGER IF EXISTS messages_search_update ON messages")
|
|
op.execute("DROP TRIGGER IF EXISTS messages_search_insert ON messages")
|
|
op.execute("DROP FUNCTION IF EXISTS message_search_sync()")
|
|
op.execute("DROP TABLE IF EXISTS message_search")
|
|
op.execute("DROP FUNCTION IF EXISTS search_vector(text)")
|