From 4db700bd4d5b1a561bd7151befa083bb13c53eca Mon Sep 17 00:00:00 2001 From: Mikhail Date: Fri, 3 Jul 2026 21:18:32 +0300 Subject: [PATCH] Add SQLAlchemy persistence layer --- alembic.ini | 37 ++ docs/backlog/01_execution_backlog.md | 6 +- docs/storage/01_schema.md | 18 + migrations/env.py | 49 ++ migrations/script.py.mako | 27 + migrations/versions/0001_initial_schema.py | 143 +++++ src/ai_orchestrator/config.py | 3 +- src/ai_orchestrator/delivery/http/app.py | 20 +- .../infrastructure/storage/factory.py | 47 ++ .../infrastructure/storage/sqlalchemy.py | 517 ++++++++++++++++++ tests/integration/test_sqlalchemy_storage.py | 79 +++ 11 files changed, 930 insertions(+), 16 deletions(-) create mode 100644 alembic.ini create mode 100644 migrations/env.py create mode 100644 migrations/script.py.mako create mode 100644 migrations/versions/0001_initial_schema.py create mode 100644 src/ai_orchestrator/infrastructure/storage/factory.py create mode 100644 src/ai_orchestrator/infrastructure/storage/sqlalchemy.py create mode 100644 tests/integration/test_sqlalchemy_storage.py diff --git a/alembic.ini b/alembic.ini new file mode 100644 index 0000000..7fd9d78 --- /dev/null +++ b/alembic.ini @@ -0,0 +1,37 @@ +[alembic] +script_location = migrations +prepend_sys_path = . +sqlalchemy.url = sqlite:///runtime/ai_orchestrator.db + +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARN +handlers = console + +[logger_sqlalchemy] +level = WARN +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = console +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s + diff --git a/docs/backlog/01_execution_backlog.md b/docs/backlog/01_execution_backlog.md index b9c6d36..0f0fb0a 100644 --- a/docs/backlog/01_execution_backlog.md +++ b/docs/backlog/01_execution_backlog.md @@ -28,9 +28,9 @@ ## Package C: Persistence -- [ ] SQLAlchemy models -- [ ] repositories for PostgreSQL/SQLite -- [ ] migrations +- [x] SQLAlchemy models +- [x] repositories for PostgreSQL/SQLite +- [x] migrations - [ ] artifact metadata persistence - [ ] event store queries with pagination diff --git a/docs/storage/01_schema.md b/docs/storage/01_schema.md index c1eee2a..8338987 100644 --- a/docs/storage/01_schema.md +++ b/docs/storage/01_schema.md @@ -5,6 +5,19 @@ - production: PostgreSQL - local development: SQLite +## Persistence Path In Code + +SQLAlchemy adapter реализован в: + +- [src/ai_orchestrator/infrastructure/storage/sqlalchemy.py](/Z:/codex/ai_orchestrator/src/ai_orchestrator/infrastructure/storage/sqlalchemy.py) +- [src/ai_orchestrator/infrastructure/storage/factory.py](/Z:/codex/ai_orchestrator/src/ai_orchestrator/infrastructure/storage/factory.py) + +Migration skeleton: + +- [alembic.ini](/Z:/codex/ai_orchestrator/alembic.ini) +- [migrations/env.py](/Z:/codex/ai_orchestrator/migrations/env.py) +- [migrations/versions/0001_initial_schema.py](/Z:/codex/ai_orchestrator/migrations/versions/0001_initial_schema.py) + ## Tables ### projects @@ -156,3 +169,8 @@ - `confirmations(status, created_at)` - `worker_sessions(status, last_heartbeat_at)` +## Active Adapter Strategy + +- `memory` backend остается для быстрых unit/integration прогонов; +- `sqlite` backend предназначен для локального development без Docker; +- `postgres` backend является production target. diff --git a/migrations/env.py b/migrations/env.py new file mode 100644 index 0000000..3e29565 --- /dev/null +++ b/migrations/env.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +from logging.config import fileConfig + +from alembic import context +from sqlalchemy import engine_from_config, pool + +from ai_orchestrator.infrastructure.storage.sqlalchemy import Base + +config = context.config + +if config.config_file_name is not None: + fileConfig(config.config_file_name) + +target_metadata = Base.metadata + + +def run_migrations_offline() -> None: + url = config.get_main_option("sqlalchemy.url") + context.configure( + url=url, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + ) + + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online() -> None: + connectable = engine_from_config( + config.get_section(config.config_ini_section, {}), + prefix="sqlalchemy.", + poolclass=pool.NullPool, + ) + + with connectable.connect() as connection: + context.configure(connection=connection, target_metadata=target_metadata) + + with context.begin_transaction(): + context.run_migrations() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() + diff --git a/migrations/script.py.mako b/migrations/script.py.mako new file mode 100644 index 0000000..08e3f5f --- /dev/null +++ b/migrations/script.py.mako @@ -0,0 +1,27 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} +""" +from __future__ import annotations + +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + + +# revision identifiers, used by Alembic. +revision = ${repr(up_revision)} +down_revision = ${repr(down_revision)} +branch_labels = ${repr(branch_labels)} +depends_on = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} + diff --git a/migrations/versions/0001_initial_schema.py b/migrations/versions/0001_initial_schema.py new file mode 100644 index 0000000..6f91ab4 --- /dev/null +++ b/migrations/versions/0001_initial_schema.py @@ -0,0 +1,143 @@ +"""Initial schema. + +Revision ID: 0001_initial_schema +Revises: +Create Date: 2026-07-03 00:00:00 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "0001_initial_schema" +down_revision = None +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.create_table( + "tasks", + sa.Column("task_id", sa.String(length=64), primary_key=True), + sa.Column("project_id", sa.String(length=128), nullable=False), + sa.Column("conversation_id", sa.String(length=128), nullable=True), + sa.Column("goal", sa.Text(), nullable=False), + sa.Column("inputs_json", sa.JSON(), nullable=False), + sa.Column("status", sa.String(length=64), nullable=False), + sa.Column("requested_mode", sa.String(length=64), nullable=False), + sa.Column("effective_mode", sa.String(length=64), nullable=False), + sa.Column("current_node_id", sa.String(length=64), nullable=True), + sa.Column("result_summary_json", sa.JSON(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + ) + op.create_index("ix_tasks_project_id", "tasks", ["project_id"]) + op.create_index("ix_tasks_conversation_id", "tasks", ["conversation_id"]) + op.create_index("ix_tasks_status", "tasks", ["status"]) + op.create_index("ix_tasks_updated_at", "tasks", ["updated_at"]) + + op.create_table( + "task_graphs", + sa.Column("task_id", sa.String(length=64), primary_key=True), + sa.Column("graph_version", sa.Integer(), nullable=False), + sa.Column("graph_metadata_json", sa.JSON(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + ) + + op.create_table( + "task_nodes", + sa.Column("node_id", sa.String(length=64), primary_key=True), + sa.Column("task_id", sa.String(length=64), nullable=False), + sa.Column("node_type", sa.String(length=64), nullable=False), + sa.Column("status", sa.String(length=64), nullable=False), + sa.Column("input_json", sa.JSON(), nullable=False), + sa.Column("output_json", sa.JSON(), nullable=False), + sa.Column("dependencies_json", sa.JSON(), nullable=False), + sa.Column("assigned_runner", sa.String(length=128), nullable=True), + sa.Column("attempts", sa.Integer(), nullable=False), + sa.Column("retryable", sa.Boolean(), nullable=False), + sa.Column("timeout_ms", sa.Integer(), nullable=True), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + ) + op.create_index("ix_task_nodes_task_id", "task_nodes", ["task_id"]) + op.create_index("ix_task_nodes_status", "task_nodes", ["status"]) + + op.create_table( + "confirmations", + sa.Column("confirmation_id", sa.String(length=64), primary_key=True), + sa.Column("task_id", sa.String(length=64), nullable=False), + sa.Column("node_id", sa.String(length=64), nullable=False), + sa.Column("status", sa.String(length=64), nullable=False), + sa.Column("scope", sa.String(length=32), nullable=False), + sa.Column("decision_json", sa.JSON(), nullable=False), + sa.Column("preview_json", sa.JSON(), nullable=False), + sa.Column("comment", sa.Text(), nullable=True), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("resolved_at", sa.DateTime(timezone=True), nullable=True), + ) + op.create_index("ix_confirmations_task_id", "confirmations", ["task_id"]) + op.create_index("ix_confirmations_node_id", "confirmations", ["node_id"]) + op.create_index("ix_confirmations_status", "confirmations", ["status"]) + + op.create_table( + "worker_sessions", + sa.Column("session_id", sa.String(length=64), primary_key=True), + sa.Column("worker_id", sa.String(length=128), nullable=False), + sa.Column("name", sa.String(length=255), nullable=False), + sa.Column("machine", sa.String(length=255), nullable=False), + sa.Column("os", sa.String(length=64), nullable=False), + sa.Column("version", sa.String(length=64), nullable=False), + sa.Column("status", sa.String(length=64), nullable=False), + sa.Column("capabilities_json", sa.JSON(), nullable=False), + sa.Column("last_heartbeat_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("current_task_id", sa.String(length=64), nullable=True), + sa.Column("connected_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("disconnected_at", sa.DateTime(timezone=True), nullable=True), + ) + op.create_index("ix_worker_sessions_worker_id", "worker_sessions", ["worker_id"]) + op.create_index("ix_worker_sessions_status", "worker_sessions", ["status"]) + op.create_index( + "ix_worker_sessions_last_heartbeat_at", + "worker_sessions", + ["last_heartbeat_at"], + ) + + op.create_table( + "task_events", + sa.Column("event_id", sa.String(length=64), primary_key=True), + sa.Column("task_id", sa.String(length=64), nullable=False), + sa.Column("conversation_id", sa.String(length=128), nullable=True), + sa.Column("node_id", sa.String(length=64), nullable=True), + sa.Column("event_type", sa.String(length=128), nullable=False), + sa.Column("correlation_id", sa.String(length=64), nullable=True), + sa.Column("causation_id", sa.String(length=64), nullable=True), + sa.Column("payload_json", sa.JSON(), nullable=False), + sa.Column("occurred_at", sa.DateTime(timezone=True), nullable=False), + ) + op.create_index("ix_task_events_task_id", "task_events", ["task_id"]) + op.create_index("ix_task_events_occurred_at", "task_events", ["occurred_at"]) + + +def downgrade() -> None: + op.drop_index("ix_task_events_occurred_at", table_name="task_events") + op.drop_index("ix_task_events_task_id", table_name="task_events") + op.drop_table("task_events") + op.drop_index("ix_worker_sessions_last_heartbeat_at", table_name="worker_sessions") + op.drop_index("ix_worker_sessions_status", table_name="worker_sessions") + op.drop_index("ix_worker_sessions_worker_id", table_name="worker_sessions") + op.drop_table("worker_sessions") + op.drop_index("ix_confirmations_status", table_name="confirmations") + op.drop_index("ix_confirmations_node_id", table_name="confirmations") + op.drop_index("ix_confirmations_task_id", table_name="confirmations") + op.drop_table("confirmations") + op.drop_index("ix_task_nodes_status", table_name="task_nodes") + op.drop_index("ix_task_nodes_task_id", table_name="task_nodes") + op.drop_table("task_nodes") + op.drop_table("task_graphs") + op.drop_index("ix_tasks_updated_at", table_name="tasks") + op.drop_index("ix_tasks_status", table_name="tasks") + op.drop_index("ix_tasks_conversation_id", table_name="tasks") + op.drop_index("ix_tasks_project_id", table_name="tasks") + op.drop_table("tasks") diff --git a/src/ai_orchestrator/config.py b/src/ai_orchestrator/config.py index 2e0f6eb..5314f23 100644 --- a/src/ai_orchestrator/config.py +++ b/src/ai_orchestrator/config.py @@ -78,5 +78,6 @@ class AppSettings(BaseSettings): env: str = "dev" host: str = "127.0.0.1" port: int = 8080 + storage_backend: str = "memory" + database_url: str = "sqlite:///runtime/ai_orchestrator.db" default_project_config: str = "configs/examples/local_only.yaml" - diff --git a/src/ai_orchestrator/delivery/http/app.py b/src/ai_orchestrator/delivery/http/app.py index 2ed6f94..3526106 100644 --- a/src/ai_orchestrator/delivery/http/app.py +++ b/src/ai_orchestrator/delivery/http/app.py @@ -29,23 +29,18 @@ from ai_orchestrator.delivery.http.schemas import ( ) from ai_orchestrator.infrastructure.config_loader import load_project_config from ai_orchestrator.infrastructure.policy import StaticProjectPolicyEvaluator -from ai_orchestrator.infrastructure.storage.memory import ( - InMemoryConfirmationRepository, - InMemoryEventStore, - InMemoryGraphRepository, - InMemoryTaskRepository, - InMemoryWorkerRepository, -) +from ai_orchestrator.infrastructure.storage.factory import create_storage_bundle def create_app(settings: AppSettings | None = None) -> FastAPI: settings = settings or AppSettings() project_config = load_project_config(settings.default_project_config) - task_repository = InMemoryTaskRepository() - graph_repository = InMemoryGraphRepository() - confirmation_repository = InMemoryConfirmationRepository() - worker_repository = InMemoryWorkerRepository() - event_store = InMemoryEventStore() + storage = create_storage_bundle(settings) + task_repository = storage.task_repository + graph_repository = storage.graph_repository + confirmation_repository = storage.confirmation_repository + worker_repository = storage.worker_repository + event_store = storage.event_store policy_evaluator = StaticProjectPolicyEvaluator( projects={project_config.project.id: project_config} ) @@ -77,6 +72,7 @@ def create_app(settings: AppSettings | None = None) -> FastAPI: app.state.confirmation_repository = confirmation_repository app.state.worker_repository = worker_repository app.state.event_store = event_store + app.state.storage = storage app.state.orchestrator = orchestrator app.state.execution_engine = execution_engine app.state.lifecycle_service = lifecycle_service diff --git a/src/ai_orchestrator/infrastructure/storage/factory.py b/src/ai_orchestrator/infrastructure/storage/factory.py new file mode 100644 index 0000000..b44bb6e --- /dev/null +++ b/src/ai_orchestrator/infrastructure/storage/factory.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from ai_orchestrator.config import AppSettings +from ai_orchestrator.infrastructure.storage.memory import ( + InMemoryConfirmationRepository, + InMemoryEventStore, + InMemoryGraphRepository, + InMemoryTaskRepository, + InMemoryWorkerRepository, +) +from ai_orchestrator.infrastructure.storage.sqlalchemy import create_sqlalchemy_storage + + +@dataclass(slots=True) +class StorageBundle: + task_repository: object + graph_repository: object + confirmation_repository: object + worker_repository: object + event_store: object + engine: object | None = None + session_factory: object | None = None + + +def create_storage_bundle(settings: AppSettings) -> StorageBundle: + if settings.storage_backend == "memory": + return StorageBundle( + task_repository=InMemoryTaskRepository(), + graph_repository=InMemoryGraphRepository(), + confirmation_repository=InMemoryConfirmationRepository(), + worker_repository=InMemoryWorkerRepository(), + event_store=InMemoryEventStore(), + ) + if settings.storage_backend in {"sqlite", "postgres"}: + bundle = create_sqlalchemy_storage(settings.database_url) + return StorageBundle( + engine=bundle.engine, + session_factory=bundle.session_factory, + task_repository=bundle.task_repository, + graph_repository=bundle.graph_repository, + confirmation_repository=bundle.confirmation_repository, + worker_repository=bundle.worker_repository, + event_store=bundle.event_store, + ) + raise ValueError(f"Unsupported storage backend: {settings.storage_backend}") diff --git a/src/ai_orchestrator/infrastructure/storage/sqlalchemy.py b/src/ai_orchestrator/infrastructure/storage/sqlalchemy.py new file mode 100644 index 0000000..183cabf --- /dev/null +++ b/src/ai_orchestrator/infrastructure/storage/sqlalchemy.py @@ -0,0 +1,517 @@ +from __future__ import annotations + +from dataclasses import dataclass +from datetime import UTC, datetime +from pathlib import Path + +from sqlalchemy import JSON, Boolean, DateTime, Integer, String, Text, create_engine, select +from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column, sessionmaker + +from ai_orchestrator.application.ports import ( + ConfirmationRepository, + EventStore, + GraphRepository, + TaskRepository, + WorkerRepository, +) +from ai_orchestrator.domain.enums import ( + ConfirmationScope, + ConfirmationStatus, + NodeStatus, + NodeType, + PolicyDecisionType, + TaskStatus, + WorkerSessionStatus, +) +from ai_orchestrator.domain.events import DomainEvent +from ai_orchestrator.domain.models import ( + ConfirmationRequest, + ExecutionGraph, + ExecutionNode, + PolicyDecision, + Task, + WorkerSession, +) + + +class Base(DeclarativeBase): + pass + + +class TaskRecord(Base): + __tablename__ = "tasks" + + task_id: Mapped[str] = mapped_column(String(64), primary_key=True) + project_id: Mapped[str] = mapped_column(String(128), index=True) + conversation_id: Mapped[str | None] = mapped_column(String(128), nullable=True, index=True) + goal: Mapped[str] = mapped_column(Text) + inputs_json: Mapped[dict] = mapped_column(JSON) + status: Mapped[str] = mapped_column(String(64), index=True) + requested_mode: Mapped[str] = mapped_column(String(64)) + effective_mode: Mapped[str] = mapped_column(String(64)) + current_node_id: Mapped[str | None] = mapped_column(String(64), nullable=True) + result_summary_json: Mapped[dict] = mapped_column(JSON) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True)) + updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), index=True) + + +class TaskGraphRecord(Base): + __tablename__ = "task_graphs" + + task_id: Mapped[str] = mapped_column(String(64), primary_key=True) + graph_version: Mapped[int] = mapped_column(Integer, default=1) + graph_metadata_json: Mapped[dict] = mapped_column(JSON) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True)) + + +class TaskNodeRecord(Base): + __tablename__ = "task_nodes" + + node_id: Mapped[str] = mapped_column(String(64), primary_key=True) + task_id: Mapped[str] = mapped_column(String(64), index=True) + node_type: Mapped[str] = mapped_column(String(64)) + status: Mapped[str] = mapped_column(String(64), index=True) + input_json: Mapped[dict] = mapped_column(JSON) + output_json: Mapped[dict] = mapped_column(JSON) + dependencies_json: Mapped[list] = mapped_column(JSON) + assigned_runner: Mapped[str | None] = mapped_column(String(128), nullable=True) + attempts: Mapped[int] = mapped_column(Integer, default=0) + retryable: Mapped[bool] = mapped_column(Boolean, default=True) + timeout_ms: Mapped[int | None] = mapped_column(Integer, nullable=True) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True)) + updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True)) + + +class ConfirmationRecord(Base): + __tablename__ = "confirmations" + + confirmation_id: Mapped[str] = mapped_column(String(64), primary_key=True) + task_id: Mapped[str] = mapped_column(String(64), index=True) + node_id: Mapped[str] = mapped_column(String(64), index=True) + status: Mapped[str] = mapped_column(String(64), index=True) + scope: Mapped[str] = mapped_column(String(32)) + decision_json: Mapped[dict] = mapped_column(JSON) + preview_json: Mapped[dict] = mapped_column(JSON) + comment: Mapped[str | None] = mapped_column(Text, nullable=True) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True)) + resolved_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) + + +class WorkerSessionRecord(Base): + __tablename__ = "worker_sessions" + + session_id: Mapped[str] = mapped_column(String(64), primary_key=True) + worker_id: Mapped[str] = mapped_column(String(128), index=True) + name: Mapped[str] = mapped_column(String(255)) + machine: Mapped[str] = mapped_column(String(255)) + os: Mapped[str] = mapped_column(String(64)) + version: Mapped[str] = mapped_column(String(64)) + status: Mapped[str] = mapped_column(String(64), index=True) + capabilities_json: Mapped[list] = mapped_column(JSON) + last_heartbeat_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), index=True) + current_task_id: Mapped[str | None] = mapped_column(String(64), nullable=True) + connected_at: Mapped[datetime] = mapped_column(DateTime(timezone=True)) + disconnected_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) + + +class TaskEventRecord(Base): + __tablename__ = "task_events" + + event_id: Mapped[str] = mapped_column(String(64), primary_key=True) + task_id: Mapped[str] = mapped_column(String(64), index=True) + conversation_id: Mapped[str | None] = mapped_column(String(128), nullable=True) + node_id: Mapped[str | None] = mapped_column(String(64), nullable=True) + event_type: Mapped[str] = mapped_column(String(128)) + correlation_id: Mapped[str | None] = mapped_column(String(64), nullable=True) + causation_id: Mapped[str | None] = mapped_column(String(64), nullable=True) + payload_json: Mapped[dict] = mapped_column(JSON) + occurred_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), index=True) + + +def build_engine(database_url: str): + if database_url.startswith("sqlite:///"): + db_path = database_url.removeprefix("sqlite:///") + if db_path not in {":memory:", ""}: + Path(db_path).parent.mkdir(parents=True, exist_ok=True) + return create_engine(database_url, future=True) + + +def build_session_factory(database_url: str): + engine = build_engine(database_url) + return engine, sessionmaker(bind=engine, autoflush=False, autocommit=False, future=True) + + +def create_all_tables(engine) -> None: + Base.metadata.create_all(engine) + + +def drop_all_tables(engine) -> None: + Base.metadata.drop_all(engine) + + +def _task_to_record(task: Task) -> TaskRecord: + return TaskRecord( + task_id=task.task_id, + project_id=task.project_id, + conversation_id=task.conversation_id, + goal=task.goal, + inputs_json=task.inputs, + status=task.status.value, + requested_mode=task.requested_mode, + effective_mode=task.effective_mode, + current_node_id=task.current_node_id, + result_summary_json=task.result_summary, + created_at=task.created_at, + updated_at=task.updated_at, + ) + + +def _record_to_task(record: TaskRecord) -> Task: + return Task( + project_id=record.project_id, + goal=record.goal, + inputs=record.inputs_json, + conversation_id=record.conversation_id, + requested_mode=record.requested_mode, + effective_mode=record.effective_mode, + task_id=record.task_id, + status=TaskStatus(record.status), + current_node_id=record.current_node_id, + result_summary=record.result_summary_json, + created_at=record.created_at, + updated_at=record.updated_at, + ) + + +def _node_to_record(node: ExecutionNode) -> TaskNodeRecord: + return TaskNodeRecord( + node_id=node.node_id, + task_id=node.task_id, + node_type=node.node_type.value, + status=node.status.value, + input_json=node.input_data, + output_json=node.output_data, + dependencies_json=node.dependencies, + assigned_runner=node.assigned_runner, + attempts=node.attempts, + retryable=node.retryable, + timeout_ms=node.timeout_ms, + created_at=node.created_at, + updated_at=node.updated_at, + ) + + +def _record_to_node(record: TaskNodeRecord) -> ExecutionNode: + return ExecutionNode( + task_id=record.task_id, + node_type=NodeType(record.node_type), + input_data=record.input_json, + dependencies=list(record.dependencies_json or []), + node_id=record.node_id, + status=NodeStatus(record.status), + output_data=record.output_json or {}, + assigned_runner=record.assigned_runner, + attempts=record.attempts, + retryable=record.retryable, + timeout_ms=record.timeout_ms, + created_at=record.created_at, + updated_at=record.updated_at, + ) + + +def _confirmation_to_record(item: ConfirmationRequest) -> ConfirmationRecord: + return ConfirmationRecord( + confirmation_id=item.confirmation_id, + task_id=item.task_id, + node_id=item.node_id, + status=item.status.value, + scope=item.scope.value, + decision_json={ + "decision": item.decision.decision.value, + "reason": item.decision.reason, + "requires_confirmation": item.decision.requires_confirmation, + }, + preview_json=item.decision.preview, + comment=item.comment, + created_at=item.created_at, + resolved_at=item.resolved_at, + ) + + +def _record_to_confirmation(record: ConfirmationRecord) -> ConfirmationRequest: + return ConfirmationRequest( + task_id=record.task_id, + node_id=record.node_id, + decision=PolicyDecision( + decision=PolicyDecisionType(record.decision_json["decision"]), + reason=record.decision_json["reason"], + requires_confirmation=record.decision_json["requires_confirmation"], + preview=record.preview_json or {}, + ), + scope=ConfirmationScope(record.scope), + confirmation_id=record.confirmation_id, + status=ConfirmationStatus(record.status), + comment=record.comment, + created_at=record.created_at, + resolved_at=record.resolved_at, + ) + + +def _worker_to_record(worker: WorkerSession) -> WorkerSessionRecord: + return WorkerSessionRecord( + session_id=worker.session_id, + worker_id=worker.worker_id, + name=worker.name, + machine=worker.machine, + os=worker.os, + version=worker.version, + status=worker.status.value, + capabilities_json=worker.capabilities, + last_heartbeat_at=worker.last_heartbeat_at, + current_task_id=worker.current_task_id, + connected_at=worker.connected_at, + disconnected_at=worker.disconnected_at, + ) + + +def _record_to_worker(record: WorkerSessionRecord) -> WorkerSession: + return WorkerSession( + worker_id=record.worker_id, + name=record.name, + machine=record.machine, + os=record.os, + version=record.version, + capabilities=list(record.capabilities_json or []), + session_id=record.session_id, + status=WorkerSessionStatus(record.status), + current_task_id=record.current_task_id, + last_heartbeat_at=record.last_heartbeat_at, + connected_at=record.connected_at, + disconnected_at=record.disconnected_at, + ) + + +def _event_to_record(event: DomainEvent) -> TaskEventRecord: + return TaskEventRecord( + event_id=event.event_id, + task_id=event.task_id, + conversation_id=event.conversation_id, + node_id=event.node_id, + event_type=event.event_type, + correlation_id=event.correlation_id, + causation_id=event.causation_id, + payload_json=event.payload, + occurred_at=event.occurred_at, + ) + + +def _record_to_event(record: TaskEventRecord) -> DomainEvent: + return DomainEvent( + event_type=record.event_type, + task_id=record.task_id, + payload=record.payload_json or {}, + conversation_id=record.conversation_id, + node_id=record.node_id, + correlation_id=record.correlation_id, + causation_id=record.causation_id, + event_id=record.event_id, + occurred_at=record.occurred_at, + ) + + +@dataclass(slots=True) +class SqlAlchemyTaskRepository(TaskRepository): + session_factory: sessionmaker[Session] + + def create(self, task: Task) -> Task: + with self.session_factory() as session: + session.add(_task_to_record(task)) + session.commit() + return task + + def save(self, task: Task) -> Task: + with self.session_factory() as session: + record = session.get(TaskRecord, task.task_id) + if record is None: + record = _task_to_record(task) + session.add(record) + else: + record.project_id = task.project_id + record.conversation_id = task.conversation_id + record.goal = task.goal + record.inputs_json = task.inputs + record.status = task.status.value + record.requested_mode = task.requested_mode + record.effective_mode = task.effective_mode + record.current_node_id = task.current_node_id + record.result_summary_json = task.result_summary + record.created_at = task.created_at + record.updated_at = task.updated_at + session.commit() + return task + + def get(self, task_id: str) -> Task | None: + with self.session_factory() as session: + record = session.get(TaskRecord, task_id) + return _record_to_task(record) if record else None + + +@dataclass(slots=True) +class SqlAlchemyGraphRepository(GraphRepository): + session_factory: sessionmaker[Session] + + def save(self, graph: ExecutionGraph) -> ExecutionGraph: + with self.session_factory() as session: + graph_record = session.get(TaskGraphRecord, graph.task_id) + if graph_record is None: + graph_record = TaskGraphRecord( + task_id=graph.task_id, + graph_version=1, + graph_metadata_json={"node_count": len(graph.nodes)}, + created_at=max( + (node.created_at for node in graph.nodes), + default=datetime.now(UTC), + ), + ) + session.add(graph_record) + else: + graph_record.graph_version += 1 + graph_record.graph_metadata_json = {"node_count": len(graph.nodes)} + existing = session.scalars( + select(TaskNodeRecord).where(TaskNodeRecord.task_id == graph.task_id) + ).all() + for record in existing: + session.delete(record) + for node in graph.nodes: + session.add(_node_to_record(node)) + session.commit() + return graph + + def get(self, task_id: str) -> ExecutionGraph | None: + with self.session_factory() as session: + graph_record = session.get(TaskGraphRecord, task_id) + if graph_record is None: + return None + node_records = session.scalars( + select(TaskNodeRecord).where(TaskNodeRecord.task_id == task_id) + ).all() + nodes = [_record_to_node(record) for record in node_records] + return ExecutionGraph(task_id=task_id, nodes=nodes) + + +@dataclass(slots=True) +class SqlAlchemyConfirmationRepository(ConfirmationRepository): + session_factory: sessionmaker[Session] + + def create(self, confirmation: ConfirmationRequest) -> ConfirmationRequest: + with self.session_factory() as session: + session.add(_confirmation_to_record(confirmation)) + session.commit() + return confirmation + + def get(self, confirmation_id: str) -> ConfirmationRequest | None: + with self.session_factory() as session: + record = session.get(ConfirmationRecord, confirmation_id) + return _record_to_confirmation(record) if record else None + + def save(self, confirmation: ConfirmationRequest) -> ConfirmationRequest: + with self.session_factory() as session: + record = session.get(ConfirmationRecord, confirmation.confirmation_id) + if record is None: + record = _confirmation_to_record(confirmation) + session.add(record) + else: + updated = _confirmation_to_record(confirmation) + record.task_id = updated.task_id + record.node_id = updated.node_id + record.status = updated.status + record.scope = updated.scope + record.decision_json = updated.decision_json + record.preview_json = updated.preview_json + record.comment = updated.comment + record.created_at = updated.created_at + record.resolved_at = updated.resolved_at + session.commit() + return confirmation + + +@dataclass(slots=True) +class SqlAlchemyWorkerRepository(WorkerRepository): + session_factory: sessionmaker[Session] + + def save(self, worker: WorkerSession) -> WorkerSession: + with self.session_factory() as session: + record = session.get(WorkerSessionRecord, worker.session_id) + if record is None: + record = _worker_to_record(worker) + session.add(record) + else: + updated = _worker_to_record(worker) + record.worker_id = updated.worker_id + record.name = updated.name + record.machine = updated.machine + record.os = updated.os + record.version = updated.version + record.status = updated.status + record.capabilities_json = updated.capabilities_json + record.last_heartbeat_at = updated.last_heartbeat_at + record.current_task_id = updated.current_task_id + record.connected_at = updated.connected_at + record.disconnected_at = updated.disconnected_at + session.commit() + return worker + + def get(self, session_id: str) -> WorkerSession | None: + with self.session_factory() as session: + record = session.get(WorkerSessionRecord, session_id) + return _record_to_worker(record) if record else None + + def list_active(self) -> list[WorkerSession]: + with self.session_factory() as session: + records = session.scalars(select(WorkerSessionRecord)).all() + return [_record_to_worker(record) for record in records] + + +@dataclass(slots=True) +class SqlAlchemyEventStore(EventStore): + session_factory: sessionmaker[Session] + + def append(self, event: DomainEvent) -> DomainEvent: + with self.session_factory() as session: + session.add(_event_to_record(event)) + session.commit() + return event + + def list_by_task(self, task_id: str) -> list[DomainEvent]: + with self.session_factory() as session: + records = session.scalars( + select(TaskEventRecord) + .where(TaskEventRecord.task_id == task_id) + .order_by(TaskEventRecord.occurred_at.asc()) + ).all() + return [_record_to_event(record) for record in records] + + +@dataclass(slots=True) +class SqlAlchemyStorageBundle: + engine: object + session_factory: sessionmaker[Session] + task_repository: SqlAlchemyTaskRepository + graph_repository: SqlAlchemyGraphRepository + confirmation_repository: SqlAlchemyConfirmationRepository + worker_repository: SqlAlchemyWorkerRepository + event_store: SqlAlchemyEventStore + + +def create_sqlalchemy_storage(database_url: str) -> SqlAlchemyStorageBundle: + engine, session_factory = build_session_factory(database_url) + create_all_tables(engine) + return SqlAlchemyStorageBundle( + engine=engine, + session_factory=session_factory, + task_repository=SqlAlchemyTaskRepository(session_factory), + graph_repository=SqlAlchemyGraphRepository(session_factory), + confirmation_repository=SqlAlchemyConfirmationRepository(session_factory), + worker_repository=SqlAlchemyWorkerRepository(session_factory), + event_store=SqlAlchemyEventStore(session_factory), + ) diff --git a/tests/integration/test_sqlalchemy_storage.py b/tests/integration/test_sqlalchemy_storage.py new file mode 100644 index 0000000..618eaeb --- /dev/null +++ b/tests/integration/test_sqlalchemy_storage.py @@ -0,0 +1,79 @@ +from pathlib import Path + +from ai_orchestrator.config import AppSettings +from ai_orchestrator.domain.enums import NodeType, PolicyDecisionType +from ai_orchestrator.domain.events import DomainEvent +from ai_orchestrator.domain.models import ( + ConfirmationRequest, + ExecutionGraph, + ExecutionNode, + PolicyDecision, + Task, + WorkerSession, +) +from ai_orchestrator.infrastructure.storage.factory import create_storage_bundle + + +def test_sqlite_storage_persists_task_graph_and_events(tmp_path: Path) -> None: + db_path = tmp_path / "orchestrator.sqlite3" + settings = AppSettings( + storage_backend="sqlite", + database_url=f"sqlite:///{db_path.as_posix()}", + ) + storage = create_storage_bundle(settings) + + task = Task(project_id="default", goal="Persist me", inputs={"a": 1}) + storage.task_repository.create(task) + node = ExecutionNode( + task_id=task.task_id, + node_type=NodeType.PLANNER, + input_data={"goal": "Persist me"}, + ) + graph = ExecutionGraph(task_id=task.task_id, nodes=[node]) + storage.graph_repository.save(graph) + storage.event_store.append( + DomainEvent(event_type="task_created", task_id=task.task_id, payload={"goal": task.goal}) + ) + + stored_task = storage.task_repository.get(task.task_id) + stored_graph = storage.graph_repository.get(task.task_id) + stored_events = storage.event_store.list_by_task(task.task_id) + + assert stored_task is not None + assert stored_task.goal == "Persist me" + assert stored_graph is not None + assert stored_graph.nodes[0].node_type == NodeType.PLANNER + assert stored_events[0].event_type == "task_created" + + +def test_sqlite_storage_persists_confirmation_and_worker(tmp_path: Path) -> None: + db_path = tmp_path / "orchestrator.sqlite3" + settings = AppSettings( + storage_backend="sqlite", + database_url=f"sqlite:///{db_path.as_posix()}", + ) + storage = create_storage_bundle(settings) + confirmation = ConfirmationRequest( + task_id="task_1", + node_id="node_1", + decision=PolicyDecision( + decision=PolicyDecisionType.CONFIRM, + reason="Need approval", + requires_confirmation=True, + ), + ) + worker = WorkerSession( + worker_id="worker_home_pc", + name="Home PC", + machine="DESKTOP-1", + os="windows", + version="0.1.0", + capabilities=["file.read"], + ) + + storage.confirmation_repository.create(confirmation) + storage.worker_repository.save(worker) + + assert storage.confirmation_repository.get(confirmation.confirmation_id) is not None + assert storage.worker_repository.get(worker.session_id) is not None + assert len(storage.worker_repository.list_active()) == 1