ITWorx_MCP_Hub-Public/services/control_api/persistence/repositories.py

"""Workspace/project-scoped repositories.

Every repository constructor requires a ``Scope``.  Project-owned repositories
reject workspace-only scopes, and every query applies the scope predicate before
loading or mutating a row.  There is intentionally no generic unrestricted
query or audit update method.
"""

from __future__ import annotations

import re
from datetime import UTC, datetime
from typing import Any, TypeVar, cast
from uuid import UUID

from pydantic import BaseModel
from sqlalchemy import Select, select, update
from sqlalchemy.orm import Session

from packages.hub_core.domain import (
    Approval,
    AuditEvent,
    Capability,
    Client,
    ClientBinding,
    Connector,
    Environment,
    KnowledgeChunk,
    KnowledgeDocument,
    KnowledgeSource,
    Operation,
    Policy,
    PolicyDecision,
    Project,
    ReconcileRun,
    RiskClass,
    Scope,
    ToolDefinition,
    VirtualServer,
    VirtualServerCapability,
    Workspace,
)
from packages.hub_core.errors import (
    AuditMutationError,
    ConcurrentUpdateError,
    DomainError,
    InvalidTransition,
    MissingProjectScope,
)
from packages.hub_core.state_machines import transition

from .models import (
    ApprovalRow,
    AuditEventRow,
    CapabilityRow,
    ClientBindingRow,
    ClientRow,
    ConnectorRow,
    EnvironmentRow,
    KnowledgeChunkRow,
    KnowledgeDocumentRow,
    KnowledgeSourceRow,
    OperationRow,
    PolicyDecisionRow,
    PolicyRow,
    ProjectRow,
    ReconcileRunRow,
    RiskClassRow,
    ToolDefinitionRow,
    VirtualServerCapabilityRow,
    VirtualServerRow,
    WorkspaceRow,
)

RowT = TypeVar("RowT")
DomainT = TypeVar("DomainT", bound=BaseModel)


def _now() -> datetime:
    return datetime.now(UTC)


def _normalized_row_data(row: Any, rename_fields: dict[str, str] | None = None) -> dict[str, Any]:
    reverse = {target: source for source, target in (rename_fields or {}).items()}
    values: dict[str, Any] = {}
    for column in row.__table__.columns:
        value = getattr(row, column.name)
        if isinstance(value, datetime) and value.tzinfo is None:
            value = value.replace(tzinfo=UTC)
        values[reverse.get(column.name, column.name)] = value
    return values


def _atomic_transition(
    session: Session,
    row_type: type[Any],
    scope: Scope,
    record_id: UUID,
    current: str,
    target: str,
    expected_version: int,
    machine: str,
    state_field: str,
    *,
    project_required: bool,
) -> None:
    next_state = transition(machine, current, target)
    statement = (
        update(row_type)
        .where(row_type.id == record_id)
        .where(row_type.workspace_id == scope.workspace_id)
        .where(row_type.revision == expected_version)
    )
    if project_required:
        statement = statement.where(row_type.project_id == scope.project_id)
    result = session.execute(
        statement.values(
            **{
                state_field: next_state,
                "revision": expected_version + 1,
                "updated_at": _now(),
            }
        )
    )
    if result.rowcount != 1:
        raise ConcurrentUpdateError(f"stale version for {row_type.__name__} {record_id}")


