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

aethergraph/storage/blob/s3_blob.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import os
+
+import boto3
+from botocore.config import Config
+
+from aethergraph.contracts.storage.blob_store import BlobStore
+from aethergraph.services.artifacts.utils import to_thread
+
+"""
+S3BlobStore: BlobStore implementation using AWS S3 as the backend.
+NOTE: This is a stub implementation; not fully tested.
+"""
+
+
+class S3BlobStore(BlobStore):
+    def __init__(self, bucket: str, prefix: str = ""):
+        self.bucket = bucket
+        self.prefix = prefix.strip("/")
+
+        # Create S3 client with default config; can be customized as needed
+        self.s3_client = boto3.client(
+            "s3", config=Config(signature_version="s3v4", max_pool_connections=50)
+        )
+
+    @property
+    def base_uri(self) -> str:
+        if self.prefix:
+            return f"s3://{self.bucket}/{self.prefix}"
+        return f"s3://{self.bucket}"
+
+    def _resolve_key(self, key: str | None, ext: str | None) -> str:
+        if key is None:
+            import uuid
+
+            key = uuid.uuid4().hex + (ext or "")
+        if self.prefix:
+            return f"{self.prefix}/{key.lstrip('/')}"
+        return key.lstrip("/")
+
+    async def put_bytes(
+        self,
+        data: bytes,
+        *,
+        key: str | None = None,
+        ext: str | None = None,
+        mime: str | None = None,
+        keep_source: bool = False,  # added for interface consistency; not used here
+    ) -> str:
+        key = self._resolve_key(key, ext)
+
+        def _upload():
+            extra = {}
+            if mime:
+                extra["ContentType"] = mime
+            self._client.put_object(
+                Bucket=self.bucket,
+                Key=key,
+                Body=data,
+                **extra,
+            )
+            return f"s3://{self.bucket}/{key}"
+
+        return await to_thread(_upload)
+
+    async def put_file(
+        self,
+        path: str,
+        *,
+        key: str | None = None,
+        mime: str | None = None,
+    ) -> str:
+        ext = os.path.splitext(path)[1]
+        key = self._resolve_key(key, ext)
+
+        def _upload_file():
+            extra = {}
+            if mime:
+                extra["ContentType"] = mime
+            self._client.upload_file(
+                Filename=os.path.abspath(path),
+                Bucket=self.bucket,
+                Key=key,
+                ExtraArgs=extra or None,
+            )
+            return f"s3://{self.bucket}/{key}"
+
+        return await to_thread(_upload_file)
+
+    async def load_bytes(self, uri: str) -> bytes:
+        # assume s3://bucket/prefix...
+        # you can parse, or just compute key from uri by stripping base_uri
+        if uri.startswith("s3://"):
+            _, rest = uri.split("://", 1)
+            bucket, key = rest.split("/", 1)
+        else:
+            bucket = self.bucket
+            key = uri
+
+        def _download():
+            obj = self._client.get_object(Bucket=bucket, Key=key)
+            return obj["Body"].read()
+
+        return await to_thread(_download)
+
+    async def load_text(
+        self,
+        uri: str,
+        *,
+        encoding: str = "utf-8",
+        errors: str = "strict",
+    ) -> str:
+        data = await self.load_bytes(uri)
+        return data.decode(encoding, errors)

aethergraph/storage/continuation_store/fs_cont.py

