sovereign 1.0.0b123__py3-none-any.whl → 1.0.0b134__py3-none-any.whl

sovereign/app.py

@@ -101,7 +101,7 @@ def init_app() -> FastAPI:
 
     @application.get("/static/{filename}", summary="Return a static asset")
     async def static(filename: str) -> Response:
-        return FileResponse(get_package_file("sovereign", f"static/{filename}"))  # type: ignore[arg-type]
+        return FileResponse(get_package_file("sovereign_files", f"static/{filename}"))  # type: ignore[arg-type]
 
     @application.get(
         "/admin/xds_dump",

sovereign/cache/__init__.py

@@ -0,0 +1,182 @@
+"""
+Sovereign Cache Module
+
+This module provides an extensible cache backend system that allows clients
+to configure their own remote cache backends through entry points.
+"""
+
+import asyncio
+import threading
+from typing import Any
+from typing_extensions import final
+
+import requests
+
+from sovereign import stats, application_logger as log
+from sovereign.schemas import config, DiscoveryRequest, RegisterClientRequest
+from sovereign.cache.types import Entry, CacheResult
+from sovereign.cache.backends import CacheBackend, get_backend
+from sovereign.cache.filesystem import FilesystemCache
+
+CLIENTS_LOCK = "sovereign_clients_lock"
+CLIENTS_KEY = "sovereign_clients"
+CACHE_READ_TIMEOUT = config.cache.read_timeout
+WORKER_URL = "http://localhost:9080/client"
+
+
+@final
+class DualCache:
+    """Cache that writes to both filesystem and remote backend, reads filesystem first with remote fallback"""
+
+    def __init__(self, fs_cache: FilesystemCache, remote_cache: CacheBackend | None):
+        self.fs_cache = fs_cache
+        self.remote_cache = remote_cache
+
+    def get(self, key: str) -> CacheResult | None:
+        # Try filesystem first
+        if value := self.fs_cache.get(key):
+            stats.increment("cache.fs.hit")
+            return CacheResult(value=value, from_remote=False)
+
+        # Fallback to remote cache if available
+        if self.remote_cache:
+            try:
+                if value := self.remote_cache.get(key):
+                    stats.increment("cache.remote.hit")
+                    # Write back to filesystem
+                    self.fs_cache.set(key, value)
+                    return CacheResult(value=value, from_remote=True)
+            except Exception as e:
+                log.warning(f"Failed to read from remote cache: {e}")
+                stats.increment("cache.remote.error")
+
+        return None
+
+    def set(self, key, value, timeout=None):
+        self.fs_cache.set(key, value, timeout)
+        if self.remote_cache:
+            try:
+                self.remote_cache.set(key, value, timeout)
+                stats.increment("cache.remote.write.success")
+            except Exception as e:
+                log.warning(f"Failed to write to remote cache: {e}")
+                stats.increment("cache.remote.write.error")
+
+    def register(self, id: str, req: DiscoveryRequest):
+        self.fs_cache.register(id, req)
+
+    def registered(self, id: str) -> bool:
+        if value := self.fs_cache.registered(id):
+            return value
+        return False
+
+    def get_registered_clients(self) -> list[tuple[str, Any]]:
+        if value := self.fs_cache.get_registered_clients():
+            return value
+        return []
+
+
+class CacheManager:
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls):
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+                    cls._instance._initialized = False
+        return cls._instance
+
+    def __init__(self):
+        if self._initialized:
+            return
+
+        filesystem_cache = FilesystemCache()
+        remote_cache = get_backend()
+
+        if remote_cache is None:
+            log.info("Cache initialized with filesystem backend only")
+        else:
+            log.info("Cache initialized with filesystem and remote backends")
+        self._cache = DualCache(filesystem_cache, remote_cache)
+        self._initialized = True
+
+    def get(self, req: DiscoveryRequest) -> Entry | None:
+        id = client_id(req)
+        if result := self._cache.get(id):
+            if result.from_remote:
+                self.register(req)
+            stats.increment("cache.hit")
+            return result.value
+        stats.increment("cache.miss")
+        return None
+
+    def set(self, key: str, value: Entry, timeout: int | None = None) -> None:
+        return self._cache.set(key, value, timeout)
+
+    def register(self, req: DiscoveryRequest) -> tuple[str, DiscoveryRequest]:
+        """Register a client using the cache backend"""
+        id = client_id(req)
+        log.debug(f"Registering client {id}")
+        self._cache.register(id, req)
+        return id, req
+
+    def registered(self, req: DiscoveryRequest) -> bool:
+        """Check if a client is registered using the cache backend"""
+        id = client_id(req)
+        is_registered = self._cache.registered(id)
+        log.debug(f"Client {id} registered={is_registered}")
+        return is_registered
+
+    def get_registered_clients(self) -> list[tuple[str, Any]]:
+        """Get all registered clients using the cache backend"""
+        return self._cache.get_registered_clients()
+
+
+@stats.timed("cache.read_ms")
+async def blocking_read(
+    req: DiscoveryRequest, timeout_s=CACHE_READ_TIMEOUT, poll_interval_s=0.5
+) -> Entry | None:
+    metric = "client.registration"
+    if entry := read(req):
+        return entry
+
+    registered = False
+    registration = RegisterClientRequest(request=req)
+    start = asyncio.get_event_loop().time()
+    attempt = 1
+    while (asyncio.get_event_loop().time() - start) < timeout_s:
+        if not registered:
+            try:
+                response = requests.put(WORKER_URL, json=registration.model_dump())
+                match response.status_code:
+                    case 200 | 202:
+                        stats.increment(metric, tags=["status:registered"])
+                        registered = True
+                    case 429:
+                        stats.increment(metric, tags=["status:ratelimited"])
+                        await asyncio.sleep(min(attempt, CACHE_READ_TIMEOUT))
+                        attempt *= 2
+            except Exception as e:
+                stats.increment(metric, tags=["status:failed"])
+                log.exception(f"Tried to register client but failed: {e}")
+        if entry := read(req):
+            return entry
+        await asyncio.sleep(poll_interval_s)
+
+    return None
+
+
+def client_id(req: DiscoveryRequest) -> str:
+    return req.cache_key(config.cache.hash_rules)
+
+
+manager = CacheManager()
+
+# Old APIs
+write = manager.set
+read = manager.get
+clients = manager.get_registered_clients
+register = manager.register
+registered = manager.registered

