feat: add agent zones
This commit is contained in:
lussierj committed
1 parent
d8cd7ad058
commit
82309db79a
29 files changed
+2242
-255
No files matched your search
@@ -0,0 +1,365 @@
|
||||
"""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")
|
||||
Reference in new issue
Block a user