govoplan-dataflow/src/govoplan_dataflow/backend/db/models.py

from __future__ import annotations

import uuid
from datetime import datetime
from typing import Any

from sqlalchemy import (
    Boolean,
    DateTime,
    ForeignKey,
    Index,
    Integer,
    JSON,
    String,
    Text,
    UniqueConstraint,
)
from sqlalchemy.orm import Mapped, mapped_column, relationship

from govoplan_core.core.concurrency import strong_resource_etag
from govoplan_core.db.base import Base, TimestampMixin


def new_uuid() -> str:
    return str(uuid.uuid4())


class DataflowPipeline(Base, TimestampMixin):
    __tablename__ = "dataflow_pipelines"
    __table_args__ = (
        Index("ix_dataflow_pipelines_tenant_status", "tenant_id", "status"),
        Index("ix_dataflow_pipelines_tenant_updated", "tenant_id", "updated_at"),
    )

    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_uuid)
    tenant_id: Mapped[str | None] = mapped_column(
        String(36),
        nullable=True,
        index=True,
    )
    scope_type: Mapped[str] = mapped_column(
        String(20),
        default="tenant",
        nullable=False,
        index=True,
    )
    scope_id: Mapped[str | None] = mapped_column(
        String(36),
        nullable=True,
        index=True,
    )
    definition_kind: Mapped[str] = mapped_column(
        String(20),
        default="flow",
        nullable=False,
        index=True,
    )
    inherit_to_lower_scopes: Mapped[bool] = mapped_column(
        Boolean,
        default=False,
        nullable=False,
    )
    allow_run: Mapped[bool] = mapped_column(
        Boolean,
        default=True,
        nullable=False,
    )
    allow_reuse: Mapped[bool] = mapped_column(
        Boolean,
        default=False,
        nullable=False,
    )
    allow_automation: Mapped[bool] = mapped_column(
        Boolean,
        default=False,
        nullable=False,
    )
    derived_from_pipeline_id: Mapped[str | None] = mapped_column(
        String(36),
        nullable=True,
        index=True,
    )
    derived_from_revision: Mapped[int | None] = mapped_column(
        Integer,
        nullable=True,
    )
    derived_from_hash: Mapped[str | None] = mapped_column(
        String(64),
        nullable=True,
    )
    derivation_provenance: Mapped[dict[str, Any]] = mapped_column(
        JSON,
        default=dict,
        nullable=False,
    )
    name: Mapped[str] = mapped_column(String(300), nullable=False)
    description: Mapped[str | None] = mapped_column(Text)
    status: Mapped[str] = mapped_column(String(32), default="draft", nullable=False, index=True)
    current_revision: Mapped[int] = mapped_column(Integer, default=1, nullable=False)
    created_by: Mapped[str | None] = mapped_column(String(255), nullable=True, index=True)
    updated_by: Mapped[str | None] = mapped_column(String(255), nullable=True, index=True)
    deleted_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)
    metadata_: Mapped[dict[str, Any]] = mapped_column("metadata", JSON, default=dict, nullable=False)

    revisions: Mapped[list["DataflowPipelineRevision"]] = relationship(
        back_populates="pipeline",
        cascade="all, delete-orphan",
        order_by="DataflowPipelineRevision.revision",
    )
    runs: Mapped[list["DataflowRun"]] = relationship(
        back_populates="pipeline",
        cascade="all, delete-orphan",
        order_by="DataflowRun.created_at",
    )
    deployments: Mapped[list["DataflowPipelineDeployment"]] = relationship(
        back_populates="pipeline",
        cascade="all, delete-orphan",
        order_by="DataflowPipelineDeployment.environment",
    )
    triggers: Mapped[list["DataflowTrigger"]] = relationship(
        back_populates="pipeline",
        cascade="all, delete-orphan",
        order_by="DataflowTrigger.created_at",
    )
    decision_sets: Mapped[list["DataflowReconciliationDecisionSet"]] = relationship(
        back_populates="pipeline",
        cascade="all, delete-orphan",
        order_by="DataflowReconciliationDecisionSet.name",
    )


class DataflowPipelineRevision(Base, TimestampMixin):
    __tablename__ = "dataflow_pipeline_revisions"
    __table_args__ = (
        UniqueConstraint("pipeline_id", "revision", name="uq_dataflow_pipeline_revision"),
        Index("ix_dataflow_revisions_tenant_pipeline", "tenant_id", "pipeline_id"),
        Index("ix_dataflow_revisions_content_hash", "tenant_id", "content_hash"),
    )

    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_uuid)
    tenant_id: Mapped[str | None] = mapped_column(
        String(36),
        nullable=True,
        index=True,
    )
    pipeline_id: Mapped[str] = mapped_column(
        ForeignKey("dataflow_pipelines.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    revision: Mapped[int] = mapped_column(Integer, nullable=False)
    schema_version: Mapped[int] = mapped_column(Integer, default=1, nullable=False)
    graph: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False)
    sql_text: Mapped[str | None] = mapped_column(Text)
    editor_mode: Mapped[str] = mapped_column(String(20), default="graph", nullable=False)
    content_hash: Mapped[str] = mapped_column(String(64), nullable=False)
    created_by: Mapped[str | None] = mapped_column(String(255), nullable=True, index=True)

    pipeline: Mapped[DataflowPipeline] = relationship(back_populates="revisions")


class DataflowReconciliationDecisionSet(Base, TimestampMixin):
    __tablename__ = "dataflow_reconciliation_decision_sets"
    __table_args__ = (
        UniqueConstraint(
            "tenant_id",
            "pipeline_id",
            "name",
            name="uq_dataflow_decision_sets_pipeline_name",
        ),
        Index(
            "ix_dataflow_decision_sets_tenant_pipeline",
            "tenant_id",
            "pipeline_id",
        ),
    )

    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_uuid)
    tenant_id: Mapped[str] = mapped_column(String(36), nullable=False, index=True)
    pipeline_id: Mapped[str] = mapped_column(
        ForeignKey("dataflow_pipelines.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    name: Mapped[str] = mapped_column(String(300), nullable=False)
    node_id: Mapped[str | None] = mapped_column(String(100), nullable=True)
    resource_revision: Mapped[int] = mapped_column(
        Integer,
        default=1,
        nullable=False,
    )
    created_by: Mapped[str | None] = mapped_column(String(255), nullable=True)
    updated_by: Mapped[str | None] = mapped_column(String(255), nullable=True)

    pipeline: Mapped[DataflowPipeline] = relationship(back_populates="decision_sets")
    decisions: Mapped[list["DataflowReconciliationDecision"]] = relationship(
        back_populates="decision_set",
        cascade="all, delete-orphan",
        order_by="DataflowReconciliationDecision.revision",
    )

    @property
    def strong_etag(self) -> str:
        return strong_resource_etag(
            "dataflow_reconciliation_decision_set",
            self.id,
            self.resource_revision,
        )


class DataflowReconciliationDecision(Base, TimestampMixin):
    __tablename__ = "dataflow_reconciliation_decisions"
    __table_args__ = (
        UniqueConstraint(
            "decision_set_id",
            "revision",
            name="uq_dataflow_reconciliation_decision_revision",
        ),
        Index(
            "ix_dataflow_reconciliation_decisions_current",
            "tenant_id",
            "decision_set_id",
            "key_hash",
            "revision",
        ),
    )

    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_uuid)
    tenant_id: Mapped[str] = mapped_column(String(36), nullable=False, index=True)
    decision_set_id: Mapped[str] = mapped_column(
        ForeignKey(
            "dataflow_reconciliation_decision_sets.id",
            ondelete="CASCADE",
        ),
        nullable=False,
        index=True,
    )
    revision: Mapped[int] = mapped_column(Integer, nullable=False)
    key_hash: Mapped[str] = mapped_column(String(64), nullable=False, index=True)
    input_hash: Mapped[str] = mapped_column(String(64), nullable=False)
    action: Mapped[str] = mapped_column(String(20), nullable=False)
    reason: Mapped[str] = mapped_column(Text, nullable=False)
    correction: Mapped[dict[str, Any] | None] = mapped_column(JSON, nullable=True)
    actor_ref: Mapped[str] = mapped_column(String(255), nullable=False)
    decided_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        nullable=False,
    )

    decision_set: Mapped[DataflowReconciliationDecisionSet] = relationship(
        back_populates="decisions"
    )


