first_contas/legacy_reference/workflows/execution_store.py

from __future__ import annotations

from contextlib import contextmanager
from datetime import datetime
from threading import Lock
from typing import Any, Iterator, Protocol, runtime_checkable

from agente_contas_tim.workflows.exceptions import (
    WorkflowExecutionNotFoundError,
    WorkflowExecutionStateError,
)
from agente_contas_tim.workflows.runtime_types import (
    ExecutionStatus,
    WorkflowExecutionRecord,
)


@runtime_checkable
class ExecutionStore(Protocol):
    def create(
        self,
        execution_id: str,
        workflow_name: str,
        workflow_version: int,
        *,
        session_id: str | None = None,
        started_by_message_id: str | None = None,
    ) -> WorkflowExecutionRecord: ...

    def get(self, execution_id: str) -> WorkflowExecutionRecord: ...

    def mark_status(
        self,
        execution_id: str,
        *,
        status: ExecutionStatus,
        current_node: str | None,
        resume_from: str | None,
        expected_input_key: str | None,
    ) -> WorkflowExecutionRecord: ...

    def claim_resume(
        self,
        execution_id: str,
        workflow_name: str,
        workflow_version: int | None,
        *,
        session_id: str | None = None,
        message_id: str | None = None,
    ) -> WorkflowExecutionRecord: ...

    def close(self) -> None: ...


class InMemoryExecutionStore:
    """Execution store em memória (sem persistência, para dev/teste)."""

    def __init__(self) -> None:
        self._records: dict[str, WorkflowExecutionRecord] = {}
        self._lock = Lock()

    def create(
        self,
        execution_id: str,
        workflow_name: str,
        workflow_version: int,
        *,
        session_id: str | None = None,
        started_by_message_id: str | None = None,
    ) -> WorkflowExecutionRecord:
        now = datetime.now()
        record = WorkflowExecutionRecord(
            execution_id=execution_id,
            workflow_name=workflow_name,
            workflow_version=workflow_version,
            status="RUNNING",
            current_node=None,
            resume_from=None,
            expected_input_key=None,
            created_at=now,
            updated_at=now,
        )
        with self._lock:
            self._records[execution_id] = record
        return record

    def get(self, execution_id: str) -> WorkflowExecutionRecord:
        with self._lock:
            record = self._records.get(execution_id)
        if record is None:
            raise WorkflowExecutionNotFoundError(
                f"execution_id={execution_id!r} nao encontrado"
            )
        return record

    def mark_status(
        self,
        execution_id: str,
        *,
        status: ExecutionStatus,
        current_node: str | None,
        resume_from: str | None,
        expected_input_key: str | None,
    ) -> WorkflowExecutionRecord:
        with self._lock:
            record = self._records.get(execution_id)
            if record is None:
                raise WorkflowExecutionNotFoundError(
                    f"execution_id={execution_id!r} nao encontrado"
                )
            updated = WorkflowExecutionRecord(
                execution_id=record.execution_id,
                workflow_name=record.workflow_name,
                workflow_version=record.workflow_version,
                status=status,
                current_node=current_node,
                resume_from=resume_from,
                expected_input_key=expected_input_key,
                created_at=record.created_at,
                updated_at=datetime.now(),
            )
            self._records[execution_id] = updated
        return updated

    def claim_resume(
        self,
        execution_id: str,
        workflow_name: str,
        workflow_version: int | None,
        *,
        session_id: str | None = None,
        message_id: str | None = None,
    ) -> WorkflowExecutionRecord:
        with self._lock:
            record = self._records.get(execution_id)
            if record is None:
                raise WorkflowExecutionNotFoundError(
                    f"execution_id={execution_id!r} nao encontrado"
                )
            if record.workflow_name != workflow_name:
                raise WorkflowExecutionStateError(
                    "workflow_name informado nao corresponde ao execution_id"
                )
            if (
                workflow_version is not None
                and record.workflow_version != workflow_version
            ):
                raise WorkflowExecutionStateError(
                    "version informada nao corresponde a execucao existente"
                )
            if record.status != "WAITING_INPUT":
                raise WorkflowExecutionStateError(
                    f"Execucao {execution_id} nao esta aguardando input"
                )
            updated = WorkflowExecutionRecord(
                execution_id=record.execution_id,
                workflow_name=record.workflow_name,
                workflow_version=record.workflow_version,
                status="RUNNING",
                current_node=record.current_node,
                resume_from=record.resume_from,
                expected_input_key=record.expected_input_key,
                created_at=record.created_at,
                updated_at=datetime.now(),
            )
            self._records[execution_id] = updated
        return record

    def close(self) -> None:
        pass


