217 lines
8.5 KiB
Python
217 lines
8.5 KiB
Python
"""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")
|