sovereign 0.14.2__py3-none-any.whl → 1.0.0a4__py3-none-any.whl

sovereign/cache/backends/s3.py

@@ -0,0 +1,161 @@
+import pickle
+import time
+from datetime import datetime, timedelta, timezone
+from importlib.util import find_spec
+from typing import Any
+from urllib.parse import quote
+
+from typing_extensions import override
+
+from sovereign import application_logger as log
+from sovereign.cache.backends import CacheBackend
+
+try:
+    # noinspection PyUnusedImports
+    import boto3
+
+    # noinspection PyUnusedImports
+    from botocore.exceptions import ClientError
+except ImportError:
+    pass
+
+BOTO_AVAILABLE = find_spec("boto3") is not None
+
+
+class S3Client:
+    def __init__(self, role_arn: str | None, client_args: dict[str, Any]):
+        self.role_arn = role_arn
+        self.client_args = client_args
+        self._client = None
+        self._credentials_expiry = None
+        self._base_session = boto3.Session()
+        self._make_client()
+
+    def _make_client(self) -> None:
+        if self.role_arn:
+            log.debug(f"Refreshing credentials for role: {self.role_arn}")
+            sts = self._base_session.client("sts")
+            duration_seconds = 3600  # 4 hours
+            response = sts.assume_role(
+                RoleArn=self.role_arn,
+                RoleSessionName="sovereign-s3-cache",
+                DurationSeconds=duration_seconds,
+            )
+            credentials = response["Credentials"]
+            session = boto3.Session(
+                aws_access_key_id=credentials["AccessKeyId"],
+                aws_secret_access_key=credentials["SecretAccessKey"],
+                aws_session_token=credentials["SessionToken"],
+            )
+            self._credentials_expiry = credentials["Expiration"]
+        else:
+            session = self._base_session
+            self._credentials_expiry = None
+        self._client = session.client("s3", **self.client_args)
+
+    def _session_expiring_soon(self) -> bool:
+        if not self.role_arn or self._credentials_expiry is None:
+            return False
+        refresh_threshold = timedelta(minutes=30).seconds
+        time_until_expiry = (
+            self._credentials_expiry - datetime.now(timezone.utc)
+        ).total_seconds()
+        return time_until_expiry <= refresh_threshold
+
+    def __getattr__(self, name):
+        if self._session_expiring_soon():
+            self._make_client()
+        return getattr(self._client, name)
+
+
+class S3Backend(CacheBackend):
+    """S3 cache backend implementation"""
+
+    # noinspection PyMissingConstructor,PyProtocol
+    @override
+    def __init__(self, config: dict[str, Any], attempts: int = 0) -> None:  # pyright: ignore[reportMissingSuperCall]
+        """Initialise S3 backend
+
+        Args:
+            config: Configuration dictionary containing S3 connection parameters
+                   Expected keys: bucket_name, prefix
+                   Optional keys: assume_role, endpoint_url
+        """
+        if not BOTO_AVAILABLE:
+            raise ImportError("boto3 not installed")
+
+        self.bucket_name = config.get("bucket_name")
+        if not self.bucket_name:
+            raise ValueError("bucket_name is required for S3 cache backend")
+
+        self.prefix = config.get("prefix", "sovereign-cache")
+        self.registration_prefix = config.get("registration_prefix", "registrations-")
+        self.role = config.get("assume_role")
+        self.remote_cache_key_suffix = config.get("remote_cache_key_suffix")
+
+        client_args: dict[str, Any] = {}
+        if endpoint_url := config.get("endpoint_url"):
+            client_args["endpoint_url"] = endpoint_url
+
+        self.client_args = client_args
+        self.s3 = S3Client(self.role, self.client_args)
+
+        try:
+            self.s3.head_bucket(Bucket=self.bucket_name)
+            log.info(f"S3 cache backend connected to bucket '{self.bucket_name}'")
+        except Exception as e:
+            if attempts == 0:
+                # wait 5 seconds and then try to access the bucket again
+                time.sleep(5)
+                self.__init__(config, attempts + 1)
+                return
+
+            log.error(
+                f"Failed to access S3 bucket '{self.bucket_name}' with current credentials: {e}"
+            )
+            raise
+
+    def _make_key(self, key: str) -> str:
+        # noop
+        if self.remote_cache_key_suffix:
+            encoded_key = quote(f"{key}-{self.remote_cache_key_suffix}", safe="")
+            log.info(f"Using remote cache key: {encoded_key}")
+        else:
+            encoded_key = quote(key, safe="")
+        return f"{self.prefix}/{encoded_key}"
+
+    def get(self, key: str) -> Any | None:
+        try:
+            log.debug(f"Retrieving object {key} from bucket")
+            response = self.s3.get_object(
+                Bucket=self.bucket_name, Key=self._make_key(key)
+            )
+            data = response["Body"].read()
+            unpickled = pickle.loads(data)
+            log.debug(f"Successfully obtained object {key} from bucket")
+            return unpickled
+        except self.s3.exceptions.NoSuchKey:
+            log.debug(f"{key} not in bucket")
+            return None
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "404":
+                log.debug(f"{key} not in bucket")
+                return None
+            log.warning(f"Failed to get key '{key}' from S3: {e}")
+            return None
+        except Exception as e:
+            log.warning(f"Failed to get key '{key}' from S3: {e}")
+            return None
+
+    @override
+    def set(self, key: str, value: Any, timeout: int | None = None) -> None:
+        try:
+            log.debug(f"Putting new object {key} into bucket")
+            self.s3.put_object(
+                Bucket=self.bucket_name,
+                Key=self._make_key(key),
+                Body=pickle.dumps(value),
+            )
+        except Exception as e:
+            log.warning(f"Failed to set key '{key}' in S3: {e}")
+            raise

