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

sovereign/v2/data/repositories.py

@@ -0,0 +1,90 @@
+import time
+
+from sovereign import stats
+from sovereign.v2.data.data_store import ComparisonOperator, DataStoreProtocol, DataType
+from sovereign.v2.types import Context, DiscoveryEntry, WorkerNode
+
+
+class ContextRepository:
+    def __init__(self, data_store: DataStoreProtocol):
+        self.data_store: DataStoreProtocol = data_store
+
+    @stats.timed("repository.context.get_ms")
+    def get(self, name: str) -> Context | None:
+        return self.data_store.get(DataType.Context, name)
+
+    @stats.timed("v2.repository.context.get_hash_ms")
+    def get_hash(self, name: str) -> int | None:
+        return self.data_store.get_property(DataType.Context, name, "data_hash")
+
+    def get_refresh_after(self, name: str) -> int | None:
+        return self.data_store.get_property(DataType.Context, name, "refresh_after")
+
+    @stats.timed("v2.repository.context.save_ms")
+    def save(self, context: Context) -> bool:
+        return self.data_store.set(DataType.Context, context.name, context)
+
+    @stats.timed("v2.repository.context.update_refresh_after_ms")
+    def update_refresh_after(self, name: str, refresh_after: int) -> bool:
+        return self.data_store.set_property(
+            DataType.Context, name, "refresh_after", refresh_after
+        )
+
+
+class DiscoveryEntryRepository:
+    def __init__(self, data_store: DataStoreProtocol):
+        self.data_store = data_store
+
+    @stats.timed("v2.repository.discovery_entry.get_ms")
+    def get(self, request_hash: str) -> DiscoveryEntry | None:
+        return self.data_store.get(DataType.DiscoveryEntry, request_hash)
+
+    @stats.timed("v2.repository.discovery_entry.find_by_template_ms")
+    def find_all_request_hashes_by_template(self, template: str) -> list[str]:
+        return self.data_store.find_all_matching_property(
+            DataType.DiscoveryEntry,
+            "template",
+            ComparisonOperator.EqualTo,
+            template,
+            "request_hash",
+        )
+
+    @stats.timed("v2.repository.discovery_entry.save_ms")
+    def save(self, entry: DiscoveryEntry) -> bool:
+        return self.data_store.set(DataType.DiscoveryEntry, entry.request_hash, entry)
+
+
+class WorkerNodeRepository:
+    def __init__(self, data_store: DataStoreProtocol):
+        self.data_store = data_store
+
+    @stats.timed("v2.repository.worker_node.heartbeat_ms")
+    def send_heartbeat(self, node_id: str) -> bool:
+        now = int(time.time())
+        return self.data_store.set(
+            DataType.WorkerNode,
+            node_id,
+            WorkerNode(node_id=node_id, last_heartbeat=now),
+        )
+
+    @stats.timed("v2.repository.worker_node.get_leader_ms")
+    def get_leader_node_id(self) -> str | None:
+        node: WorkerNode | None = self.data_store.min_by_property(
+            DataType.WorkerNode, "node_id"
+        )
+        if node:
+            return node.node_id
+        return None
+
+    @stats.timed("v2.repository.worker_node.prune_ms")
+    def prune_dead_nodes(self) -> bool:
+        """
+        Remove any nodes that have not sent a heartbeat in the last 10 minutes.
+        """
+        now = int(time.time())
+        return self.data_store.delete_matching(
+            DataType.WorkerNode,
+            "last_heartbeat",
+            ComparisonOperator.LessThanOrEqualTo,
+            now - 600,
+        )

sovereign/v2/data/utils.py

@@ -0,0 +1,33 @@
+from sovereign import config
+from sovereign.utils.entry_point_loader import EntryPointLoader
+from sovereign.v2.data.data_store import DataStoreProtocol
+from sovereign.v2.data.worker_queue import QueueProtocol
+
+
+def get_data_store() -> DataStoreProtocol:
+    entry_points = EntryPointLoader("data_stores")
+    data_store: DataStoreProtocol | None = None
+
+    for entry_point in entry_points.groups["data_stores"]:
+        if entry_point.name == config.worker_v2_data_store_provider:
+            data_store = entry_point.load()()
+            break
+
+    if data_store is None:
+        raise ValueError(
+            f"Data store '{config.worker_v2_data_store_provider}' not found in entry points"
+        )
+
+    return data_store
+
+
+def get_queue() -> QueueProtocol:
+    entry_points = EntryPointLoader("queues")
+
+    for entry_point in entry_points.groups["queues"]:
+        if entry_point.name == config.worker_v2_queue_provider:
+            return entry_point.load()()
+
+    raise ValueError(
+        f"Queue '{config.worker_v2_queue_provider}' not found in entry points"
+    )

sovereign/v2/data/worker_queue.py

@@ -0,0 +1,273 @@
+import logging
+import sqlite3
+import time
+import uuid
+from dataclasses import dataclass
+from typing import Protocol, runtime_checkable
+
+from structlog.typing import FilteringBoundLogger
+
+from sovereign import config
+from sovereign.v2.logging import get_named_logger
+from sovereign.v2.types import QueueJob, queue_job_type_adapter
+
+if config.worker_v2_queue_invsibility_time is None:
+    DEFAULT_VISIBILITY_TIMEOUT_SECONDS = int(config.cache.read_timeout) + 30
+else:
+    DEFAULT_VISIBILITY_TIMEOUT_SECONDS = config.worker_v2_queue_invsibility_time
+
+
+@dataclass
+class QueueMessage:
+    """A message retrieved from the queue, containing the job and a receipt handle for acknowledgement."""
+
+    job: QueueJob
+    receipt_handle: str
+
+
+@runtime_checkable
+class QueueProtocol(Protocol):
+    def put(self, job: QueueJob) -> str | None: ...
+
+    def get(self) -> QueueMessage | None: ...
+
+    def ack(self, receipt_handle: str) -> bool: ...
+
+
+class InMemoryQueue(QueueProtocol):
+    """
+    Messages become invisible when retrieved and must be acknowledged within the
+    visibility timeout, otherwise they become visible again for other workers.
+    """
+
+    def __init__(
+        self, visibility_timeout: int | None = DEFAULT_VISIBILITY_TIMEOUT_SECONDS
+    ) -> None:
+        self.logger: FilteringBoundLogger = get_named_logger(
+            f"{self.__class__.__module__}.{self.__class__.__qualname__}",
+            level=logging.DEBUG,
+        )
+
+        self.visibility_timeout: int = (
+            visibility_timeout
+            if visibility_timeout is not None
+            else DEFAULT_VISIBILITY_TIMEOUT_SECONDS
+        )
+
+        # storage for messages: message_id -> (job, invisible_until, receipt_handle)
+        self._messages: dict[str, tuple[QueueJob, float | None, str | None]] = {}
+
+    def put(self, job: QueueJob) -> str | None:
+        message_id = str(uuid.uuid4())
+        self._messages[message_id] = (job, None, None)  # visible, no receipt handle
+        self.logger.debug(
+            "Putting job in queue",
+            job=job,
+            message_id=message_id,
+            queue_size=len(self._messages),
+        )
+        return message_id
+
+    def get(self) -> QueueMessage | None:
+        timeout = 30
+        start_time = int(time.time())
+        poll_interval_seconds = 0.5
+
+        while int(time.time()) - start_time < timeout:
+            now = int(time.time())
+            # find first visible message (not invisible, or invisibility expired)
+            for message_id, (job, invisible_until, _) in self._messages.items():
+                # check if message is visible
+                if invisible_until is None or invisible_until <= now:
+                    # make it invisible and generate new receipt handle
+                    receipt_handle = str(uuid.uuid4())
+                    new_invisible_until = now + self.visibility_timeout
+                    self._messages[message_id] = (
+                        job,
+                        new_invisible_until,
+                        receipt_handle,
+                    )
+
+                    self.logger.debug(
+                        "Retrieved job from queue",
+                        message_id=message_id,
+                        receipt_handle=receipt_handle,
+                        invisible_until=new_invisible_until,
+                    )
+                    return QueueMessage(job=job, receipt_handle=receipt_handle)
+
+            time.sleep(poll_interval_seconds)
+
+        return None
+
+    def ack(self, receipt_handle: str) -> bool:
+        """
+        Acknowledge a message, permanently removing it from the queue.
+
+        Returns True if the message was successfully acknowledged, False if the
+        receipt handle was invalid (message doesn't exist or was redelivered).
+        """
+        for message_id, (job, invisible_until, stored_receipt) in list(
+            self._messages.items()
+        ):
+            if stored_receipt == receipt_handle:
+                del self._messages[message_id]
+                self.logger.debug(
+                    "Acknowledged job",
+                    message_id=message_id,
+                    receipt_handle=receipt_handle,
+                )
+                return True
+
+        self.logger.warning(
+            "Failed to acknowledge job, invalid receipt handle",
+            receipt_handle=receipt_handle,
+        )
+        return False
+
+    def is_empty(self) -> bool:
+        return not self._messages
+
+
+class SqliteQueue(QueueProtocol):
+    """
+    SQLite-backed queue with visibility timeout support.
+
+    Messages become invisible when retrieved and must be acknowledged within the
+    visibility timeout, otherwise they become visible again for other workers.
+    """
+
+    def __init__(self, visibility_timeout: int = DEFAULT_VISIBILITY_TIMEOUT_SECONDS):
+        self.logger: FilteringBoundLogger = get_named_logger(
+            f"{self.__class__.__module__}.{self.__class__.__qualname__}",
+            level=logging.DEBUG,
+        )
+        self.visibility_timeout = visibility_timeout
+        self.db_path = config.worker_v2_queue_path
+        self._init_db()
+
+    def _get_connection(self) -> sqlite3.Connection:
+        # check_same_thread=False allows SQLite connections to be shared across threads
+        # and means that we need to ensure thread safety ourselves.
+        # isolation_level=None uses autocommit mode,
+        # which prevents "cannot commit - no transaction is active" errors in multi-threaded contexts.
+        conn = sqlite3.connect(
+            self.db_path, check_same_thread=False, isolation_level=None
+        )
+        conn.row_factory = sqlite3.Row
+        return conn
+
+    def _init_db(self):
+        try:
+            with self._get_connection() as conn:
+                conn.execute("""
+                             CREATE TABLE IF NOT EXISTS queue
+                             (
+                                 id INTEGER PRIMARY KEY AUTOINCREMENT,
+                                 data TEXT NOT NULL,
+                                 invisible_until INT,
+                                 receipt_handle TEXT
+                             )
+                             """)
+                conn.execute(
+                    "CREATE INDEX IF NOT EXISTS idx_invisible_until ON queue (invisible_until)"
+                )
+                conn.execute(
+                    "CREATE INDEX IF NOT EXISTS idx_receipt_handle ON queue (receipt_handle)"
+                )
+                conn.commit()
+        except Exception:
+            self.logger.exception("Failed to initialise SQLite queue database")
+            raise
+
+    def put(self, job: QueueJob) -> str | None:
+        try:
+            with self._get_connection() as conn:
+                cursor = conn.execute(
+                    "INSERT INTO queue (data, invisible_until, receipt_handle) VALUES (?, NULL, NULL)",
+                    (job.model_dump_json(),),
+                )
+                job_id = str(cursor.lastrowid)
+                self.logger.debug("Put job in SQLite queue", job=job, job_id=job_id)
+                return str(job_id)
+        except Exception:
+            self.logger.exception("Failed to put job in SQLite queue", job=job)
+            return None
+
+    def get(self) -> QueueMessage | None:
+        timeout = 30
+        start_time = time.time()
+        poll_interval_seconds = 0.5
+
+        while time.time() - start_time < timeout:
+            try:
+                with self._get_connection() as conn:
+                    now = int(time.time())
+                    # find first visible message (invisible_until is NULL or expired)
+                    cursor = conn.execute(
+                        """
+                        SELECT id, data
+                        FROM queue
+                        WHERE invisible_until IS NULL
+                           OR invisible_until <= ? LIMIT 1
+                        """,
+                        (now,),
+                    )
+                    row = cursor.fetchone()
+                    if row:
+                        # generate receipt handle and make message invisible
+                        receipt_handle = str(uuid.uuid4())
+                        invisible_until = now + self.visibility_timeout
+                        conn.execute(
+                            "UPDATE queue SET invisible_until = ?, receipt_handle = ? WHERE id = ?",
+                            (invisible_until, receipt_handle, row["id"]),
+                        )
+                        conn.commit()
+                        self.logger.debug(
+                            "Retrieved job from queue",
+                            job_id=row["id"],
+                            receipt_handle=receipt_handle,
+                            invisible_until=invisible_until,
+                        )
+                        job = queue_job_type_adapter.validate_json(row["data"])
+                        return QueueMessage(job=job, receipt_handle=receipt_handle)
+            except Exception:
+                self.logger.exception("Failed to get job from SQLite queue")
+                return None
+
+            time.sleep(poll_interval_seconds)
+
+        return None
+
+    def ack(self, receipt_handle: str) -> bool:
+        """
+        Acknowledge a message, permanently removing it from the queue.
+
+        Returns True if the message was successfully acknowledged, False if the
+        receipt handle was invalid (message doesn't exist or was redelivered).
+        """
+        try:
+            with self._get_connection() as conn:
+                cursor = conn.execute(
+                    "DELETE FROM queue WHERE receipt_handle = ?",
+                    (receipt_handle,),
+                )
+                conn.commit()
+                if cursor.rowcount > 0:
+                    self.logger.debug(
+                        "Acknowledged job",
+                        receipt_handle=receipt_handle,
+                    )
+                    return True
+                else:
+                    self.logger.warning(
+                        "Failed to acknowledge job, invalid receipt handle",
+                        receipt_handle=receipt_handle,
+                    )
+                    return False
+        except Exception:
+            self.logger.exception(
+                "Failed to acknowledge job in SQLite queue",
+                receipt_handle=receipt_handle,
+            )
+            return False