class ScopedRepository[RowT, DomainT]:
    """Small typed CRUD boundary with mandatory workspace/project filtering."""

    row_type: Any
    domain_type: Any
    project_required = True
    project_filter_if_present = True
    rename_fields: dict[str, str] = {}

    def __init__(self, session: Session, scope: Scope) -> None:
        if not isinstance(scope, Scope):
            raise TypeError("repository scope must be a Scope")
        self.session = session
        self.scope = scope
        if self.project_required and scope.project_id is None:
            raise MissingProjectScope(f"{type(self).__name__} requires project scope")

    def _statement(self) -> Select[tuple[RowT]]:
        row = self.row_type
        statement = select(row).where(row.workspace_id == self.scope.workspace_id)
        if self.project_required or (
            self.project_filter_if_present and self.scope.project_id is not None
        ):
            statement = statement.where(row.project_id == self.scope.project_id)
        return statement.order_by(row.id)

    def _row_to_domain(self, row: RowT) -> DomainT:
        data = _normalized_row_data(row, self.rename_fields)
        data = {key: value for key, value in data.items() if key in self.domain_type.model_fields}
        return cast(DomainT, self.domain_type.model_validate(data))

    def _model_data(self, model: DomainT) -> dict[str, Any]:
        values = cast(Any, model).model_dump()
        columns = set(self.row_type.__table__.columns.keys())
        data: dict[str, Any] = {}
        for field, value in values.items():
            target = self.rename_fields.get(field, field)
            if target in columns:
                data[target] = value
        return data

    def list(self, *, limit: int = 100, offset: int = 0) -> list[DomainT]:
        if not 1 <= limit <= 500:
            raise ValueError("limit must be between 1 and 500")
        if offset < 0:
            raise ValueError("offset must be non-negative")
        rows = self.session.scalars(self._statement().limit(limit).offset(offset)).all()
        return [self._row_to_domain(row) for row in rows]

    def get(self, record_id: UUID) -> DomainT | None:
        row = self.session.scalars(self._statement().where(self.row_type.id == record_id)).first()
        return self._row_to_domain(row) if row is not None else None

    def create(self, model: DomainT) -> DomainT:
        data = self._model_data(model)
        if data.get("workspace_id") != self.scope.workspace_id:
            raise ValueError("model workspace does not match repository scope")
        if self.project_required and data.get("project_id") != self.scope.project_id:
            raise ValueError("model project does not match repository scope")
        row = self.row_type(**data)
        self.session.add(row)
        self.session.flush()
        return self._row_to_domain(row)

    def update(self, model: DomainT, *, expected_version: int) -> DomainT:
        record_id = cast(Any, model).id
        data = self._model_data(model)
        data.pop("id", None)
        data.pop("created_at", None)
        data.pop("revision", None)
        data["updated_at"] = _now()
        data["revision"] = expected_version + 1
        statement = (
            update(self.row_type)
            .where(self.row_type.id == record_id)
            .where(self.row_type.workspace_id == self.scope.workspace_id)
            .where(self.row_type.revision == expected_version)
        )
        if self.project_required or (
            self.project_filter_if_present and self.scope.project_id is not None
        ):
            statement = statement.where(self.row_type.project_id == self.scope.project_id)
        result = self.session.execute(statement.values(**data))
        if result.rowcount != 1:
            raise ConcurrentUpdateError(f"stale version for {type(self).__name__} {record_id}")
        self.session.flush()
        return self.get(record_id)  # type: ignore[return-value]


class WorkspaceRepository(ScopedRepository[WorkspaceRow, Workspace]):
    row_type = WorkspaceRow
    domain_type = Workspace
    project_required = False
    project_filter_if_present = False

    def create(self, model: Workspace) -> Workspace:
        if model.id != self.scope.workspace_id:
            raise ValueError("workspace model does not match repository scope")
        row = WorkspaceRow(**model.model_dump())
        self.session.add(row)
        self.session.flush()
        return self._row_to_domain(row)

    def _statement(self) -> Select[tuple[WorkspaceRow]]:
        return (
            select(self.row_type)
            .where(self.row_type.id == self.scope.workspace_id)
            .order_by(self.row_type.id)
        )


