aethergraph 0.1.0a1__py3-none-any.whl → 0.1.0a3__py3-none-any.whl

aethergraph/storage/runs/sqlite_run_store.py

@@ -0,0 +1,403 @@
+# aethergraph/storage/sqlite_run_store.py
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import asdict, is_dataclass
+from datetime import datetime
+import json
+from pathlib import Path
+import sqlite3
+import threading
+from typing import Any
+
+from aethergraph.contracts.services.runs import RunStore
+from aethergraph.core.runtime.run_types import RunRecord, RunStatus
+
+
+def _dt_to_ts(dt: datetime | None) -> float | None:
+    if dt is None:
+        return None
+    if dt.tzinfo is None:
+        # assume UTC if naive
+        return dt.replace(tzinfo=datetime.timezone.utc).timestamp()
+    return dt.timestamp()
+
+
+def _encode_run(record: RunRecord) -> dict[str, Any]:
+    """Convert RunRecord -> plain dict with JSON-safe types."""
+    if is_dataclass(record):  # noqa: SIM108
+        data = asdict(record)
+    else:
+        # fallback; should not really happen
+        data = dict(record.__dict__)
+
+    for k, v in list(data.items()):
+        if isinstance(v, datetime):
+            data[k] = v.isoformat()
+    return data
+
+
+def _decode_run(data: dict[str, Any]) -> RunRecord:
+    """Convert dict from JSON back into RunRecord."""
+
+    # Best-effort datetime parsing for common fields
+    def _parse_dt(val: Any) -> datetime | None:
+        if val is None:
+            return None
+        if isinstance(val, datetime):
+            return val
+        if isinstance(val, str):
+            try:
+                return datetime.fromisoformat(val)
+            except Exception:
+                return None
+        return None
+
+    for key in (
+        "created_at",
+        "updated_at",
+        "started_at",
+        "finished_at",
+        "first_artifact_at",
+        "last_artifact_at",
+    ):
+        if key in data:
+            parsed = _parse_dt(data[key])
+            if parsed is not None:
+                data[key] = parsed
+
+    return RunRecord(**data)
+
+
+class SQLiteRunStoreSync:
+    """
+    SQLite-backed RunStore.
+
+    - Stores full RunRecord as JSON in `data_json`
+    - Promotes a few fields to columns for fast filtering:
+        run_id, graph_id, status, user_id, org_id, session_id,
+        started_at, finished_at
+    """
+
+    def __init__(self, path: str):
+        path_obj = Path(path)
+        path_obj.parent.mkdir(parents=True, exist_ok=True)
+
+        self._db = sqlite3.connect(
+            str(path_obj),
+            check_same_thread=False,
+            isolation_level=None,  # autocommit
+        )
+        self._db.execute("PRAGMA journal_mode=WAL;")
+        self._db.execute("PRAGMA synchronous=NORMAL;")
+
+        # Base table
+        self._db.execute(
+            """
+            CREATE TABLE IF NOT EXISTS runs (
+                run_id       TEXT PRIMARY KEY,
+                data_json    TEXT NOT NULL,
+                graph_id     TEXT,
+                status       TEXT,
+                user_id      TEXT,
+                org_id       TEXT,
+                session_id   TEXT,
+                started_at   REAL,
+                finished_at  REAL
+            )
+            """
+        )
+
+        # Indices for common queries
+        self._db.execute(
+            "CREATE INDEX IF NOT EXISTS idx_runs_graph_started ON runs(graph_id, started_at DESC)"
+        )
+        self._db.execute(
+            "CREATE INDEX IF NOT EXISTS idx_runs_status_started ON runs(status, started_at DESC)"
+        )
+        self._db.execute(
+            "CREATE INDEX IF NOT EXISTS idx_runs_user_started ON runs(user_id, started_at DESC)"
+        )
+        self._db.execute(
+            "CREATE INDEX IF NOT EXISTS idx_runs_org_started ON runs(org_id, started_at DESC)"
+        )
+        self._db.execute(
+            "CREATE INDEX IF NOT EXISTS idx_runs_session_started ON runs(session_id, started_at DESC)"
+        )
+
+        self._lock = threading.RLock()
+
+    # --- core ops ---
+
+    def create(self, record: RunRecord) -> None:
+        data = _encode_run(record)
+        payload = json.dumps(data, ensure_ascii=False)
+        started_ts = _dt_to_ts(getattr(record, "started_at", None))
+        finished_ts = _dt_to_ts(getattr(record, "finished_at", None))
+
+        with self._lock:
+            self._db.execute(
+                """
+                INSERT INTO runs (
+                    run_id, data_json,
+                    graph_id, status,
+                    user_id, org_id, session_id,
+                    started_at, finished_at
+                )
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+                """,
+                (
+                    record.run_id,
+                    payload,
+                    getattr(record, "graph_id", None),
+                    record.status.value
+                    if isinstance(record.status, RunStatus)
+                    else str(record.status),
+                    getattr(record, "user_id", None),
+                    getattr(record, "org_id", None),
+                    getattr(record, "session_id", None),
+                    started_ts,
+                    finished_ts,
+                ),
+            )
+
+    def update_status(
+        self,
+        run_id: str,
+        status: RunStatus,
+        *,
+        finished_at: datetime | None = None,
+        error: str | None = None,
+    ) -> None:
+        with self._lock:
+            row = self._db.execute(
+                "SELECT data_json FROM runs WHERE run_id = ?",
+                (run_id,),
+            ).fetchone()
+            if not row:
+                return
+
+            data = json.loads(row[0])
+            data["status"] = status.value if isinstance(status, RunStatus) else str(status)
+            if finished_at is not None:
+                data["finished_at"] = finished_at.isoformat()
+            if error is not None:
+                data["error"] = error
+
+            payload = json.dumps(data, ensure_ascii=False)
+            finished_ts = _dt_to_ts(finished_at)
+
+            self._db.execute(
+                """
+                UPDATE runs
+                SET data_json = ?, status = ?, finished_at = ?
+                WHERE run_id = ?
+                """,
+                (payload, status.value, finished_ts, run_id),
+            )
+
+    def get(self, run_id: str) -> RunRecord | None:
+        with self._lock:
+            row = self._db.execute(
+                "SELECT data_json FROM runs WHERE run_id = ?",
+                (run_id,),
+            ).fetchone()
+        if not row:
+            return None
+        data = json.loads(row[0])
+        return _decode_run(data)
+
+    def list(
+        self,
+        *,
+        graph_id: str | None = None,
+        status: RunStatus | None = None,
+        user_id: str | None = None,
+        org_id: str | None = None,
+        session_id: str | None = None,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> list[RunRecord]:
+        """
+        List runs ordered by started_at DESC.
+
+        NOTE: session_id is optional; you can ignore it if you want to keep
+        the signature 100% identical to your current RunStore, or add it
+        and update RunManager accordingly.
+        """
+        where: list[str] = []
+        params: list[Any] = []
+
+        if graph_id is not None:
+            where.append("graph_id = ?")
+            params.append(graph_id)
+
+        if status is not None:
+            where.append("status = ?")
+            status_val = status.value if isinstance(status, RunStatus) else str(status)
+            params.append(status_val)
+
+        if org_id is not None:
+            where.append("org_id = ?")
+            params.append(org_id)
+
+        if user_id is not None:
+            where.append("user_id = ?")
+            params.append(user_id)
+
+        if session_id is not None:
+            where.append("session_id = ?")
+            params.append(session_id)
+
+        sql = "SELECT data_json FROM runs"
+        if where:
+            sql += " WHERE " + " AND ".join(where)
+        sql += " ORDER BY started_at DESC"
+
+        if limit is not None:
+            sql += " LIMIT ? OFFSET ?"
+            params.extend([limit, offset])
+
+        with self._lock:
+            rows = self._db.execute(sql, params).fetchall()
+
+        return [_decode_run(json.loads(r[0])) for r in rows]
+
+    def record_artifact(
+        self,
+        run_id: str,
+        *,
+        artifact_id: str,
+        created_at: datetime | None = None,
+        max_recent: int = 10,
+    ) -> None:
+        """
+        Optional API used by ArtifactFacade._record via getattr(..., 'record_artifact', None).
+
+        Updates artifact-related metadata:
+
+          - artifact_count
+          - first_artifact_at
+          - last_artifact_at
+          - recent_artifact_ids (bounded to `max_recent`)
+
+        No-op if the run does not exist.
+        """
+        with self._lock:
+            row = self._db.execute(
+                "SELECT data_json FROM runs WHERE run_id = ?",
+                (run_id,),
+            ).fetchone()
+
+            if not row:
+                return
+
+            # Decode current RunRecord from JSON
+            data = json.loads(row[0])
+            record = _decode_run(data)
+
+            # Choose timestamp
+            ts = created_at or datetime.utcnow()
+
+            # Update stats
+            record.artifact_count = (record.artifact_count or 0) + 1
+
+            if record.first_artifact_at is None or ts < record.first_artifact_at:
+                record.first_artifact_at = ts
+
+            if record.last_artifact_at is None or ts > record.last_artifact_at:
+                record.last_artifact_at = ts
+
+            # Maintain a small rolling window of recent IDs
+            if artifact_id:
+                recent = list(record.recent_artifact_ids or [])
+                recent.append(artifact_id)
+                record.recent_artifact_ids = recent[-max_recent:]
+
+            # Re-encode and persist JSON
+            new_data = _encode_run(record)
+            payload = json.dumps(new_data, ensure_ascii=False)
+
+            self._db.execute(
+                """
+                UPDATE runs
+                SET data_json = ?
+                WHERE run_id = ?
+                """,
+                (payload, run_id),
+            )
+
+
+class SQLiteRunStore(RunStore):
+    """
+    Async RunStore implementation that delegates to SQLiteRunStoreSync
+    using asyncio.to_thread for I/O.
+    """
+
+    def __init__(self, path: str):
+        self._sync = SQLiteRunStoreSync(path)
+
+    async def create(self, record: RunRecord) -> None:
+        await asyncio.to_thread(self._sync.create, record)
+
+    async def update_status(
+        self,
+        run_id: str,
+        status: RunStatus,
+        *,
+        finished_at: datetime | None = None,
+        error: str | None = None,
+    ) -> None:
+        await asyncio.to_thread(
+            self._sync.update_status,
+            run_id,
+            status,
+            finished_at=finished_at,
+            error=error,
+        )
+
+    async def get(self, run_id: str) -> RunRecord | None:
+        return await asyncio.to_thread(self._sync.get, run_id)
+
+    async def list(
+        self,
+        *,
+        graph_id: str | None = None,
+        status: RunStatus | None = None,
+        user_id: str | None = None,
+        org_id: str | None = None,
+        session_id: str | None = None,
+        limit: int = 100,
+        offset: int = 0,
+        # If you decide to expose session_id here, add it and thread it down.
+    ) -> list[RunRecord]:
+        # For now we only use graph_id/status; session_id can be added later
+        return await asyncio.to_thread(
+            self._sync.list,
+            graph_id=graph_id,
+            status=status,
+            user_id=user_id,
+            org_id=org_id,
+            session_id=session_id,
+            limit=limit,
+            offset=offset,
+        )
+
+    async def record_artifact(
+        self,
+        run_id: str,
+        *,
+        artifact_id: str,
+        created_at: datetime | None = None,
+    ) -> None:
+        """
+        Async façade for artifact stats update.
+        Called from ArtifactFacade._record via getattr(..., 'record_artifact', None).
+        """
+        await asyncio.to_thread(
+            self._sync.record_artifact,
+            run_id,
+            artifact_id=artifact_id,
+            created_at=created_at,
+        )

aethergraph/storage/sessions/doc_store.py

@@ -0,0 +1,183 @@
+from __future__ import annotations
+
+import asyncio
+from datetime import datetime, timezone
+from typing import Any
+import uuid
+
+from aethergraph.api.v1.schemas import Session
+from aethergraph.contracts.services.sessions import SessionStore
+from aethergraph.contracts.storage.doc_store import (
+    DocStore,  # wherever your DocStore Protocol lives
+)
+from aethergraph.core.runtime.run_types import SessionKind
+
+
+def _encode_dt(dt: datetime | None) -> str | None:
+    return dt.isoformat() if dt else None
+
+
+def _decode_dt(s: str | None) -> datetime | None:
+    if not s:
+        return None
+    if s.endswith("Z"):
+        s = s[:-1] + "+00:00"
+    return datetime.fromisoformat(s)
+
+
+def _session_to_doc(s: Session) -> dict[str, Any]:
+    d = s.model_dump() if hasattr(s, "model_dump") else dict(s)  # pydantic v2/v1 tolerant
+    d["created_at"] = _encode_dt(getattr(s, "created_at", None))
+    d["updated_at"] = _encode_dt(getattr(s, "updated_at", None))
+    return d
+
+
+def _doc_to_session(doc: dict[str, Any]) -> Session:
+    doc = dict(doc)
+    doc["created_at"] = _decode_dt(doc.get("created_at"))
+    doc["updated_at"] = _decode_dt(doc.get("updated_at"))
+    return Session(**doc)
+
+
+class DocSessionStore(SessionStore):
+    """
+    SessionStore backed by an arbitrary DocStore.
+
+    - Uses doc IDs like "<prefix><session_id>" (prefix defaults to "session:").
+    - Persists Session as JSON-friendly dicts (ISO datetimes).
+    - Supports FS-backed or SQLite-backed DocStore transparently.
+
+    The only requirement is that the underlying DocStore implements `list()`
+    if you want list_for_user() to work.
+    """
+
+    def __init__(self, doc_store: DocStore, *, prefix: str = "session:") -> None:
+        self._ds = doc_store
+        self._prefix = prefix
+        self._lock = asyncio.Lock()
+
+    def _doc_id(self, session_id: str) -> str:
+        return f"{self._prefix}{session_id}"
+
+    async def create(
+        self,
+        *,
+        kind: SessionKind,
+        user_id: str | None,
+        org_id: str | None,
+        title: str | None = None,
+        source: str = "webui",
+        external_ref: str | None = None,
+    ) -> Session:
+        now = datetime.now(timezone.utc)
+        session_id = f"sess_{uuid.uuid4().hex[:8]}"
+        sess = Session(
+            session_id=session_id,
+            kind=kind,
+            title=title,
+            user_id=user_id,
+            org_id=org_id,
+            source=source,
+            external_ref=external_ref,
+            created_at=now,
+            updated_at=now,
+        )
+
+        async with self._lock:
+            await self._ds.put(self._doc_id(session_id), _session_to_doc(sess))
+        return sess
+
+    async def get(self, session_id: str) -> Session | None:
+        doc_id = self._doc_id(session_id)
+        async with self._lock:
+            doc = await self._ds.get(doc_id)
+        if not doc:
+            return None
+        return _doc_to_session(doc)
+
+    async def list_for_user(
+        self,
+        *,
+        user_id: str | None,
+        org_id: str | None = None,
+        kind: SessionKind | None = None,
+        limit: int = 50,
+        offset: int = 0,
+    ) -> list[Session]:
+        # Same tradeoff as DocRunStore.list(): scan all, filter in Python
+        if not hasattr(self._ds, "list"):
+            raise RuntimeError(
+                "Underlying DocStore does not implement list(); "
+                "cannot support SessionStore.list_for_user()."
+            )
+
+        async with self._lock:
+            doc_ids: list[str] = await self._ds.list()  # type: ignore[attr-defined]
+            doc_ids = [d for d in doc_ids if d.startswith(self._prefix)]
+
+            records: list[Session] = []
+            for doc_id in doc_ids:
+                doc = await self._ds.get(doc_id)
+                if not doc:
+                    continue
+                sess = _doc_to_session(doc)
+
+                if user_id is not None and sess.user_id != user_id:
+                    continue
+                if org_id is not None and sess.org_id != org_id:
+                    continue
+                if kind is not None and sess.kind != kind:
+                    continue
+
+                records.append(sess)
+
+        records.sort(key=lambda s: s.created_at, reverse=True)
+
+        if offset > 0:
+            records = records[offset:]
+        if limit is not None:
+            records = records[:limit]
+        return records
+
+    async def touch(
+        self,
+        session_id: str,
+        *,
+        updated_at: datetime | None = None,
+    ) -> None:
+        doc_id = self._doc_id(session_id)
+        async with self._lock:
+            doc = await self._ds.get(doc_id)
+            if doc is None:
+                return
+            doc["updated_at"] = _encode_dt(updated_at or datetime.now(timezone.utc))
+            await self._ds.put(doc_id, doc)
+
+    async def update(
+        self,
+        session_id: str,
+        *,
+        title: str | None = None,
+        external_ref: str | None = None,
+    ) -> Session | None:
+        doc_id = self._doc_id(session_id)
+        async with self._lock:
+            doc = await self._ds.get(doc_id)
+            if doc is None:
+                return None
+
+            if title is not None:
+                doc["title"] = title
+            if external_ref is not None:
+                doc["external_ref"] = external_ref
+
+            # Always bump updated_at
+            doc["updated_at"] = _encode_dt(datetime.now(timezone.utc))
+
+            await self._ds.put(doc_id, doc)
+
+        return _doc_to_session(doc)
+
+    async def delete(self, session_id: str) -> None:
+        async with self._lock:
+            await self._ds.delete(self._doc_id(session_id))

aethergraph/storage/sessions/inmem_store.py

@@ -0,0 +1,110 @@
+import asyncio
+from datetime import datetime, timezone
+import uuid
+
+from aethergraph.api.v1.schemas import Session
+from aethergraph.contracts.services.sessions import SessionStore
+from aethergraph.core.runtime.run_types import SessionKind
+
+
+class InMemorySessionStore(SessionStore):
+    def __init__(self) -> None:
+        self._sessions: dict[str, Session] = {}
+        self._lock = asyncio.Lock()  # TODO: confirm async lock is fine bc this will only be used inside uvicorn process with UI.
+
+    async def create(
+        self,
+        *,
+        kind: SessionKind,
+        user_id: str | None,
+        org_id: str | None,
+        title: str | None = None,
+        source: str = "webui",
+        external_ref: str | None = None,
+    ) -> Session:
+        async with self._lock:
+            now = datetime.now(timezone.utc)
+            session_id = f"sess_{uuid.uuid4().hex[:8]}"
+            sess = Session(
+                session_id=session_id,
+                kind=kind,
+                title=title,
+                user_id=user_id,
+                org_id=org_id,
+                source=source,
+                external_ref=external_ref,
+                created_at=now,
+                updated_at=now,
+            )
+            self._sessions[session_id] = sess
+            return sess
+
+    async def get(self, session_id: str) -> Session | None:
+        async with self._lock:
+            return self._sessions.get(session_id)
+
+    async def list_for_user(
+        self,
+        *,
+        user_id: str | None,
+        org_id: str | None = None,
+        kind: SessionKind | None = None,
+        limit: int = 50,
+        offset: int = 0,
+    ) -> list[Session]:
+        async with self._lock:
+            records = list(self._sessions.values())
+
+            if user_id is not None:
+                records = [s for s in records if s.user_id == user_id]
+            if org_id is not None:
+                records = [s for s in records if s.org_id == org_id]
+            if kind is not None:
+                records = [s for s in records if s.kind == kind]
+
+            records.sort(key=lambda s: s.created_at, reverse=True)
+
+            if offset:
+                records = records[offset:]
+            if limit:
+                records = records[:limit]
+
+            return records
+
+    async def touch(
+        self,
+        session_id: str,
+        *,
+        updated_at: datetime | None = None,
+    ) -> None:
+        async with self._lock:
+            sess = self._sessions.get(session_id)
+            if not sess:
+                return
+            sess.updated_at = updated_at or datetime.now(timezone.utc)
+
+    async def update(
+        self,
+        session_id: str,
+        *,
+        title: str | None = None,
+        external_ref: str | None = None,
+    ) -> Session | None:
+        async with self._lock:
+            sess = self._sessions.get(session_id)
+            if not sess:
+                return None
+
+            # Mutate in-place (Session is a Pydantic model or similar)
+            if title is not None:
+                sess.title = title
+            if external_ref is not None:
+                sess.external_ref = external_ref
+
+            sess.updated_at = datetime.now(timezone.utc)
+            self._sessions[session_id] = sess
+            return sess
+
+    async def delete(self, session_id: str) -> None:
+        async with self._lock:
+            self._sessions.pop(session_id, None)