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

sovereign/v2/logging.py

@@ -0,0 +1,81 @@
+import logging
+import os
+from typing import Any, MutableMapping
+
+import structlog
+from pydantic import BaseModel
+
+# noinspection PyProtectedMember
+from structlog.dev import RichTracebackFormatter
+from structlog.typing import FilteringBoundLogger
+
+
+def get_named_logger(name: str, level: int = logging.INFO) -> FilteringBoundLogger:
+    """
+    Gets a structured logger with a specific name to allow us to control log levels separately.
+
+    Set LOG_FORMAT=human for pretty-printed, colourful output.
+    """
+
+    # noinspection PyUnusedLocal
+    def filter_by_level(
+        logger: Any, method_name: str, event_dict: MutableMapping[str, Any]
+    ) -> MutableMapping[str, Any]:
+        level_map = {
+            "debug": logging.DEBUG,
+            "info": logging.INFO,
+            "warn": logging.WARN,
+            "warning": logging.WARNING,
+            "error": logging.ERROR,
+            "critical": logging.CRITICAL,
+            "exception": logging.ERROR,
+        }
+        method_level = level_map.get(method_name, logging.INFO)
+        if method_level < level:
+            raise structlog.DropEvent
+        return event_dict
+
+    # noinspection PyUnusedLocal
+    def serialise_pydantic_models(
+        logger: FilteringBoundLogger,
+        method_name: str,
+        event_dict: MutableMapping[str, Any],
+    ) -> MutableMapping[str, Any]:
+        for key, value in event_dict.items():
+            if isinstance(value, BaseModel):
+                event_dict[key] = value.model_dump()
+        return event_dict
+
+    log_format = os.environ.get("LOG_FORMAT", "json").lower()
+    is_human_format = log_format == "human"
+
+    base_processors = [
+        filter_by_level,
+        structlog.stdlib.add_log_level,
+        serialise_pydantic_models,
+        structlog.processors.TimeStamper(
+            fmt="iso" if not is_human_format else "%Y-%m-%d %H:%M:%S"
+        ),
+    ]
+
+    if is_human_format:
+        # human-readable format with colours
+        final_processors = base_processors + [
+            structlog.processors.UnicodeDecoder(),
+            structlog.dev.ConsoleRenderer(
+                colors=True,
+                exception_formatter=RichTracebackFormatter(show_locals=False),
+                pad_event=30,
+                sort_keys=False,
+            ),
+        ]
+    else:
+        # JSON format for production/machine consumption
+        current_processors = list(structlog.get_config()["processors"])
+        final_processors = base_processors + current_processors
+
+    return structlog.wrap_logger(
+        structlog.PrintLogger(),
+        final_processors,
+        context_class=dict,
+    ).bind(logger_name=name)

sovereign/v2/types.py

@@ -0,0 +1,41 @@
+from typing import Any
+
+import pydantic
+from pydantic import TypeAdapter
+
+from sovereign.types import DiscoveryRequest, DiscoveryResponse
+
+
+class Context(pydantic.BaseModel):
+    name: str
+    data: Any
+    data_hash: int
+    last_refreshed_at: int | None = None
+    refresh_after: int
+
+
+class DiscoveryEntry(pydantic.BaseModel):
+    request_hash: str
+    template: str
+    request: DiscoveryRequest
+    response: DiscoveryResponse | None
+    last_rendered_at: int | None = None
+
+
+class RefreshContextJob(pydantic.BaseModel):
+    context_name: str
+
+
+class RenderDiscoveryJob(pydantic.BaseModel):
+    request_hash: str
+
+
+QueueJob = RefreshContextJob | RenderDiscoveryJob
+
+
+queue_job_type_adapter = TypeAdapter(QueueJob)
+
+
+class WorkerNode(pydantic.BaseModel):
+    node_id: str
+    last_heartbeat: int

sovereign/v2/web.py