class ProjectRepository(ScopedRepository[ProjectRow, Project]):
    row_type = ProjectRow
    domain_type = Project
    project_required = False
    project_filter_if_present = False

    def _statement(self) -> Select[tuple[ProjectRow]]:
        statement = select(self.row_type).where(
            self.row_type.workspace_id == self.scope.workspace_id
        )
        if self.scope.project_id is not None:
            statement = statement.where(self.row_type.id == self.scope.project_id)
        return statement.order_by(self.row_type.id)

    def create(self, model: Project) -> Project:
        if self.scope.project_id is not None and model.id != self.scope.project_id:
            raise ValueError("project model does not match repository scope")
        return super().create(model)


class RiskClassRepository:
    """Read-only catalog boundary for safe global risk classes."""

    def __init__(self, session: Session, scope: Scope) -> None:
        if not isinstance(scope, Scope):
            raise TypeError("repository scope must be a Scope")
        self.session = session
        self.scope = scope

    def list(self) -> list[RiskClass]:
        rows = self.session.scalars(select(RiskClassRow).order_by(RiskClassRow.code)).all()
        return [RiskClass.model_validate(row) for row in rows]


class CapabilityRepository:
    """Read/create boundary for global capability definitions, always scoped by caller."""

    def __init__(self, session: Session, scope: Scope) -> None:
        if not isinstance(scope, Scope):
            raise TypeError("repository scope must be a Scope")
        self.session = session
        self.scope = scope

    def list(self, *, limit: int = 100, offset: int = 0) -> list[Capability]:
        rows = self.session.scalars(
            select(CapabilityRow)
            .order_by(CapabilityRow.canonical_name, CapabilityRow.version)
            .limit(limit)
            .offset(offset)
        ).all()
        return [Capability.model_validate(_normalized_row_data(row)) for row in rows]

    def create(self, model: Capability) -> Capability:
        row = CapabilityRow(**model.model_dump())
        self.session.add(row)
        self.session.flush()
        return Capability.model_validate(_normalized_row_data(row))


class ToolDefinitionRepository(ScopedRepository[ToolDefinitionRow, ToolDefinition]):
    row_type = ToolDefinitionRow
    domain_type = ToolDefinition


class VirtualServerCapabilityRepository(
    ScopedRepository[VirtualServerCapabilityRow, VirtualServerCapability]
):
    row_type = VirtualServerCapabilityRow
    domain_type = VirtualServerCapability


class ClientBindingRepository(ScopedRepository[ClientBindingRow, ClientBinding]):
    row_type = ClientBindingRow
    domain_type = ClientBinding


class EnvironmentRepository(ScopedRepository[EnvironmentRow, Environment]):
    row_type = EnvironmentRow
    domain_type = Environment


class ConnectorRepository(ScopedRepository[ConnectorRow, Connector]):
    row_type = ConnectorRow
    domain_type = Connector

    def _model_data(self, model: Connector) -> dict[str, Any]:
        data = super()._model_data(model)
        data["environment_scope"] = (
            str(model.environment_id) if model.environment_id else "__project__"
        )
        return data

    def transition_health(
        self, connector_id: UUID, target: str, *, expected_version: int
    ) -> Connector:
        row = self.session.scalars(
            self._statement().where(self.row_type.id == connector_id)
        ).first()
        if row is None:
            raise LookupError(f"connector not found in scope: {connector_id}")
        _atomic_transition(
            self.session,
            self.row_type,
            self.scope,
            connector_id,
            row.health_state,
            target,
            expected_version,
            "connector_health",
            "health_state",
            project_required=True,
        )
        self.session.flush()
        return self.get(connector_id)  # type: ignore[return-value]


class VirtualServerRepository(ScopedRepository[VirtualServerRow, VirtualServer]):
    row_type = VirtualServerRow
    domain_type = VirtualServer
    rename_fields = {"metadata": "record_metadata"}


