hyperion-sdk 0.2.0.dev1741815359__py3-none-any.whl

hyperion/infrastructure/geo/location.py

@@ -0,0 +1,186 @@
+"""Common location classes and functions."""
+
+import math
+from collections import defaultdict
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import ClassVar, Generic, TypeVar, cast
+
+import haversine
+import numpy as np
+import numpy.typing as npt
+from haversine.haversine import get_avg_earth_radius
+
+from hyperion.infrastructure.cache import Cache
+from hyperion.logging import get_logger
+
+LATITUDE_DEGREE_TO_METERS = 111_000
+EARTH_RADIUS_METERS = get_avg_earth_radius(haversine.Unit.METERS)
+
+logger = get_logger("hyperion-geo")
+
+AnyLocation = TypeVar("AnyLocation", bound="Location")
+
+
+def meters_to_degrees(meters: float, at_latitude: float) -> tuple[float, float]:
+    """Convert meters to degrees at a given latitude.
+
+    Args:
+        meters (float): The distance in meters.
+        at_latitude (float): The latitude at which the conversion should be done.
+
+    Returns:
+        tuple[float, float]: The distance in degrees for latitude and longitude.
+    """
+    return meters / LATITUDE_DEGREE_TO_METERS, meters / (
+        LATITUDE_DEGREE_TO_METERS * math.cos(math.radians(at_latitude))
+    )
+
+
+class SpatialKMeans(Generic[AnyLocation]):
+    """K-means clustering for geographical locations."""
+
+    def __init__(self, locations: Iterable[AnyLocation]) -> None:
+        """Initialize the K-means clustering with the given locations.
+
+        Args:
+            locations (Iterable[Location]): The locations to cluster.
+        """
+        self.locations = list(locations)
+        self.locations_array = np.array(
+            [[location.latitude, location.longitude] for location in self.locations],
+        )
+
+    def _get_distances_from_centroids(self, centroids: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+        return np.array(
+            [
+                [haversine.haversine((point[0], point[1]), (centroid[0], centroid[1])) for centroid in centroids]
+                for point in self.locations_array
+            ]
+        )
+
+    def fit(self, k: int, max_iters: int = 100) -> dict["Location", list[AnyLocation]]:
+        """Fit the K-means model with k clusters and return the clusters.
+
+        Args:
+            k (int): The desired number of clusters (must be less than the number of locations).
+            max_iters (int, optional): The maximum number of iterations. Defaults to 100.
+
+        Returns:
+            dict[Location, list[Location]]: The clusters with the centroids as keys.
+        """
+        centroids = self.locations_array[np.random.choice(len(self.locations_array), k, replace=False), :2]
+
+        for _ in range(max_iters):
+            distances = self._get_distances_from_centroids(centroids)
+
+            cluster_assignments = np.argmin(distances, axis=1)
+
+            new_centroids_list = []
+            for cluster_idx in range(k):
+                cluster_points = self.locations_array[cluster_assignments == cluster_idx, :2]
+                if len(cluster_points) > 0:
+                    new_centroids_list.append(cluster_points.mean(axis=0))
+                else:
+                    new_centroids_list.append(centroids[cluster_idx])  # Keep old centroid for empty clusters
+            new_centroids = np.array(new_centroids_list)
+
+            if np.allclose(centroids, new_centroids):
+                break
+            centroids = new_centroids
+
+        clusters: dict[Location, list[AnyLocation]] = defaultdict(list)
+        centroid_locations = [Location(float(lat), float(lon)) for lat, lon in centroids]
+
+        for location_id, centroid_id in enumerate(cluster_assignments):
+            clusters[centroid_locations[centroid_id]].append(self.locations[location_id])
+        return dict(clusters)
+
+
+@dataclass(frozen=True, eq=True)
+class Location:
+    """A geographical location."""
+
+    _cache: ClassVar[Cache | None] = None
+
+    latitude: float
+    longitude: float
+    title: str | None = None
+    address: str | None = None
+
+    @classmethod
+    def _get_cache(cls) -> Cache:
+        if cls._cache is None:
+            cls._cache = Cache.from_config()
+        return cls._cache
+
+    def _get_distance_haversine(self, other: "Location") -> float:
+        return cast(
+            float,
+            haversine.haversine(
+                (self.latitude, self.longitude), (other.latitude, other.longitude), haversine.Unit.METERS
+            ),
+        )
+
+    def _get_distance_euclidean(self, other: "Location") -> float:
+        lat_diff = math.radians(self.latitude - other.latitude)
+        lon_diff = math.radians(self.longitude - other.longitude)
+        x = lon_diff * math.cos(math.radians((self.latitude + other.latitude) / 2))
+        y = lat_diff
+        return cast(float, EARTH_RADIUS_METERS) * math.sqrt(x**2 + y**2)
+
+    def get_distance(self, other: "Location", approximate: bool = False) -> float:
+        """Get the distance to another location in meters.
+        If approximate is False (default), will use haversine. Otherwise uses Euclidean to approximate the distance.
+
+        Args:
+            other (Location): The other location.
+            approximate (bool, optional): Whether to approximate the distance. Defaults to False.
+
+        Returns:
+            float: The distance in meters.
+        """
+        cache_key = str((self.latitude, self.longitude, other.latitude, other.longitude, approximate))
+        cache = self._get_cache()
+        if cached := cache.get(cache_key):
+            return float(cached)
+        if approximate:
+            return self._get_distance_euclidean(other)
+        return self._get_distance_haversine(other)
+
+    def get_nearest(
+        self, others: Iterable[AnyLocation], threshold: float | None = None, approximate: bool = False
+    ) -> AnyLocation:
+        """Get the closest location from the iterable of locations.
+
+        If threshold is given and all locations are further than the threshold in meters, an ValueError is raised.
+
+        Args:
+            others (Iterable[Location]): The other locations.
+            threshold (float, optional): The maximum distance in meters. Defaults to None.
+            approximate (bool, optional): Whether to approximate the distance. Defaults to False.
+
+        Returns:
+            Location: The nearest location.
+        """
+        nearest: tuple[AnyLocation, float] | None = None
+        for other in others:
+            distance = self.get_distance(other, approximate=approximate)
+            if nearest is None or nearest[1] > distance:
+                nearest = (other, distance)
+        if nearest is None:
+            raise ValueError(f"None of the given locations is close enough to {self!r}.")
+        logger.debug("Found nearest location.", this=self, other=nearest[0], distance=nearest[1])
+        return nearest[0]
+
+
+@dataclass(frozen=True, eq=True)
+class NamedLocation:
+    location: Location
+    route: str | None = None
+    neighborhood: str | None = None
+    sublocality: str | None = None
+    administrative_area: str | None = None
+    administrative_area_level_2: str | None = None
+    country: str | None = None
+    address: str | None = None

hyperion/infrastructure/http.py

@@ -0,0 +1,62 @@
+import asyncio
+from types import TracebackType
+from urllib.parse import urlparse
+
+import httpx
+
+from hyperion.config import http_config
+from hyperion.logging import get_logger
+
+logger = get_logger("hyperion-http")
+
+
+def redact_url(url: str, replace: str = "***") -> str:
+    """Replace password from {url} (if any) with {replace}.
+    If the url contains no password, url is returned unchanged.
+    """
+    parsed = urlparse(url)
+    if parsed.password:
+        return parsed._replace(netloc=f"{parsed.username}:{replace}@{parsed.hostname}").geturl()
+    return url
+
+
+class AsyncHTTPClientWrapper:
+    def __init__(self) -> None:
+        self._client: httpx.AsyncClient | None = None
+        self._stacklevel = 0
+        self._lock = asyncio.Lock()
+
+    async def __aenter__(self) -> httpx.AsyncClient:
+        async with self._lock:
+            self._stacklevel += 1
+            if self._client is None or self._client.is_closed:
+                logger.debug("Creating and entering httpx async client.")
+                self._client = httpx.AsyncClient(mounts=self._get_proxy_mounts())
+                await self._client.__aenter__()
+        return self._client
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None = None,
+        exc_value: BaseException | None = None,
+        traceback: TracebackType | None = None,
+    ) -> None:
+        async with self._lock:
+            self._stacklevel -= 1
+            if self._stacklevel == 0 and self._client is not None:
+                logger.info("Closing httpx async client.")
+                await self._client.__aexit__(exc_type, exc_value, traceback)
+                self._client = None
+
+    @staticmethod
+    def _get_proxy_mounts() -> dict[str, httpx.AsyncHTTPTransport]:
+        proxy_mounts: dict[str, httpx.AsyncHTTPTransport] = {}
+        if http_config.proxy_http:
+            redacted_url = redact_url(http_config.proxy_http)
+            logger.info("Configuring HTTP proxy for http://", proxy_url=redacted_url)
+            proxy_mounts["http://"] = httpx.AsyncHTTPTransport(proxy=http_config.proxy_http)
+        if http_config.proxy_https:
+            redacted_url = redact_url(http_config.proxy_https)
+            logger.info("Configuring HTTP proxy for https://", proxy_url=redacted_url)
+            proxy_mounts["https://"] = httpx.AsyncHTTPTransport(proxy=http_config.proxy_https)
+        return proxy_mounts

