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

aethergraph/storage/artifacts/cas_store.py

@@ -0,0 +1,422 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+import datetime
+import json
+import logging
+import os
+from pathlib import Path
+import shutil
+import tempfile
+from typing import Any, BinaryIO
+
+from aethergraph.contracts.services.artifacts import Artifact
+from aethergraph.contracts.storage.artifact_store import AsyncArtifactStore
+from aethergraph.contracts.storage.blob_store import BlobStore
+
+from .utils import (
+    _now_iso,
+    _sha256_file,
+    _tree_manifest_and_hash,
+    to_thread,
+)
+
+logger = logging.getLogger("aethergraph.services.artifacts.cas_store")
+
+
+class CASArtifactStore(AsyncArtifactStore):
+    """
+    Content-addressed artifact store built on top of a BlobStore.
+
+    - Uses local staging_dir for temp files/dirs.
+    - Stores blobs via BlobStore with keys derived from SHA-256 hashes.
+    - Persists minimal manifest/metadata as blobs too (for directories).
+    """
+
+    def __init__(self, blob: BlobStore, staging_dir: str):
+        self._blob = blob
+        self._staging_dir = os.path.abspath(staging_dir)
+        os.makedirs(self._staging_dir, exist_ok=True)
+        self.last_artifact: Artifact | None = None
+
+    @property
+    def base_uri(self) -> str:
+        return self._blob.base_uri
+
+    def _augment_labels_with_filename(
+        self,
+        labels: dict | None,
+        *,
+        suggested_uri: str | None = None,
+        path: str | None = None,
+    ) -> dict:
+        """
+        Ensure labels contains a stable 'filename' key when we can infer one.
+
+        - Prefer an explicit suggested_uri basename.
+        - Fallback to the local path basename.
+        - Do NOT override an existing 'filename' or 'name' key.
+        """
+        out: dict[str, Any] = dict(labels or {})
+
+        # Don't stomp on explicit naming
+        if "filename" in out or "name" in out:
+            return out
+
+        candidate: str | None = None
+        if suggested_uri:
+            candidate = os.path.basename(suggested_uri.rstrip("/"))
+        elif path:
+            candidate = os.path.basename(path.rstrip(os.sep))
+
+        if candidate:
+            out["filename"] = candidate
+
+        return out
+
+    # ---------- staging utils ----------
+    async def plan_staging_path(self, planned_ext: str = "") -> str:
+        def _mk():
+            fd, p = tempfile.mkstemp(suffix=planned_ext, dir=self._staging_dir)
+            os.close(fd)
+            return p
+
+        return await to_thread(_mk)
+
+    async def plan_staging_dir(self, suffix: str = "") -> str:
+        def _mkd():
+            return tempfile.mkdtemp(prefix="dir_", suffix=suffix, dir=self._staging_dir)
+
+        return await to_thread(_mkd)
+
+    # ---------- basic save / ingest ----------
+    async def save_file(
+        self,
+        *,
+        path: str,
+        kind: str,
+        run_id: str,
+        graph_id: str,
+        node_id: str,
+        tool_name: str,
+        tool_version: str,
+        suggested_uri: str | None = None,  # NOTE: only metadata / pretty; impl may ignore
+        pin: bool = False,
+        labels: dict | None = None,
+        metrics: dict | None = None,
+        preview_uri: str | None = None,  # NOTE: only metadata / pretty; impl may ignore
+        cleanup: bool = True,
+    ) -> Artifact:
+        sha, nbytes = await to_thread(_sha256_file, path)
+        ext = os.path.splitext(path)[1]
+        key = os.path.join("cas", "blobs", f"{sha}{ext}")
+
+        blob_uri = await self._blob.put_file(path, key=key, mime=None, keep_source=not cleanup)
+
+        eff_labels = self._augment_labels_with_filename(
+            labels,
+            suggested_uri=suggested_uri,
+            path=path,
+        )
+
+        a = Artifact(
+            artifact_id=sha,
+            uri=blob_uri,
+            kind=kind,
+            bytes=nbytes,
+            sha256=sha,
+            mime=None,  # callers can fill in if desired
+            run_id=run_id,
+            graph_id=graph_id,
+            node_id=node_id,
+            tool_name=tool_name,
+            tool_version=tool_version,
+            created_at=_now_iso(),
+            labels=eff_labels,
+            metrics=metrics or {},
+            preview_uri=preview_uri,
+            pinned=pin,
+        )
+        self.last_artifact = a
+        return a
+
+    # ---------- streaming writer ----------
+    @asynccontextmanager
+    async def open_writer(
+        self,
+        *,
+        kind: str,
+        run_id: str,
+        graph_id: str,
+        node_id: str,
+        tool_name: str,
+        tool_version: str,
+        planned_ext: str | None = None,
+        pin: bool = False,
+    ) -> AsyncIterator[Any]:
+        staged_path = await self.plan_staging_path(planned_ext or "")
+
+        class _Writer:
+            """Helper class for streaming writes to a temp file."""
+
+            def __init__(self, path: str, f: BinaryIO):
+                self.tmp_path = path
+                self._f = f
+                self._labels: dict[str, str] = {}
+                self._metrics: dict[str, float] = {}
+                self.artifact: Artifact | None = None  # filled after finalize
+
+            def write(self, chunk: bytes) -> None:
+                self._f.write(chunk)
+
+            def add_labels(self, labels: dict[str, str]) -> None:
+                self._labels.update(labels or {})
+
+            def add_metrics(self, metrics: dict[str, float]) -> None:
+                self._metrics.update(metrics or {})
+
+        writer: _Writer | None = None
+
+        try:
+            # Ruff-friendly: file is opened via a context manager, and kept
+            # open for the duration of the user’s writes.
+            with open(staged_path, "wb") as f:
+                writer = _Writer(staged_path, f)
+                # Yield to caller; they can await inside and call writer.write(...)
+                yield writer
+            # <-- file is closed here when the with-block exits
+
+            # Now ingest the staged file into CAS and create the Artifact
+            if writer is not None:
+                a = await self.ingest_staged_file(
+                    staged_path=staged_path,
+                    kind=kind,
+                    run_id=run_id,
+                    graph_id=graph_id,
+                    node_id=node_id,
+                    tool_name=tool_name,
+                    tool_version=tool_version,
+                    pin=pin,
+                    labels=writer._labels,
+                    metrics=writer._metrics,
+                )
+                writer.artifact = a
+
+        except Exception:
+            # Best-effort cleanup of staged file on error
+            try:
+                if os.path.exists(staged_path):
+                    os.remove(staged_path)
+            finally:
+                raise
+
+    async def ingest_staged_file(
+        self,
+        *,
+        staged_path: str,
+        kind: str,
+        run_id: str,
+        graph_id: str,
+        node_id: str,
+        tool_name: str,
+        tool_version: str,
+        pin: bool = False,
+        labels: dict | None = None,
+        metrics: dict | None = None,
+        preview_uri: str | None = None,
+        suggested_uri: str | None = None,
+    ) -> Artifact:
+        # just delegate to save_file (same semantics)
+        a = await self.save_file(
+            path=staged_path,
+            kind=kind,
+            run_id=run_id,
+            graph_id=graph_id,
+            node_id=node_id,
+            tool_name=tool_name,
+            tool_version=tool_version,
+            suggested_uri=suggested_uri,
+            pin=pin,
+            labels=labels,
+            metrics=metrics,
+            preview_uri=preview_uri,
+        )
+        try:
+            os.remove(staged_path)
+        except Exception:
+            logger.warning("ingest_staged_file: failed to delete staged file %s", staged_path)
+        return a
+
+    async def ingest_directory(
+        self,
+        *,
+        staged_dir: str,
+        kind: str = "dataset",
+        run_id: str,
+        graph_id: str,
+        node_id: str,
+        tool_name: str,
+        tool_version: str,
+        include: list[str] | None = None,
+        exclude: list[str] | None = None,
+        index_children: bool = False,  # TODO: use later for per-file artifacts
+        pin: bool = False,
+        labels: dict | None = None,
+        metrics: dict | None = None,
+        suggested_uri: str | None = None,
+        archive: bool = False,
+        archive_name: str = "bundle.tar.gz",
+        cleanup: bool = True,
+        store: str | None = None,  # "archive" | "manifest"
+    ) -> Artifact:
+        if not os.path.isdir(staged_dir):
+            raise ValueError(f"ingest_directory: not a directory: {staged_dir}")
+
+        if store is None:
+            store = "archive" if archive else "manifest"
+
+        manifest_entries, tree_sha = await to_thread(
+            _tree_manifest_and_hash, staged_dir, include, exclude
+        )
+
+        # manifest blob
+        manifest_key = os.path.join("cas", "trees", tree_sha, "manifest.json")
+
+        def _dump_manifest() -> bytes:
+            return json.dumps(
+                {
+                    "files": manifest_entries,
+                    "created_at": _now_iso(),
+                    "tool_name": tool_name,
+                    "tool_version": tool_version,
+                },
+                indent=2,
+            ).encode("utf-8")
+
+        # manifest URI -> use it in future if needed
+        _ = await self._blob.put_bytes(
+            _dump_manifest(),
+            key=manifest_key,
+            ext=".json",
+            mime="application/json",
+        )
+
+        archive_uri: str | None = None
+        if store == "archive":
+            # build tar.gz locally, then upload
+            archive_path = os.path.join(self._staging_dir, f"{tree_sha}.tar.gz")
+
+            def _make_tar():
+                import tarfile
+
+                with tarfile.open(archive_path, mode="w:gz") as tar:
+                    for e in sorted(manifest_entries, key=lambda x: x["path"]):
+                        abs_file = os.path.join(staged_dir, e["path"])
+                        tar.add(abs_file, arcname=e["path"])
+                return archive_path
+
+            archive_path = await to_thread(_make_tar)
+            archive_key = os.path.join("cas", "trees", tree_sha, archive_name)
+            archive_uri = await self._blob.put_file(
+                archive_path,
+                key=archive_key,
+                mime="application/gzip",
+            )
+
+        elif store == "manifest":
+            if cleanup:
+                # we will delete staged_dir; only OK if user accepts that artifacts
+                # are now represented by manifest (+ optional archive)
+                pass
+        else:
+            raise ValueError(f"unknown store mode: {store}")
+
+        # Directory "handle" URI: base_uri + prefix
+        dir_prefix = os.path.join("cas", "trees", tree_sha)
+        # NOTE: we don't require an actual object at dir_prefix; it's a logical handle.
+        dir_uri = self.base_uri.rstrip("/") + "/" + dir_prefix.replace(os.sep, "/")
+
+        total_bytes = sum(e["bytes"] for e in manifest_entries)
+
+        eff_labels = self._augment_labels_with_filename(
+            labels,
+            suggested_uri=suggested_uri or archive_name,
+            path=staged_dir,
+        )
+
+        a = Artifact(
+            artifact_id=tree_sha,
+            uri=dir_uri,
+            kind=kind,
+            bytes=total_bytes,
+            sha256=tree_sha,
+            mime="application/vnd.aethergraph.bundle+dir",
+            run_id=run_id,
+            graph_id=graph_id,
+            node_id=node_id,
+            tool_name=tool_name,
+            tool_version=tool_version,
+            created_at=_now_iso(),
+            labels=eff_labels,
+            metrics=metrics or {},
+            preview_uri=archive_uri,
+            pinned=pin,
+        )
+        self.last_artifact = a
+
+        if cleanup:
+            try:
+                shutil.rmtree(staged_dir, ignore_errors=True)
+            except Exception:
+                logger.warning("ingest_directory: failed to cleanup staged dir %s", staged_dir)
+
+        return a
+
+    # ---------- load ----------
+    async def load_bytes(self, uri):
+        return await self._blob.load_bytes(uri)
+
+    async def load_text(self, uri: str, *, encoding: str = "utf-8", errors: str = "strict") -> str:
+        return await self._blob.load_text(uri, encoding=encoding, errors=errors)
+
+    async def load_artifact_bytes(self, uri: str) -> bytes:
+        return await self._blob.load_bytes(uri)
+
+    async def load_artifact_dir(self, uri):
+        """
+        Normalize a directory artifact to a local path.
+
+        FS backend can simply return the directory; S3 backend
+        will download files described by manifest into a temp dir.
+        For now, implement generic: if it's already a file:// path,
+        just return as-is; otherwise, ArtifactFacade can add a helper
+        `as_local_dir(artifact)` that handles S3 download.
+        """
+        return uri
+
+    async def load_artifact(self, uri):
+        # Compatibility: if direcotry URI, return as-is, else load blob content
+        if uri.endswith("/"):
+            # directory handle URI
+            return await self.load_artifact_dir(uri)
+        # else, blob URI
+        return await self._blob.load_bytes(uri)
+
+    # ---------- cleanup ----------
+    async def cleanup_tmp(self, max_age_hours: int = 24) -> None:
+        now = datetime.datetime.now(datetime.timezone.utc).timestamp()
+
+        def _cleanup():
+            for p in Path(self._staging_dir).rglob("*"):
+                try:
+                    age_h = (now - p.stat().st_mtime) / 3600.0
+                    if age_h > max_age_hours:
+                        if p.is_file():
+                            p.unlink(missing_ok=True)
+                        else:
+                            shutil.rmtree(p, ignore_errors=True)
+                except Exception:
+                    pass
+
+        await to_thread(_cleanup)

