292 lines
13 KiB
Python
292 lines
13 KiB
Python
"""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}'
|
|
))
|