445 lines
14 KiB
Python
445 lines
14 KiB
Python
|
"""add structural tenant and platform scope to Agent assets
|
||
|
|
|
||
|
|
Revision ID: 20260717_0026
|
||
|
|
Revises: 20260717_0025
|
||
|
|
Create Date: 2026-07-17 14:30:00
|
||
|
|
"""
|
||
|
|
|
||
|
|
from collections.abc import Sequence
|
||
|
|
|
||
|
|
import sqlalchemy as sa
|
||
|
|
|
||
|
|
from alembic import op
|
||
|
|
|
||
|
|
revision: str = "20260717_0026"
|
||
|
|
down_revision: str | None = "20260717_0025"
|
||
|
|
branch_labels: str | Sequence[str] | None = None
|
||
|
|
depends_on: str | Sequence[str] | None = None
|
||
|
|
|
||
|
|
_ASSET_TABLE = "agent_assets"
|
||
|
|
_ASSET_CHILD_TABLES = (
|
||
|
|
"agent_asset_versions",
|
||
|
|
"agent_asset_reviews",
|
||
|
|
"agent_asset_test_runs",
|
||
|
|
"agent_asset_rule_feedback",
|
||
|
|
)
|
||
|
|
_OWNERSHIP_CHILD_TABLES = (
|
||
|
|
"agent_asset_versions",
|
||
|
|
"agent_asset_reviews",
|
||
|
|
)
|
||
|
|
_ONLYOFFICE_SESSION_TABLE = "agent_asset_onlyoffice_sessions"
|
||
|
|
|
||
|
|
|
||
|
|
def _require_postgresql() -> None:
|
||
|
|
dialect_name = op.get_bind().dialect.name
|
||
|
|
if dialect_name != "postgresql":
|
||
|
|
raise RuntimeError(
|
||
|
|
"20260717_0026 only supports PostgreSQL; "
|
||
|
|
f"refusing to mutate {dialect_name} without transactional constraint DDL"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _has_complete_agent_asset_schema() -> bool:
|
||
|
|
inspector = sa.inspect(op.get_bind())
|
||
|
|
table_names = (_ASSET_TABLE, *_ASSET_CHILD_TABLES)
|
||
|
|
existing = {table_name for table_name in table_names if inspector.has_table(table_name)}
|
||
|
|
if not existing:
|
||
|
|
return False
|
||
|
|
missing = set(table_names) - existing
|
||
|
|
if missing:
|
||
|
|
raise RuntimeError(
|
||
|
|
"cannot migrate partial Agent asset schema; missing tables: "
|
||
|
|
+ ", ".join(sorted(missing))
|
||
|
|
)
|
||
|
|
return True
|
||
|
|
|
||
|
|
|
||
|
|
def _create_onlyoffice_session_table() -> None:
|
||
|
|
if sa.inspect(op.get_bind()).has_table(_ONLYOFFICE_SESSION_TABLE):
|
||
|
|
return
|
||
|
|
op.create_table(
|
||
|
|
_ONLYOFFICE_SESSION_TABLE,
|
||
|
|
sa.Column("jti", sa.String(length=36), nullable=False),
|
||
|
|
sa.Column("tenant_id", sa.String(length=64), nullable=False),
|
||
|
|
sa.Column("resource_scope", sa.String(length=16), nullable=False),
|
||
|
|
sa.Column("asset_id", sa.String(length=100), nullable=False),
|
||
|
|
sa.Column("document_key", sa.String(length=200), nullable=False),
|
||
|
|
sa.Column("document_version", sa.String(length=30), nullable=False),
|
||
|
|
sa.Column("document_fingerprint", sa.String(length=160), nullable=False),
|
||
|
|
sa.Column("audience", sa.String(length=80), nullable=False),
|
||
|
|
sa.Column("writable", sa.Boolean(), nullable=False),
|
||
|
|
sa.Column("status", sa.String(length=16), nullable=False),
|
||
|
|
sa.Column("actor", sa.String(length=160), nullable=False),
|
||
|
|
sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False),
|
||
|
|
sa.Column("claimed_at", sa.DateTime(timezone=True), nullable=True),
|
||
|
|
sa.Column("consumed_at", sa.DateTime(timezone=True), nullable=True),
|
||
|
|
sa.Column("failure_reason", sa.Text(), nullable=False),
|
||
|
|
sa.Column(
|
||
|
|
"created_at",
|
||
|
|
sa.DateTime(timezone=True),
|
||
|
|
nullable=False,
|
||
|
|
server_default=sa.func.now(),
|
||
|
|
),
|
||
|
|
sa.CheckConstraint(
|
||
|
|
"(resource_scope = 'platform' AND tenant_id = 'platform') OR "
|
||
|
|
"(resource_scope = 'tenant' AND tenant_id <> 'platform')",
|
||
|
|
name="ck_agent_asset_onlyoffice_sessions_scope_tenant",
|
||
|
|
),
|
||
|
|
sa.CheckConstraint(
|
||
|
|
"status IN ('active', 'processing', 'consumed', 'failed', 'revoked')",
|
||
|
|
name="ck_agent_asset_onlyoffice_sessions_status",
|
||
|
|
),
|
||
|
|
sa.CheckConstraint(
|
||
|
|
"(status = 'active' AND claimed_at IS NULL AND consumed_at IS NULL) OR "
|
||
|
|
"(status IN ('processing', 'failed') AND claimed_at IS NOT NULL "
|
||
|
|
"AND consumed_at IS NULL) OR "
|
||
|
|
"(status = 'consumed' AND claimed_at IS NOT NULL AND consumed_at IS NOT NULL) OR "
|
||
|
|
"(status = 'revoked' AND consumed_at IS NULL)",
|
||
|
|
name="ck_agent_asset_onlyoffice_sessions_lifecycle",
|
||
|
|
),
|
||
|
|
sa.ForeignKeyConstraint(
|
||
|
|
["tenant_id"],
|
||
|
|
["tenants.tenant_id"],
|
||
|
|
name="fk_agent_asset_onlyoffice_sessions_tenant",
|
||
|
|
ondelete="CASCADE",
|
||
|
|
),
|
||
|
|
sa.PrimaryKeyConstraint("jti"),
|
||
|
|
)
|
||
|
|
op.create_index(
|
||
|
|
"ix_agent_asset_onlyoffice_sessions_tenant_asset",
|
||
|
|
_ONLYOFFICE_SESSION_TABLE,
|
||
|
|
["tenant_id", "resource_scope", "asset_id", "created_at"],
|
||
|
|
)
|
||
|
|
op.create_index(
|
||
|
|
"ix_agent_asset_onlyoffice_sessions_status_expiry",
|
||
|
|
_ONLYOFFICE_SESSION_TABLE,
|
||
|
|
["status", "expires_at"],
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _add_scope_columns(table_name: str) -> None:
|
||
|
|
op.add_column(table_name, sa.Column("tenant_id", sa.String(length=64), nullable=True))
|
||
|
|
op.add_column(table_name, sa.Column("scope", sa.String(length=16), nullable=True))
|
||
|
|
|
||
|
|
|
||
|
|
def _assert_known_asset_tenants() -> None:
|
||
|
|
unknown = int(
|
||
|
|
op.get_bind().scalar(
|
||
|
|
sa.text(
|
||
|
|
"""
|
||
|
|
SELECT COUNT(*)
|
||
|
|
FROM agent_assets AS asset
|
||
|
|
LEFT JOIN tenants AS tenant
|
||
|
|
ON tenant.tenant_id = NULLIF(BTRIM(asset.config_json ->> 'tenant_id'), '')
|
||
|
|
WHERE NULLIF(BTRIM(asset.config_json ->> 'tenant_id'), '') IS NOT NULL
|
||
|
|
AND tenant.tenant_id IS NULL
|
||
|
|
"""
|
||
|
|
)
|
||
|
|
)
|
||
|
|
or 0
|
||
|
|
)
|
||
|
|
if unknown:
|
||
|
|
raise RuntimeError(
|
||
|
|
"cannot migrate Agent assets: config_json contains tenant ids absent from tenants "
|
||
|
|
f"(unknown_assets={unknown})"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _backfill_scope() -> None:
|
||
|
|
_assert_known_asset_tenants()
|
||
|
|
op.execute(
|
||
|
|
"""
|
||
|
|
UPDATE agent_assets
|
||
|
|
SET tenant_id = COALESCE(
|
||
|
|
NULLIF(BTRIM(config_json ->> 'tenant_id'), ''),
|
||
|
|
'platform'
|
||
|
|
),
|
||
|
|
scope = CASE
|
||
|
|
WHEN NULLIF(BTRIM(config_json ->> 'tenant_id'), '') IS NULL
|
||
|
|
THEN 'platform'
|
||
|
|
ELSE 'tenant'
|
||
|
|
END
|
||
|
|
"""
|
||
|
|
)
|
||
|
|
for table_name in _OWNERSHIP_CHILD_TABLES:
|
||
|
|
op.execute(
|
||
|
|
f"""
|
||
|
|
UPDATE {table_name} AS child
|
||
|
|
SET tenant_id = asset.tenant_id,
|
||
|
|
scope = asset.scope
|
||
|
|
FROM agent_assets AS asset
|
||
|
|
WHERE child.asset_id = asset.id
|
||
|
|
"""
|
||
|
|
)
|
||
|
|
op.execute(
|
||
|
|
"""
|
||
|
|
UPDATE agent_asset_test_runs AS child
|
||
|
|
SET tenant_id = COALESCE(
|
||
|
|
NULLIF(BTRIM(child.input_json ->> 'target_tenant_id'), ''),
|
||
|
|
NULLIF(BTRIM(asset.config_json ->> 'tenant_id'), ''),
|
||
|
|
asset.tenant_id
|
||
|
|
),
|
||
|
|
scope = CASE
|
||
|
|
WHEN COALESCE(
|
||
|
|
NULLIF(BTRIM(child.input_json ->> 'target_tenant_id'), ''),
|
||
|
|
NULLIF(BTRIM(asset.config_json ->> 'tenant_id'), ''),
|
||
|
|
asset.tenant_id
|
||
|
|
) = 'platform' THEN 'platform'
|
||
|
|
ELSE 'tenant'
|
||
|
|
END
|
||
|
|
FROM agent_assets AS asset
|
||
|
|
WHERE child.asset_id = asset.id
|
||
|
|
"""
|
||
|
|
)
|
||
|
|
op.execute(
|
||
|
|
"""
|
||
|
|
UPDATE agent_asset_rule_feedback AS child
|
||
|
|
SET tenant_id = COALESCE(
|
||
|
|
NULLIF(BTRIM(asset.config_json ->> 'tenant_id'), ''),
|
||
|
|
asset.tenant_id
|
||
|
|
),
|
||
|
|
scope = CASE
|
||
|
|
WHEN COALESCE(
|
||
|
|
NULLIF(BTRIM(asset.config_json ->> 'tenant_id'), ''),
|
||
|
|
asset.tenant_id
|
||
|
|
) = 'platform' THEN 'platform'
|
||
|
|
ELSE 'tenant'
|
||
|
|
END
|
||
|
|
FROM agent_assets AS asset
|
||
|
|
WHERE child.asset_id = asset.id
|
||
|
|
"""
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _set_scope_not_null_and_defaults(table_name: str) -> None:
|
||
|
|
op.alter_column(
|
||
|
|
table_name,
|
||
|
|
"tenant_id",
|
||
|
|
existing_type=sa.String(length=64),
|
||
|
|
nullable=False,
|
||
|
|
server_default="platform",
|
||
|
|
)
|
||
|
|
op.alter_column(
|
||
|
|
table_name,
|
||
|
|
"scope",
|
||
|
|
existing_type=sa.String(length=16),
|
||
|
|
nullable=False,
|
||
|
|
server_default="platform",
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _assert_known_scoped_tenants() -> None:
|
||
|
|
bind = op.get_bind()
|
||
|
|
for table_name in (_ASSET_TABLE, *_ASSET_CHILD_TABLES):
|
||
|
|
unknown = int(
|
||
|
|
bind.scalar(
|
||
|
|
sa.text(
|
||
|
|
f"""
|
||
|
|
SELECT COUNT(*)
|
||
|
|
FROM {table_name} AS scoped
|
||
|
|
LEFT JOIN tenants AS tenant ON tenant.tenant_id = scoped.tenant_id
|
||
|
|
WHERE tenant.tenant_id IS NULL
|
||
|
|
"""
|
||
|
|
)
|
||
|
|
)
|
||
|
|
or 0
|
||
|
|
)
|
||
|
|
if unknown:
|
||
|
|
raise RuntimeError(
|
||
|
|
f"cannot migrate {table_name}: scoped tenant is absent from tenants "
|
||
|
|
f"(unknown_rows={unknown})"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _create_scope_check(table_name: str) -> None:
|
||
|
|
op.create_check_constraint(
|
||
|
|
f"ck_{table_name}_scope_tenant",
|
||
|
|
table_name,
|
||
|
|
"(scope = 'platform' AND tenant_id = 'platform') OR "
|
||
|
|
"(scope = 'tenant' AND tenant_id <> 'platform')",
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _drop_legacy_code_uniqueness() -> None:
|
||
|
|
op.execute("ALTER TABLE agent_assets DROP CONSTRAINT IF EXISTS agent_assets_code_key")
|
||
|
|
op.execute("DROP INDEX IF EXISTS ix_agent_assets_code")
|
||
|
|
op.create_index("ix_agent_assets_code", _ASSET_TABLE, ["code"], unique=False)
|
||
|
|
|
||
|
|
|
||
|
|
def upgrade() -> None:
|
||
|
|
_require_postgresql()
|
||
|
|
has_agent_asset_schema = _has_complete_agent_asset_schema()
|
||
|
|
_create_onlyoffice_session_table()
|
||
|
|
if not has_agent_asset_schema:
|
||
|
|
return
|
||
|
|
_add_scope_columns(_ASSET_TABLE)
|
||
|
|
for table_name in _ASSET_CHILD_TABLES:
|
||
|
|
_add_scope_columns(table_name)
|
||
|
|
_backfill_scope()
|
||
|
|
_set_scope_not_null_and_defaults(_ASSET_TABLE)
|
||
|
|
for table_name in _ASSET_CHILD_TABLES:
|
||
|
|
_set_scope_not_null_and_defaults(table_name)
|
||
|
|
_assert_known_scoped_tenants()
|
||
|
|
|
||
|
|
_drop_legacy_code_uniqueness()
|
||
|
|
op.create_unique_constraint(
|
||
|
|
"uq_agent_assets_tenant_scope_id",
|
||
|
|
_ASSET_TABLE,
|
||
|
|
["tenant_id", "scope", "id"],
|
||
|
|
)
|
||
|
|
op.create_unique_constraint(
|
||
|
|
"uq_agent_assets_tenant_scope_code",
|
||
|
|
_ASSET_TABLE,
|
||
|
|
["tenant_id", "scope", "code"],
|
||
|
|
)
|
||
|
|
op.create_foreign_key(
|
||
|
|
"fk_agent_assets_tenant",
|
||
|
|
_ASSET_TABLE,
|
||
|
|
"tenants",
|
||
|
|
["tenant_id"],
|
||
|
|
["tenant_id"],
|
||
|
|
ondelete="RESTRICT",
|
||
|
|
)
|
||
|
|
_create_scope_check(_ASSET_TABLE)
|
||
|
|
op.create_index(
|
||
|
|
"ix_agent_assets_scope_tenant",
|
||
|
|
_ASSET_TABLE,
|
||
|
|
["scope", "tenant_id"],
|
||
|
|
)
|
||
|
|
|
||
|
|
for table_name in _ASSET_CHILD_TABLES:
|
||
|
|
op.create_foreign_key(
|
||
|
|
f"fk_{table_name}_tenant",
|
||
|
|
table_name,
|
||
|
|
"tenants",
|
||
|
|
["tenant_id"],
|
||
|
|
["tenant_id"],
|
||
|
|
ondelete="RESTRICT",
|
||
|
|
)
|
||
|
|
_create_scope_check(table_name)
|
||
|
|
op.create_index(
|
||
|
|
f"ix_{table_name}_tenant_asset",
|
||
|
|
table_name,
|
||
|
|
["tenant_id", "scope", "asset_id"],
|
||
|
|
)
|
||
|
|
for table_name in _OWNERSHIP_CHILD_TABLES:
|
||
|
|
op.create_foreign_key(
|
||
|
|
f"fk_{table_name}_tenant_asset",
|
||
|
|
table_name,
|
||
|
|
_ASSET_TABLE,
|
||
|
|
["tenant_id", "scope", "asset_id"],
|
||
|
|
["tenant_id", "scope", "id"],
|
||
|
|
ondelete="CASCADE",
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _require_lossless_downgrade() -> None:
|
||
|
|
bind = op.get_bind()
|
||
|
|
tenant_assets = int(
|
||
|
|
bind.scalar(
|
||
|
|
sa.text("SELECT COUNT(*) FROM agent_assets WHERE scope = 'tenant'")
|
||
|
|
)
|
||
|
|
or 0
|
||
|
|
)
|
||
|
|
tenant_evidence = sum(
|
||
|
|
int(
|
||
|
|
bind.scalar(
|
||
|
|
sa.text(f"SELECT COUNT(*) FROM {table_name} WHERE scope = 'tenant'")
|
||
|
|
)
|
||
|
|
or 0
|
||
|
|
)
|
||
|
|
for table_name in ("agent_asset_test_runs", "agent_asset_rule_feedback")
|
||
|
|
)
|
||
|
|
if tenant_assets or tenant_evidence:
|
||
|
|
raise RuntimeError(
|
||
|
|
"cannot downgrade Agent asset tenant security: tenant-owned facts exist "
|
||
|
|
f"(tenant_assets={tenant_assets}, tenant_evidence={tenant_evidence})"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _require_no_onlyoffice_sessions() -> None:
|
||
|
|
if not sa.inspect(op.get_bind()).has_table(_ONLYOFFICE_SESSION_TABLE):
|
||
|
|
return
|
||
|
|
session_count = int(
|
||
|
|
op.get_bind().scalar(
|
||
|
|
sa.text(f"SELECT COUNT(*) FROM {_ONLYOFFICE_SESSION_TABLE}")
|
||
|
|
)
|
||
|
|
or 0
|
||
|
|
)
|
||
|
|
if session_count:
|
||
|
|
raise RuntimeError(
|
||
|
|
"cannot downgrade Agent asset tenant security: ONLYOFFICE session evidence "
|
||
|
|
f"exists (sessions={session_count})"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _drop_onlyoffice_session_table() -> None:
|
||
|
|
if not sa.inspect(op.get_bind()).has_table(_ONLYOFFICE_SESSION_TABLE):
|
||
|
|
return
|
||
|
|
op.drop_index(
|
||
|
|
"ix_agent_asset_onlyoffice_sessions_status_expiry",
|
||
|
|
table_name=_ONLYOFFICE_SESSION_TABLE,
|
||
|
|
)
|
||
|
|
op.drop_index(
|
||
|
|
"ix_agent_asset_onlyoffice_sessions_tenant_asset",
|
||
|
|
table_name=_ONLYOFFICE_SESSION_TABLE,
|
||
|
|
)
|
||
|
|
op.drop_table(_ONLYOFFICE_SESSION_TABLE)
|
||
|
|
|
||
|
|
|
||
|
|
def _drop_constraint_if_exists(
|
||
|
|
table_name: str,
|
||
|
|
constraint_name: str,
|
||
|
|
) -> None:
|
||
|
|
"""删除本迁移负责的约束，兼容模型建表产生的不同外键名称。"""
|
||
|
|
preparer = op.get_bind().dialect.identifier_preparer
|
||
|
|
op.execute(
|
||
|
|
sa.text(
|
||
|
|
f"ALTER TABLE {preparer.quote(table_name)} "
|
||
|
|
f"DROP CONSTRAINT IF EXISTS {preparer.quote(constraint_name)}"
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def downgrade() -> None:
|
||
|
|
_require_postgresql()
|
||
|
|
has_agent_asset_schema = _has_complete_agent_asset_schema()
|
||
|
|
_require_no_onlyoffice_sessions()
|
||
|
|
_drop_onlyoffice_session_table()
|
||
|
|
if not has_agent_asset_schema:
|
||
|
|
return
|
||
|
|
_require_lossless_downgrade()
|
||
|
|
for table_name in reversed(_OWNERSHIP_CHILD_TABLES):
|
||
|
|
_drop_constraint_if_exists(
|
||
|
|
table_name,
|
||
|
|
f"fk_{table_name}_tenant_asset",
|
||
|
|
)
|
||
|
|
for table_name in reversed(_ASSET_CHILD_TABLES):
|
||
|
|
_drop_constraint_if_exists(
|
||
|
|
table_name,
|
||
|
|
f"fk_{table_name}_tenant",
|
||
|
|
)
|
||
|
|
op.drop_index(f"ix_{table_name}_tenant_asset", table_name=table_name)
|
||
|
|
_drop_constraint_if_exists(
|
||
|
|
table_name,
|
||
|
|
f"ck_{table_name}_scope_tenant",
|
||
|
|
)
|
||
|
|
op.drop_index("ix_agent_assets_scope_tenant", table_name=_ASSET_TABLE)
|
||
|
|
_drop_constraint_if_exists(_ASSET_TABLE, "ck_agent_assets_scope_tenant")
|
||
|
|
_drop_constraint_if_exists(_ASSET_TABLE, "fk_agent_assets_tenant")
|
||
|
|
_drop_constraint_if_exists(
|
||
|
|
_ASSET_TABLE,
|
||
|
|
"uq_agent_assets_tenant_scope_code",
|
||
|
|
)
|
||
|
|
_drop_constraint_if_exists(
|
||
|
|
_ASSET_TABLE,
|
||
|
|
"uq_agent_assets_tenant_scope_id",
|
||
|
|
)
|
||
|
|
op.drop_index("ix_agent_assets_code", table_name=_ASSET_TABLE)
|
||
|
|
op.create_index("ix_agent_assets_code", _ASSET_TABLE, ["code"], unique=True)
|
||
|
|
for table_name in reversed(_ASSET_CHILD_TABLES):
|
||
|
|
op.drop_column(table_name, "scope")
|
||
|
|
op.drop_column(table_name, "tenant_id")
|
||
|
|
op.drop_column(_ASSET_TABLE, "scope")
|
||
|
|
op.drop_column(_ASSET_TABLE, "tenant_id")
|