aethergraph/storage/artifacts/fs_cas.py

@@ -0,0 +1,18 @@
+import os
+
+from aethergraph.storage.artifacts.cas_store import CASArtifactStore
+from aethergraph.storage.blob.fs_blob import FSBlobStore
+
+
+class FSArtifactStore(CASArtifactStore):
+    # Initialize with a base directory for storing artifacts
+
+    def __init__(self, base_dir: str):
+        base_dir = os.path.abspath(base_dir)
+        blob = FSBlobStore(os.path.join(base_dir, "blobs"))
+        staging_dir = os.path.join(base_dir, "staging")
+        super().__init__(blob=blob, staging_dir=staging_dir)
+
+    # TODO: Add any FS-specific optimizations if needed
+    # Optionally override load_artifact_dir to return actual local dir path if uri is file://cas/trees/...
+    # and implement FS-only "pretty" symlinks.

aethergraph/storage/artifacts/s3_cas.py

@@ -0,0 +1,14 @@
+from aethergraph.storage.artifacts.cas_store import CASArtifactStore
+from aethergraph.storage.blob.s3_blob import S3BlobStore
+
+
+class S3ArtifactStore(CASArtifactStore):
+    # Initialize with S3 bucket and optional prefix for storing artifacts
+    def __init__(self, bucket: str, prefix: str, staging_dir: str):
+        blob = S3BlobStore(bucket=bucket, prefix=prefix)
+        super().__init__(blob=blob, staging_dir=staging_dir)
+
+    # TODO: Optionally add any S3-specific optimizations if needed
+    # - parse tree_sha from uri
+    # - download files listed in manifest.json into a local temp dir
+    # - return that path

