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

sovereign/utils/eds.py

@@ -1,9 +1,11 @@
 import random
-from typing import Dict, Any, Optional, List
 from copy import deepcopy
+from typing import Any, Dict, List, Optional
+
 from starlette.exceptions import HTTPException
-from sovereign import config
-from sovereign.schemas import DiscoveryRequest
+
+from sovereign.configuration import config
+from sovereign.types import DiscoveryRequest
 from sovereign.utils.templates import resolve
 
 HARD_FAIL_ON_DNS_FAILURE = config.legacy_fields.dns_hard_fail
@@ -26,12 +28,16 @@ def _upstream_kwargs(
         if hard_fail:
             raise
         ip_addresses = [upstream["address"]]
-    return {
+    ret = {
         "addrs": ip_addresses,
         "port": upstream["port"],
         "region": default_region or upstream.get("region", "unknown"),
         "zone": proxy_region,
     }
+    if "health_check_config" in upstream:
+        ret["health_check_config"] = upstream["health_check_config"]
+
+    return ret
 
 
 def total_zones(endpoints: List[Dict[str, Dict[str, Any]]]) -> int:
@@ -84,14 +90,21 @@ def locality_lb_endpoints(
     return ret
 
 
-def lb_endpoints(addrs: List[str], port: int, region: str, zone: str) -> Dict[str, Any]:
+def lb_endpoints(
+    addrs: List[str],
+    port: int,
+    region: str,
+    zone: str,
+    health_check_config: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
     """
     Creates an envoy endpoint.LbEndpoints proto
 
-    :param addrs:  The IP addresses or hostname(s) of the upstream.
-    :param port:   The port that the upstream should be accessed on.
-    :param region: The region of the upstream.
-    :param zone:   The region of the proxy asking for the endpoint configuration.
+    :param addrs:                   The IP addresses or hostname(s) of the upstream.
+    :param port:                    The port that the upstream should be accessed on.
+    :param region:                  The region of the upstream.
+    :param zone:                    The region of the proxy asking for the endpoint configuration.
+    :param health_check_config:     Optional health check config for the upstream.
     """
     if PRIORITY_MAPPING is None:
         raise RuntimeError(
@@ -100,20 +113,25 @@ def lb_endpoints(addrs: List[str], port: int, region: str, zone: str) -> Dict[st
         )
     node_priorities = PRIORITY_MAPPING.get(zone, {})
     priority = node_priorities.get(region, 10)
+
+    endpoints = []
+    for addr in addrs:
+        endpoint = {
+            "endpoint": {
+                "address": {
+                    "socket_address": {
+                        "address": addr,
+                        "port_value": port,
+                    }
+                },
+            }
+        }
+        if health_check_config:
+            endpoint["endpoint"]["health_check_config"] = health_check_config
+        endpoints.append(endpoint)
+
     return {
         "priority": priority,
         "locality": {"zone": region},
-        "lb_endpoints": [
-            {
-                "endpoint": {
-                    "address": {
-                        "socket_address": {
-                            "address": addr,
-                            "port_value": port,
-                        }
-                    }
-                }
-            }
-            for addr in addrs
-        ],
+        "lb_endpoints": endpoints,
     }

sovereign/utils/entry_point_loader.py

@@ -0,0 +1,18 @@
+from functools import cached_property
+from importlib.metadata import EntryPoints, entry_points
+from typing import Dict
+
+
+class EntryPointLoader:
+    ENTRY_POINTS = entry_points()
+
+    def __init__(self, *args: str) -> None:
+        self._groups = args
+
+    @classmethod
+    def _select(cls, group: str) -> EntryPoints:
+        return cls.ENTRY_POINTS.select(group=f"sovereign.{group}")
+
+    @cached_property
+    def groups(self) -> Dict[str, EntryPoints]:
+        return {group: self._select(group) for group in self._groups}

sovereign/utils/mock.py

@@ -1,32 +1,75 @@
-from typing import Optional, Dict, List
+import ast
+import re
 from random import randint
-from sovereign.schemas import DiscoveryRequest, Node, Locality, Resources
+from typing import Dict, List, Optional
+
+from sovereign.types import DiscoveryRequest, Locality, Node, Status
+
+scrub = re.compile(r"[^a-zA-Z_\.]")
+
+
+class NodeExpressionError(Exception):
+    pass
 
 
 def mock_discovery_request(
-    service_cluster: Optional[str] = None,
-    resource_names: Optional[List[str]] = None,
-    region: str = "none",
-    version: str = "1.11.1",
+    api_version: Optional[str] = "V3",
+    resource_type: Optional[str] = None,
+    resource_names: Optional[List[str] | str] = None,
+    region: Optional[str] = "none",
+    version: Optional[str] = "<envoy_version>",
     metadata: Optional[Dict[str, str]] = None,
+    error_message: Optional[str] = None,
+    expressions: Optional[list[str]] = None,
 ) -> DiscoveryRequest:
     if resource_names is None:
-        resource_names = Resources()
-    else:
-        resource_names = Resources(resource_names)
+        resource_names = []
+    if isinstance(resource_names, str):
+        resource_names = [resource_names]
+    if expressions is None:
+        expressions = []
+    base_node = Node(
+        id="sovereign-interface",
+        cluster="*",
+        build_version=f"<randomHash>/{version}/Clean/RELEASE",
+        locality=Locality(zone=region),
+    ).model_dump()
+    set_node_expressions(base_node, expressions)
     request = DiscoveryRequest(
         type_url=None,
-        node=Node(
-            id="mock",
-            cluster=service_cluster or "",
-            build_version=f"e5f864a82d4f27110359daa2fbdcb12d99e415b9/{version}/Clean/RELEASE",
-            locality=Locality(zone=region, region="example", sub_zone="a"),
-        ),
+        node=Node.model_validate(base_node),
         version_info=str(randint(100000, 1000000000)),
         resource_names=resource_names,
-        hide_private_keys=True,
-        desired_controlplane="controlplane",
+        is_internal_request=True,
+        desired_controlplane="__sovereign__",
+        error_detail=Status(code=200, message="None", details=["None"]),
+        api_version=api_version,
+        resource_type=resource_type,
     )
     if isinstance(metadata, dict):
         request.node.metadata = metadata
+    if error_message:
+        request.error_detail = Status(code=666, message=error_message, details=["foo"])
     return request
+
+
+def set_node_expressions(node, expressions):
+    for expr in expressions:
+        try:
+            field, value = re.split(r"\s*=\s*", expr, maxsplit=1)
+            value = f'"{value}"'
+        except ValueError:
+            raise NodeExpressionError(f"Invalid node filter format: {expr}")
+
+        field = scrub.sub("", field)
+        parts = field.split(".")
+
+        try:
+            value = ast.literal_eval(value)
+        except Exception as e:
+            raise NodeExpressionError(f"Invalid node filter value: {value}") from e
+
+        current = node
+        for part in parts[:-1]:
+            current = current.setdefault(part, {})
+        current[parts[-1]] = value

sovereign/utils/resources.py

@@ -0,0 +1,17 @@
+import importlib.resources as res
+from functools import cache
+from importlib.resources.abc import Traversable
+
+
+@cache
+def get_package(name: str) -> Traversable:
+    return res.files(name)
+
+
+def get_package_file(package_name: str, filename: str) -> Traversable:
+    return get_package(package_name).joinpath(filename)
+
+
+def get_package_file_bytes(package_name: str, filename: str) -> bytes:
+    file = get_package_file(package_name, filename)
+    return file.read_bytes()

sovereign/utils/templates.py

@@ -1,7 +1,9 @@
-from typing import List, Optional, Any, Dict
-from socket import gethostbyname_ex
 from socket import gaierror as dns_error
+from socket import gethostbyname_ex
+from typing import Any, Dict, List, Optional
+
 from starlette.exceptions import HTTPException
+
 from sovereign import config, stats
 
 REGIONS = config.legacy_fields.regions

sovereign/utils/timer.py

@@ -6,9 +6,11 @@ from typing import Any, Callable, Coroutine, NoReturn
 from croniter import croniter
 
 
-def wait_until(dt: datetime) -> float:
-    now = datetime.now()
-    sleep_time = (dt - now).total_seconds()
+def wait_until(target_dt: datetime) -> float:
+    now = datetime.now(dt.timezone.utc)
+    if target_dt.tzinfo is None:
+        target_dt = target_dt.replace(tzinfo=dt.timezone.utc)
+    sleep_time = (target_dt - now).total_seconds()
     return sleep_time
 
 

sovereign/utils/version_info.py

@@ -11,3 +11,11 @@ def compute_hash(*args: Any) -> str:
         zlib.crc32(data) & 0xFFFFFFFF
     )  # same numeric value across all py versions & platforms
     return str(version_info)
+
+
+def compute_hash_int(*args: Any) -> int:
+    data: bytes = repr(args).encode()
+    version_info = (
+        zlib.crc32(data) & 0xFFFFFFFF
+    )  # same numeric value across all py versions & platforms
+    return version_info

sovereign/utils/weighted_clusters.py

@@ -1,5 +1,6 @@
 from __future__ import division
-from typing import List, Any, Dict, Generator
+
+from typing import Any, Dict, Generator, List
 
 
 def _round_to_n(sequence: List[int], n: int = 100) -> List[int]: