366 lines
11 KiB
Python
366 lines
11 KiB
Python
"""add agent zones
|
|
|
|
Revision ID: 0016_agent_zones
|
|
Revises: 0015_miner_control_foundation
|
|
Create Date: 2026-08-25
|
|
"""
|
|
|
|
from typing import Sequence, Union
|
|
|
|
import sqlalchemy as sa
|
|
from alembic import op
|
|
from sqlalchemy.dialects import postgresql
|
|
|
|
|
|
revision: str = "0016_agent_zones"
|
|
down_revision: Union[str, None] = "0015_miner_control_foundation"
|
|
branch_labels: Union[str, Sequence[str], None] = None
|
|
depends_on: Union[str, Sequence[str], None] = None
|
|
|
|
DEFAULT_ZONE_NAME = "Default Zone"
|
|
|
|
|
|
def _format_rows(rows: list[tuple]) -> str:
|
|
return ", ".join("/".join(str(value) for value in row) for row in rows)
|
|
|
|
|
|
def upgrade() -> None:
|
|
op.create_table(
|
|
"agent_zones",
|
|
sa.Column(
|
|
"id",
|
|
postgresql.UUID(as_uuid=True),
|
|
primary_key=True,
|
|
server_default=sa.text("gen_random_uuid()"),
|
|
),
|
|
sa.Column(
|
|
"site_id",
|
|
postgresql.UUID(as_uuid=True),
|
|
sa.ForeignKey("sites.id", ondelete="RESTRICT"),
|
|
nullable=False,
|
|
),
|
|
sa.Column("name", sa.Text(), nullable=False),
|
|
sa.Column("description", sa.Text(), nullable=True),
|
|
sa.Column(
|
|
"is_enabled", sa.Boolean(), nullable=False, server_default=sa.text("true")
|
|
),
|
|
sa.Column(
|
|
"created_at",
|
|
sa.DateTime(timezone=True),
|
|
nullable=False,
|
|
server_default=sa.text("now()"),
|
|
),
|
|
sa.Column(
|
|
"updated_at",
|
|
sa.DateTime(timezone=True),
|
|
nullable=False,
|
|
server_default=sa.text("now()"),
|
|
),
|
|
sa.CheckConstraint(
|
|
"name = trim(name) AND trim(name) <> '' AND char_length(name) <= 120",
|
|
name="ck_agent_zones_name",
|
|
),
|
|
sa.UniqueConstraint("id", "site_id", name="uq_agent_zones_id_site_id"),
|
|
)
|
|
op.create_index("idx_agent_zones_site_id", "agent_zones", ["site_id"])
|
|
op.create_index(
|
|
"uq_agent_zones_active_site_name",
|
|
"agent_zones",
|
|
["site_id", sa.text("lower(name)")],
|
|
unique=True,
|
|
postgresql_where=sa.text("is_enabled"),
|
|
)
|
|
|
|
op.create_table(
|
|
"agent_zone_network_ranges",
|
|
sa.Column(
|
|
"id",
|
|
postgresql.UUID(as_uuid=True),
|
|
primary_key=True,
|
|
server_default=sa.text("gen_random_uuid()"),
|
|
),
|
|
sa.Column(
|
|
"agent_zone_id",
|
|
postgresql.UUID(as_uuid=True),
|
|
sa.ForeignKey("agent_zones.id", ondelete="CASCADE"),
|
|
nullable=False,
|
|
),
|
|
sa.Column("network", postgresql.CIDR(), nullable=False),
|
|
sa.Column(
|
|
"created_at",
|
|
sa.DateTime(timezone=True),
|
|
nullable=False,
|
|
server_default=sa.text("now()"),
|
|
),
|
|
sa.UniqueConstraint(
|
|
"agent_zone_id", "network", name="uq_agent_zone_network_ranges_zone_network"
|
|
),
|
|
)
|
|
op.create_index(
|
|
"idx_agent_zone_network_ranges_zone_id",
|
|
"agent_zone_network_ranges",
|
|
["agent_zone_id"],
|
|
)
|
|
|
|
op.add_column(
|
|
"miners",
|
|
sa.Column("agent_zone_id", postgresql.UUID(as_uuid=True), nullable=True),
|
|
)
|
|
op.create_check_constraint(
|
|
"ck_miners_agent_zone_requires_site",
|
|
"miners",
|
|
"agent_zone_id IS NULL OR site_id IS NOT NULL",
|
|
)
|
|
op.create_foreign_key(
|
|
"fk_miners_agent_zone_site",
|
|
"miners",
|
|
"agent_zones",
|
|
["agent_zone_id", "site_id"],
|
|
["id", "site_id"],
|
|
ondelete="RESTRICT",
|
|
)
|
|
op.create_index("idx_miners_agent_zone_id", "miners", ["agent_zone_id"])
|
|
|
|
op.add_column(
|
|
"minerops_site_agents",
|
|
sa.Column("agent_zone_id", postgresql.UUID(as_uuid=True), nullable=True),
|
|
)
|
|
op.create_check_constraint(
|
|
"ck_minerops_site_agents_zone_requires_site",
|
|
"minerops_site_agents",
|
|
"agent_zone_id IS NULL OR site_id IS NOT NULL",
|
|
)
|
|
op.create_foreign_key(
|
|
"fk_minerops_site_agents_zone_site",
|
|
"minerops_site_agents",
|
|
"agent_zones",
|
|
["agent_zone_id", "site_id"],
|
|
["id", "site_id"],
|
|
ondelete="RESTRICT",
|
|
)
|
|
|
|
op.add_column(
|
|
"minerops_agent_snapshots",
|
|
sa.Column("agent_zone_id", postgresql.UUID(as_uuid=True), nullable=True),
|
|
)
|
|
op.create_foreign_key(
|
|
"fk_minerops_agent_snapshots_agent_zone_id",
|
|
"minerops_agent_snapshots",
|
|
"agent_zones",
|
|
["agent_zone_id"],
|
|
["id"],
|
|
ondelete="RESTRICT",
|
|
)
|
|
op.create_index(
|
|
"idx_minerops_agent_snapshots_agent_zone_id",
|
|
"minerops_agent_snapshots",
|
|
["agent_zone_id"],
|
|
)
|
|
|
|
op.add_column(
|
|
"minerops_agent_discoveries",
|
|
sa.Column("agent_zone_id", postgresql.UUID(as_uuid=True), nullable=True),
|
|
)
|
|
op.create_foreign_key(
|
|
"fk_minerops_agent_discoveries_agent_zone_id",
|
|
"minerops_agent_discoveries",
|
|
"agent_zones",
|
|
["agent_zone_id"],
|
|
["id"],
|
|
ondelete="RESTRICT",
|
|
)
|
|
op.create_index(
|
|
"idx_minerops_agent_discoveries_agent_zone_id",
|
|
"minerops_agent_discoveries",
|
|
["agent_zone_id"],
|
|
)
|
|
|
|
connection = op.get_bind()
|
|
connection.execute(
|
|
sa.text(
|
|
"""
|
|
INSERT INTO agent_zones (site_id, name)
|
|
SELECT sites.id, :default_name
|
|
FROM sites
|
|
ORDER BY sites.id
|
|
"""
|
|
),
|
|
{"default_name": DEFAULT_ZONE_NAME},
|
|
)
|
|
connection.execute(
|
|
sa.text(
|
|
"""
|
|
UPDATE miners AS miner
|
|
SET agent_zone_id = zone.id
|
|
FROM agent_zones AS zone
|
|
WHERE miner.site_id = zone.site_id
|
|
AND zone.name = :default_name
|
|
AND miner.agent_zone_id IS NULL
|
|
"""
|
|
),
|
|
{"default_name": DEFAULT_ZONE_NAME},
|
|
)
|
|
connection.execute(
|
|
sa.text(
|
|
"""
|
|
UPDATE minerops_site_agents AS agent
|
|
SET agent_zone_id = zone.id
|
|
FROM agent_zones AS zone
|
|
WHERE agent.site_id = zone.site_id
|
|
AND zone.name = :default_name
|
|
AND agent.agent_zone_id IS NULL
|
|
"""
|
|
),
|
|
{"default_name": DEFAULT_ZONE_NAME},
|
|
)
|
|
|
|
connection.execute(
|
|
sa.text(
|
|
"""
|
|
UPDATE minerops_agent_snapshots AS snapshot
|
|
SET agent_zone_id = agent.agent_zone_id
|
|
FROM minerops_site_agents AS agent
|
|
WHERE snapshot.agent_id = agent.agent_id
|
|
AND agent.agent_zone_id IS NOT NULL
|
|
AND snapshot.agent_zone_id IS NULL
|
|
"""
|
|
)
|
|
)
|
|
connection.execute(
|
|
sa.text(
|
|
"""
|
|
UPDATE minerops_agent_discoveries AS discovery
|
|
SET agent_zone_id = agent.agent_zone_id
|
|
FROM minerops_site_agents AS agent
|
|
WHERE discovery.agent_id = agent.agent_id
|
|
AND agent.agent_zone_id IS NOT NULL
|
|
AND discovery.agent_zone_id IS NULL
|
|
"""
|
|
)
|
|
)
|
|
|
|
op.create_index(
|
|
"uq_minerops_site_agents_agent_zone_id",
|
|
"minerops_site_agents",
|
|
["agent_zone_id"],
|
|
unique=True,
|
|
postgresql_where=sa.text("agent_zone_id IS NOT NULL"),
|
|
)
|
|
op.drop_index(
|
|
"uq_minerops_site_agents_site_id", table_name="minerops_site_agents"
|
|
)
|
|
|
|
|
|
def downgrade() -> None:
|
|
connection = op.get_bind()
|
|
|
|
duplicate_site_agents = connection.execute(
|
|
sa.text(
|
|
"""
|
|
SELECT site_id::text, count(*)
|
|
FROM minerops_site_agents
|
|
WHERE site_id IS NOT NULL
|
|
GROUP BY site_id
|
|
HAVING count(*) > 1
|
|
ORDER BY site_id
|
|
"""
|
|
)
|
|
).all()
|
|
if duplicate_site_agents:
|
|
raise RuntimeError(
|
|
"Cannot downgrade 0016 to 0015: multiple Agent installations are "
|
|
"assigned within Site(s): "
|
|
f"{_format_rows(list(duplicate_site_agents))}. Move or unassign Agents "
|
|
"until each Site has at most one Agent before retrying; no Agent Zone "
|
|
"data was changed."
|
|
)
|
|
|
|
non_reducible_zones = connection.execute(
|
|
sa.text(
|
|
"""
|
|
SELECT site_id::text, count(*)
|
|
FROM agent_zones
|
|
GROUP BY site_id
|
|
HAVING count(*) > 1
|
|
OR bool_or(name <> :default_name OR description IS NOT NULL OR NOT is_enabled)
|
|
ORDER BY site_id
|
|
"""
|
|
),
|
|
{"default_name": DEFAULT_ZONE_NAME},
|
|
).all()
|
|
range_count = connection.execute(
|
|
sa.text("SELECT count(*) FROM agent_zone_network_ranges")
|
|
).scalar_one()
|
|
if non_reducible_zones or range_count:
|
|
detail = _format_rows(list(non_reducible_zones)) or "network ranges exist"
|
|
raise RuntimeError(
|
|
"Cannot downgrade 0016 to 0015: Agent Zone configuration cannot be "
|
|
f"represented by the one-Agent-per-Site model ({detail}). Remove "
|
|
"additional/custom Zones and Zone network ranges before retrying; no "
|
|
"Agent Zone data was changed."
|
|
)
|
|
|
|
op.create_index(
|
|
"uq_minerops_site_agents_site_id",
|
|
"minerops_site_agents",
|
|
["site_id"],
|
|
unique=True,
|
|
postgresql_where=sa.text("site_id IS NOT NULL"),
|
|
)
|
|
op.drop_index(
|
|
"uq_minerops_site_agents_agent_zone_id",
|
|
table_name="minerops_site_agents",
|
|
)
|
|
|
|
op.drop_index(
|
|
"idx_minerops_agent_discoveries_agent_zone_id",
|
|
table_name="minerops_agent_discoveries",
|
|
)
|
|
op.drop_constraint(
|
|
"fk_minerops_agent_discoveries_agent_zone_id",
|
|
"minerops_agent_discoveries",
|
|
type_="foreignkey",
|
|
)
|
|
op.drop_column("minerops_agent_discoveries", "agent_zone_id")
|
|
|
|
op.drop_index(
|
|
"idx_minerops_agent_snapshots_agent_zone_id",
|
|
table_name="minerops_agent_snapshots",
|
|
)
|
|
op.drop_constraint(
|
|
"fk_minerops_agent_snapshots_agent_zone_id",
|
|
"minerops_agent_snapshots",
|
|
type_="foreignkey",
|
|
)
|
|
op.drop_column("minerops_agent_snapshots", "agent_zone_id")
|
|
|
|
op.drop_constraint(
|
|
"fk_minerops_site_agents_zone_site",
|
|
"minerops_site_agents",
|
|
type_="foreignkey",
|
|
)
|
|
op.drop_constraint(
|
|
"ck_minerops_site_agents_zone_requires_site",
|
|
"minerops_site_agents",
|
|
type_="check",
|
|
)
|
|
op.drop_column("minerops_site_agents", "agent_zone_id")
|
|
|
|
op.drop_index("idx_miners_agent_zone_id", table_name="miners")
|
|
op.drop_constraint(
|
|
"fk_miners_agent_zone_site", "miners", type_="foreignkey"
|
|
)
|
|
op.drop_constraint(
|
|
"ck_miners_agent_zone_requires_site", "miners", type_="check"
|
|
)
|
|
op.drop_column("miners", "agent_zone_id")
|
|
|
|
op.drop_index(
|
|
"idx_agent_zone_network_ranges_zone_id",
|
|
table_name="agent_zone_network_ranges",
|
|
)
|
|
op.drop_table("agent_zone_network_ranges")
|
|
op.drop_index("uq_agent_zones_active_site_name", table_name="agent_zones")
|
|
op.drop_index("idx_agent_zones_site_id", table_name="agent_zones")
|
|
op.drop_table("agent_zones")
|