Fix: DB Migration
This commit is contained in:
@@ -0,0 +1,216 @@
|
||||
"""Schema cleanup: id UniqueConstraints, organization_api_keys index/timestamp fixes.
|
||||
|
||||
Revision ID: 026_schema_cleanup
|
||||
Revises: 025_fix_zt_timestamps
|
||||
Create Date: 2026-03-23
|
||||
|
||||
Addresses all `db check` differences after 025 on a database upgraded from
|
||||
production (021_merge_heads):
|
||||
|
||||
1. Add UniqueConstraint on `id` for all pre-existing tables that inherit
|
||||
BaseModel (which declares id with unique=True). The ZeroTier tables
|
||||
already got these in 024_fix_zerotier_schema; this covers the rest.
|
||||
|
||||
2. organization_api_keys — fix schema drift vs. the current model:
|
||||
- TIMESTAMPTZ → TIMESTAMP WITHOUT TIME ZONE (align with rest of codebase)
|
||||
- Drop legacy unique constraint 'organization_api_keys_key_hash_key'
|
||||
and replace with named index 'ix_organization_api_keys_key_hash'
|
||||
- Drop extra index 'idx_org_api_key_org_id' (superseded by
|
||||
'ix_organization_api_keys_organization_id')
|
||||
- Add 'ix_organization_api_keys_organization_id' and
|
||||
'ix_organization_api_keys_is_revoked' named indexes expected by model
|
||||
|
||||
3. Drop 'idx_dept_can_sudo' index from departments — created by an old
|
||||
migration but not declared in the current Department model.
|
||||
|
||||
All operations are guarded so the migration is safe to re-run.
|
||||
"""
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
|
||||
revision = "026_schema_cleanup"
|
||||
down_revision = "025_fix_zt_timestamps"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _constraint_exists(conn, name: str) -> bool:
|
||||
row = conn.execute(sa.text(
|
||||
"SELECT 1 FROM information_schema.table_constraints "
|
||||
"WHERE constraint_name = :n"
|
||||
), {"n": name}).first()
|
||||
return row is not None
|
||||
|
||||
|
||||
def _index_exists(conn, table: str, index: str) -> bool:
|
||||
row = conn.execute(sa.text(
|
||||
"SELECT 1 FROM pg_indexes "
|
||||
"WHERE tablename = :t AND indexname = :i"
|
||||
), {"t": table, "i": index}).first()
|
||||
return row is not None
|
||||
|
||||
|
||||
def _col_is_timestamptz(conn, table: str, column: str) -> bool:
|
||||
row = conn.execute(sa.text(
|
||||
"SELECT data_type FROM information_schema.columns "
|
||||
"WHERE table_name = :t AND column_name = :c"
|
||||
), {"t": table, "c": column}).first()
|
||||
return row is not None and row[0] == "timestamp with time zone"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tables that inherit BaseModel and need an id UniqueConstraint.
|
||||
# ZeroTier tables were handled in 024; all others are listed here.
|
||||
# ---------------------------------------------------------------------------
|
||||
_LEGACY_TABLES = [
|
||||
"application_provider_configs",
|
||||
"audit_logs",
|
||||
"authentication_methods",
|
||||
"ca_permissions",
|
||||
"cas",
|
||||
"certificate_audit_logs",
|
||||
"department_cert_policies",
|
||||
"department_memberships",
|
||||
"department_principals",
|
||||
"departments",
|
||||
"email_verification_tokens",
|
||||
"external_provider_configs",
|
||||
"mfa_policy_compliance",
|
||||
"oauth_states",
|
||||
"oidc_audit_logs",
|
||||
"oidc_authorization_codes",
|
||||
"oidc_clients",
|
||||
"oidc_refresh_tokens",
|
||||
"oidc_sessions",
|
||||
# oidc_token_metadata intentionally excluded: its id column overrides
|
||||
# BaseModel without unique=True (JTI is the PK but not separately unique)
|
||||
"org_invite_tokens",
|
||||
"organization_api_keys",
|
||||
"organization_members",
|
||||
"organization_provider_overrides",
|
||||
"organization_security_policies",
|
||||
"organizations",
|
||||
"password_reset_tokens",
|
||||
"principal_memberships",
|
||||
"principals",
|
||||
"sessions",
|
||||
"ssh_certificates",
|
||||
"ssh_keys",
|
||||
"user_security_policies",
|
||||
"users",
|
||||
]
|
||||
|
||||
|
||||
def upgrade():
|
||||
conn = op.get_bind()
|
||||
|
||||
# ── 1. Add id UniqueConstraint to all legacy BaseModel tables ─────────
|
||||
for tbl in _LEGACY_TABLES:
|
||||
cname = f"{tbl}_id_key"
|
||||
if not _constraint_exists(conn, cname):
|
||||
op.create_unique_constraint(cname, tbl, ["id"])
|
||||
|
||||
# Drop the wrongly-added constraint on oidc_token_metadata if present
|
||||
# (its id column overrides BaseModel without unique=True)
|
||||
if _constraint_exists(conn, "oidc_token_metadata_id_key"):
|
||||
op.drop_constraint("oidc_token_metadata_id_key", "oidc_token_metadata", type_="unique")
|
||||
|
||||
# ── 2. organization_api_keys: timestamp columns TIMESTAMPTZ → TIMESTAMP
|
||||
for col in ("created_at", "updated_at", "deleted_at", "last_used_at", "revoked_at"):
|
||||
if _col_is_timestamptz(conn, "organization_api_keys", col):
|
||||
conn.execute(sa.text(
|
||||
f'ALTER TABLE organization_api_keys ALTER COLUMN "{col}" '
|
||||
f'TYPE TIMESTAMP WITHOUT TIME ZONE '
|
||||
f'USING CASE WHEN "{col}" IS NULL THEN NULL '
|
||||
f'ELSE "{col}" AT TIME ZONE \'UTC\' END'
|
||||
))
|
||||
|
||||
# ── 3. organization_api_keys: replace legacy unique constraint + indexes
|
||||
# Drop the anonymous unique constraint on key_hash (created by
|
||||
# sa.UniqueConstraint('key_hash') in the original migration)
|
||||
if _constraint_exists(conn, "organization_api_keys_key_hash_key"):
|
||||
op.drop_constraint(
|
||||
"organization_api_keys_key_hash_key",
|
||||
"organization_api_keys",
|
||||
type_="unique",
|
||||
)
|
||||
# Add named unique index for key_hash expected by the model
|
||||
if not _index_exists(conn, "organization_api_keys", "ix_organization_api_keys_key_hash"):
|
||||
op.create_index(
|
||||
"ix_organization_api_keys_key_hash",
|
||||
"organization_api_keys",
|
||||
["key_hash"],
|
||||
unique=True,
|
||||
)
|
||||
|
||||
# Drop the legacy plain org-id index (superseded by the named one below)
|
||||
if _index_exists(conn, "organization_api_keys", "idx_org_api_key_org_id"):
|
||||
op.drop_index("idx_org_api_key_org_id", table_name="organization_api_keys")
|
||||
|
||||
# Add named org-id index expected by the model
|
||||
if not _index_exists(conn, "organization_api_keys", "ix_organization_api_keys_organization_id"):
|
||||
op.create_index(
|
||||
"ix_organization_api_keys_organization_id",
|
||||
"organization_api_keys",
|
||||
["organization_id"],
|
||||
)
|
||||
|
||||
# Add named is_revoked index expected by the model
|
||||
if not _index_exists(conn, "organization_api_keys", "ix_organization_api_keys_is_revoked"):
|
||||
op.create_index(
|
||||
"ix_organization_api_keys_is_revoked",
|
||||
"organization_api_keys",
|
||||
["is_revoked"],
|
||||
)
|
||||
|
||||
# ── 4. Drop orphan idx_dept_can_sudo from departments ─────────────────
|
||||
if _index_exists(conn, "departments", "idx_dept_can_sudo"):
|
||||
op.drop_index("idx_dept_can_sudo", table_name="departments")
|
||||
|
||||
# NOTE: ix_ssh_certificates_serial uniqueness is handled in
|
||||
# 027_fix_cert_serial_uniqueness (composite unique per CA).
|
||||
|
||||
|
||||
def downgrade():
|
||||
conn = op.get_bind()
|
||||
|
||||
# Restore idx_dept_can_sudo
|
||||
if not _index_exists(conn, "departments", "idx_dept_can_sudo"):
|
||||
op.create_index("idx_dept_can_sudo", "departments", ["organization_id", "can_sudo"])
|
||||
|
||||
# Restore organization_api_keys indexes
|
||||
if _index_exists(conn, "organization_api_keys", "ix_organization_api_keys_is_revoked"):
|
||||
op.drop_index("ix_organization_api_keys_is_revoked", table_name="organization_api_keys")
|
||||
if _index_exists(conn, "organization_api_keys", "ix_organization_api_keys_organization_id"):
|
||||
op.drop_index("ix_organization_api_keys_organization_id", table_name="organization_api_keys")
|
||||
if _index_exists(conn, "organization_api_keys", "ix_organization_api_keys_key_hash"):
|
||||
op.drop_index("ix_organization_api_keys_key_hash", table_name="organization_api_keys")
|
||||
if not _constraint_exists(conn, "organization_api_keys_key_hash_key"):
|
||||
op.create_unique_constraint(
|
||||
"organization_api_keys_key_hash_key",
|
||||
"organization_api_keys",
|
||||
["key_hash"],
|
||||
)
|
||||
if not _index_exists(conn, "organization_api_keys", "idx_org_api_key_org_id"):
|
||||
op.create_index("idx_org_api_key_org_id", "organization_api_keys", ["organization_id"])
|
||||
|
||||
# Restore TIMESTAMPTZ on organization_api_keys
|
||||
for col in ("created_at", "updated_at", "deleted_at", "last_used_at", "revoked_at"):
|
||||
conn.execute(sa.text(
|
||||
f'ALTER TABLE organization_api_keys ALTER COLUMN "{col}" '
|
||||
f'TYPE TIMESTAMP WITH TIME ZONE '
|
||||
f'USING CASE WHEN "{col}" IS NULL THEN NULL '
|
||||
f'ELSE "{col}" AT TIME ZONE \'UTC\' END'
|
||||
))
|
||||
|
||||
# Drop id UniqueConstraints from legacy tables
|
||||
for tbl in reversed(_LEGACY_TABLES):
|
||||
cname = f"{tbl}_id_key"
|
||||
if _constraint_exists(conn, cname):
|
||||
op.drop_constraint(cname, tbl, type_="unique")
|
||||
Reference in New Issue
Block a user