"""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")