Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -0,0 +1,164 @@
"""Add tenant scaffolding and targets schema.

Additive, behavior-preserving migration:

- Adds opaque string ``tenant_id`` columns to ``agents``, ``controls``,
``policies``, ``agent_controls``, and ``agent_policies``. Existing rows are
backfilled to ``default-tenant`` and columns are then made NOT NULL. A
DB-level ``server_default`` keeps writes that omit a tenant working.
- Creates new tables ``targets`` and ``target_controls``. Uniqueness on
``targets`` covers ``(tenant_id, target_type, external_id)`` and on
``target_controls`` covers ``(target_id, control_id)``.
``target_controls.target_id`` uses ``ON DELETE CASCADE`` because the
attachment has no meaning without its target; ``control_id`` uses the
default restrictive behavior so control deletion does not silently cascade
into attachment cleanup.
- Intentionally omitted from this migration (to be addressed separately):
* ``policy_controls.tenant_id`` (tenant scope inherited transitively
through ``policy_id`` and ``control_id``).
* ``control_execution_events.tenant_id`` (observability tables out of
scope here).
* ``updated_at`` columns (no established auto-maintenance pattern in the
repo yet).
* Indexes on the new ``tenant_id`` columns (read paths do not filter on
tenant yet, so unused indexes would just add write cost).

Revision ID: 7c1a9b4e2d30
Revises: 5f2b5f4e1a90
Create Date: 2026-04-20 00:00:00.000000

"""
from alembic import op
import sqlalchemy as sa


revision = "7c1a9b4e2d30"
down_revision = "5f2b5f4e1a90"
branch_labels = None
depends_on = None


DEFAULT_TENANT_ID = "default-tenant"

_TENANT_SCOPED_TABLES = (
"agents",
"controls",
"policies",
"agent_controls",
"agent_policies",
)


def upgrade() -> None:
# Step 1: add tenant_id as nullable on all affected tables.
for table in _TENANT_SCOPED_TABLES:
op.add_column(
table,
sa.Column("tenant_id", sa.String(length=64), nullable=True),
)

# Step 2: backfill existing rows to the synthetic default tenant.
for table in _TENANT_SCOPED_TABLES:
op.execute(
sa.text(
f"UPDATE {table} SET tenant_id = :tenant WHERE tenant_id IS NULL"
).bindparams(tenant=DEFAULT_TENANT_ID)
)

# Step 3: make tenant_id NOT NULL and install the DB-level default so
# unscoped OSS writes continue to land in the default tenant automatically.
for table in _TENANT_SCOPED_TABLES:
op.alter_column(
table,
"tenant_id",
existing_type=sa.String(length=64),
nullable=False,
server_default=sa.text(f"'{DEFAULT_TENANT_ID}'"),
)

# Step 4: create the new target schema objects.
op.create_table(
"targets",
sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
sa.Column(
"tenant_id",
sa.String(length=64),
nullable=False,
server_default=sa.text(f"'{DEFAULT_TENANT_ID}'"),
),
sa.Column("target_type", sa.String(length=64), nullable=False),
sa.Column("external_id", sa.String(length=255), nullable=False),
sa.Column("name", sa.String(length=255), nullable=True),
sa.Column(
"data",
sa.dialects.postgresql.JSONB(astext_type=sa.Text()),
server_default=sa.text("'{}'::jsonb"),
nullable=False,
),
sa.Column(
"created_at",
sa.DateTime(timezone=True),
server_default=sa.text("CURRENT_TIMESTAMP"),
nullable=False,
),
sa.PrimaryKeyConstraint("id"),
sa.UniqueConstraint(
"tenant_id",
"target_type",
"external_id",
name="uq_targets_tenant_type_external_id",
),
)

op.create_table(
"target_controls",
sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
sa.Column("target_id", sa.Integer(), nullable=False),
sa.Column("control_id", sa.Integer(), nullable=False),
sa.Column(
"enabled",
sa.Boolean(),
nullable=False,
server_default=sa.text("true"),
),
sa.Column(
"created_at",
sa.DateTime(timezone=True),
server_default=sa.text("CURRENT_TIMESTAMP"),
nullable=False,
),
sa.ForeignKeyConstraint(
["target_id"], ["targets.id"], ondelete="CASCADE"
),
sa.ForeignKeyConstraint(["control_id"], ["controls.id"]),
sa.PrimaryKeyConstraint("id"),
sa.UniqueConstraint(
"target_id", "control_id", name="uq_target_controls_target_control"
),
)
op.create_index(
op.f("ix_target_controls_target_id"),
"target_controls",
["target_id"],
unique=False,
)
op.create_index(
op.f("ix_target_controls_control_id"),
"target_controls",
["control_id"],
unique=False,
)


def downgrade() -> None:
op.drop_index(
op.f("ix_target_controls_control_id"), table_name="target_controls"
)
op.drop_index(
op.f("ix_target_controls_target_id"), table_name="target_controls"
)
op.drop_table("target_controls")
op.drop_table("targets")

for table in _TENANT_SCOPED_TABLES:
op.drop_column(table, "tenant_id")
119 changes: 118 additions & 1 deletion server/src/agent_control_server/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from agent_control_models.server import EvaluatorSchema
from pydantic import Field
from sqlalchemy import (
Boolean,
CheckConstraint,
Column,
DateTime,
Expand All @@ -14,13 +15,19 @@
Integer,
String,
Table,
UniqueConstraint,
text,
)
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Mapped, mapped_column, relationship, validates

from .db import Base

# Synthetic tenant used when no explicit tenant is resolved for a request.
# In this initial rollout tenant_id is inert metadata on existing tables:
# writes stamp it via an ORM/DB default, but read paths do not filter on it.
DEFAULT_TENANT_ID = "default-tenant"