@@ -0,0 +1,101 @@
+import asyncio
+import logging
+import os
+import threading
+
+from structlog.typing import FilteringBoundLogger
+
+from sovereign import config
+from sovereign.types import DiscoveryRequest, DiscoveryResponse
+from sovereign.v2.data.repositories import DiscoveryEntryRepository
+from sovereign.v2.data.utils import get_data_store, get_queue
+from sovereign.v2.logging import get_named_logger
+from sovereign.v2.types import DiscoveryEntry, RenderDiscoveryJob
+
+
+async def wait_for_discovery_response(
+    request: DiscoveryRequest,
+) -> DiscoveryResponse | None:
+    # 1 - check if the entry already exists in the database with a non-empty response
+    # 2 - if it does, return it
+    # 3 - if it doesn't, enqueue a new job to render it
+    # 4 - poll for up to CACHE_READ_TIMEOUT seconds, if we find a response, return it
+
+    request_hash = request.cache_key(config.cache.hash_rules)
+
+    logger: FilteringBoundLogger = get_named_logger(
+        f"{__name__}.{wait_for_discovery_response.__qualname__} ({__file__})",
+        level=logging.DEBUG,
+    ).bind(
+        request_hash=request_hash,
+        template_resource_type=request.template.resource_type,
+        process_id=os.getpid(),
+        thread_id=threading.get_ident(),
+    )
+
+    logger.debug("Starting lookup for discovery response")
+
+    data_store = get_data_store()
+    discovery_entry_repository = DiscoveryEntryRepository(data_store)
+
+    queue = get_queue()
+
+    discovery_entry = discovery_entry_repository.get(request_hash)
+
+    if not discovery_entry:
+        logger.debug(
+            "No existing discovery entry found, creating new entry and enqueuing job"
+        )
+
+        # we need to save this request to the database
+        discovery_entry = DiscoveryEntry(
+            request_hash=request_hash,
+            template=request.template.resource_type,
+            request=request,
+            response=None,
+        )
+        discovery_entry_repository.save(discovery_entry)
+
+    if not discovery_entry.response:
+        # enqueue a job to render this discovery request (duplicates handled in the worker)
+        job = RenderDiscoveryJob(request_hash=request_hash)
+        queue.put(job)
+
+    if discovery_entry.response:
+        logger.debug("Returning cached response immediately")
+        return discovery_entry.response
+
+    # wait for up to CACHE_READ_TIMEOUT seconds for the response to be populated
+    logger.debug(
+        "Polling for response",
+        timeout=config.cache.read_timeout,
+        poll_interval=config.cache.poll_interval_secs,
+    )
+
+    start_time = asyncio.get_event_loop().time()
+    attempts = 0
+
+    while (
+        asyncio.get_event_loop().time() - start_time
+    ) < config.cache.read_timeout and discovery_entry.response is None:
+        attempts += 1
+        discovery_entry = discovery_entry_repository.get(request_hash)
+        if discovery_entry is None:
+            logger.error("No discovery entry found while polling for response")
+            return None
+        await asyncio.sleep(config.cache.poll_interval_secs)
+
+    elapsed_time = asyncio.get_event_loop().time() - start_time
+
+    if discovery_entry.response:
+        logger.debug(
+            "Response received after polling",
+            attempts=attempts,
+            elapsed_time=elapsed_time,
+        )
+    else:
+        logger.error(
+            "Timeout waiting for response", attempts=attempts, elapsed_time=elapsed_time
+        )
+
+    return discovery_entry.response

sovereign/v2/worker.py