class ClientRepository(ScopedRepository[ClientRow, Client]):
    row_type = ClientRow
    domain_type = Client
    project_required = False
    project_filter_if_present = False
    rename_fields = {"metadata": "record_metadata"}

    def transition_status(self, client_id: UUID, target: str, *, expected_version: int) -> Client:
        row = self.session.scalars(self._statement().where(self.row_type.id == client_id)).first()
        if row is None:
            raise LookupError(f"client not found in workspace: {client_id}")
        _atomic_transition(
            self.session,
            self.row_type,
            self.scope,
            client_id,
            row.status,
            target,
            expected_version,
            "client",
            "status",
            project_required=False,
        )
        if target == "revoked":
            row.revoked_at = _now()
        self.session.flush()
        return self.get(client_id)  # type: ignore[return-value]


class PolicyRepository(ScopedRepository[PolicyRow, Policy]):
    row_type = PolicyRow
    domain_type = Policy
    project_required = False
    rename_fields = {"match": "match_json", "conditions": "conditions_json"}


class PolicyDecisionRepository:
    """Short-retention decision record boundary with mandatory workspace scope."""

    def __init__(self, session: Session, scope: Scope) -> None:
        if not isinstance(scope, Scope):
            raise TypeError("repository scope must be a Scope")
        self.session = session
        self.scope = scope

    def create(self, model: PolicyDecision) -> PolicyDecision:
        if model.workspace_id != self.scope.workspace_id:
            raise ValueError("policy decision workspace does not match scope")
        if self.scope.project_id is not None and model.project_id != self.scope.project_id:
            raise ValueError("policy decision project does not match scope")
        row = PolicyDecisionRow(**model.model_dump())
        self.session.add(row)
        self.session.flush()
        return PolicyDecision.model_validate(_normalized_row_data(row))

    def list(self, *, limit: int = 100, offset: int = 0) -> list[PolicyDecision]:
        statement = select(PolicyDecisionRow).where(
            PolicyDecisionRow.workspace_id == self.scope.workspace_id
        )
        if self.scope.project_id is not None:
            statement = statement.where(PolicyDecisionRow.project_id == self.scope.project_id)
        statement = (
            statement.order_by(PolicyDecisionRow.created_at, PolicyDecisionRow.id)
            .limit(limit)
            .offset(offset)
        )
        return [
            PolicyDecision.model_validate(_normalized_row_data(row))
            for row in self.session.scalars(statement).all()
        ]


class ReconcileRunRepository(ScopedRepository[ReconcileRunRow, ReconcileRun]):
    row_type = ReconcileRunRow
    domain_type = ReconcileRun

    def transition_status(
        self, run_id: UUID, target: str, *, expected_version: int
    ) -> ReconcileRun:
        row = self.session.scalars(self._statement().where(self.row_type.id == run_id)).first()
        if row is None:
            raise LookupError(f"reconcile run not found in scope: {run_id}")
        _atomic_transition(
            self.session,
            self.row_type,
            self.scope,
            run_id,
            row.status,
            target,
            expected_version,
            "reconcile",
            "status",
            project_required=True,
        )
        if target in {"succeeded", "failed", "invalid"}:
            row.completed_at = _now()
        self.session.flush()
        return self.get(run_id)  # type: ignore[return-value]


class OperationRepository(ScopedRepository[OperationRow, Operation]):
    row_type = OperationRow
    domain_type = Operation

    def transition_state(
        self, operation_id: UUID, target: str, *, expected_version: int
    ) -> Operation:
        row = self.session.scalars(
            self._statement().where(self.row_type.id == operation_id)
        ).first()
        if row is None:
            raise LookupError(f"operation not found in scope: {operation_id}")
        _atomic_transition(
            self.session,
            self.row_type,
            self.scope,
            operation_id,
            row.state,
            target,
            expected_version,
            "operation",
            "state",
            project_required=True,
        )
        self.session.flush()
        return self.get(operation_id)  # type: ignore[return-value]


