refactor: change alelmbic migration number and made migrations idempotent

This commit is contained in:
Anish Sarkar 2026-02-06 12:32:55 +05:30
parent 0fdd194d92
commit b41c22842f
3 changed files with 63 additions and 31 deletions

View file

@ -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

View file

@ -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$$;
"""
)

View file

@ -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