@@ -0,0 +1,199 @@
+import logging
+import os
+import random
+import threading
+import time
+
+from structlog.typing import FilteringBoundLogger
+
+from sovereign import stats
+from sovereign.configuration import config
+from sovereign.dynamic_config import Loadable
+from sovereign.v2.data.data_store import DataStoreProtocol
+from sovereign.v2.data.repositories import (
+    ContextRepository,
+    DiscoveryEntryRepository,
+    WorkerNodeRepository,
+)
+from sovereign.v2.data.utils import get_data_store, get_queue
+from sovereign.v2.data.worker_queue import QueueProtocol
+from sovereign.v2.jobs.refresh_context import get_refresh_after, refresh_context
+from sovereign.v2.jobs.render_discovery_job import render_discovery_response
+from sovereign.v2.logging import get_named_logger
+from sovereign.v2.types import (
+    QueueJob,
+    RefreshContextJob,
+    RenderDiscoveryJob,
+)
+
+
+class Worker:
+    context_repository: ContextRepository
+    discovery_entry_repository: DiscoveryEntryRepository
+    worker_node_repository: WorkerNodeRepository
+
+    queue: QueueProtocol
+
+    def __init__(
+        self,
+        data_store: DataStoreProtocol | None = None,
+        node_id: str | None = None,
+        queue: QueueProtocol | None = None,
+    ) -> None:
+        self.logger: FilteringBoundLogger = get_named_logger(
+            f"{self.__class__.__module__}.{self.__class__.__qualname__}",
+            level=logging.INFO,
+        )
+
+        self.node_id = (
+            node_id
+            if node_id is not None
+            else f"{time.time()}{random.randint(0, 1000000)}"
+        )
+
+        data_store = data_store if data_store is not None else get_data_store()
+
+        self.context_repository = ContextRepository(data_store)
+        self.discovery_entry_repository = DiscoveryEntryRepository(data_store)
+        self.worker_node_repository = WorkerNodeRepository(data_store)
+
+        self.queue = queue if queue is not None else get_queue()
+
+    def start(self):
+        # start the context refresh loop and daemonise it
+        threading.Thread(daemon=True, target=self.context_refresh_loop).start()
+
+        # pull from the queue for eternity and process the messages
+        while True:
+            try:
+                if message := self.queue.get():
+                    job_type = type(message.job).__name__
+                    stats.increment(
+                        "v2.worker.queue.message_received",
+                        tags=[f"job_type:{job_type}"],
+                    )
+                    self.process_job(message.job)
+                    self.queue.ack(message.receipt_handle)
+                    stats.increment(
+                        "v2.worker.queue.message_acked", tags=[f"job_type:{job_type}"]
+                    )
+            except Exception:
+                stats.increment("v2.worker.queue.error")
+                self.logger.exception("Error while processing job")
+
+    def process_job(self, job: QueueJob):
+        self.logger.info(
+            "Processing job from queue",
+            job_type=type(job),
+            job=job,
+            node_id=self.node_id,
+            process_id=os.getpid(),
+            thread_id=threading.get_ident(),
+        )
+
+        match job:
+            case RefreshContextJob():
+                refresh_context(
+                    job.context_name,
+                    self.node_id,
+                    config,
+                    self.context_repository,
+                    self.discovery_entry_repository,
+                    self.queue,
+                )
+            case RenderDiscoveryJob():
+                render_discovery_response(
+                    job.request_hash,
+                    self.context_repository,
+                    self.discovery_entry_repository,
+                    self.node_id,
+                )
+
+    def context_refresh_loop(self):
+        self.logger.info(
+            "Starting context refresh loop",
+            node_id=self.node_id,
+            process_id=os.getpid(),
+            thread_id=threading.get_ident(),
+        )
+
+        is_leader = False
+
+        while True:
+            try:
+                self.worker_node_repository.send_heartbeat(self.node_id)
+                self.worker_node_repository.prune_dead_nodes()
+
+                if not self.worker_node_repository.get_leader_node_id() == self.node_id:
+                    is_leader = False
+                    self.logger.info(
+                        "This node is not the leader, checking again in 60 seconds",
+                        node_id=self.node_id,
+                        process_id=os.getpid(),
+                        thread_id=threading.get_ident(),
+                    )
+                    time.sleep(60)
+                    continue
+
+                # I am the leader
+                if not is_leader:
+                    is_leader = True
+                    self.logger.info(
+                        "This node is the leader, begin refreshing contexts",
+                        node_id=self.node_id,
+                        process_id=os.getpid(),
+                        thread_id=threading.get_ident(),
+                    )
+
+                name: str
+                loadable: Loadable
+                for name, loadable in config.template_context.context.items():
+                    refresh_after: int | None = (
+                        self.context_repository.get_refresh_after(name)
+                    )
+
+                    time_now = time.time()
+
+                    if refresh_after is None or refresh_after <= time.time():
+                        job = RefreshContextJob(context_name=name)
+
+                        self.logger.info(
+                            "Queuing context refresh",
+                            node_id=self.node_id,
+                            process_id=os.getpid(),
+                            thread_id=threading.get_ident(),
+                            name=name,
+                            refresh_after=refresh_after,
+                            refresh_after_seconds=(refresh_after or time_now)
+                            - time_now,
+                        )
+                        self.queue.put(job)
+                        stats.increment(
+                            "v2.worker.context_refresh.queued", tags=[f"context:{name}"]
+                        )
+
+                        # update refresh_after to ensure that, at most, we refresh once per interval
+                        self.context_repository.update_refresh_after(
+                            name, get_refresh_after(config, loadable)
+                        )
+                    else:
+                        stats.increment(
+                            "v2.worker.context_refresh.skipped",
+                            tags=[f"context:{name}"],
+                        )
+                        self.logger.debug(
+                            "Skipping context refresh",
+                            node_id=self.node_id,
+                            process_id=os.getpid(),
+                            thread_id=threading.get_ident(),
+                            name=name,
+                            refresh_after=refresh_after,
+                            refresh_after_seconds=(refresh_after or time_now)
+                            - time_now,
+                        )
+
+                time.sleep(1)
+            except Exception:
+                stats.increment("v2.worker.context_refresh.error")
+                self.logger.exception("Error while refreshing context")
+                time.sleep(5)