class DataflowRun(Base, TimestampMixin):
    __tablename__ = "dataflow_runs"
    __table_args__ = (
        UniqueConstraint(
            "tenant_id",
            "pipeline_id",
            "idempotency_key",
            name="uq_dataflow_run_idempotency",
        ),
        Index("ix_dataflow_runs_tenant_status", "tenant_id", "status"),
        Index("ix_dataflow_runs_pipeline_created", "pipeline_id", "created_at"),
        Index(
            "ix_dataflow_runs_worker_queue",
            "status",
            "available_at",
            "created_at",
        ),
    )

    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_uuid)
    tenant_id: Mapped[str] = mapped_column(String(36), nullable=False, index=True)
    pipeline_id: Mapped[str] = mapped_column(
        ForeignKey("dataflow_pipelines.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    pipeline_revision_id: Mapped[str] = mapped_column(
        ForeignKey("dataflow_pipeline_revisions.id", ondelete="RESTRICT"),
        nullable=False,
        index=True,
    )
    run_type: Mapped[str] = mapped_column(String(30), default="preview", nullable=False, index=True)
    status: Mapped[str] = mapped_column(String(30), nullable=False, index=True)
    execution_backend: Mapped[str] = mapped_column(
        String(30),
        default="auto",
        nullable=False,
        index=True,
    )
    environment: Mapped[str] = mapped_column(
        String(30),
        default="development",
        nullable=False,
        index=True,
    )
    executor_version: Mapped[str] = mapped_column(String(40), nullable=False)
    definition_hash: Mapped[str] = mapped_column(String(64), nullable=False)
    idempotency_key: Mapped[str | None] = mapped_column(
        String(255),
        nullable=True,
        index=True,
    )
    request_hash: Mapped[str | None] = mapped_column(String(64), nullable=True)
    request_: Mapped[dict[str, Any]] = mapped_column(
        "request",
        JSON,
        default=dict,
        nullable=False,
    )
    invocation_kind: Mapped[str] = mapped_column(
        String(30),
        default="manual",
        nullable=False,
        index=True,
    )
    trigger_id: Mapped[str | None] = mapped_column(
        String(36),
        nullable=True,
        index=True,
    )
    trigger_delivery_id: Mapped[str | None] = mapped_column(
        String(36),
        nullable=True,
        index=True,
    )
    correlation_id: Mapped[str | None] = mapped_column(
        String(128),
        nullable=True,
        index=True,
    )
    causation_id: Mapped[str | None] = mapped_column(
        String(128),
        nullable=True,
        index=True,
    )
    source_fingerprints: Mapped[list[dict[str, Any]]] = mapped_column(JSON, default=list, nullable=False)
    result_schema: Mapped[list[dict[str, Any]]] = mapped_column(JSON, default=list, nullable=False)
    diagnostics: Mapped[list[dict[str, Any]]] = mapped_column(JSON, default=list, nullable=False)
    input_row_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
    output_row_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
    output_publication_ref: Mapped[str | None] = mapped_column(
        String(500),
        nullable=True,
    )
    output_datasource_ref: Mapped[str | None] = mapped_column(
        String(500),
        nullable=True,
    )
    output_materialization_ref: Mapped[str | None] = mapped_column(
        String(500),
        nullable=True,
    )
    attempts: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
    max_attempts: Mapped[int] = mapped_column(Integer, default=3, nullable=False)
    queued_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True),
        nullable=True,
    )
    available_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True),
        nullable=True,
        index=True,
    )
    claimed_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True),
        nullable=True,
    )
    lease_expires_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True),
        nullable=True,
        index=True,
    )
    heartbeat_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True),
        nullable=True,
    )
    worker_id: Mapped[str | None] = mapped_column(
        String(255),
        nullable=True,
        index=True,
    )
    cancellation_requested_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True),
        nullable=True,
    )
    progress_percent: Mapped[int] = mapped_column(
        Integer,
        default=0,
        nullable=False,
    )
    progress_phase: Mapped[str] = mapped_column(
        String(80),
        default="queued",
        nullable=False,
    )
    retention_until: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True),
        nullable=True,
        index=True,
    )
    purged_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True),
        nullable=True,
    )
    authorization_: Mapped[dict[str, Any]] = mapped_column(
        "authorization",
        JSON,
        default=dict,
        nullable=False,
    )
    resource_budget: Mapped[dict[str, Any]] = mapped_column(
        JSON,
        default=dict,
        nullable=False,
    )
    started_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True),
        nullable=True,
    )
    finished_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
    error: Mapped[str | None] = mapped_column(Text)
    created_by: Mapped[str | None] = mapped_column(String(255), nullable=True, index=True)

    pipeline: Mapped[DataflowPipeline] = relationship(back_populates="runs")


