loki-mode 7.79.0 → 7.80.1

package/lokistore/cloud.py

@@ -0,0 +1,305 @@
+"""
+Optional cloud LokiStore backends: S3, GCS, Azure Blob.
+
+Each backend's SDK is imported ONLY inside its __init__, so:
+- The SDK is never imported at module import time.
+- Local-only users never touch these classes and never pay an import cost.
+- Selecting a cloud backend without its SDK installed raises a clear,
+  actionable BackendNotAvailableError naming the missing package.
+
+Backends are intentionally thin. They do not add retry/credential frameworks;
+they rely on each SDK's default credential chain (env vars, instance metadata,
+etc.). Keys are normalized identically to LocalStore (traversal rejected) and
+an optional store-level prefix is prepended to every object key.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import List, Optional, Union
+
+from .base import (
+    BackendNotAvailableError,
+    LokiStore,
+    normalize_key,
+    read_source_bytes,
+)
+
+
+def _join_prefix(prefix: str, key: str) -> str:
+    """Join a normalized store prefix and a normalized key into an object key."""
+    key = normalize_key(key)
+    if not prefix:
+        return key
+    return f"{prefix.rstrip('/')}/{key}"
+
+
+def _normalize_prefix(prefix: Optional[str]) -> str:
+    """Normalize an optional store prefix to a clean POSIX-style path (or '')."""
+    if not prefix:
+        return ""
+    # Reuse the key guard so a prefix cannot smuggle traversal either.
+    return normalize_key(prefix)
+
+
+def _dir_prefix(obj_prefix: str) -> str:
+    """
+    Coerce a non-empty object-key prefix to its directory form (slash
+    terminated) so cloud list() matches at slash boundaries like LocalStore's
+    directory walk, per the LokiStore.list contract. An empty prefix (list
+    everything) is returned unchanged.
+    """
+    if obj_prefix and not obj_prefix.endswith("/"):
+        return obj_prefix + "/"
+    return obj_prefix
+
+
+class S3Store(LokiStore):
+    """Amazon S3 (or S3-compatible) backend (requires boto3)."""
+
+    def __init__(
+        self,
+        bucket: str,
+        prefix: Optional[str] = None,
+        region: Optional[str] = None,
+    ):
+        if not bucket:
+            raise ValueError("S3Store requires a bucket name")
+        try:
+            import boto3  # noqa: F401  (imported here so missing dep is lazy)
+        except ImportError as exc:
+            raise BackendNotAvailableError(
+                "The s3 storage backend requires the 'boto3' package. "
+                "Install it with: pip install boto3"
+            ) from exc
+
+        self._bucket = bucket
+        self._prefix = _normalize_prefix(prefix)
+        # boto3 picks up credentials from its default chain (env, shared
+        # config, instance/role metadata). region_name is optional.
+        kwargs = {"region_name": region} if region else {}
+        self._client = boto3.client("s3", **kwargs)
+
+    def _object_key(self, key: str) -> str:
+        return _join_prefix(self._prefix, key)
+
+    def put(self, key: str, data: Union[bytes, bytearray, str, os.PathLike]) -> None:
+        payload = read_source_bytes(data)
+        self._client.put_object(
+            Bucket=self._bucket, Key=self._object_key(key), Body=payload
+        )
+
+    def get(self, key: str) -> bytes:
+        try:
+            resp = self._client.get_object(
+                Bucket=self._bucket, Key=self._object_key(key)
+            )
+        except self._client.exceptions.NoSuchKey as exc:
+            raise FileNotFoundError(f"no such key: {key!r}") from exc
+        return resp["Body"].read()
+
+    def get_to(self, key: str, dest_path: Union[str, os.PathLike]) -> None:
+        dest = Path(dest_path).expanduser()
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        self._client.download_file(self._bucket, self._object_key(key), str(dest))
+
+    def exists(self, key: str) -> bool:
+        try:
+            self._client.head_object(Bucket=self._bucket, Key=self._object_key(key))
+            return True
+        except Exception:
+            # head_object raises ClientError (404) when absent; treat any
+            # lookup failure as "not present" for the boolean contract.
+            return False
+
+    def list(self, prefix: str = "") -> List[str]:
+        # A caller prefix is matched at slash boundaries (directory form); the
+        # store-level prefix, when listing everything, is left as-is.
+        obj_prefix = _dir_prefix(self._object_key(prefix)) if prefix else self._prefix
+        paginator = self._client.get_paginator("list_objects_v2")
+        strip = (self._prefix + "/") if self._prefix else ""
+        keys: List[str] = []
+        for page in paginator.paginate(Bucket=self._bucket, Prefix=obj_prefix):
+            for item in page.get("Contents", []):
+                full = item["Key"]
+                rel = full[len(strip):] if strip and full.startswith(strip) else full
+                keys.append(rel)
+        keys.sort()
+        return keys
+
+    def delete(self, key: str) -> bool:
+        # S3 delete is idempotent and does not signal whether an object
+        # existed. Probe first so the boolean contract is honest.
+        existed = self.exists(key)
+        self._client.delete_object(Bucket=self._bucket, Key=self._object_key(key))
+        return existed
+
+
+class GCSStore(LokiStore):
+    """Google Cloud Storage backend (requires google-cloud-storage)."""
+
+    def __init__(
+        self,
+        bucket: str,
+        prefix: Optional[str] = None,
+        region: Optional[str] = None,  # accepted for signature parity, unused
+    ):
+        if not bucket:
+            raise ValueError("GCSStore requires a bucket name")
+        try:
+            from google.cloud import storage  # type: ignore
+        except ImportError as exc:
+            raise BackendNotAvailableError(
+                "The gcs storage backend requires the 'google-cloud-storage' "
+                "package. Install it with: pip install google-cloud-storage"
+            ) from exc
+
+        self._prefix = _normalize_prefix(prefix)
+        # The client uses Application Default Credentials (env, gcloud, or
+        # workload identity).
+        self._client = storage.Client()
+        self._bucket = self._client.bucket(bucket)
+
+    def _object_key(self, key: str) -> str:
+        return _join_prefix(self._prefix, key)
+
+    def put(self, key: str, data: Union[bytes, bytearray, str, os.PathLike]) -> None:
+        payload = read_source_bytes(data)
+        blob = self._bucket.blob(self._object_key(key))
+        blob.upload_from_string(payload)
+
+    def get(self, key: str) -> bytes:
+        blob = self._bucket.blob(self._object_key(key))
+        if not blob.exists():
+            raise FileNotFoundError(f"no such key: {key!r}")
+        return blob.download_as_bytes()
+
+    def get_to(self, key: str, dest_path: Union[str, os.PathLike]) -> None:
+        dest = Path(dest_path).expanduser()
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        blob = self._bucket.blob(self._object_key(key))
+        if not blob.exists():
+            raise FileNotFoundError(f"no such key: {key!r}")
+        blob.download_to_filename(str(dest))
+
+    def exists(self, key: str) -> bool:
+        return self._bucket.blob(self._object_key(key)).exists()
+
+    def list(self, prefix: str = "") -> List[str]:
+        obj_prefix = _dir_prefix(self._object_key(prefix)) if prefix else self._prefix
+        strip = (self._prefix + "/") if self._prefix else ""
+        keys: List[str] = []
+        for blob in self._client.list_blobs(self._bucket, prefix=obj_prefix):
+            full = blob.name
+            rel = full[len(strip):] if strip and full.startswith(strip) else full
+            keys.append(rel)
+        keys.sort()
+        return keys
+
+    def delete(self, key: str) -> bool:
+        blob = self._bucket.blob(self._object_key(key))
+        if not blob.exists():
+            return False
+        blob.delete()
+        return True
+
+
+class AzureBlobStore(LokiStore):
+    """Azure Blob Storage backend (requires azure-storage-blob)."""
+
+    def __init__(
+        self,
+        bucket: str,
+        prefix: Optional[str] = None,
+        region: Optional[str] = None,  # accepted for signature parity, unused
+    ):
+        # In Azure terms `bucket` is the container name.
+        if not bucket:
+            raise ValueError("AzureBlobStore requires a container name (bucket)")
+        try:
+            from azure.storage.blob import ContainerClient  # type: ignore
+        except ImportError as exc:
+            raise BackendNotAvailableError(
+                "The azure-blob storage backend requires the "
+                "'azure-storage-blob' package. "
+                "Install it with: pip install azure-storage-blob"
+            ) from exc
+
+        self._prefix = _normalize_prefix(prefix)
+
+        # Credentials come from the SDK's default chain. We accept either a
+        # full connection string (AZURE_STORAGE_CONNECTION_STRING) or an
+        # account URL (AZURE_STORAGE_ACCOUNT_URL) + DefaultAzureCredential.
+        conn = os.environ.get("AZURE_STORAGE_CONNECTION_STRING")
+        account_url = os.environ.get("AZURE_STORAGE_ACCOUNT_URL")
+        if conn:
+            self._container = ContainerClient.from_connection_string(
+                conn, container_name=bucket
+            )
+        elif account_url:
+            try:
+                from azure.identity import DefaultAzureCredential  # type: ignore
+            except ImportError as exc:
+                raise BackendNotAvailableError(
+                    "Azure account-URL auth requires the 'azure-identity' "
+                    "package. Install it with: pip install azure-identity"
+                ) from exc
+            self._container = ContainerClient(
+                account_url=account_url,
+                container_name=bucket,
+                credential=DefaultAzureCredential(),
+            )
+        else:
+            raise BackendNotAvailableError(
+                "The azure-blob backend needs AZURE_STORAGE_CONNECTION_STRING "
+                "or AZURE_STORAGE_ACCOUNT_URL to be set."
+            )
+
+    def _object_key(self, key: str) -> str:
+        return _join_prefix(self._prefix, key)
+
+    def put(self, key: str, data: Union[bytes, bytearray, str, os.PathLike]) -> None:
+        payload = read_source_bytes(data)
+        self._container.upload_blob(
+            name=self._object_key(key), data=payload, overwrite=True
+        )
+
+    def get(self, key: str) -> bytes:
+        from azure.core.exceptions import ResourceNotFoundError  # type: ignore
+
+        try:
+            downloader = self._container.download_blob(self._object_key(key))
+        except ResourceNotFoundError as exc:
+            raise FileNotFoundError(f"no such key: {key!r}") from exc
+        return downloader.readall()
+
+    def get_to(self, key: str, dest_path: Union[str, os.PathLike]) -> None:
+        payload = self.get(key)
+        dest = Path(dest_path).expanduser()
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        with open(dest, "wb") as f:
+            f.write(payload)
+
+    def exists(self, key: str) -> bool:
+        return self._container.get_blob_client(self._object_key(key)).exists()
+
+    def list(self, prefix: str = "") -> List[str]:
+        obj_prefix = _dir_prefix(self._object_key(prefix)) if prefix else self._prefix
+        strip = (self._prefix + "/") if self._prefix else ""
+        keys: List[str] = []
+        for blob in self._container.list_blobs(name_starts_with=obj_prefix or None):
+            full = blob.name
+            rel = full[len(strip):] if strip and full.startswith(strip) else full
+            keys.append(rel)
+        keys.sort()
+        return keys
+
+    def delete(self, key: str) -> bool:
+        from azure.core.exceptions import ResourceNotFoundError  # type: ignore
+
+        try:
+            self._container.delete_blob(self._object_key(key))
+            return True
+        except ResourceNotFoundError:
+            return False

package/lokistore/factory.py

@@ -0,0 +1,187 @@
+"""
+LokiStore factory: select and construct a backend from config or env.
+
+Defaults to local. With nothing configured, get_store() returns a LocalStore
+rooted at the project `.loki/` directory, honoring LOKI_DIR / TARGET_DIR
+exactly like the rest of the codebase. Cloud backends are constructed only when
+explicitly selected, and only then is their SDK imported.
+
+Config precedence: an explicit config dict overrides env, env overrides the
+local default. Recognized config keys mirror the env vars (without the
+LOKI_STORAGE_ prefix): backend, bucket, prefix, region, base_dir.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Dict, Optional
+
+from .base import LokiStore, StoreError
+from .local import LocalStore
+
+
+def resolve_local_base(explicit_base: Optional[str] = None) -> str:
+    """
+    Resolve the local `.loki/` base directory exactly like run.sh and the
+    dashboard do.
+
+    Resolution order, matching autonomy/run.sh
+    (`${LOKI_DIR:-${TARGET_DIR:-.}/.loki}`):
+      1. explicit_base argument, if given
+      2. $LOKI_DIR, if set (used verbatim, like the shell)
+      3. $TARGET_DIR/.loki, if TARGET_DIR is set
+      4. ./.loki  (current working directory)
+
+    Returns:
+        The resolved base directory as a string path. It is NOT created here;
+        LocalStore creates parents lazily on first write.
+    """
+    if explicit_base:
+        return explicit_base
+    loki_dir = os.environ.get("LOKI_DIR")
+    if loki_dir:
+        return loki_dir
+    target_dir = os.environ.get("TARGET_DIR", ".")
+    return os.path.join(target_dir, ".loki")
+
+
+def _config_from_env() -> Dict[str, Any]:
+    """Read the LOKI_STORAGE_* env vars into a config dict (omitting unset)."""
+    cfg: Dict[str, Any] = {}
+    backend = os.environ.get("LOKI_STORAGE_BACKEND")
+    if backend:
+        cfg["backend"] = backend
+    bucket = os.environ.get("LOKI_STORAGE_BUCKET")
+    if bucket:
+        cfg["bucket"] = bucket
+    prefix = os.environ.get("LOKI_STORAGE_PREFIX")
+    if prefix:
+        cfg["prefix"] = prefix
+    region = os.environ.get("LOKI_STORAGE_REGION")
+    if region:
+        cfg["region"] = region
+    return cfg
+
+
+def build_store(config: Optional[Dict[str, Any]] = None) -> LokiStore:
+    """
+    Construct a LokiStore from an explicit config dict (no env fallback).
+
+    Use this when a caller has already assembled config from its own source.
+    Prefer get_store() for the standard env-aware path.
+
+    Recognized keys:
+      backend : "local" (default) | "s3" | "gcs" | "azure-blob"
+      bucket  : bucket/container name (cloud backends)
+      prefix  : key prefix within the bucket (optional)
+      region  : region (s3, optional)
+      base_dir: local base directory (local backend only; overrides resolution)
+    """
+    config = dict(config or {})
+    backend = (config.get("backend") or "local").strip().lower()
+
+    if backend in ("local", "", "file", "filesystem"):
+        return LocalStore(resolve_local_base(config.get("base_dir")))
+
+    bucket = config.get("bucket")
+    prefix = config.get("prefix")
+    region = config.get("region")
+
+    if backend in ("s3", "aws", "aws-s3"):
+        from .cloud import S3Store
+
+        return S3Store(bucket=bucket, prefix=prefix, region=region)
+
+    if backend in ("gcs", "gcp", "google", "google-cloud-storage"):
+        from .cloud import GCSStore
+
+        return GCSStore(bucket=bucket, prefix=prefix, region=region)
+
+    if backend in ("azure-blob", "azure", "azureblob", "azure_blob"):
+        from .cloud import AzureBlobStore
+
+        return AzureBlobStore(bucket=bucket, prefix=prefix, region=region)
+
+    raise StoreError(
+        f"unknown storage backend: {backend!r} "
+        "(expected one of: local, s3, gcs, azure-blob)"
+    )
+
+
+def get_store(config: Optional[Dict[str, Any]] = None) -> LokiStore:
+    """
+    Return a LokiStore, defaulting to local.
+
+    Resolution:
+      1. Start from the LOKI_STORAGE_* env vars.
+      2. Overlay any explicit config dict (config keys win over env).
+      3. If no backend is selected, default to local.
+
+    With a clean environment and no config, this returns a LocalStore rooted at
+    the project `.loki/` -- zero new deps, zero behavior change.
+    """
+    merged = _config_from_env()
+    if config:
+        merged.update({k: v for k, v in config.items() if v is not None})
+    return build_store(merged)
+
+
+# ---------------------------------------------------------------------------
+# Metadata backend selector (stub for this release).
+#
+# Blobs (state, checkpoints, artifacts) go through LokiStore above. Structured
+# metadata (the dashboard's task/run index) lives in a relational store. Today
+# that is the existing async-SQLite database in dashboard/database.py at
+# $LOKI_DATA_DIR/dashboard.db. A postgres backend can be added later for shared
+# multi-instance fleets; the selector below is the seam for that, defaulting to
+# sqlite so there is no behavior change now.
+# ---------------------------------------------------------------------------
+
+def get_metadata_backend(config: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+    """
+    Resolve the metadata (relational) backend selection.
+
+    This is intentionally a thin, documented descriptor rather than a live
+    connection: the existing dashboard/database.py owns the sqlite engine, and
+    a postgres implementation is deferred. Returning a descriptor lets future
+    callers branch on the choice without this module importing SQLAlchemy or
+    any database driver.
+
+    Selection (config["metadata_backend"] overrides
+    $LOKI_METADATA_BACKEND; default "sqlite"):
+      sqlite   -> reuse $LOKI_DATA_DIR/dashboard.db (default, implemented today
+                  in dashboard/database.py)
+      postgres -> read $LOKI_METADATA_URL (NOT YET IMPLEMENTED; selecting it
+                  returns the descriptor with implemented=False so callers can
+                  fail with a clear message until the impl lands)
+
+    Returns:
+        A descriptor dict: {backend, implemented, dsn, note}.
+    """
+    config = dict(config or {})
+    backend = (
+        config.get("metadata_backend")
+        or os.environ.get("LOKI_METADATA_BACKEND")
+        or "sqlite"
+    ).strip().lower()
+
+    if backend in ("sqlite", "", "default"):
+        data_dir = os.environ.get("LOKI_DATA_DIR", os.path.expanduser("~/.loki"))
+        return {
+            "backend": "sqlite",
+            "implemented": True,
+            "dsn": os.path.join(data_dir, "dashboard.db"),
+            "note": "reuses dashboard/database.py async-sqlite engine",
+        }
+
+    if backend in ("postgres", "postgresql", "pg"):
+        return {
+            "backend": "postgres",
+            "implemented": False,
+            "dsn": config.get("metadata_url") or os.environ.get("LOKI_METADATA_URL"),
+            "note": "postgres metadata backend is planned, not yet implemented",
+        }
+
+    raise StoreError(
+        f"unknown metadata backend: {backend!r} (expected: sqlite, postgres)"
+    )