sovereign/views/__init__.py

@@ -0,0 +1,7 @@
+from sovereign import cache, config
+
+if not config.worker_v2_enabled:
+    reader = cache.CacheReader()
+else:
+    # don't use the cache reader when in v2 worker mode
+    reader = None

sovereign/views/api.py

@@ -0,0 +1,82 @@
+import json
+from typing import Annotated, Optional
+
+from fastapi import APIRouter, Path, Query
+from fastapi.responses import Response
+
+from sovereign.cache import Entry
+from sovereign.configuration import ConfiguredResourceTypes, config
+from sovereign.utils.mock import mock_discovery_request
+from sovereign.v2.web import wait_for_discovery_response
+from sovereign.views import reader
+
+router = APIRouter()
+
+
+def _traverse(data, prefix, expressions):
+    for key, value in data.items():
+        path = f"{prefix}.{key}" if prefix else key
+        if isinstance(value, dict):
+            yield from _traverse(value, path, expressions)
+        else:
+            yield f"{path}={value}"
+
+
+def expand_metadata_to_expr(m):
+    exprs = []
+    yield from _traverse(m, "", exprs)
+
+
+# noinspection DuplicatedCode
+@router.get("/resources/{resource_type}", summary="Get resources for a given type")
+async def resource(
+    resource_type: Annotated[ConfiguredResourceTypes, Path(title="xDS Resource type")],
+    resource_name: Optional[str] = Query(None, title="Resource name"),
+    api_version: Optional[str] = Query("v3", title="Envoy API version"),
+    service_cluster: Optional[str] = Query("*", title="Envoy Service cluster"),
+    region: Optional[str] = Query(None, title="Locality Zone"),
+    version: Optional[str] = Query(None, title="Envoy Semantic Version"),
+    metadata: Optional[str] = Query(None, title="Envoy node metadata to filter by"),
+) -> Response:
+    # todo: rewrite for worker v2
+
+    expressions = [f"cluster={service_cluster}"]
+    try:
+        data = {"metadata": json.loads(metadata or "{}")}
+        for expr in expand_metadata_to_expr(data):
+            expressions.append(expr)
+    except Exception:
+        pass
+    kwargs = {
+        "api_version": api_version,
+        "resource_type": ConfiguredResourceTypes(resource_type).value,
+        "resource_names": resource_name,
+        "version": version,
+        "region": region,
+        "expressions": expressions,
+    }
+    req = mock_discovery_request(**{k: v for k, v in kwargs.items() if v is not None})  # type: ignore
+
+    entry: Entry | None = None
+
+    if config.worker_v2_enabled:
+        # we're set up to use v2 of the worker
+        discovery_response = await wait_for_discovery_response(req)
+        if discovery_response is not None:
+            entry = Entry(
+                text=discovery_response.model_dump_json(indent=None),
+                len=len(discovery_response.resources),
+                version=discovery_response.version_info,
+                node=req.node,
+            )
+
+    else:
+        entry = await reader.blocking_read(req)  # ty: ignore[possibly-missing-attribute]
+
+    if content := getattr(entry, "text", None):
+        return Response(content, media_type="application/json")
+    else:
+        return Response(
+            json.dumps({"title": "No resources found", "status": 404}),
+            media_type="application/json+problem",
+        )

sovereign/views/crypto.py