sovereign/v2/jobs/refresh_context.py

@@ -0,0 +1,117 @@
+import datetime
+import logging
+import os
+import threading
+import time
+import zlib
+from typing import Any
+
+from croniter import croniter
+from structlog.typing import FilteringBoundLogger
+
+from sovereign.configuration import SovereignConfigv2
+from sovereign.context import CronInterval, SecondsInterval, TaskInterval, stats
+from sovereign.dynamic_config import Loadable
+from sovereign.utils.timer import wait_until
+from sovereign.v2.data.repositories import ContextRepository, DiscoveryEntryRepository
+from sovereign.v2.data.worker_queue import QueueProtocol
+from sovereign.v2.logging import get_named_logger
+from sovereign.v2.types import Context, RenderDiscoveryJob
+
+
+def refresh_context(
+    name: str,
+    node_id: str,
+    config: SovereignConfigv2,
+    context_repository: ContextRepository,
+    discovery_job_repository: DiscoveryEntryRepository,
+    queue: QueueProtocol,
+):
+    with stats.timed("v2.worker.job.refresh_context_ms", context=name):
+        loadable = config.template_context.context[name]
+
+        logger: FilteringBoundLogger = get_named_logger(
+            f"{__name__}.{refresh_context.__qualname__} ({__file__})",
+            level=logging.DEBUG,
+        ).bind(
+            name=name,
+            node_id=node_id,
+            process_id=os.getpid(),
+            thread_id=threading.get_ident(),
+        )
+
+        try:
+            value: Any = loadable.load()
+            context_hash = _get_hash(value)
+
+            if context_repository.get_hash(name) != context_hash:
+                context = Context(
+                    name=name,
+                    data=value,
+                    data_hash=context_hash,
+                    last_refreshed_at=int(time.time()),
+                    refresh_after=get_refresh_after(config, loadable),
+                )
+                context_repository.save(context)
+
+                request_hashes: set[str] = set()
+
+                for version, version_templates in (
+                    {"default": config.templates.default} | config.templates.versions
+                ).items():
+                    for template in version_templates:
+                        if name in template.depends_on:
+                            for request_hash in discovery_job_repository.find_all_request_hashes_by_template(
+                                template.type
+                            ):
+                                request_hashes.add(request_hash)
+
+                for request_hash in request_hashes:
+                    logger.info(
+                        "Queuing render for discovery request because context changed",
+                        request_hash=request_hash,
+                        context=name,
+                    )
+                    queue.put(RenderDiscoveryJob(request_hash=request_hash))
+        except Exception:
+            # if loadable.retry_policy is not None:
+            # print(loadable.retry_policy)
+            # todo: handle exceptions/retries
+            # todo: use the default retry logic instead
+            logger.exception("Failed to load context")
+
+
+def _get_hash(value: Any) -> int:
+    data: bytes = repr(value).encode()
+    return zlib.adler32(data) & 0xFFFFFFFF
+
+
+# noinspection PyUnreachableCode
+def _seconds_til_next_run(task_interval: TaskInterval) -> int:
+    match task_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 0
+
+
+def get_refresh_after(config: SovereignConfigv2, loadable: Loadable) -> int:
+    interval = loadable.interval
+
+    # get the default interval from config if not specified in loadable
+    if interval is None:
+        template_context_config = config.template_context
+        if template_context_config.refresh_rate is not None:
+            interval = str(template_context_config.refresh_rate)
+        elif template_context_config.refresh_cron is not None:
+            interval = template_context_config.refresh_cron
+        else:
+            interval = "60"
+
+    task_interval = TaskInterval.from_str(interval)
+
+    return int(time.time() + _seconds_til_next_run(task_interval))