sovereign/cache/backends/__init__.py

@@ -0,0 +1,110 @@
+"""
+Cache backends module
+
+This module provides the protocol definition for cache backends and
+the loading mechanism for extensible cache backends via entry points.
+"""
+
+from collections.abc import Sequence
+from importlib.metadata import EntryPoints
+from typing import Protocol, Any, runtime_checkable
+
+from sovereign import application_logger as log
+from sovereign.utils.entry_point_loader import EntryPointLoader
+
+
+@runtime_checkable
+class CacheBackend(Protocol):
+    def __init__(self, config: dict[str, Any]) -> None:
+        """Initialize the cache backend with generic configuration
+
+        Args:
+            config: Dictionary containing backend-specific configuration
+        """
+        ...
+
+    def get(self, key: str) -> Any | None:
+        """Get a value from the cache
+
+        Args:
+            key: The cache key
+
+        Returns:
+            The cached value or None if not found
+        """
+        ...
+
+    def set(self, key: str, value: Any, timeout: int | None = None) -> None:
+        """Set a value in the cache
+
+        Args:
+            key: The cache key
+            value: The value to cache
+            timeout: Optional timeout in seconds
+        """
+        ...
+
+
+def get_backend() -> CacheBackend | None:
+    from sovereign import config
+
+    cache_config = config.cache.remote_backend
+    if not cache_config:
+        log.info("No remote cache backend configured, using filesystem only")
+        return None
+
+    backend_type = cache_config.type
+
+    loader = EntryPointLoader("cache.backends")
+    entry_points: EntryPoints | Sequence[Any] = loader.groups.get("cache.backends", [])
+
+    backend = None
+    for ep in entry_points:
+        if ep.name == backend_type:
+            backend = ep.load()
+            break
+
+    if not backend:
+        raise KeyError(
+            (
+                f"Cache backend '{backend_type}' not found. "
+                f"Available backends: {[ep.name for ep in entry_points]}"
+            )
+        )
+
+    backend_config = _process_loadable_config(cache_config.config)
+    instance = backend(backend_config)
+
+    if not isinstance(instance, CacheBackend):
+        raise TypeError(
+            (f"Cache backend '{backend_type}' does not implement CacheBackend protocol")
+        )
+
+    log.info(f"Successfully initialized cache backend: {backend_type}")
+    return instance
+
+
+def _process_loadable_config(config: dict[str, Any]) -> dict[str, Any]:
+    from sovereign.dynamic_config import Loadable
+
+    processed = {}
+    for key, value in config.items():
+        try:
+            if isinstance(value, str):
+                loadable = Loadable.from_legacy_fmt(value)
+                processed[key] = loadable.load()
+            elif isinstance(value, dict):
+                loadable = Loadable(**value)
+                processed[key] = loadable.load()
+            else:
+                processed[key] = value
+            continue
+        except Exception as e:
+            log.warning(f"Failed to load value for {key}: {e}")
+
+        if isinstance(value, dict):
+            processed[key] = _process_loadable_config(value)
+        else:
+            processed[key] = value
+
+    return processed