@@ -1,31 +1,38 @@
-from typing import Dict
-from pydantic import BaseModel, Field
+from typing import Any, Dict, Optional
+
 from fastapi import APIRouter, Body
 from fastapi.responses import JSONResponse
-from sovereign import json_response_class, cipher_suite
-from sovereign.utils.crypto import generate_key
+from pydantic import BaseModel, Field
+
+from sovereign import logs, server_cipher_container
+from sovereign.configuration import EncryptionConfig
+from sovereign.response_class import json_response_class
+from sovereign.utils.crypto.crypto import CipherContainer
+from sovereign.utils.crypto.suites import EncryptionType
 
 router = APIRouter()
 
 
 class EncryptionRequest(BaseModel):
     data: str = Field(..., title="Text to be encrypted", min_length=1, max_length=65535)
-    key: str = Field(
+    key: Optional[str] = Field(
         None,
-        title="Optional Fernet encryption key to use to encrypt",
+        title="Optional encryption key to use to encrypt",
         min_length=44,
         max_length=44,
     )
+    encryption_type: str = Field(default="fernet", title="Encryption type to be used")
 
 
 class DecryptionRequest(BaseModel):
     data: str = Field(..., title="Text to be decrypted", min_length=1, max_length=65535)
     key: str = Field(
         ...,
-        title="Fernet encryption key to use to decrypt",
+        title="Encryption key to use to decrypt",
         min_length=44,
         max_length=44,
     )
+    encryption_type: str = Field(default="fernet", title="Encryption type to be used")
 
 
 class DecryptableRequest(BaseModel):
@@ -37,17 +44,37 @@ class DecryptableRequest(BaseModel):
     summary="Decrypt provided encrypted data using a provided key",
     response_class=json_response_class,
 )
-async def _decrypt(request: DecryptionRequest = Body(None)) -> Dict[str, str]:
-    return {"result": cipher_suite.decrypt(request.data, request.key)}
+async def _decrypt(request: DecryptionRequest = Body(None)) -> dict[str, Any]:
+    user_cipher_container = CipherContainer.from_encryption_configs(
+        encryption_configs=[
+            EncryptionConfig(
+                encryption_key=request.key,
+                encryption_type=EncryptionType(request.encryption_type),
+            )
+        ],
+        logger=logs.application_logger.logger,
+    )
+    return {**user_cipher_container.decrypt_with_type(request.data)}
 
 
 @router.post(
     "/encrypt",
-    summary="Encrypt provided data using this servers key",
+    summary="Encrypt provided data using this servers key or provided key",
     response_class=json_response_class,
 )
-async def _encrypt(request: EncryptionRequest = Body(None)) -> Dict[str, str]:
-    return {"result": cipher_suite.encrypt(data=request.data, key=request.key)}
+async def _encrypt(request: EncryptionRequest = Body(None)) -> dict[str, Any]:
+    if request.key:
+        user_cipher_container = CipherContainer.from_encryption_configs(
+            encryption_configs=[
+                EncryptionConfig(
+                    encryption_key=request.key,
+                    encryption_type=EncryptionType(request.encryption_type),
+                )
+            ],
+            logger=logs.application_logger.logger,
+        )
+        return {**user_cipher_container.encrypt(request.data)}
+    return {**server_cipher_container.encrypt(request.data)}
 
 
 @router.post(
@@ -56,7 +83,7 @@ async def _encrypt(request: EncryptionRequest = Body(None)) -> Dict[str, str]:
     response_class=json_response_class,
 )
 async def _decryptable(request: DecryptableRequest = Body(None)) -> JSONResponse:
-    cipher_suite.decrypt(request.data)
+    server_cipher_container.decrypt(request.data)
     return json_response_class({})
 
 
@@ -65,5 +92,9 @@ async def _decryptable(request: DecryptableRequest = Body(None)) -> JSONResponse
     summary="Generate a new asymmetric encryption key",
     response_class=json_response_class,
 )
-def _generate_key() -> Dict[str, str]:
-    return {"result": generate_key()}
+def _generate_key(encryption_type: str = "fernet") -> Dict[str, str]:
+    cipher_suite = CipherContainer.get_cipher_suite(EncryptionType(encryption_type))
+    return {
+        "key": cipher_suite.generate_key().decode(),
+        "encryption_type": encryption_type,
+    }