hyperion/infrastructure/keyval.py

@@ -0,0 +1,264 @@
+"""Key-Value stores.
+
+If you need a store with a TTL, look at Cache instead.
+"""
+
+import gzip
+from abc import ABC, abstractmethod
+from collections.abc import Iterable, Iterator
+from fnmatch import fnmatch
+from typing import Literal, TypeGuard, cast
+
+import boto3
+import snappy
+
+CompressionType = Literal["snappy", "gzip"]
+
+
+def is_valid_compression_type(compression_type: str) -> TypeGuard[CompressionType]:
+    return compression_type in ("snappy", "gzip")
+
+
+class KeyValueStore(ABC, Iterable[str]):
+    """A key-value store."""
+
+    def __init__(self, prefix: str | None = None, compression: CompressionType | None = None) -> None:
+        """Initialize the store.
+
+        Args:
+            prefix: A prefix to add to all keys.
+            compression: The compression algorithm to use.
+        """
+        self.prefix = prefix
+        self.compression = compression
+
+    def _key(self, key: str) -> str:
+        """Return the key with the prefix added."""
+        prefix = f"{self.prefix}:" if self.prefix else ""
+        return f"{prefix}{key}"
+
+    def _compress(self, value: str) -> bytes:
+        """Compress a value using configured compression.
+
+        If no compression is selected, the string is encoded as utf-8 into bytes.
+
+        Args:
+            value: The value to compress.
+
+        Returns:
+            bytes: The compressed value.
+
+        Raises:
+            ValueError: If an unsupported compression is selected.
+        """
+        value_bytes = value.encode("utf-8")
+        match self.compression:
+            case None:
+                return value_bytes
+            case "snappy":
+                return cast(bytes, snappy.compress(value_bytes))
+            case "gzip":
+                return gzip.compress(value_bytes)
+            case _:
+                raise ValueError(f"Unsupported compression {self.compression!r}.")
+
+    def _decompress(self, value: bytes) -> str:
+        """Decompress a value using configured compression.
+
+        If no compression is selected, the bytes are decoded as utf-8 into a string.
+
+        Args:
+            value: The value to decompress.
+
+        Returns:
+            str: The decompressed value.
+
+        Raises:
+            ValueError: If an unsupported compression is selected.
+            TypeError: If the decompression returns an unexpected type.
+        """
+        match self.compression:
+            case None:
+                return value.decode("utf-8")
+            case "snappy":
+                value_decompressed = snappy.decompress(value)
+                if isinstance(value_decompressed, str):
+                    return value_decompressed
+                elif isinstance(value_decompressed, bytes):
+                    return value_decompressed.decode("utf-8")
+                raise TypeError(
+                    "Unexpected value returned from snappy decompression, "
+                    f"expected str | bytes, got {type(value_decompressed)!r}."
+                )
+            case "gzip":
+                return gzip.decompress(value).decode("utf-8")
+            case _:
+                raise ValueError(f"Unsupported compression {self.compression!r}.")
+
+    @abstractmethod
+    def _get_raw(self, hashed_key: str) -> bytes | None:
+        """Get a raw uncompressed value from the store.
+
+        The key passed here must already be hashed and prefixed.
+
+        Args:
+            hashed_key: The key to get.
+
+        Returns:
+            bytes | None: The raw value, or None if not found.
+        """
+
+    def get(self, key: str) -> str | None:
+        """Gets a value from the store."""
+        if value := self._get_raw(self._key(key)):
+            return self._decompress(value)
+        return None
+
+    @abstractmethod
+    def _set_raw(self, hashed_key: str, compresed_value: bytes) -> None:
+        """Set a raw, compressed bytes value in the store.
+
+        The key passed here must already by hashed and prefixed.
+
+        Args:
+            hashed_key: The key to set.
+            compresed_value: The compressed value to set.
+        """
+
+    def set(self, key: str, value: str) -> None:
+        """Sets a value in the store."""
+        return self._set_raw(self._key(key), self._compress(value))
+
+    @abstractmethod
+    def _delete_raw(self, hashed_key: str) -> None:
+        """Delete a value from the store.
+
+        The key passed here must already by hashed and prefixed.
+
+        Args:
+            hashed_key: The key to delete.
+        """
+
+    def __iter__(self) -> Iterator[str]:
+        """Iterate all keys available in the store."""
+        return iter(self._iter_all_keys())
+
+    @abstractmethod
+    def _iter_all_keys(self) -> Iterable[str]:
+        """Iterate all keys available in the store."""
+
+    def keys(self, match: str = "*") -> Iterable[str]:
+        """Iterate all keys available in the store that match the given expression.
+
+        Nothing fancy is supported, UNIX-style fnmatch is performed.
+
+        Args:
+            match: The expression to match.
+
+        Returns:
+            Iterable[str]: The keys that match the expression
+        """
+        for key in self._iter_all_keys():
+            if fnmatch(key, match):
+                yield key
+
+    def delete(self, key: str) -> None:
+        """Delete a value from the store.
+
+        Args:
+            key: The key to delete.
+        """
+        return self._delete_raw(self._key(key))
+
+    def exists(self, key: str) -> bool:
+        """Returns whether a key exists in the store.
+
+        Args:
+            key: The key to check.
+        """
+        return key in self.keys()
+
+    def __contains__(self, key: str) -> bool:
+        return self.exists(key)
+
+    def __getitem__(self, key: str) -> str | None:
+        return self.get(key)
+
+    def __setitem__(self, key: str, value: str) -> None:
+        return self.set(key, value)
+
+    def __delitem__(self, key: str) -> None:
+        return self.delete(key)
+
+
+class DynamoDBStore(KeyValueStore):
+    """A key-value store using DynamoDB.
+
+    The table must have a key attribute and a value attribute.
+    """
+
+    def __init__(
+        self,
+        prefix: str | None = None,
+        compression: CompressionType | None = None,
+        table_name: str | None = None,
+        key_attribute: str = "key",
+        value_attribute: str = "value",
+    ):
+        super().__init__(prefix, compression)
+        self.client = boto3.resource("dynamodb")
+        if table_name is None:
+            raise ValueError("No table name provided for DynamoDBStore.")
+        self.table_name = table_name
+        self.table = self.client.Table(self.table_name)
+        self.key_attribute = key_attribute
+        self.value_attribute = value_attribute
+
+    def _get_raw(self, hashed_key: str) -> bytes | None:
+        response = self.table.get_item(Key={self.key_attribute: hashed_key})
+        item = response.get("Item")
+        if not item:
+            return None
+        return bytes(item[self.value_attribute])
+
+    def _set_raw(self, hashed_key: str, compresed_value: bytes) -> None:
+        self.table.put_item(Item={self.key_attribute: hashed_key, self.value_attribute: compresed_value})
+
+    def _delete_raw(self, hashed_key: str) -> None:
+        self.table.delete_item(Key={self.key_attribute: hashed_key})
+
+    def _iter_all_keys(self) -> Iterable[str]:
+        scan_kwargs = {"ProjectionExpression": "#k", "ExpressionAttributeNames": {"#k": self.key_attribute}}
+        while True:
+            response = self.table.scan(**scan_kwargs)
+            for item in response.get("Items", []):
+                key = str(item[self.key_attribute])
+                if self.prefix:
+                    key = key.replace(f"{self.prefix}:", "", 1)
+                yield key
+            if "LastEvaluatedKey" not in response:
+                break
+            scan_kwargs["ExclusiveStartKey"] = response["LastEvaluatedKey"]
+
+
+class InMemoryStore(KeyValueStore):
+    """An in-memory key-value store.
+
+    This store is not persistent and is not shared between instances.
+    """
+
+    def __init__(self, prefix: str | None = None, compression: CompressionType | None = None):
+        super().__init__(prefix, compression)
+        self.store: dict[str, bytes] = {}
+
+    def _get_raw(self, hashed_key: str) -> bytes | None:
+        return self.store.get(hashed_key)
+
+    def _set_raw(self, hashed_key: str, compresed_value: bytes) -> None:
+        self.store[hashed_key] = compresed_value
+
+    def _delete_raw(self, hashed_key: str) -> None:
+        self.store.pop(hashed_key, None)
+
+    def _iter_all_keys(self) -> Iterable[str]:
+        yield from self.store

hyperion/infrastructure/queue.py

@@ -0,0 +1,151 @@
+import abc
+import datetime
+from collections.abc import Iterator
+from enum import Enum
+from typing import ClassVar
+
+from aws_lambda_typing.events import SQSEvent
+from boto3 import client
+
+# After a long consideration, I decided to use pydantic to serialize / deserialize the messages.
+# The convenience outweighs the performance hit.
+from pydantic import BaseModel, Field
+
+from hyperion.config import config, queue_config
+from hyperion.dateutils import utcnow
+from hyperion.entities.catalog import DataLakeAsset
+from hyperion.logging import get_logger
+
+logger = get_logger("hyperion-queue")
+
+
+class ArrivalEvent(str, Enum):
+    ARRIVED = "arrived"
+
+
+class Message(BaseModel):
+    _subclasses: ClassVar[dict[str, type["Message"]]] = {}
+
+    def __init_subclass__(cls: type["Message"]) -> None:
+        cls._subclasses[cls.__name__] = cls
+
+    @staticmethod
+    def deserialize(json_str: str, message_type: str, receipt_handle: str | None = None) -> "Message":
+        msg = Message._subclasses[message_type].model_validate_json(json_str)
+        if receipt_handle and not msg.receipt_handle:
+            msg.receipt_handle = receipt_handle
+        return msg
+
+    created: datetime.datetime = Field(default_factory=utcnow)
+    sender: str = config.service_name
+    receipt_handle: str | None = None
+
+
+class DataLakeArrivalMessage(Message):
+    asset: DataLakeAsset
+    event: ArrivalEvent
+
+
+class SourceBackfillMessage(Message):
+    source: str
+    start_date: datetime.datetime | datetime.date
+    end_date: datetime.datetime | datetime.date | None = None
+    notify: bool = True
+
+
+def iter_messages_from_sqs_event(event: SQSEvent) -> Iterator[Message]:
+    if not isinstance(event, dict) or "Records" not in event or not isinstance(event["Records"], list):
+        raise ValueError("Provided event is not a valid SQS Event.")
+    for record in event["Records"]:
+        logger.info("Deserializing message.", message_id=record["messageId"])
+        if "MessageType" not in record["messageAttributes"]:
+            logger.warning("Message has no type and probably does not come from Hyperion.", message=record)
+            continue
+        message_type = record["messageAttributes"]["MessageType"]["stringValue"]
+        logger.debug("Attempting to deserialize message.", message=record, type=message_type)
+        yield Message.deserialize(record["body"], message_type, record["receiptHandle"])
+
+
+def create_backfill_event(
+    message: SourceBackfillMessage, message_id: str = "test", receipt_handle: str = ""
+) -> SQSEvent:
+    return {
+        "Records": [
+            {
+                "body": message.model_dump_json(),
+                "messageId": message_id,
+                "receiptHandle": receipt_handle,
+                "messageAttributes": {
+                    "MessageType": {
+                        "binaryListValues": [],
+                        "binaryValue": None,
+                        "dataType": "String",
+                        "stringListValues": [],
+                        "stringValue": "SourceBackfillMessage",
+                    }
+                },
+            }
+        ]
+    }
+
+
+class Queue(abc.ABC):
+    @staticmethod
+    def from_config() -> "Queue":
+        if queue_config.url is None:
+            logger.info("Using in-memory queue.")
+            return InMemoryQueue()
+        logger.info("Using SQS queue.")
+        return SQSQueue(queue_config.url)
+
+    @abc.abstractmethod
+    def send(self, message: Message) -> None:
+        """Send a message to the queue."""
+
+    @abc.abstractmethod
+    def delete(self, receipt_handle: str) -> None:
+        """Delete a message from the queue."""
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}>"
+
+
+class InMemoryQueue(Queue):
+    def __init__(self) -> None:
+        super().__init__()
+        self._messages: list[Message] = []
+
+    def send(self, message: Message) -> None:
+        self._messages.append(message)
+
+    def delete(self, receipt_handle: str) -> None:
+        """Delete the message using the created as isoformat."""
+        for message in self._messages:
+            if message.created.isoformat() == receipt_handle:
+                self._messages.remove(message)
+                break
+
+
+class SQSQueue(Queue):
+    def __init__(self, queue_url: str) -> None:
+        super().__init__()
+        self._queue_url = queue_url
+        self._client = client("sqs")
+
+    def send(self, message: Message) -> None:
+        self._client.send_message(
+            QueueUrl=self._queue_url,
+            MessageBody=message.model_dump_json(),
+            MessageAttributes={
+                "MessageType": {
+                    "DataType": "String",
+                    "StringValue": message.__class__.__name__,
+                }
+            },
+        )
+
+    def delete(self, receipt_handle: str) -> None:
+        self._client.delete_message(QueueUrl=self._queue_url, ReceiptHandle=receipt_handle)
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__} {self._queue_url}>"