sovereign/cache/backends/s3.py

@@ -0,0 +1,139 @@
+import pickle
+from datetime import datetime, timezone, timedelta
+from typing import Any
+from typing_extensions import override
+from urllib.parse import quote
+from importlib.util import find_spec
+
+from sovereign import application_logger as log
+from sovereign.cache.backends import CacheBackend
+
+try:
+    import boto3
+    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"""
+
+    @override
+    def __init__(self, config: dict[str, Any]) -> None:  # pyright: ignore[reportMissingSuperCall]
+        """Initialize S3 backend
+
+        Args:
+            config: Configuration dictionary containing S3 connection parameters
+                   Expected keys: bucket_name, key
+                   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.key = config.get("key", "sovereign-cache/")
+        self.registration_prefix = config.get("registration_prefix", "registrations-")
+        self.role = config.get("assume_role")
+
+        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:
+            log.error(
+                f"Failed to access S3 bucket '{self.bucket_name}' with current credentials: {e}"
+            )
+            raise
+
+    def _make_key(self, key: str) -> str:
+        encoded_key = quote(key, safe="")
+        return f"{self.key}{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()
+            return pickle.loads(data)
+        except self.s3.exceptions.NoSuchKey:
+            return None
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "404":
+                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,42 @@
+from hashlib import sha256
+from cachelib import FileSystemCache
+
+from sovereign import config
+from sovereign.schemas import DiscoveryRequest
+
+
+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
+        self._cache = FileSystemCache(
+            cache_dir=self.cache_path,
+            default_timeout=self.default_timeout,
+            hash_method=sha256,
+        )
+
+    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:
+        clients = self.get_registered_clients()
+        if (id, req) in clients:
+            return
+        clients.append((id, req))
+        _ = self._cache.set("_registered_clients", clients)
+
+    def registered(self, id: str) -> bool:
+        clients = self.get_registered_clients()
+        return any(cid == id for cid, _ in clients)
+
+    def get_registered_clients(self) -> list[tuple[str, DiscoveryRequest]]:
+        return self._cache.get("_registered_clients") or []

sovereign/cache/types.py

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

sovereign/context.py

@@ -6,6 +6,7 @@ import asyncio
 import inspect
 from enum import Enum
 from typing import Any, Callable, Optional, Union
+from typing_extensions import final, override
 
 import pydantic
 from croniter import croniter
@@ -14,14 +15,15 @@ from sovereign.schemas import DiscoveryRequest, config
 from sovereign.statistics import configure_statsd
 from sovereign.utils.timer import wait_until
 from sovereign.dynamic_config import Loadable
+from sovereign.events import bus, Event, Topic
 
 
 stats = configure_statsd()
 DEFAULT_RETRY_INTERVAL = config.template_context.refresh_retry_interval_secs
 DEFAULT_NUM_RETRIES = config.template_context.refresh_num_retries
-NEW_CONTEXT = asyncio.Event()
 
 
+@final
 class ScheduledTask:
     def __init__(self, task: "ContextTask"):
         self.task = task
@@ -37,22 +39,24 @@ class ScheduledTask:
         self.due = time.monotonic() + self.task.seconds_til_next_run
         heapq.heappush(tasks, self)
 
+    @override
     def __str__(self) -> str:
         return f"ScheduledTask({self.task.name})"
 
 