class AgentData(BaseModel):
"""Agent metadata stored in JSONB."""
Expand All @@ -30,7 +37,10 @@ class AgentData(BaseModel):
evaluators: list[EvaluatorSchema] = Field(default_factory=list)


# Association table for Policy <> Control many-to-many relationship
# Association table for Policy <> Control many-to-many relationship.
# ``policy_controls`` deliberately does not carry tenant_id: tenant scope is
# inherited transitively through policy_id and control_id, both of which
# already point to tenant-owned rows.
policy_controls: Table = Table(
"policy_controls",
Base.metadata,
Expand All @@ -44,6 +54,13 @@ class AgentData(BaseModel):
Base.metadata,
Column("agent_name", ForeignKey("agents.name"), primary_key=True, index=True),
Column("policy_id", ForeignKey("policies.id"), primary_key=True, index=True),
Column(
"tenant_id",
String(64),
nullable=False,
server_default=text(f"'{DEFAULT_TENANT_ID}'"),
default=DEFAULT_TENANT_ID,
),
)

# Association table for Agent <> Control many-to-many direct relationship
Expand All @@ -52,6 +69,13 @@ class AgentData(BaseModel):
Base.metadata,
Column("agent_name", ForeignKey("agents.name"), primary_key=True, index=True),
Column("control_id", ForeignKey("controls.id"), primary_key=True, index=True),
Column(
"tenant_id",
String(64),
nullable=False,
server_default=text(f"'{DEFAULT_TENANT_ID}'"),
default=DEFAULT_TENANT_ID,
),
)


Expand All @@ -60,6 +84,12 @@ class Policy(Base):

id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
name: Mapped[str] = mapped_column(String(255), nullable=False, unique=True)
tenant_id: Mapped[str] = mapped_column(
String(64),
nullable=False,
server_default=text(f"'{DEFAULT_TENANT_ID}'"),
default=DEFAULT_TENANT_ID,
)
agents: Mapped[list["Agent"]] = relationship(
"Agent", secondary=lambda: agent_policies, back_populates="policies"
)
Expand All @@ -74,6 +104,12 @@ class Control(Base):

id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
name: Mapped[str] = mapped_column(String(255), nullable=False, unique=True)
tenant_id: Mapped[str] = mapped_column(
String(64),
nullable=False,
server_default=text(f"'{DEFAULT_TENANT_ID}'"),
default=DEFAULT_TENANT_ID,
)
# JSONB payload describing control specifics
data: Mapped[dict[str, Any]] = mapped_column(
JSONB, server_default=text("'{}'::jsonb"), nullable=False
Expand All @@ -96,6 +132,12 @@ class Agent(Base):
)

name: Mapped[str] = mapped_column(String(255), primary_key=True)
tenant_id: Mapped[str] = mapped_column(
String(64),
nullable=False,
server_default=text(f"'{DEFAULT_TENANT_ID}'"),
default=DEFAULT_TENANT_ID,
)
data: Mapped[dict[str, Any]] = mapped_column(
JSONB, server_default=text("'{}'::jsonb"), nullable=False
)
Expand All @@ -114,6 +156,81 @@ def _normalize_name(self, _key: str, value: str) -> str:
return normalize_agent_name(value)


# =============================================================================
# Target Models
# =============================================================================
#
# Targets are typed, tenant-scoped attachable objects. The schema is introduced
# here without being wired into runtime control resolution or management APIs;
# both are added in follow-up changes. ``target_controls`` inherits tenant
# scope transitively through ``target_id``.


class Target(Base):
"""A typed, tenant-scoped attachable object (e.g. ``environment``).

The column is named ``target_type`` rather than ``type`` to avoid
shadowing Python's builtin and to keep greps for the field specific.
"""

__tablename__ = "targets"
__table_args__ = (
UniqueConstraint(
"tenant_id",
"target_type",
"external_id",
name="uq_targets_tenant_type_external_id",
),
)

id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
tenant_id: Mapped[str] = mapped_column(
String(64),
nullable=False,
server_default=text(f"'{DEFAULT_TENANT_ID}'"),
default=DEFAULT_TENANT_ID,
)
target_type: Mapped[str] = mapped_column(String(64), nullable=False)
external_id: Mapped[str] = mapped_column(String(255), nullable=False)
name: Mapped[str | None] = mapped_column(String(255), nullable=True)
data: Mapped[dict[str, Any]] = mapped_column(
JSONB, server_default=text("'{}'::jsonb"), nullable=False
)
created_at: Mapped[dt.datetime] = mapped_column(
DateTime(timezone=True),
server_default=text("CURRENT_TIMESTAMP"),
nullable=False,
)


class TargetControl(Base):
"""Attachment of a control to a target with per-target enablement."""

__tablename__ = "target_controls"
__table_args__ = (
UniqueConstraint("target_id", "control_id", name="uq_target_controls_target_control"),
)

id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
# CASCADE on target_id: a target_control row has no meaning without its target.
target_id: Mapped[int] = mapped_column(
Integer, ForeignKey("targets.id", ondelete="CASCADE"), nullable=False, index=True
)
# RESTRICT (default) on control_id: do not silently fan control deletes into
# attachment cleanup; callers must remove attachments explicitly.
control_id: Mapped[int] = mapped_column(
Integer, ForeignKey("controls.id"), nullable=False, index=True
)
enabled: Mapped[bool] = mapped_column(
Boolean, nullable=False, server_default=text("true"), default=True
)
created_at: Mapped[dt.datetime] = mapped_column(
DateTime(timezone=True),
server_default=text("CURRENT_TIMESTAMP"),
nullable=False,
)


# =============================================================================
# Observability Models
# =============================================================================
Expand Down
Loading
Loading