"""Fix ZeroTier table schema: enum types, unique constraints, indexes, drop cert_token. Revision ID: 024_fix_zerotier_schema Revises: 023_zerotier_drop_legacy Create Date: 2026-03-22 Addresses all `db check` differences after 023: - Cast VARCHAR(40) enum columns to their proper PostgreSQL enum types (guarded — skipped if columns are already native enum, e.g. on a fresh DB where 020_zerotier created them correctly) - Replace partial unique indexes with named UniqueConstraints - Fix devices.node_id partial index -> plain index - Add UniqueConstraint on `id` for all new ZeroTier tables (BaseModel.unique=True) - Drop orphan cert_token column and its index from ssh_certificates """ from alembic import op import sqlalchemy as sa revision = "024_fix_zerotier_schema" down_revision = "023_zerotier_drop_legacy" branch_labels = None depends_on = None # --------------------------------------------------------------------------- # helpers # --------------------------------------------------------------------------- def _col_data_type(conn, table: str, column: str) -> str | None: """Return the PostgreSQL data_type string for a column, or None.""" 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[0] if row else None def _column_exists(conn, table: str, column: str) -> bool: return _col_data_type(conn, table, column) is not None def _index_exists(conn, table: str, index: str) -> bool: from sqlalchemy.engine.reflection import Inspector insp = Inspector.from_engine(conn) return any(i["name"] == index for i in insp.get_indexes(table)) def _constraint_exists(conn, constraint: str) -> bool: row = conn.execute(sa.text( "SELECT 1 FROM information_schema.table_constraints " "WHERE constraint_name = :c" ), {"c": constraint}).first() return row is not None def upgrade(): conn = op.get_bind() # ------------------------------------------------------------------------- # 1. Cast VARCHAR(40) enum columns to proper PostgreSQL enum types. # GUARDED: On a fresh DB, 020_zerotier already created these as native # enum types. We only cast if the column is currently 'character varying'. # ------------------------------------------------------------------------- enum_casts = [ ("portal_networks", "environment", "network_environment", None), ("portal_networks", "request_mode", "network_request_mode", None), ("devices", "status", "device_status", "'active'::device_status"), ("device_network_memberships", "state", "membership_state", "'pending_device_registration'::membership_state"), ("user_network_approvals", "grant_type", "approval_grant_type", "'requested'::approval_grant_type"), ("user_network_approvals", "state", "approval_state", "'pending'::approval_state"), ("activation_sessions", "end_reason", "activation_end_reason", None), ("kill_switch_events", "scope", "kill_switch_scope", "'organization'::kill_switch_scope"), ] for table, col, enum_type, new_default in enum_casts: dtype = _col_data_type(conn, table, col) if dtype == "character varying": conn.execute(sa.text(f'ALTER TABLE "{table}" ALTER COLUMN "{col}" DROP DEFAULT')) conn.execute(sa.text( f'ALTER TABLE "{table}" ALTER COLUMN "{col}" TYPE {enum_type} ' f'USING "{col}"::text::{enum_type}' )) if new_default: conn.execute(sa.text( f'ALTER TABLE "{table}" ALTER COLUMN "{col}" SET DEFAULT {new_default}' )) elif dtype == "USER-DEFINED" and new_default: # Already native enum (fresh DB path). Ensure server_default is set # if 020 used `default=` (Python-side) instead of `server_default=`. # This is harmless — SET DEFAULT is idempotent. conn.execute(sa.text( f'ALTER TABLE "{table}" ALTER COLUMN "{col}" SET DEFAULT {new_default}' )) # ------------------------------------------------------------------------- # 2. portal_networks: drop partial unique index, add named UniqueConstraint # ------------------------------------------------------------------------- if _index_exists(conn, "portal_networks", "ix_portal_networks_org_zt"): op.drop_index("ix_portal_networks_org_zt", table_name="portal_networks") if not _constraint_exists(conn, "uix_org_zt_network_id"): op.create_unique_constraint( "uix_org_zt_network_id", "portal_networks", ["organization_id", "zerotier_network_id"], ) # ------------------------------------------------------------------------- # 3. device_network_memberships: drop partial unique index, add named UC # ------------------------------------------------------------------------- if _index_exists(conn, "device_network_memberships", "ix_device_network_memberships_device_network"): op.drop_index("ix_device_network_memberships_device_network", table_name="device_network_memberships") if not _constraint_exists(conn, "uix_device_network"): op.create_unique_constraint( "uix_device_network", "device_network_memberships", ["device_id", "portal_network_id", "deleted_at"], ) # ------------------------------------------------------------------------- # 4. user_network_approvals: drop partial unique index, add named UC # ------------------------------------------------------------------------- if _index_exists(conn, "user_network_approvals", "ix_user_network_approvals_user_network"): op.drop_index("ix_user_network_approvals_user_network", table_name="user_network_approvals") if not _constraint_exists(conn, "uix_user_network_approval"): op.create_unique_constraint( "uix_user_network_approval", "user_network_approvals", ["user_id", "portal_network_id", "deleted_at"], ) # ------------------------------------------------------------------------- # 5. zerotier_memberships: drop index, add named UniqueConstraint # ------------------------------------------------------------------------- if _index_exists(conn, "zerotier_memberships", "ix_zerotier_memberships_network_node"): op.drop_index("ix_zerotier_memberships_network_node", table_name="zerotier_memberships") if not _constraint_exists(conn, "uix_zt_network_node"): op.create_unique_constraint( "uix_zt_network_node", "zerotier_memberships", ["zerotier_network_id", "node_id"], ) # ------------------------------------------------------------------------- # 6. devices.node_id: drop partial unique index, add plain non-unique index # ------------------------------------------------------------------------- if _index_exists(conn, "devices", "ix_devices_node_id_active"): op.drop_index("ix_devices_node_id_active", table_name="devices") if not _index_exists(conn, "devices", "ix_devices_node_id"): op.create_index("ix_devices_node_id", "devices", ["node_id"]) # ------------------------------------------------------------------------- # 7. Add UniqueConstraint on `id` for all ZeroTier tables # BaseModel defines id with unique=True → separate _id_key constraint. # ------------------------------------------------------------------------- zt_tables = [ "portal_networks", "devices", "device_network_memberships", "user_network_approvals", "activation_sessions", "zerotier_memberships", "kill_switch_events", ] for tbl in zt_tables: cname = f"{tbl}_id_key" if not _constraint_exists(conn, cname): op.create_unique_constraint(cname, tbl, ["id"]) # ------------------------------------------------------------------------- # 8. Drop orphan cert_token column and its index from ssh_certificates. # cert_token was created by 3de11c5dc2d5 but the SSHCertificate model # never uses it. Guarded in case a future revision removes it first. # ------------------------------------------------------------------------- if _index_exists(conn, "ssh_certificates", "ix_ssh_certificates_cert_token"): op.drop_index("ix_ssh_certificates_cert_token", table_name="ssh_certificates") if _column_exists(conn, "ssh_certificates", "cert_token"): op.drop_column("ssh_certificates", "cert_token") def downgrade(): conn = op.get_bind() # Restore cert_token if it was dropped if not _column_exists(conn, "ssh_certificates", "cert_token"): op.add_column( "ssh_certificates", sa.Column("cert_token", sa.String(64), nullable=True), ) if not _index_exists(conn, "ssh_certificates", "ix_ssh_certificates_cert_token"): op.create_index( "ix_ssh_certificates_cert_token", "ssh_certificates", ["cert_token"], unique=True, ) # Drop id unique constraints on ZeroTier tables zt_tables = [ "portal_networks", "devices", "device_network_memberships", "user_network_approvals", "activation_sessions", "zerotier_memberships", "kill_switch_events", ] for tbl in zt_tables: cname = f"{tbl}_id_key" if _constraint_exists(conn, cname): op.drop_constraint(cname, tbl, type_="unique") # Restore devices node_id index if _index_exists(conn, "devices", "ix_devices_node_id"): op.drop_index("ix_devices_node_id", table_name="devices") if not _index_exists(conn, "devices", "ix_devices_node_id_active"): op.create_index( "ix_devices_node_id_active", "devices", ["node_id"], unique=True, postgresql_where=sa.text("deleted_at IS NULL"), ) # Restore zerotier_memberships index if _constraint_exists(conn, "uix_zt_network_node"): op.drop_constraint("uix_zt_network_node", "zerotier_memberships", type_="unique") if not _index_exists(conn, "zerotier_memberships", "ix_zerotier_memberships_network_node"): op.create_index( "ix_zerotier_memberships_network_node", "zerotier_memberships", ["zerotier_network_id", "node_id"], unique=True, ) # Restore user_network_approvals partial unique index if _constraint_exists(conn, "uix_user_network_approval"): op.drop_constraint("uix_user_network_approval", "user_network_approvals", type_="unique") if not _index_exists(conn, "user_network_approvals", "ix_user_network_approvals_user_network"): op.create_index( "ix_user_network_approvals_user_network", "user_network_approvals", ["user_id", "portal_network_id"], unique=True, postgresql_where=sa.text("deleted_at IS NULL"), ) # Restore device_network_memberships partial unique index if _constraint_exists(conn, "uix_device_network"): op.drop_constraint("uix_device_network", "device_network_memberships", type_="unique") if not _index_exists(conn, "device_network_memberships", "ix_device_network_memberships_device_network"): op.create_index( "ix_device_network_memberships_device_network", "device_network_memberships", ["device_id", "portal_network_id"], unique=True, postgresql_where=sa.text("deleted_at IS NULL"), ) # Restore portal_networks partial unique index if _constraint_exists(conn, "uix_org_zt_network_id"): op.drop_constraint("uix_org_zt_network_id", "portal_networks", type_="unique") if not _index_exists(conn, "portal_networks", "ix_portal_networks_org_zt"): op.create_index( "ix_portal_networks_org_zt", "portal_networks", ["organization_id", "zerotier_network_id"], unique=True, postgresql_where=sa.text("deleted_at IS NULL"), ) # Cast enum columns back to VARCHAR(40) — only if currently native enum enum_casts = [ ("portal_networks", "environment", "'development'::character varying"), ("portal_networks", "request_mode", "'approval_required'::character varying"), ("devices", "status", "'active'::character varying"), ("device_network_memberships", "state", "'pending_device_registration'::character varying"), ("user_network_approvals", "grant_type", "'requested'::character varying"), ("user_network_approvals", "state", "'pending'::character varying"), ("activation_sessions", "end_reason", None), ("kill_switch_events", "scope", "'organization'::character varying"), ] for table, col, old_default in enum_casts: conn.execute(sa.text(f'ALTER TABLE "{table}" ALTER COLUMN "{col}" DROP DEFAULT')) conn.execute(sa.text( f'ALTER TABLE "{table}" ALTER COLUMN "{col}" TYPE VARCHAR(40) ' f'USING "{col}"::text' )) if old_default: conn.execute(sa.text( f'ALTER TABLE "{table}" ALTER COLUMN "{col}" SET DEFAULT {old_default}' ))