+@final
 class TemplateContext:
     def __init__(
         self,
-        middleware: Optional[list[Callable[[DiscoveryRequest, dict], None]]] = None,
+        middleware: list[Callable[[DiscoveryRequest, dict[str, Any]], None]]
+        | None = None,
     ) -> None:
         self.tasks: dict[str, ContextTask] = dict()
         self.results: dict[str, ContextResult] = dict()
         self.hashes: dict[str, int] = dict()
         self.scheduled: list[ScheduledTask] = list()
-        self.middleware = middleware or list()
         self.running: set[str] = set()
-        self.notify_consumers: Optional[asyncio.Task] = None
+        self.middleware = middleware or list()
 
     @classmethod
     def from_config(cls) -> "TemplateContext":
@@ -68,7 +72,7 @@ class TemplateContext:
     def register_task_from_loadable(self, name: str, loadable: Loadable) -> None:
         self.register_task(ContextTask.from_loadable(name, loadable))
 
-    def update_hash(self, task: "ContextTask"):
+    async def update_hash(self, task: "ContextTask"):
         name = task.name
         result = self.results.get(name)
         old = self.hashes.get(name)
@@ -77,18 +81,7 @@ class TemplateContext:
         if old != new:
             stats.increment("context.updated", tags=[f"context:{name}"])
             self.hashes[name] = new
-            # Debounced event notification to the worker
-            if self.notify_consumers:
-                # cancel existing and reset timer
-                self.notify_consumers.cancel()
-            self.notify_consumers = asyncio.create_task(self.publish_event())
-
-    async def publish_event(self):
-        try:
-            await asyncio.sleep(3.0)
-            NEW_CONTEXT.set()
-        except asyncio.CancelledError:
-            pass
+            await task.notify()
 
     def get_context(self, req: DiscoveryRequest) -> dict[str, Any]:
         ret = {r.name: r.data for r in self.results.values()}
@@ -107,7 +100,7 @@ class TemplateContext:
         self.running.add(task.name)
         try:
             await task.refresh(self.results)
-            self.update_hash(task)
+            await self.update_hash(task)
         finally:
             self.running.remove(task.name)
 
@@ -161,6 +154,15 @@ class ContextTask(pydantic.BaseModel):
     interval: "TaskInterval"
     retry_policy: Optional["TaskRetryPolicy"] = None
 
+    async def notify(self):
+        await bus.publish(
+            Topic.CONTEXT,
+            Event(
+                message=f"Context {self.name} updated",
+                metadata={"name": self.name},
+            ),
+        )
+
     async def refresh(self, output: dict[str, "ContextResult"]) -> None:
         output[self.name] = await self.try_load()
 

sovereign/events.py

@@ -0,0 +1,49 @@
+import pydantic
+from enum import IntEnum
+from asyncio import Queue, gather
+from collections import defaultdict
+from typing import final, Sequence
+
+
+Primitives = str | int | float | bool | Sequence[str]
+
+
+class Topic(IntEnum):
+    CONTEXT = 1
+
+
+class Event(pydantic.BaseModel):
+    message: str
+    metadata: dict[str, Primitives] = pydantic.Field(default_factory=dict)
+
+
+@final
+class EventBus:
+    def __init__(self, maxsize: int = 0):
+        self._topics: dict[Topic, list[Queue[Event]]] = defaultdict(list)
+        self._maxsize = maxsize
+
+    def subscribe(self, topic: Topic) -> Queue[Event]:
+        q: Queue[Event] = Queue(self._maxsize)
+        self._topics[topic].append(q)
+        return q
+
+    def unsubscribe(self, topic: Topic, q: Queue[Event]) -> None:
+        qs = self._topics.get(topic)
+        if not qs:
+            return
+        try:
+            qs.remove(q)
+        except ValueError:
+            pass
+        if not qs:
+            _ = self._topics.pop(topic, None)
+
+    async def publish(self, topic: Topic, msg: Event) -> None:
+        qs = self._topics.get(topic, [])
+        if not qs:
+            return
+        _ = await gather(*(q.put(msg) for q in qs))
+
+
+bus = EventBus()

sovereign/middlewares.py

@@ -19,7 +19,7 @@ class RequestContextLogMiddleware(BaseHTTPMiddleware):
             response = await call_next(request)
         finally:
             req_id = get_request_id()
-            response.headers["X-Request-ID"] = req_id
+            req_id = response.headers.setdefault("X-Request-Id", get_request_id())
             logs.access_logger.queue_log_fields(REQUEST_ID=req_id)
             _request_id_ctx_var.reset(token)
         return response