@@ -0,0 +1,283 @@
+from __future__ import annotations
+
+import asyncio
+from dataclasses import asdict
+from datetime import datetime, timezone
+import hashlib
+import hmac
+import json
+import logging
+import os
+from pathlib import Path
+import re
+import threading
+from typing import Any
+
+from aethergraph.services.continuations.continuation import Continuation, Correlator
+
+
+class FSContinuationStore:  # implements AsyncContinuationStore
+    def __init__(self, root: str | Path, secret: bytes):
+        self.root = Path(root)
+        self.secret = secret
+        self._lock = threading.RLock()
+
+    # ---------- helpers ----------
+    def _cont_path(self, run_id: str, node_id: str) -> Path:
+        return self.root / "runs" / run_id / "nodes" / node_id / "continuation.json"
+
+    def _token_idx_path(self, token: str) -> Path:
+        return self.root / "index" / "tokens" / f"{token}.json"
+
+    def _rev_idx_path(self, token: str) -> Path:
+        return self.root / "index" / "rev" / f"{token}.json"
+
+    def _safe(self, s: str) -> str:
+        s = (s or "").replace("\\", "_").replace("/", "_").replace(":", "_")
+        s = re.sub(r"[^A-Za-z0-9._@-]", "_", s)
+        s = re.sub(r"_+", "_", s).strip("._ ") or "x"
+        if len(s) > 100:
+            h = hashlib.sha1(s.encode()).hexdigest()[:8]
+            s = f"{s[:92]}_{h}"
+        return s
+
+    def _corr_dir(self, scheme: str, channel: str, thread: str, message: str) -> Path:
+        return (
+            self.root
+            / "index"
+            / "corr"
+            / self._safe(scheme)
+            / self._safe(channel)
+            / self._safe(thread)
+            / self._safe(message)
+        )
+
+    def _corr_tokens_path(self, scheme: str, channel: str, thread: str, message: str) -> Path:
+        return self._corr_dir(scheme, channel, thread, message) / "tokens.json"
+
+    def _hmac(self, *parts: str) -> str:
+        h = hmac.new(self.secret, digestmod=hashlib.sha256)
+        for p in parts:
+            h.update(p.encode("utf-8"))
+        return h.hexdigest()
+
+    async def mint_token(self, run_id: str, node_id: str, attempts: int) -> str:
+        return self._hmac(run_id, node_id, str(attempts), os.urandom(8).hex())
+
+    # ---------- core ----------
+    async def save(self, cont: Continuation) -> None:
+        def _write():
+            path = self._cont_path(cont.run_id, cont.node_id)
+            path.parent.mkdir(parents=True, exist_ok=True)
+            payload = cont.to_dict() if hasattr(cont, "to_dict") else asdict(cont)
+            for k in ("deadline", "next_wakeup_at", "created_at"):
+                v = payload.get(k)
+                if isinstance(v, datetime):
+                    payload[k] = v.astimezone(timezone.utc).isoformat()
+            tmp = path.with_suffix(".tmp")
+            tmp.write_text(json.dumps(payload, indent=2), encoding="utf-8")
+            tmp.replace(path)
+            self._write_token_index(cont.run_id, cont.node_id, cont.token)
+
+        await asyncio.to_thread(_write)
+
+    async def get(self, run_id: str, node_id: str) -> Continuation | None:
+        def _read():
+            path = self._cont_path(run_id, node_id)
+            if not path.exists():
+                return None
+            raw = json.loads(path.read_text(encoding="utf-8"))
+            for k in ("deadline", "next_wakeup_at", "created_at"):
+                if raw.get(k):
+                    raw[k] = datetime.fromisoformat(raw[k])
+            raw["closed"] = bool(raw.get("closed", False))
+            return Continuation(**raw)
+
+        return await asyncio.to_thread(_read)
+
+    async def list_cont_by_run(self, run_id: str) -> list[Continuation]:
+        def _list():
+            out = []
+            run_path = self.root / "runs" / run_id / "nodes"
+            if not run_path.exists():
+                return out
+            for node_dir in run_path.iterdir():
+                cont_path = node_dir / "continuation.json"
+                if cont_path.exists():
+                    raw = json.loads(cont_path.read_text(encoding="utf-8"))
+                    for k in ("deadline", "next_wakeup_at", "created_at"):
+                        if raw.get(k):
+                            raw[k] = datetime.fromisoformat(raw[k])
+                    raw["closed"] = bool(raw.get("closed", False))
+                    out.append(Continuation(**raw))
+            return out
+
+        return await asyncio.to_thread(_list)
+
+    async def delete(self, run_id: str, node_id: str) -> None:
+        def _del():
+            p = self._cont_path(run_id, node_id)
+            if p.exists():
+                p.unlink()
+
+        await asyncio.to_thread(_del)
+
+    # ---------- token helpers ----------
+    def _write_token_index(self, run_id: str, node_id: str, token: str) -> None:
+        with self._lock:
+            p = self._token_idx_path(token)
+            p.parent.mkdir(parents=True, exist_ok=True)
+            p.write_text(
+                json.dumps({"run_id": run_id, "node_id": node_id}, indent=2), encoding="utf-8"
+            )
+
+    async def get_by_token(self, token: str) -> Continuation | None:
+        def _lookup():
+            p = self._token_idx_path(token)
+            if not p.exists():
+                return None
+            ref = json.loads(p.read_text(encoding="utf-8"))
+            return ref["run_id"], ref["node_id"]
+
+        ref = await asyncio.to_thread(_lookup)
+        if not ref:
+            return None
+        run_id, node_id = ref
+        return await self.get(run_id, node_id)
+
+    async def mark_closed(self, token: str) -> None:
+        c = await self.get_by_token(token)
+        if not c:
+            return
+        if not c.closed:
+            c.closed = True
+            await self.save(c)
+
+    async def verify_token(self, run_id: str, node_id: str, token: str) -> bool:
+        c = await self.get(run_id, node_id)
+        return bool(c and hmac.compare_digest(token, c.token))
+
+    # ---------- correlators ----------
+    async def bind_correlator(self, *, token: str, corr: Correlator) -> None:
+        def _bind():
+            scheme, channel, thread, message = corr.key()
+            tokens_path = self._corr_tokens_path(scheme, channel, thread, message)
+            tokens_path.parent.mkdir(parents=True, exist_ok=True)
+            toks: list[str] = []
+            if tokens_path.exists():
+                try:
+                    toks = json.loads(tokens_path.read_text(encoding="utf-8"))
+                except Exception:
+                    toks = []
+            if token not in toks:
+                toks.append(token)
+                tokens_path.write_text(json.dumps(toks, indent=2), encoding="utf-8")
+                # reverse index
+                r = self._rev_idx_path(token)
+                r.parent.mkdir(parents=True, exist_ok=True)
+                paths = []
+                if r.exists():
+                    try:
+                        paths = json.loads(r.read_text(encoding="utf-8"))
+                    except Exception:
+                        paths = []
+                key_path = str(tokens_path.relative_to(self.root))
+                if key_path not in paths:
+                    paths.append(key_path)
+                    r.write_text(json.dumps(paths, indent=2), encoding="utf-8")
+
+        await asyncio.to_thread(_bind)
+
+    async def find_by_correlator(self, *, corr: Correlator) -> Continuation | None:
+        def _read_toks():
+            scheme, channel, thread, message = corr.key()
+            p = self._corr_tokens_path(scheme, channel, thread, message)
+            if not p.exists():
+                return []
+            try:
+                return json.loads(p.read_text(encoding="utf-8")) or []
+            except Exception:
+                return []
+
+        toks = await asyncio.to_thread(_read_toks)
+        for tok in reversed(toks):
+            c = await self.get_by_token(tok)
+            if c and not c.closed:
+                if c.deadline and datetime.now(timezone.utc) > c.deadline.astimezone(timezone.utc):
+                    continue
+                return c
+        return None
+
+    async def last_open(self, *, channel: str, kind: str) -> Continuation | None:
+        # Optional slow scan (dev only)
+        def _scan():
+            waits = []
+            runs_path = self.root / "runs"
+            if not runs_path.exists():
+                return waits
+            for run_dir in runs_path.iterdir():
+                nodes_dir = run_dir / "nodes"
+                if not nodes_dir.exists():
+                    continue
+                for node_dir in nodes_dir.iterdir():
+                    cont_path = node_dir / "continuation.json"
+                    if cont_path.exists():
+                        waits.append(json.loads(cont_path.read_text(encoding="utf-8")))
+            return waits
+
+        waits = await asyncio.to_thread(_scan)
+        for raw in reversed(waits):
+            if raw.get("closed"):
+                continue
+            if raw.get("channel") == channel and raw.get("kind") == kind:
+                for k in ("deadline", "next_wakeup_at", "created_at"):
+                    if raw.get(k):
+                        raw[k] = datetime.fromisoformat(raw[k])
+                return Continuation(**raw)
+        return None
+
+    async def list_waits(self) -> list[dict[str, Any]]:
+        def _scan():
+            out = []
+            runs_path = self.root / "runs"
+            if not runs_path.exists():
+                return out
+            for run_dir in runs_path.iterdir():
+                nodes_dir = run_dir / "nodes"
+                if not nodes_dir.exists():
+                    continue
+                for node_dir in nodes_dir.iterdir():
+                    cont_path = node_dir / "continuation.json"
+                    if cont_path.exists():
+                        out.append(json.loads(cont_path.read_text(encoding="utf-8")))
+            return out
+
+        return await asyncio.to_thread(_scan)
+
+    async def clear(self) -> None:
+        def _clear():
+            for sub in ("runs", "index"):
+                p = self.root / sub
+                if p.exists():
+                    for root, dirs, files in os.walk(p, topdown=False):
+                        for f in files:
+                            try:
+                                os.remove(Path(root) / f)
+                            except Exception:
+                                logger = logging.getLogger(
+                                    "aethergraph.services.continuations.stores.fs_store"
+                                )
+                                logger.warning("Failed to remove file: %s", Path(root) / f)
+                        for d in dirs:
+                            try:
+                                os.rmdir(Path(root) / d)
+                            except Exception:
+                                logger = logging.getLogger(
+                                    "aethergraph.services.continuations.stores.fs_store"
+                                )
+                                logger.warning("Failed to remove dir: %s", Path(root) / d)
+
+        await asyncio.to_thread(_clear)
+
+    async def alias_for(self, token: str) -> str | None:
+        return token[:24]