class ApprovalRepository(ScopedRepository[ApprovalRow, Approval]):
    """Approval boundary with plan-hash and state-machine enforcement."""

    row_type = ApprovalRow
    domain_type = Approval

    def transition_state(
        self, approval_id: UUID, target: str, *, expected_version: int
    ) -> Approval:
        row = self.session.scalars(self._statement().where(self.row_type.id == approval_id)).first()
        if row is None:
            raise LookupError(f"approval not found in scope: {approval_id}")
        if target == "approved":
            operation = self.session.scalars(
                select(OperationRow)
                .where(OperationRow.id == row.operation_id)
                .where(OperationRow.workspace_id == self.scope.workspace_id)
                .where(OperationRow.project_id == self.scope.project_id)
            ).first()
            if operation is None or operation.plan_hash != row.plan_hash:
                raise DomainError("approval plan hash does not match operation")
            if operation.state != "approval_required":
                raise InvalidTransition("operation is not waiting for approval")
        _atomic_transition(
            self.session,
            self.row_type,
            self.scope,
            approval_id,
            row.state,
            target,
            expected_version,
            "approval",
            "state",
            project_required=True,
        )
        if target == "approved":
            row.approved_at = _now()
        self.session.flush()
        return self.get(approval_id)  # type: ignore[return-value]


class AuditEventRepository:
    """Append/query-only audit boundary; update and delete are not exposed."""

    _sensitive_key = re.compile(
        r"(?:password|secret|token|private|authorization|cookie|credential|connection)", re.I
    )

    def __init__(self, session: Session, scope: Scope) -> None:
        if not isinstance(scope, Scope):
            raise TypeError("repository scope must be a Scope")
        self.session = session
        self.scope = scope

    @classmethod
    def _assert_sanitized(cls, value: Any) -> None:
        if isinstance(value, dict):
            for key, item in value.items():
                if cls._sensitive_key.search(str(key)):
                    raise AuditMutationError(f"sensitive audit field is prohibited: {key}")
                cls._assert_sanitized(item)
        elif isinstance(value, list):
            for item in value:
                cls._assert_sanitized(item)

    def create(self, event: AuditEvent) -> AuditEvent:
        if event.workspace_id != self.scope.workspace_id:
            raise ValueError("audit workspace does not match repository scope")
        if self.scope.project_id is not None and event.project_id != self.scope.project_id:
            raise ValueError("audit project does not match repository scope")
        self._assert_sanitized(event.sanitized_arguments)
        row = AuditEventRow(**event.model_dump())
        self.session.add(row)
        self.session.flush()
        return AuditEvent.model_validate(row)

    def list(self, *, limit: int = 100, offset: int = 0) -> list[AuditEvent]:
        if not 1 <= limit <= 500 or offset < 0:
            raise ValueError("invalid pagination")
        statement = select(AuditEventRow).where(
            AuditEventRow.workspace_id == self.scope.workspace_id
        )
        if self.scope.project_id is not None:
            statement = statement.where(AuditEventRow.project_id == self.scope.project_id)
        statement = (
            statement.order_by(AuditEventRow.occurred_at, AuditEventRow.id)
            .limit(limit)
            .offset(offset)
        )
        return [
            AuditEvent.model_validate(_normalized_row_data(row))
            for row in self.session.scalars(statement).all()
        ]


class KnowledgeSourceRepository(ScopedRepository[KnowledgeSourceRow, KnowledgeSource]):
    row_type = KnowledgeSourceRow
    domain_type = KnowledgeSource


class KnowledgeDocumentRepository(ScopedRepository[KnowledgeDocumentRow, KnowledgeDocument]):
    row_type = KnowledgeDocumentRow
    domain_type = KnowledgeDocument
    rename_fields = {"metadata": "record_metadata"}


class KnowledgeChunkRepository(ScopedRepository[KnowledgeChunkRow, KnowledgeChunk]):
    row_type = KnowledgeChunkRow
    domain_type = KnowledgeChunk
    rename_fields = {"metadata": "record_metadata"}