class DataflowPipelineDeployment(Base, TimestampMixin):
    __tablename__ = "dataflow_pipeline_deployments"
    __table_args__ = (
        UniqueConstraint(
            "tenant_id",
            "pipeline_id",
            "environment",
            name="uq_dataflow_pipeline_deployment_environment",
        ),
        Index(
            "ix_dataflow_deployments_tenant_environment",
            "tenant_id",
            "environment",
        ),
    )

    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_uuid)
    tenant_id: Mapped[str] = mapped_column(String(36), nullable=False, index=True)
    pipeline_id: Mapped[str] = mapped_column(
        ForeignKey("dataflow_pipelines.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    pipeline_revision_id: Mapped[str] = mapped_column(
        ForeignKey("dataflow_pipeline_revisions.id", ondelete="RESTRICT"),
        nullable=False,
        index=True,
    )
    environment: Mapped[str] = mapped_column(String(30), nullable=False, index=True)
    source_environment: Mapped[str] = mapped_column(String(30), nullable=False)
    status: Mapped[str] = mapped_column(
        String(30),
        default="active",
        nullable=False,
        index=True,
    )
    provenance: Mapped[dict[str, Any]] = mapped_column(
        JSON,
        default=dict,
        nullable=False,
    )
    promoted_by: Mapped[str | None] = mapped_column(
        String(255),
        nullable=True,
        index=True,
    )

    pipeline: Mapped[DataflowPipeline] = relationship(back_populates="deployments")


class DataflowTrigger(Base, TimestampMixin):
    __tablename__ = "dataflow_triggers"
    __table_args__ = (
        Index(
            "ix_dataflow_triggers_due",
            "status",
            "next_fire_at",
        ),
        Index(
            "ix_dataflow_triggers_tenant_pipeline",
            "tenant_id",
            "pipeline_id",
        ),
    )

    id: Mapped[str] = mapped_column(
        String(36),
        primary_key=True,
        default=new_uuid,
    )
    tenant_id: Mapped[str] = mapped_column(
        String(36),
        nullable=False,
        index=True,
    )
    pipeline_id: Mapped[str] = mapped_column(
        ForeignKey("dataflow_pipelines.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    pipeline_revision_id: Mapped[str] = mapped_column(
        ForeignKey("dataflow_pipeline_revisions.id", ondelete="RESTRICT"),
        nullable=False,
        index=True,
    )
    name: Mapped[str] = mapped_column(String(300), nullable=False)
    kind: Mapped[str] = mapped_column(
        String(30),
        nullable=False,
        index=True,
    )
    status: Mapped[str] = mapped_column(
        String(30),
        default="disabled",
        nullable=False,
        index=True,
    )
    revision: Mapped[int] = mapped_column(Integer, default=1, nullable=False)
    config_: Mapped[dict[str, Any]] = mapped_column(
        "config",
        JSON,
        default=dict,
        nullable=False,
    )
    publication_: Mapped[dict[str, Any] | None] = mapped_column(
        "publication",
        JSON,
        nullable=True,
    )
    row_limit: Mapped[int] = mapped_column(
        Integer,
        default=500,
        nullable=False,
    )
    catch_up_policy: Mapped[str] = mapped_column(
        String(20),
        default="coalesce",
        nullable=False,
    )
    max_concurrent_runs: Mapped[int] = mapped_column(
        Integer,
        default=1,
        nullable=False,
    )
    next_fire_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True),
        nullable=True,
        index=True,
    )
    last_fire_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True),
        nullable=True,
    )
    last_status: Mapped[str | None] = mapped_column(String(30), nullable=True)
    last_error: Mapped[str | None] = mapped_column(Text, nullable=True)
    authorization_account_id: Mapped[str] = mapped_column(
        String(36),
        nullable=False,
    )
    authorization_membership_id: Mapped[str] = mapped_column(
        String(36),
        nullable=False,
    )
    authorization_ref: Mapped[str] = mapped_column(
        String(255),
        nullable=False,
        unique=True,
    )
    grant_scopes: Mapped[list[str]] = mapped_column(
        JSON,
        default=list,
        nullable=False,
    )
    created_by: Mapped[str | None] = mapped_column(
        String(255),
        nullable=True,
        index=True,
    )
    updated_by: Mapped[str | None] = mapped_column(
        String(255),
        nullable=True,
        index=True,
    )

    pipeline: Mapped[DataflowPipeline] = relationship(
        back_populates="triggers"
    )
    deliveries: Mapped[list["DataflowTriggerDelivery"]] = relationship(
        back_populates="trigger",
        cascade="all, delete-orphan",
        order_by="DataflowTriggerDelivery.created_at",
    )


