From b41c22842f35d19308e5fcd512dc05c57a35d103 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Fri, 6 Feb 2026 12:32:55 +0530 Subject: [PATCH] refactor: change alelmbic migration number and made migrations idempotent --- ...ications_table_and_electric_replication.py | 36 +++++++------ ...4_add_access_token_to_image_generations.py | 50 ++++++++++++++----- ...mn.py => 95_add_document_status_column.py} | 8 +-- 3 files changed, 63 insertions(+), 31 deletions(-) rename surfsense_backend/alembic/versions/{93_add_document_status_column.py => 95_add_document_status_column.py} (95%) diff --git a/surfsense_backend/alembic/versions/66_add_notifications_table_and_electric_replication.py b/surfsense_backend/alembic/versions/66_add_notifications_table_and_electric_replication.py index dc25a1edd..182bf981c 100644 --- a/surfsense_backend/alembic/versions/66_add_notifications_table_and_electric_replication.py +++ b/surfsense_backend/alembic/versions/66_add_notifications_table_and_electric_replication.py @@ -17,13 +17,6 @@ from collections.abc import Sequence from alembic import context, op -# Get Electric SQL user credentials from env.py configuration -_config = context.config -ELECTRIC_DB_USER = _config.get_main_option("electric_db_user", "electric") -ELECTRIC_DB_PASSWORD = _config.get_main_option( - "electric_db_password", "electric_password" -) - # revision identifiers, used by Alembic. revision: str = "66" down_revision: str | None = "65" @@ -31,8 +24,21 @@ branch_labels: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None +def _get_electric_credentials() -> tuple[str, str]: + """Get Electric SQL credentials from Alembic config. + + Must be called inside upgrade()/downgrade(), not at module level, + because context.config is only available during migration execution. + """ + _config = context.config + user = _config.get_main_option("electric_db_user", "electric") + password = _config.get_main_option("electric_db_password", "electric_password") + return user, password + + def upgrade() -> None: """Upgrade schema - add notifications table and Electric SQL replication.""" + electric_db_user, electric_db_password = _get_electric_credentials() # Create notifications table op.execute( """ @@ -74,8 +80,8 @@ def upgrade() -> None: f""" DO $$ BEGIN - IF NOT EXISTS (SELECT FROM pg_user WHERE usename = '{ELECTRIC_DB_USER}') THEN - CREATE USER {ELECTRIC_DB_USER} WITH REPLICATION PASSWORD '{ELECTRIC_DB_PASSWORD}'; + IF NOT EXISTS (SELECT FROM pg_user WHERE usename = '{electric_db_user}') THEN + CREATE USER {electric_db_user} WITH REPLICATION PASSWORD '{electric_db_password}'; END IF; END $$; @@ -89,19 +95,19 @@ def upgrade() -> None: DECLARE db_name TEXT := current_database(); BEGIN - EXECUTE format('GRANT CONNECT ON DATABASE %I TO {ELECTRIC_DB_USER}', db_name); + EXECUTE format('GRANT CONNECT ON DATABASE %I TO {electric_db_user}', db_name); END $$; """ ) - op.execute(f"GRANT USAGE ON SCHEMA public TO {ELECTRIC_DB_USER};") - op.execute(f"GRANT SELECT ON ALL TABLES IN SCHEMA public TO {ELECTRIC_DB_USER};") - op.execute(f"GRANT SELECT ON ALL SEQUENCES IN SCHEMA public TO {ELECTRIC_DB_USER};") + op.execute(f"GRANT USAGE ON SCHEMA public TO {electric_db_user};") + op.execute(f"GRANT SELECT ON ALL TABLES IN SCHEMA public TO {electric_db_user};") + op.execute(f"GRANT SELECT ON ALL SEQUENCES IN SCHEMA public TO {electric_db_user};") op.execute( - f"ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON TABLES TO {ELECTRIC_DB_USER};" + f"ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON TABLES TO {electric_db_user};" ) op.execute( - f"ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON SEQUENCES TO {ELECTRIC_DB_USER};" + f"ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON SEQUENCES TO {electric_db_user};" ) # Create the publication if not exists diff --git a/surfsense_backend/alembic/versions/94_add_access_token_to_image_generations.py b/surfsense_backend/alembic/versions/94_add_access_token_to_image_generations.py index 09bea2c19..92f027e00 100644 --- a/surfsense_backend/alembic/versions/94_add_access_token_to_image_generations.py +++ b/surfsense_backend/alembic/versions/94_add_access_token_to_image_generations.py @@ -10,8 +10,6 @@ SECRET_KEY rotation. from collections.abc import Sequence -import sqlalchemy as sa - from alembic import op # revision identifiers, used by Alembic. @@ -23,17 +21,45 @@ depends_on: str | Sequence[str] | None = None def upgrade() -> None: # Add access_token column (nullable so existing rows are unaffected) - op.add_column( - "image_generations", - sa.Column("access_token", sa.String(64), nullable=True), - ) - op.create_index( - "ix_image_generations_access_token", - "image_generations", - ["access_token"], + # Guard: skip entirely if image_generations table doesn't exist + op.execute( + """ + DO $$ + BEGIN + IF EXISTS ( + SELECT 1 FROM information_schema.tables + WHERE table_name = 'image_generations' + ) THEN + -- Add column if not exists + IF NOT EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_name = 'image_generations' AND column_name = 'access_token' + ) THEN + ALTER TABLE image_generations + ADD COLUMN access_token VARCHAR(64); + END IF; + + -- Create index if not exists + CREATE INDEX IF NOT EXISTS ix_image_generations_access_token + ON image_generations (access_token); + END IF; + END$$; + """ ) def downgrade() -> None: - op.drop_index("ix_image_generations_access_token", table_name="image_generations") - op.drop_column("image_generations", "access_token") + op.execute("DROP INDEX IF EXISTS ix_image_generations_access_token") + op.execute( + """ + DO $$ + BEGIN + IF EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_name = 'image_generations' AND column_name = 'access_token' + ) THEN + ALTER TABLE image_generations DROP COLUMN access_token; + END IF; + END$$; + """ + ) diff --git a/surfsense_backend/alembic/versions/93_add_document_status_column.py b/surfsense_backend/alembic/versions/95_add_document_status_column.py similarity index 95% rename from surfsense_backend/alembic/versions/93_add_document_status_column.py rename to surfsense_backend/alembic/versions/95_add_document_status_column.py index 382db6109..f5a6fa65d 100644 --- a/surfsense_backend/alembic/versions/93_add_document_status_column.py +++ b/surfsense_backend/alembic/versions/95_add_document_status_column.py @@ -1,7 +1,7 @@ """Add status column to documents table for per-document processing status -Revision ID: 93 -Revises: 92 +Revision ID: 95 +Revises: 94 Create Date: 2026-02-05 Changes: @@ -16,8 +16,8 @@ from collections.abc import Sequence from alembic import op # revision identifiers, used by Alembic. -revision: str = "93" -down_revision: str | None = "92" +revision: str = "95" +down_revision: str | None = "94" branch_labels: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None