sovereign/v2/jobs/render_discovery_job.py

@@ -0,0 +1,145 @@
+import logging
+import os
+import threading
+import time
+
+from structlog.typing import FilteringBoundLogger
+
+from sovereign import config, disabled_ciphersuite, server_cipher_container, stats
+from sovereign.rendering_common import (
+    add_type_urls,
+    deserialize_config,
+    filter_resources,
+)
+from sovereign.types import DiscoveryResponse, ProcessedTemplate
+from sovereign.utils import templates
+from sovereign.v2.data.repositories import ContextRepository, DiscoveryEntryRepository
+from sovereign.v2.logging import get_named_logger
+from sovereign.v2.types import Context, DiscoveryEntry
+
+
+# noinspection DuplicatedCode
+def render_discovery_response(
+    request_hash: str,
+    context_repository: ContextRepository,
+    discovery_entry_repository: DiscoveryEntryRepository,
+    node_id: str,
+):
+    logger: FilteringBoundLogger = get_named_logger(
+        f"{__name__}.{render_discovery_response.__qualname__} ({__file__})",
+        level=logging.DEBUG,
+    ).bind(
+        request_hash=request_hash,
+        node_id=node_id,
+        process_id=os.getpid(),
+        thread_id=threading.get_ident(),
+    )
+
+    try:
+        logger.debug("Starting rendering of discovery response")
+
+        discovery_entry = discovery_entry_repository.get(request_hash)
+
+        if discovery_entry is None:
+            logger.error("No discovery entry found for request hash")
+            return True  # don't retry this job, it won't succeed
+
+        request = discovery_entry.request
+
+        with stats.timed(
+            "v2.worker.job.render_discovery_response_ms",
+            template=discovery_entry.request.template.resource_type,
+        ):
+            logger = logger.bind(
+                template=discovery_entry.request.template.resource_type
+            )
+
+            dependencies = request.template.depends_on
+            contexts: dict[str, Context | None] = {
+                name: context_repository.get(name) for name in dependencies
+            }
+
+            missing_contexts = [
+                name
+                for name, context in contexts.items()
+                if context is None or context.last_refreshed_at is None
+            ]
+            if missing_contexts:
+                logger.error(
+                    "Cannot render template for request, required contexts not yet loaded",
+                    missing_contexts=missing_contexts,
+                )
+                return False
+
+            # in order to handle duplicate jobs for the same request_hash, check the last_rendered_at property - if this is
+            # greater than the all the last_refreshed_at values for the contexts, then we can skip rendering
+            refresh_times = [
+                context.last_refreshed_at
+                for context in contexts.values()
+                if context is not None and context.last_refreshed_at is not None
+            ]
+
+            if refresh_times:
+                latest_context_refresh = max(refresh_times)
+
+                if (
+                    discovery_entry.last_rendered_at
+                    and latest_context_refresh < discovery_entry.last_rendered_at
+                ):
+                    # the template was last rendered after all the contexts were refreshed, so we can skip rendering
+                    logger.info("Skipping rendering for duplicate job")
+                    return True
+
+            raw_contexts = {
+                name: context.data
+                for (name, context) in contexts.items()
+                if context is not None
+            }
+
+            logger.debug(
+                "Contexts loaded for rendering discovery response",
+                contexts=raw_contexts.keys(),
+                depends_on=request.template.depends_on,
+            )
+
+            if request.is_internal_request:
+                raw_contexts["__hide_from_ui"] = lambda v: "(value hidden)"
+                raw_contexts["crypto"] = disabled_ciphersuite
+            else:
+                raw_contexts["__hide_from_ui"] = lambda v: v
+                raw_contexts["crypto"] = server_cipher_container
+
+            raw_contexts["config"] = config
+
+            result = request.template.generate(
+                discovery_request=request,
+                host_header=request.desired_controlplane,
+                resource_names=request.resources,
+                utils=templates,
+                **raw_contexts,
+            )
+
+            if not request.template.is_python_source:
+                assert isinstance(result, str)
+                result = deserialize_config(result)
+
+            assert isinstance(result, dict)
+            resources = filter_resources(result["resources"], request.resources)
+            add_type_urls(request.api_version, request.resource_type, resources)
+            processed_template = ProcessedTemplate(resources=resources)
+            response = DiscoveryResponse(
+                resources=resources, version_info=processed_template.version_info
+            )
+
+            if not discovery_entry_repository.save(
+                DiscoveryEntry(
+                    request_hash=request_hash,
+                    template=request.template.resource_type,
+                    request=request,
+                    response=response,
+                    last_rendered_at=int(time.time()),
+                )
+            ):
+                logger.error("Failed to save discovery entry")
+    finally:
+        logger.debug("Finished rendering of discovery response")