aethergraph/storage/continuation_store/inmem_cont.py

@@ -0,0 +1,146 @@
+from dataclasses import asdict
+from datetime import datetime, timezone
+import hashlib
+import hmac
+import os
+import threading
+from typing import Any
+
+from aethergraph.contracts.services.continuations import AsyncContinuationStore
+from aethergraph.services.continuations.continuation import Continuation, Correlator
+
+
+class InMemoryContinuationStore(AsyncContinuationStore):
+    """
+    Process-local, in-memory continuation store.
+
+    - Thread-safe via RLock (main loop + sidecar can share it).
+    - No persistence beyond process lifetime.
+
+    NOTE: suitable for testing or single-process deployments only.
+    """
+
+    def __init__(self, *, secret: bytes):
+        self._secret = secret
+        self._by_id: dict[tuple[str, str], Continuation] = {}
+        self._by_token: dict[str, tuple[str, str]] = {}
+        self._corr_tokens: dict[str, list[str]] = {}
+        self._lock = threading.RLock()
+
+    # --- helpers ---
+
+    def _key(self, run_id: str, node_id: str) -> tuple[str, str]:
+        return (run_id, node_id)
+
+    def _corr_key(self, corr: Correlator) -> str:
+        scheme, channel, thread, message = corr.key()
+        return f"{scheme}:{channel}:{thread}:{message}"
+
+    def _hmac(self, *parts: str) -> str:
+        h = hmac.new(self._secret, digestmod=hashlib.sha256)
+        for p in parts:
+            h.update(p.encode("utf-8"))
+        return h.hexdigest()
+
+    async def mint_token(self, run_id: str, node_id: str, attempts: int) -> str:
+        return self._hmac(run_id, node_id, str(attempts), os.urandom(8).hex())
+
+    # --- core ---
+
+    async def save(self, cont: Continuation) -> None:
+        with self._lock:
+            self._by_id[self._key(cont.run_id, cont.node_id)] = cont
+            self._by_token[cont.token] = (cont.run_id, cont.node_id)
+
+    async def get(self, run_id: str, node_id: str) -> Continuation | None:
+        with self._lock:
+            return self._by_id.get(self._key(run_id, node_id))
+
+    async def delete(self, run_id: str, node_id: str) -> None:
+        with self._lock:
+            key = self._key(run_id, node_id)
+            cont = self._by_id.pop(key, None)
+            if cont:
+                self._by_token.pop(cont.token, None)
+
+    async def list_cont_by_run(self, run_id: str) -> list[Continuation]:
+        with self._lock:
+            return [c for (r, _), c in self._by_id.items() if r == run_id]
+
+    # --- token ---
+
+    async def get_by_token(self, token: str) -> Continuation | None:
+        with self._lock:
+            ref = self._by_token.get(token)
+            if not ref:
+                return None
+            run_id, node_id = ref
+            return self._by_id.get(self._key(run_id, node_id))
+
+    async def mark_closed(self, token: str) -> None:
+        c = await self.get_by_token(token)
+        if not c:
+            return
+        if not c.closed:
+            c.closed = True
+            await self.save(c)
+
+    async def verify_token(self, run_id: str, node_id: str, token: str) -> bool:
+        c = await self.get(run_id, node_id)
+        return bool(c and hmac.compare_digest(token, c.token))
+
+    # --- correlators ---
+
+    async def bind_correlator(self, *, token: str, corr: Correlator) -> None:
+        key = self._corr_key(corr)
+        with self._lock:
+            toks = self._corr_tokens.setdefault(key, [])
+            if token not in toks:
+                toks.append(token)
+
+    async def find_by_correlator(self, *, corr: Correlator) -> Continuation | None:
+        key = self._corr_key(corr)
+        with self._lock:
+            toks = list(self._corr_tokens.get(key, []))
+
+        now = datetime.now(timezone.utc)
+        for tok in reversed(toks):
+            c = await self.get_by_token(tok)
+            if not c or c.closed:
+                continue
+            if c.deadline and now > c.deadline.astimezone(timezone.utc):
+                continue
+            return c
+        return None
+
+    # --- scans ---
+
+    async def last_open(self, *, channel: str, kind: str) -> Continuation | None:
+        with self._lock:
+            candidates = [
+                c
+                for c in self._by_id.values()
+                if not c.closed and c.channel == channel and c.kind == kind
+            ]
+        if not candidates:
+            return None
+
+        # pick most recent created_at
+        def _ts(c: Continuation) -> float:
+            dt = c.created_at or datetime.min.replace(tzinfo=timezone.utc)
+            return dt.timestamp()
+
+        return max(candidates, key=_ts)
+
+    async def list_waits(self) -> list[dict[str, Any]]:
+        with self._lock:
+            return [asdict(c) for c in self._by_id.values()]
+
+    async def clear(self) -> None:
+        with self._lock:
+            self._by_id.clear()
+            self._by_token.clear()
+            self._corr_tokens.clear()
+
+    async def alias_for(self, token: str) -> str | None:
+        return token[:24]