class DataflowTriggerDelivery(Base, TimestampMixin):
    __tablename__ = "dataflow_trigger_deliveries"
    __table_args__ = (
        UniqueConstraint(
            "trigger_id",
            "source_key",
            name="uq_dataflow_trigger_delivery_source",
        ),
        Index(
            "ix_dataflow_trigger_deliveries_queue",
            "status",
            "created_at",
        ),
        Index(
            "ix_dataflow_trigger_deliveries_tenant_trigger",
            "tenant_id",
            "trigger_id",
        ),
    )

    id: Mapped[str] = mapped_column(
        String(36),
        primary_key=True,
        default=new_uuid,
    )
    tenant_id: Mapped[str] = mapped_column(
        String(36),
        nullable=False,
        index=True,
    )
    trigger_id: Mapped[str] = mapped_column(
        ForeignKey("dataflow_triggers.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    pipeline_id: Mapped[str] = mapped_column(
        ForeignKey("dataflow_pipelines.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    pipeline_revision_id: Mapped[str] = mapped_column(
        ForeignKey("dataflow_pipeline_revisions.id", ondelete="RESTRICT"),
        nullable=False,
        index=True,
    )
    source_key: Mapped[str] = mapped_column(String(500), nullable=False)
    invocation_kind: Mapped[str] = mapped_column(
        String(30),
        nullable=False,
        index=True,
    )
    status: Mapped[str] = mapped_column(
        String(30),
        default="queued",
        nullable=False,
        index=True,
    )
    scheduled_for: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True),
        nullable=True,
    )
    event_: Mapped[dict[str, Any] | None] = mapped_column(
        "event",
        JSON,
        nullable=True,
    )
    run_id: Mapped[str | None] = mapped_column(
        ForeignKey("dataflow_runs.id", ondelete="SET NULL"),
        nullable=True,
        index=True,
    )
    attempts: Mapped[int] = mapped_column(
        Integer,
        default=0,
        nullable=False,
    )
    authorization_provenance: Mapped[dict[str, Any]] = mapped_column(
        JSON,
        default=dict,
        nullable=False,
    )
    error: Mapped[str | None] = mapped_column(Text, nullable=True)
    claimed_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True),
        nullable=True,
    )
    finished_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True),
        nullable=True,
    )

    trigger: Mapped[DataflowTrigger] = relationship(
        back_populates="deliveries"
    )


__all__ = [
    "DataflowPipeline",
    "DataflowPipelineRevision",
    "DataflowRun",
    "DataflowTrigger",
    "DataflowTriggerDelivery",
    "new_uuid",
]