class PostgresExecutionStore:
    """Persistencia duravel e lock atomico para execucoes de workflow em Postgres."""

    def __init__(self, dsn: str) -> None:
        self._dsn = dsn
        self._setup_lock = Lock()
        self._setup()

    def create(
        self,
        execution_id: str,
        workflow_name: str,
        workflow_version: int,
        *,
        session_id: str | None = None,
        started_by_message_id: str | None = None,
    ) -> WorkflowExecutionRecord:
        with (
            self._connect() as conn,
            conn.cursor(row_factory=self._dict_row_factory()) as cur,
        ):
            cur.execute(
                """
                INSERT INTO workflow_execution (
                    execution_id,
                    workflow_name,
                    workflow_version,
                    status,
                    current_node,
                    resume_from,
                    expected_input_key,
                    created_at,
                    updated_at
                ) VALUES (%s, %s, %s, %s, %s, %s, %s, NOW(), NOW())
                RETURNING
                    execution_id,
                    workflow_name,
                    workflow_version,
                    status,
                    current_node,
                    resume_from,
                    expected_input_key,
                    created_at,
                    updated_at
                """,
                (
                    execution_id,
                    workflow_name,
                    workflow_version,
                    "RUNNING",
                    None,
                    None,
                    None,
                ),
            )
            row = cur.fetchone()
        return self._row_to_record(row)

    def get(self, execution_id: str) -> WorkflowExecutionRecord:
        with (
            self._connect() as conn,
            conn.cursor(row_factory=self._dict_row_factory()) as cur,
        ):
            cur.execute(
                """
                SELECT
                    execution_id,
                    workflow_name,
                    workflow_version,
                    status,
                    current_node,
                    resume_from,
                    expected_input_key,
                    created_at,
                    updated_at
                FROM workflow_execution
                WHERE execution_id = %s
                """,
                (execution_id,),
            )
            row = cur.fetchone()

        if row is None:
            raise WorkflowExecutionNotFoundError(
                f"execution_id={execution_id!r} nao encontrado"
            )
        return self._row_to_record(row)

    def mark_status(
        self,
        execution_id: str,
        *,
        status: ExecutionStatus,
        current_node: str | None,
        resume_from: str | None,
        expected_input_key: str | None,
    ) -> WorkflowExecutionRecord:
        with (
            self._connect() as conn,
            conn.cursor(row_factory=self._dict_row_factory()) as cur,
        ):
            cur.execute(
                """
                UPDATE workflow_execution
                SET status = %s,
                    current_node = %s,
                    resume_from = %s,
                    expected_input_key = %s,
                    updated_at = NOW()
                WHERE execution_id = %s
                RETURNING
                    execution_id,
                    workflow_name,
                    workflow_version,
                    status,
                    current_node,
                    resume_from,
                    expected_input_key,
                    created_at,
                    updated_at
                """,
                (
                    status,
                    current_node,
                    resume_from,
                    expected_input_key,
                    execution_id,
                ),
            )
            row = cur.fetchone()

        if row is None:
            raise WorkflowExecutionNotFoundError(
                f"execution_id={execution_id!r} nao encontrado"
            )
        return self._row_to_record(row)

    def claim_resume(
        self,
        execution_id: str,
        workflow_name: str,
        workflow_version: int | None,
        *,
        session_id: str | None = None,
        message_id: str | None = None,
    ) -> WorkflowExecutionRecord:
        with (
            self._connect() as conn,
            conn.cursor(row_factory=self._dict_row_factory()) as cur,
        ):
            cur.execute(
                """
                SELECT
                    execution_id,
                    workflow_name,
                    workflow_version,
                    status,
                    current_node,
                    resume_from,
                    expected_input_key,
                    created_at,
                    updated_at
                FROM workflow_execution
                WHERE execution_id = %s
                FOR UPDATE
                """,
                (execution_id,),
            )
            row = cur.fetchone()

            if row is None:
                raise WorkflowExecutionNotFoundError(
                    f"execution_id={execution_id!r} nao encontrado"
                )

            record = self._row_to_record(row)
            if record.workflow_name != workflow_name:
                raise WorkflowExecutionStateError(
                    "workflow_name informado nao corresponde ao execution_id"
                )
            if (
                workflow_version is not None
                and record.workflow_version != workflow_version
            ):
                raise WorkflowExecutionStateError(
                    "version informada nao corresponde a execucao existente"
                )
            if record.status != "WAITING_INPUT":
                raise WorkflowExecutionStateError(
                    f"Execucao {execution_id} nao esta aguardando input"
                )

            cur.execute(
                """
                UPDATE workflow_execution
                SET status = %s,
                    updated_at = NOW()
                WHERE execution_id = %s
                  AND status = %s
                RETURNING
                    execution_id,
                    workflow_name,
                    workflow_version,
                    status,
                    current_node,
                    resume_from,
                    expected_input_key,
                    created_at,
                    updated_at
                """,
                ("RUNNING", execution_id, "WAITING_INPUT"),
            )
            updated = cur.fetchone()
            if updated is None:
                raise WorkflowExecutionStateError(
                    f"Execucao {execution_id} foi retomada por outra requisicao"
                )

        return self._row_to_record(updated)

    def close(self) -> None:
        return None

    def _setup(self) -> None:
        with self._setup_lock:
            with self._connect() as conn, conn.cursor() as cur:
                cur.execute(
                    """
                    CREATE TABLE IF NOT EXISTS workflow_execution (
                        execution_id TEXT PRIMARY KEY,
                        workflow_name TEXT NOT NULL,
                        workflow_version INTEGER NOT NULL,
                        status TEXT NOT NULL,
                        current_node TEXT,
                        resume_from TEXT,
                        expected_input_key TEXT,
                        created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
                        updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
                    )
                    """
                )
                cur.execute(
                    """
                    CREATE INDEX IF NOT EXISTS idx_workflow_execution_status
                    ON workflow_execution (status)
                    """
                )

    @contextmanager
    def _connect(self) -> Iterator[Any]:
        psycopg = self._import_psycopg()
        conn = psycopg.connect(self._dsn, autocommit=False)
        try:
            yield conn
            conn.commit()
        except Exception:
            conn.rollback()
            raise
        finally:
            conn.close()

    @staticmethod
    def _import_psycopg() -> Any:
        try:
            import psycopg
        except ModuleNotFoundError as exc:  # pragma: no cover - depende do ambiente
            raise ModuleNotFoundError(
                "psycopg nao esta instalado. "
                "Adicione a dependencia para usar workflows em PostgreSQL."
            ) from exc
        return psycopg

    @staticmethod
    def _dict_row_factory() -> Any:
        try:
            from psycopg.rows import dict_row
        except ModuleNotFoundError as exc:  # pragma: no cover - depende do ambiente
            raise ModuleNotFoundError(
                "psycopg nao esta instalado. "
                "Adicione a dependencia para usar workflows em PostgreSQL."
            ) from exc
        return dict_row

    @staticmethod
    def _row_to_record(row: Any) -> WorkflowExecutionRecord:
        if row is None:
            raise WorkflowExecutionNotFoundError("Registro de workflow nao encontrado")

        return WorkflowExecutionRecord(
            execution_id=str(row["execution_id"]),
            workflow_name=str(row["workflow_name"]),
            workflow_version=int(row["workflow_version"]),
            status=str(row["status"]),  # type: ignore[arg-type]
            current_node=row["current_node"],
            resume_from=row["resume_from"],
            expected_input_key=row["expected_input_key"],
            created_at=PostgresExecutionStore._as_datetime(row["created_at"]),
            updated_at=PostgresExecutionStore._as_datetime(row["updated_at"]),
        )

    @staticmethod
    def _as_datetime(value: Any) -> datetime:
        if isinstance(value, datetime):
            return value
        return datetime.fromisoformat(str(value))