aethergraph/storage/artifacts/utils.py

@@ -0,0 +1,124 @@
+import asyncio
+from datetime import datetime, timezone
+from fnmatch import fnmatch
+import hashlib
+import json
+import os
+from pathlib import Path
+
+
+def now_iso() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+async def to_thread(fn, *a, **k):
+    return await asyncio.to_thread(fn, *a, **k)
+
+
+# ----- helpers ----- NOTE: we have multiple copies of these in different places, consider centralizing -----
+def _now_iso():
+    return datetime.now(timezone.utc).isoformat()
+
+
+def _sha256_file(path: str, chunk=1024 * 1024) -> tuple[str, int]:
+    """Return (sha256 hex, size in bytes) of a file."""
+    h = hashlib.sha256()
+    total = 0
+    with open(path, "rb") as f:
+        while True:
+            b = f.read(chunk)
+            if not b:
+                break
+            h.update(b)
+            total += len(b)
+    return h.hexdigest(), total
+
+
+def _content_addr_path(base_dir: str, sha256: str, ext: str | None) -> str:
+    """Return a content-addressed path under base_dir for a given sha256 and optional extension.
+    Creates subdirectories as needed.
+
+    It works as follows:
+    - Takes the first 4 characters of the sha256 hash to create two levels of subdirectories.
+    - The first two characters form the first subdirectory (sub1).
+    - The next two characters form the second subdirectory (sub2).
+    - The full sha256 hash, optionally followed by the provided file extension, is used as the filename.
+    - Ensures that the target directory exists by creating it if necessary.
+    - Returns the full path to the content-addressed file.
+
+    The final path structure will look like:
+    base_dir/sub1/sub2/sha256[.ext]
+    """
+    sub1, sub2 = sha256[:2], sha256[2:4]
+    fname = sha256 + (ext or "")
+    target_dir = os.path.join(base_dir, sub1, sub2)
+    os.makedirs(target_dir, exist_ok=True)
+    return os.path.join(target_dir, fname)
+
+
+def _walk_dir(root: str, include: list[str] | None, exclude: list[str] | None):
+    """Yield (relpath, abspath) for files under root honoring include/exclude globs."""
+    root_p = Path(root)
+    for p in root_p.rglob("*"):
+        if not p.is_file():
+            continue
+        rel = str(p.relative_to(root_p)).replace("\\", "/")
+        if exclude and any(fnmatch.fnmatch(rel, pat) for pat in exclude):
+            continue
+        if include and not any(fnmatch.fnmatch(rel, pat) for pat in include):
+            continue
+        yield rel, str(p)
+
+
+def _tree_manifest_and_hash(root: str, include: list[str] | None, exclude: list[str] | None):
+    """
+    Build a deterministic manifest of files: [{"path": rel, "sha256": sha, "bytes": n}, ...]
+    The tree hash is sha256 over JSON lines: "<rel> <sha> <bytes>\n" sorted by rel.
+    """
+    entries = []
+    lines = []
+    for rel, abspath in _walk_dir(root, include, exclude):
+        sha, nbytes = _sha256_file(abspath)
+        entries.append({"path": rel, "sha256": sha, "bytes": nbytes})
+        lines.append(f"{rel}\t{sha}\t{nbytes}\n")
+    # sort for determinism
+    lines.sort()
+    h = hashlib.sha256()
+    for line in lines:
+        h.update(line.encode("utf-8"))
+    tree_sha = h.hexdigest()
+    return entries, tree_sha
+
+
+def _content_addr_dir_path(base_dir: str, tree_sha: str):
+    # content-addressed folder to hold manifest (and optional archive)
+    sub1, sub2 = tree_sha[:2], tree_sha[2:4]
+    target_dir = os.path.join(base_dir, sub1, sub2, tree_sha)
+    os.makedirs(target_dir, exist_ok=True)
+    return target_dir
+
+
+def _write_json(path: str, obj: dict | list):
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(obj, f, ensure_ascii=False, separators=(",", ":"))
+
+
+def _maybe_cleanup_tmp_parent(tmp_root: str, path: str):
+    """Remove empty parent dirs strictly under tmp_root (never _tmp itself)."""
+    try:
+        parent = os.path.dirname(os.path.abspath(path))
+        tmp_root_abs = os.path.abspath(tmp_root)
+
+        # Only operate if `parent` is inside tmp_root
+        while (
+            os.path.commonpath([parent, tmp_root_abs]) == tmp_root_abs
+            and os.path.normcase(parent)
+            != os.path.normcase(tmp_root_abs)  # don't delete _tmp itself
+        ):
+            try:
+                os.rmdir(parent)  # only removes if empty
+            except OSError:
+                break
+            parent = os.path.dirname(parent)
+    except Exception:
+        pass

