Fix: DB Migration
This commit is contained in:
@@ -0,0 +1,291 @@
|
||||
"""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}'
|
||||
))
|
||||
Reference in New Issue
Block a user