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

sovereign/context.py

@@ -1,119 +1,285 @@
-from typing import Dict, Any, Generator, Iterable, NoReturn, Optional
-from copy import deepcopy
-from fastapi import HTTPException
-from sovereign.config_loader import Loadable
-from sovereign.schemas import DiscoveryRequest, XdsTemplate
-from sovereign.sources import SourcePoller
-from sovereign.utils.crypto import CipherSuite, CipherContainer
-from sovereign.utils.timer import poll_forever, poll_forever_cron
+import asyncio
+import datetime
+import heapq
+import inspect
+import time
+import zlib
+from enum import Enum
+from typing import Any, Callable, Optional, Union
 
+import pydantic
+from croniter import croniter
+from typing_extensions import final, override
 
+from sovereign import application_logger as log
+from sovereign.configuration import config
+from sovereign.dynamic_config import Loadable
+from sovereign.events import Event, Topic, bus
+from sovereign.statistics import configure_statsd
+from sovereign.types import DiscoveryRequest
+from sovereign.utils.timer import wait_until
+
+stats = configure_statsd()
+DEFAULT_RETRY_INTERVAL = config.template_context.refresh_retry_interval_secs
+DEFAULT_NUM_RETRIES = config.template_context.refresh_num_retries
+
+
+@final
+class ScheduledTask:
+    def __init__(self, task: "ContextTask"):
+        self.task = task
+        self.due = time.monotonic()
+
+    def __lt__(self, other: "ScheduledTask") -> bool:
+        return self.due < other.due
+
+    async def run(
+        self, output: dict[str, "ContextResult"], tasks: list["ScheduledTask"]
+    ):
+        await self.task.refresh(output)
+        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,
-        refresh_rate: Optional[int],
-        refresh_cron: Optional[str],
-        configured_context: Dict[str, Loadable],
-        poller: SourcePoller,
-        encryption_suite: CipherContainer,
-        disabled_suite: CipherSuite,
-        logger: Any,
-        stats: Any,
+        middleware: list[Callable[[DiscoveryRequest, dict[str, Any]], None]]
+        | None = None,
     ) -> None:
-        self.poller = poller
-        self.refresh_rate = refresh_rate
-        self.refresh_cron = refresh_cron
-        self.configured_context = configured_context
-        self.crypto = encryption_suite
-        self.disabled_suite = disabled_suite
-        self.logger = logger
-        self.stats = stats
-        # initial load
-        self.context = self.load_context_variables()
-
-    async def start_refresh_context(self) -> NoReturn:
-        if self.refresh_cron is not None:
-            await poll_forever_cron(self.refresh_cron, self.refresh_context)
-        elif self.refresh_rate is not None:
-            await poll_forever(self.refresh_rate, self.refresh_context)
-
-        raise RuntimeError("Failed to start refresh_context, this should never happen")
-
-    async def refresh_context(self) -> None:
-        self.context = self.load_context_variables()
-
-    def load_context_variables(self) -> Dict[str, Any]:
-        ret = dict()
-        for k, v in self.configured_context.items():
+        self.tasks: dict[str, ContextTask] = dict()
+        self.results: dict[str, ContextResult] = dict()
+        self.hashes: dict[str, int] = dict()
+        self.scheduled: list[ScheduledTask] = list()
+        self.running: set[str] = set()
+        self.middleware = middleware or list()
+
+    @classmethod
+    def from_config(cls) -> "TemplateContext":
+        ret = TemplateContext()
+        for name, spec in config.template_context.context.items():
+            ret.register_task_from_loadable(name, spec)
+        return ret
+
+    def register_task(self, task: "ContextTask") -> None:
+        self.tasks[task.name] = task
+        self.scheduled.append(ScheduledTask(task))
+
+    def register_task_from_loadable(self, name: str, loadable: Loadable) -> None:
+        self.register_task(ContextTask.from_loadable(name, loadable))
+
+    async def update_hash(self, task: "ContextTask"):
+        name = task.name
+        result = self.results.get(name)
+        old = self.hashes.get(name)
+        new = hash(result)
+
+        if old != new:
+            stats.increment("context.updated", tags=[f"context:{name}"])
+            self.hashes[name] = new
+            await task.notify()
+
+    def get_context(self, req: DiscoveryRequest) -> dict[str, Any]:
+        ret = {r.name: r.data for r in self.results.values()}
+        for fn in self.middleware:
+            fn(req, ret)
+        return ret
+
+    def get(self, key: str, default: Any = None) -> Any:
+        if result := self.results.get(key):
+            return result.data
+        return default
+
+    async def _run_task(self, task: "ContextTask"):
+        if task.name in self.running:
+            return
+        self.running.add(task.name)
+        try:
+            await task.refresh(self.results)
+            await self.update_hash(task)
+        finally:
+            self.running.remove(task.name)
+
+    async def run_once(self):
+        heapq.heapify(self.scheduled)
+        for next_ in self.scheduled:
+            await self._run_task(next_.task)
+
+    async def start(self):
+        if not self.scheduled:
+            # No context jobs configured
+            return
+        heapq.heapify(self.scheduled)
+        while True:
+            # Obtain next task
+            next_ = heapq.heappop(self.scheduled)
+            task = next_.task
+            # Wait for due date
+            delay = max(0, next_.due - time.monotonic())
+            await asyncio.sleep(delay)
+            # reschedule immediately (at next due date)
+            next_.due = time.monotonic() + task.seconds_til_next_run
+            heapq.heappush(self.scheduled, next_)
+            # fire and forget, task writes to mutable dict reference
+            # no data race because each task writes to its unique key
+            asyncio.create_task(self._run_task(task))
+
+
+class ContextStatus(Enum):
+    READY = "ready"
+    PENDING = "pending"
+    FAILED = "failed"
+
+
+class ContextResult(pydantic.BaseModel):
+    name: str
+    data: Any = None
+    state: ContextStatus = ContextStatus.PENDING
+
+    def __str__(self) -> str:
+        return f"ContextResult({self.name}, {self.state.value})"
+
+    def __hash__(self) -> int:
+        data: bytes = repr(self.data).encode()
+        return zlib.adler32(data) & 0xFFFFFFFF
+
+
+class ContextTask(pydantic.BaseModel):
+    name: str
+    spec: Loadable
+    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:
+        result = await self.try_load()
+        if result.state == ContextStatus.READY:
+            output[self.name] = result
+
+    async def try_load(self) -> "ContextResult":
+        attempts_remaining, retry_interval = TaskRetryPolicy.from_task(self)
+        data = None
+        state = ContextStatus.PENDING
+        while attempts_remaining > 0:
+            stats.increment("context.refresh.attempt", tags=[f"context:{self.name}"])
             try:
-                if isinstance(v, Loadable):
-                    ret[k] = v.load()
-                elif isinstance(v, str):
-                    ret[k] = Loadable.from_legacy_fmt(v).load()
-                self.stats.increment(
-                    "context.refresh.success",
-                    tags=[f"context:{k}"],
+                load_fn = self.spec.load
+                if inspect.iscoroutinefunction(load_fn):
+                    data = await load_fn()
+                else:
+                    data = load_fn()
+                stats.increment(
+                    "context.refresh.success", tags=[f"context:{self.name}"]
                 )
-            except Exception as e:  # pylint: disable=broad-exception-caught
-                self.logger(event=e)
-                self.stats.increment(
-                    "context.refresh.error",
-                    tags=[f"context:{k}"],
+                state = ContextStatus.READY
+                break
+            except Exception as e:
+                log.error(
+                    "Context failed to refresh",
+                    error=[line for line in str(e).splitlines()],
                 )
-        if "crypto" not in ret:
-            ret["crypto"] = self.crypto
-        return ret
+                state = ContextStatus.FAILED
+                stats.increment("context.refresh.error", tags=[f"context:{self.name}"])
+            attempts_remaining -= 1
+            await asyncio.sleep(retry_interval)
+        return ContextResult(
+            name=self.name,
+            data=data,
+            state=state,
+        )
 
-    def build_new_context_from_instances(self, node_value: str) -> Dict[str, Any]:
-        matches = self.poller.match_node(node_value=node_value)
-        ret = dict()
-        for key, value in self.context.items():
-            try:
-                ret[key] = deepcopy(value)
-            except TypeError:
-                ret[key] = value
-
-        to_add = dict()
-        for scope, instances in matches.scopes.items():
-            if scope in ("default", None):
-                to_add["instances"] = instances
+    @property
+    def seconds_til_next_run(self) -> int:
+        match self.interval.value:
+            case CronInterval(cron=expression):
+                cron = croniter(expression)
+                next_date = cron.get_next(datetime.datetime)
+                return int(wait_until(next_date))
+            case SecondsInterval(seconds=seconds):
+                return seconds
+            case _:
+                return 1
+
+    @classmethod
+    def from_loadable(cls, name: str, loadable: Loadable) -> "ContextTask":
+        interval = loadable.interval
+        if interval is None:
+            cfg = config.template_context
+            if cfg.refresh_rate is not None:
+                interval = str(cfg.refresh_rate)
+            elif cfg.refresh_cron is not None:
+                interval = cfg.refresh_cron
             else:
-                to_add[scope] = instances
-        if to_add == {}:
-            raise HTTPException(
-                detail=(
-                    "This node does not match any instances! ",
-                    "If node matching is enabled, check that the node "
-                    "match key aligns with the source match key. "
-                    "If you don't know what any of this is, disable "
-                    "node matching via the config",
-                ),
-                status_code=400,
-            )
-        ret.update(to_add)
-        return ret
+                interval = "60"
+        retry_policy = None
+        if policy := loadable.retry_policy:
+            retry_policy = TaskRetryPolicy(**policy)
 
-    def get_context(
-        self, request: DiscoveryRequest, template: XdsTemplate
-    ) -> Dict[str, Any]:
-        ret = self.build_new_context_from_instances(
-            node_value=self.poller.extract_node_key(request.node),
+        return ContextTask(
+            name=name,
+            spec=loadable,
+            interval=TaskInterval.from_str(interval),
+            retry_policy=retry_policy,
         )
-        if request.hide_private_keys:
-            ret["crypto"] = self.disabled_suite
-        if not template.is_python_source:
-            keys_to_remove = self.unused_variables(list(ret), template.jinja_variables)
-            for key in keys_to_remove:
-                ret.pop(key, None)
-        return ret
+
+    def __str__(self) -> str:
+        return f"ContextTask({self.name}, {self.spec})"
+
+    __repr__ = __str__
+
+
+class TaskRetryPolicy(pydantic.BaseModel):
+    num_retries: int
+    interval: int
 
     @staticmethod
-    def unused_variables(
-        keys: Iterable[str], variables: Iterable[str]
-    ) -> Generator[str, None, None]:
-        for key in keys:
-            if key not in variables:
-                yield key
-
-    def get(self, *args: Any, **kwargs: Any) -> Any:
-        return self.context.get(*args, **kwargs)
+    def from_task(t: "ContextTask") -> tuple[int, int]:
+        interval = DEFAULT_RETRY_INTERVAL
+        attempts = 1
+        if policy := t.spec.retry_policy:
+            try:
+                retry_policy = TaskRetryPolicy(**policy)
+                interval = retry_policy.interval
+                attempts += retry_policy.num_retries
+            except Exception as e:
+                # TODO: warning
+                print(f"Failed to parse retry policy of task: {t}. Error: {e}")
+        else:
+            attempts += DEFAULT_NUM_RETRIES
+        return attempts, interval
+
+
+class TaskInterval(pydantic.BaseModel):
+    value: "TaskIntervalKind"
+
+    @classmethod
+    def from_str(cls, s: str) -> "TaskInterval":
+        if s.isdigit():
+            return TaskInterval(value=SecondsInterval(seconds=int(s)))
+        if croniter.is_valid(s):
+            return TaskInterval(value=CronInterval(cron=s))
+        raise ValueError(f"Invalid interval string: {s}")
+
+
+class CronInterval(pydantic.BaseModel):
+    cron: str
+
+
+class SecondsInterval(pydantic.BaseModel):
+    seconds: int
+
+
+TaskIntervalKind = Union[CronInterval, SecondsInterval]

sovereign/dynamic_config/__init__.py

@@ -0,0 +1,112 @@
+import inspect
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from sovereign.dynamic_config.deser import ConfigDeserializer
+from sovereign.dynamic_config.loaders import CustomLoader
+from sovereign.utils.entry_point_loader import EntryPointLoader
+
+LOADERS: dict[str, CustomLoader] = {}
+DESERIALIZERS: dict[str, ConfigDeserializer] = {}
+
+
+class Loadable(BaseModel):
+    path: str = Field(alias="target")
+    protocol: str = Field(alias="loader")
+    serialization: str | None = Field(None, alias="deserialize_with")
+    interval: str | None = None
+    retry_policy: dict[str, Any] | None = None
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    def load(self, default: Any = None) -> Any:
+        global LOADERS
+        if not LOADERS:
+            init_loaders()
+
+        global DESERIALIZERS
+        if not DESERIALIZERS:
+            init_deserializers()
+
+        if self.protocol not in LOADERS:
+            raise KeyError(
+                f"Could not find CustomLoader {self.protocol}. Available: {LOADERS}"
+            )
+        loader = LOADERS[self.protocol]
+
+        ser = self.serialization
+        if ser is None:
+            ser = loader.default_deser
+        elif ser not in DESERIALIZERS:
+            raise KeyError(
+                f"Could not find Deserializer {ser}. Available: {DESERIALIZERS}"
+            )
+        deserializer = DESERIALIZERS[ser]
+
+        try:
+            data = loader.load(self.path)
+            return deserializer.deserialize(data)
+        except Exception as original_error:
+            if default is not None:
+                return default
+            raise Exception(
+                f"Could not load value. {self.__str__()}, {original_error=}"
+            )
+
+    @staticmethod
+    def from_legacy_fmt(fmt_string: str) -> "Loadable":
+        if "://" not in fmt_string:
+            return Loadable(
+                loader="inline",
+                deserialize_with="string",
+                target=fmt_string,
+            )
+        try:
+            scheme, path = fmt_string.split("://")
+        except ValueError:
+            raise ValueError(fmt_string)
+        try:
+            proto, ser = scheme.split("+")
+        except ValueError:
+            proto, ser = scheme, "yaml"
+
+        if proto in ("python", "module"):
+            ser = "raw"
+        if proto in ("http", "https"):
+            path = "://".join([proto, path])
+
+        return Loadable(
+            loader=proto,
+            deserialize_with=ser,
+            target=path,
+        )
+
+    def __str__(self) -> str:
+        return f"Loadable({self.protocol}+{self.serialization}://{self.path})"
+
+
+def init_loaders():
+    global LOADERS
+    for entry_point in EntryPointLoader("loaders").groups["loaders"]:
+        custom_loader = entry_point.load()
+        func = custom_loader()
+        method = getattr(func, "load")
+        if not inspect.ismethod(method):
+            raise AttributeError(
+                f"CustomLoader {entry_point.name} does not implement .load()"
+            )
+        LOADERS[entry_point.name] = func
+
+
+def init_deserializers():
+    global DESERIALIZERS
+    for entry_point in EntryPointLoader("deserializers").groups["deserializers"]:
+        deserializer = entry_point.load()
+        func = deserializer()
+        method = getattr(func, "deserialize")
+        if not inspect.ismethod(method):
+            raise AttributeError(
+                f"Deserializer {entry_point.name} does not implement .deserialize()"
+            )
+        DESERIALIZERS[entry_point.name] = func

sovereign/dynamic_config/deser.py

@@ -0,0 +1,78 @@
+import json
+from typing import Any, Protocol
+
+import jinja2
+import yaml
+
+try:
+    import ujson
+
+    UJSON_AVAILABLE = True
+except ImportError:
+    UJSON_AVAILABLE = False
+
+try:
+    import orjson
+
+    ORJSON_AVAILABLE = True
+except ImportError:
+    ORJSON_AVAILABLE = False
+
+jinja_env = jinja2.Environment(autoescape=True)
+
+
+class ConfigDeserializer(Protocol):
+    """
+    Deserializers can be added to sovereign by creating a subclass
+    and then specified in config:
+
+    template_context:
+      context:
+          ...:
+            protocol: ...
+            serialization: <serializer name>
+            path: ...
+    """
+
+    def deserialize(self, input: Any) -> Any: ...
+
+
+class YamlDeserializer(ConfigDeserializer):
+    def deserialize(self, input: Any) -> Any:
+        return yaml.safe_load(input)
+
+
+class JsonDeserializer(ConfigDeserializer):
+    def deserialize(self, input: Any) -> Any:
+        return json.loads(input)
+
+
+class JinjaDeserializer(ConfigDeserializer):
+    def deserialize(self, input: Any) -> Any:
+        return jinja_env.from_string(input)
+
+
+class StringDeserializer(ConfigDeserializer):
+    def deserialize(self, input: Any) -> Any:
+        return str(input)
+
+
+class PassthroughDeserializer(ConfigDeserializer):
+    def deserialize(self, input: Any) -> Any:
+        return input
+
+
+class UjsonDeserializer(ConfigDeserializer):
+    def deserialize(self, input: Any) -> Any:
+        if not UJSON_AVAILABLE:
+            raise ImportError("Configured a UJSON deserializer but it's not installed")
+        return ujson.loads(input)
+
+
+class OrjsonDeserializer(ConfigDeserializer):
+    def deserialize(self, input: Any) -> Any:
+        if not ORJSON_AVAILABLE:
+            raise ImportError(
+                "Configured an ORJSON deserializer but it's not installed"
+            )
+        return orjson.loads(input)