aethergraph/storage/blob/fs_blob.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import os
+import shutil
+
+from aethergraph.contracts.storage.blob_store import BlobStore
+from aethergraph.storage.fs_utils import _from_uri_or_path, _to_file_uri, to_thread
+
+
+class FSBlobStore(BlobStore):
+    def __init__(self, base_dir: str):
+        self.base_dir = os.path.abspath(base_dir)
+        os.makedirs(self.base_dir, exist_ok=True)
+
+    @property
+    def base_uri(self) -> str:
+        return _to_file_uri(self.base_dir)
+
+    def _resolve_key(self, key: str | None, ext: str | None) -> str:
+        if key is None:
+            # fall back to some random-ish name under "blobs/"
+            import uuid
+
+            name = uuid.uuid4().hex + (ext or "")
+            key = os.path.join("blobs", name)
+        return key
+
+    async def put_bytes(
+        self,
+        data: bytes,
+        *,
+        key: str | None = None,
+        ext: str | None = None,
+        mime: str | None = None,
+    ) -> str:
+        key = self._resolve_key(key, ext)
+        path = os.path.join(self.base_dir, key)
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+
+        def _write():
+            with open(path, "wb") as f:
+                f.write(data)
+            return _to_file_uri(path)
+
+        return await to_thread(_write)
+
+    async def put_file(
+        self,
+        path: str,
+        *,
+        key: str | None = None,
+        mime: str | None = None,
+        keep_source: bool = False,
+    ) -> str:
+        ext = os.path.splitext(path)[1]
+        key = self._resolve_key(key, ext)
+        dst = os.path.join(self.base_dir, key)
+        os.makedirs(os.path.dirname(dst), exist_ok=True)
+
+        def _move():
+            if keep_source:
+                shutil.copy2(os.path.abspath(path), dst)
+            else:
+                shutil.move(os.path.abspath(path), dst)
+            return _to_file_uri(dst)
+
+        return await to_thread(_move)
+
+    async def load_bytes(self, uri: str) -> bytes:
+        path = _from_uri_or_path(uri)
+
+        def _read():
+            with open(path, "rb") as f:
+                return f.read()
+
+        return await to_thread(_read)
+
+    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)