sovereign/cache/filesystem.py

@@ -0,0 +1,74 @@
+import json
+import sqlite3
+from hashlib import sha256
+from pathlib import Path
+
+from cachelib import FileSystemCache
+from typing_extensions import final
+
+from sovereign.configuration import config
+from sovereign.types import DiscoveryRequest
+
+INIT = """
+CREATE TABLE IF NOT EXISTS registered_clients (
+    client_id TEXT PRIMARY KEY,
+    discovery_request TEXT NOT NULL
+)
+"""
+INSERT = "INSERT OR IGNORE INTO registered_clients (client_id, discovery_request) VALUES (?, ?)"
+LIST = "SELECT client_id, discovery_request FROM registered_clients"
+SEARCH = "SELECT 1 FROM registered_clients WHERE client_id = ?"
+
+
+@final
+class FilesystemCache:
+    def __init__(self, cache_path: str | None = None, default_timeout: int = 0):
+        self.cache_path = cache_path or config.cache.local_fs_path
+        self.default_timeout = default_timeout  # 0 = infinite TTL
+
+        self._cache = FileSystemCache(
+            cache_dir=self.cache_path,
+            default_timeout=self.default_timeout,
+            hash_method=sha256,
+        )
+
+        # Initialize SQLite for client registration
+        Path(self.cache_path).mkdir(parents=True, exist_ok=True)
+        self._db_path = Path(self.cache_path) / "clients.db"
+        self._init_db()
+
+    def _init_db(self):
+        with sqlite3.connect(self._db_path) as conn:
+            _ = conn.execute(INIT)
+
+    def get(self, key):
+        return self._cache.get(key)
+
+    def set(self, key, value, timeout=None):
+        return self._cache.set(key, value, timeout)
+
+    def delete(self, key):
+        return self._cache.delete(key)
+
+    def clear(self):
+        return self._cache.clear()
+
+    def register(self, id: str, req: DiscoveryRequest) -> None:
+        with sqlite3.connect(self._db_path) as conn:
+            _ = conn.execute(INSERT, (id, json.dumps(req.model_dump())))
+
+    def registered(self, id: str) -> bool:
+        with sqlite3.connect(self._db_path) as conn:
+            cursor = conn.execute(SEARCH, (id,))
+            return cursor.fetchone() is not None
+
+    def get_registered_clients(self) -> list[tuple[str, DiscoveryRequest]]:
+        with sqlite3.connect(self._db_path) as conn:
+            cursor = conn.execute(LIST)
+            rows = cursor.fetchall()
+
+        result = []
+        for client_id, req_json in rows:
+            req = DiscoveryRequest.model_validate(json.loads(req_json))
+            result.append((client_id, req))
+        return result

sovereign/cache/types.py

@@ -0,0 +1,17 @@
+from typing import Any
+
+from pydantic import BaseModel
+
+from sovereign.types import Node
+
+
+class CacheResult(BaseModel):
+    value: Any
+    from_remote: bool
+
+
+class Entry(BaseModel):
+    text: str
+    len: int
+    version: str
+    node: Node