diff --git a/surfsense_backend/alembic/versions/117_optimize_zero_publication_column_lists.py b/surfsense_backend/alembic/versions/117_optimize_zero_publication_column_lists.py index 78a26a381..3ad5a043b 100644 --- a/surfsense_backend/alembic/versions/117_optimize_zero_publication_column_lists.py +++ b/surfsense_backend/alembic/versions/117_optimize_zero_publication_column_lists.py @@ -79,40 +79,44 @@ def _terminate_blocked_pids(conn, table: str) -> None: def upgrade() -> None: conn = op.get_bind() + # asyncpg requires LOCK TABLE inside a transaction block. Alembic already + # opened one via context.begin_transaction(), but the driver still errors + # unless we use an explicit SAVEPOINT (nested transaction) for this block. + tx = conn.begin_nested() if conn.in_transaction() else conn.begin() + with tx: + conn.execute(sa.text("SET lock_timeout = '10s'")) - conn.execute(sa.text("SET lock_timeout = '10s'")) + for tbl in sorted(TABLES_WITH_FULL_IDENTITY): + _terminate_blocked_pids(conn, tbl) + conn.execute(sa.text(f'LOCK TABLE "{tbl}" IN ACCESS EXCLUSIVE MODE')) - for tbl in sorted(TABLES_WITH_FULL_IDENTITY): - _terminate_blocked_pids(conn, tbl) - conn.execute(sa.text(f'LOCK TABLE "{tbl}" IN ACCESS EXCLUSIVE MODE')) + for tbl in TABLES_WITH_FULL_IDENTITY: + conn.execute(sa.text(f'ALTER TABLE "{tbl}" REPLICA IDENTITY DEFAULT')) - for tbl in TABLES_WITH_FULL_IDENTITY: - conn.execute(sa.text(f'ALTER TABLE "{tbl}" REPLICA IDENTITY DEFAULT')) + conn.execute(sa.text(f"DROP PUBLICATION IF EXISTS {PUBLICATION_NAME}")) - conn.execute(sa.text(f"DROP PUBLICATION IF EXISTS {PUBLICATION_NAME}")) + has_zero_ver = conn.execute( + sa.text( + "SELECT 1 FROM information_schema.columns " + "WHERE table_name = 'documents' AND column_name = '_0_version'" + ) + ).fetchone() - has_zero_ver = conn.execute( - sa.text( - "SELECT 1 FROM information_schema.columns " - "WHERE table_name = 'documents' AND column_name = '_0_version'" + cols = DOCUMENT_COLS + (['"_0_version"'] if has_zero_ver else []) + col_list = ", ".join(cols) + + conn.execute( + sa.text( + f"CREATE PUBLICATION {PUBLICATION_NAME} FOR TABLE " + f"notifications, " + f"documents ({col_list}), " + f"folders, " + f"search_source_connectors, " + f"new_chat_messages, " + f"chat_comments, " + f"chat_session_state" + ) ) - ).fetchone() - - cols = DOCUMENT_COLS + (['"_0_version"'] if has_zero_ver else []) - col_list = ", ".join(cols) - - conn.execute( - sa.text( - f"CREATE PUBLICATION {PUBLICATION_NAME} FOR TABLE " - f"notifications, " - f"documents ({col_list}), " - f"folders, " - f"search_source_connectors, " - f"new_chat_messages, " - f"chat_comments, " - f"chat_session_state" - ) - ) def downgrade() -> None: diff --git a/surfsense_backend/alembic/versions/121_add_memory_md_columns.py b/surfsense_backend/alembic/versions/121_add_memory_md_columns.py index d5ff967fd..ac248dfca 100644 --- a/surfsense_backend/alembic/versions/121_add_memory_md_columns.py +++ b/surfsense_backend/alembic/versions/121_add_memory_md_columns.py @@ -12,8 +12,6 @@ from __future__ import annotations from collections.abc import Sequence -import sqlalchemy as sa - from alembic import op revision: str = "121" @@ -23,16 +21,30 @@ depends_on: str | Sequence[str] | None = None def upgrade() -> None: - op.add_column( - "user", - sa.Column("memory_md", sa.Text(), nullable=True, server_default=""), - ) - op.add_column( - "searchspaces", - sa.Column("shared_memory_md", sa.Text(), nullable=True, server_default=""), + # Idempotent: column(s) may already exist after a failed run or manual DDL. + op.execute( + """ + DO $$ + BEGIN + IF NOT EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_schema = 'public' AND table_name = 'user' + AND column_name = 'memory_md' + ) THEN + ALTER TABLE "user" ADD COLUMN memory_md TEXT DEFAULT ''; + END IF; + IF NOT EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_schema = 'public' AND table_name = 'searchspaces' + AND column_name = 'shared_memory_md' + ) THEN + ALTER TABLE searchspaces ADD COLUMN shared_memory_md TEXT DEFAULT ''; + END IF; + END$$; + """ ) def downgrade() -> None: - op.drop_column("searchspaces", "shared_memory_md") - op.drop_column("user", "memory_md") + op.execute("ALTER TABLE searchspaces DROP COLUMN IF EXISTS shared_memory_md") + op.execute('ALTER TABLE "user" DROP COLUMN IF EXISTS memory_md')