mirror of
https://github.com/hoshikawa2/first_contas.git
synced 2026-07-09 18:24:20 +00:00
first commit
This commit is contained in:
477
legacy_reference/workflows/execution_store.py
Normal file
477
legacy_reference/workflows/execution_store.py
Normal file
@@ -0,0 +1,477 @@
|
||||
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))
|
||||
Reference in New Issue
Block a user