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

hyperion/infrastructure/secrets.py

@@ -0,0 +1,63 @@
+import abc
+import os
+import re
+from typing import cast
+
+import boto3
+
+from hyperion.config import secrets_config
+from hyperion.logging import get_logger
+
+logger = get_logger("hyperion-secrets")
+
+SECRET_PATTERN = re.compile(r"!#secret:#(?P<secret_name>.+)")
+
+
+class SecretsManager(abc.ABC):
+    _instance: "SecretsManager | None" = None
+
+    @staticmethod
+    def _create_new() -> "SecretsManager":
+        if secrets_config.backend is None:
+            logger.warning("No secrets backend is configured. Using dummy secrets manager.")
+            return DummySecretsManager()
+        if secrets_config.backend == "AWSSecretsManager":
+            logger.info("Using AWS Secrets Manager.")
+            return AWSSecretsManager()
+        raise ValueError(f"Unsupported secrets backend: {secrets_config.backend!r}.")
+
+    @staticmethod
+    def from_config() -> "SecretsManager":
+        if SecretsManager._instance is None:
+            SecretsManager._instance = SecretsManager._create_new()
+        return SecretsManager._instance
+
+    @staticmethod
+    def translate_env_vars() -> None:
+        """Loop through all env variables and replace secrets with their values.
+
+        Only variables with value pattern of `!#secret:#secret_name` will be replaced.
+        """
+        for key, value in os.environ.items():
+            if (match := SECRET_PATTERN.match(value)) is not None:
+                secret_name = match.group("secret_name")
+                logger.info("Replacing secret in environment variable.", key=key, secret_name=secret_name)
+                os.environ[key] = SecretsManager.from_config().get_secret(secret_name)
+
+    @abc.abstractmethod
+    def get_secret(self, secret_name: str) -> str:
+        """Get the secret with the given name."""
+
+
+class DummySecretsManager(SecretsManager):
+    def get_secret(self, secret_name: str) -> str:
+        logger.warning("Using dummy secrets manager, no values will be returned.", secret_name=secret_name)
+        return ""
+
+
+class AWSSecretsManager(SecretsManager):
+    def __init__(self) -> None:
+        self.client = boto3.client("secretsmanager")
+
+    def get_secret(self, secret_name: str) -> str:
+        return cast(str, self.client.get_secret_value(SecretId=secret_name)["SecretString"])

hyperion/logging.py

@@ -0,0 +1,122 @@
+import logging as _logging
+import sys
+import time
+import traceback
+from collections.abc import Callable
+from types import TracebackType
+from typing import Any, ParamSpec, TypeVar
+
+import loguru
+
+from hyperion.config import config
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+LOGURU_PRETTY_FORMAT = (
+    "<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | "
+    "<level>{level: <8}</level> | "
+    "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - "
+    "<level>{message}</level>\n<dim>{extra}</dim>"
+)
+
+
+def setup_logger_from_env() -> None:
+    loguru.logger.remove()
+    if config.log_pretty:
+        loguru.logger.add(sys.stderr, colorize=True, format=LOGURU_PRETTY_FORMAT, level=config.log_level)
+    else:
+        loguru.logger.add(sys.stderr, format="{message}", serialize=True, level=config.log_level)
+
+
+class InterceptHandler(_logging.Handler):
+    def emit(self, record: _logging.LogRecord) -> None:
+        log_level: str | int
+        try:
+            log_level = loguru.logger.level(record.levelname).name
+        except ValueError:
+            log_level = record.levelno
+        frame = _logging.currentframe()
+        depth = 2
+        while frame.f_code.co_filename == _logging.__file__:
+            if frame.f_back is None:
+                break
+            frame = frame.f_back
+            depth += 1
+        loguru.logger.opt(depth=depth, exception=record.exc_info).log(log_level, record.getMessage())
+
+
+def intercept_python_loggers() -> None:
+    _logging.root.handlers = []
+    _logging.root.setLevel(_logging.DEBUG)
+    _logging.root.addHandler(InterceptHandler())
+
+
+class LogTiming:
+    def __init__(self, action: str | None = None, logger: "loguru.Logger | None" = None) -> None:
+        self.action = action
+        self.start_time: float = 0
+        self.logger = logger or get_logger("timing")
+
+    def __enter__(self) -> None:
+        self.start_time = time.monotonic()
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None = None,
+        exc: BaseException | None = None,
+        tb: TracebackType | None = None,
+    ) -> None:
+        duration = time.monotonic() - self.start_time
+        action = self.action or "unnamed action"
+        if exc:
+            self.logger.warning(
+                f"Action {action!r} failed after {duration:0.3f} seconds.",
+                action=self.action,
+                duration=duration,
+                exception_type=exc_type.__name__ if exc_type else "unknown",
+                exception=str(exc),
+                traceback="\n".join(traceback.format_tb(tb)),
+            )
+            return
+        self.logger.info(
+            f"Action {action!r} finished after {duration:0.3f} seconds.", action=self.action, duration=duration
+        )
+
+    def __call__(self, func: Callable[P, T]) -> Callable[P, T]:
+        if self.action is None:
+            self.action = func.__name__
+
+        def _wrapped(*args: P.args, **kwargs: P.kwargs) -> T:
+            with self:
+                return func(*args, **kwargs)
+
+        return _wrapped
+
+
+setup_logger_from_env()
+intercept_python_loggers()
+
+
+def get_logger(service: str, **context: Any) -> "loguru.Logger":
+    """Get logger bound to a service with optional contextual variables."""
+    return loguru.logger.bind(service=service, **context)
+
+
+if __name__ == "__main__":
+    demo_logger = get_logger("demo")
+    demo_logger.trace("Trace message")
+    demo_logger.debug("Debug message")
+    demo_logger.info("Informational message")
+    demo_logger.warning("A warning!")
+    demo_logger.error("An error...")
+    demo_logger.critical("This is terrible.")
+    with LogTiming("This will take around 1 second.", demo_logger):
+        time.sleep(1)
+
+    try:
+        with LogTiming("This will fail after around 1 second.", demo_logger):
+            time.sleep(1)
+            raise ValueError("This failed.")
+    except Exception as error:
+        demo_logger.exception(error)

hyperion/sources/base.py

@@ -0,0 +1,105 @@
+"""Base abstract class for sources."""
+
+import abc
+import asyncio
+import datetime
+from collections.abc import AsyncIterator, Awaitable, Iterable
+from dataclasses import dataclass
+from typing import Any, ClassVar, cast
+
+from aws_lambda_typing.context import Context
+from aws_lambda_typing.events import EventBridgeEvent, SQSEvent
+
+from hyperion.asyncutils import AsyncTaskQueue, get_loop
+from hyperion.catalog import Catalog
+from hyperion.config import storage_config
+from hyperion.entities.catalog import DataLakeAsset
+from hyperion.infrastructure.queue import SourceBackfillMessage, SQSQueue, iter_messages_from_sqs_event
+from hyperion.logging import get_logger
+
+SourceEventType = EventBridgeEvent | SQSEvent
+
+logger = get_logger("hyperion-source")
+
+
+@dataclass(eq=True, frozen=True)
+class SourceAsset:
+    asset: DataLakeAsset
+    data: Iterable[dict[str, Any]]
+
+
+class Source(abc.ABC):
+    source: ClassVar[str] = NotImplemented
+
+    def __init__(self, catalog: Catalog) -> None:
+        if self.source is NotImplemented:
+            raise NotImplementedError("Cannot instantiate a source without a source name.")
+        self.catalog = catalog
+
+    @abc.abstractmethod
+    def run(
+        self,
+        start_date: datetime.datetime | datetime.date | None = None,
+        end_date: datetime.datetime | datetime.date | None = None,
+    ) -> Awaitable[Iterable[SourceAsset]] | AsyncIterator[SourceAsset]:
+        """The main coroutine that runs the source extraction."""
+
+    @classmethod
+    async def _run(
+        cls,
+        catalog: Catalog,
+        notify: bool = True,
+        start_date: datetime.datetime | datetime.date | None = None,
+        end_date: datetime.datetime | datetime.date | None = None,
+    ) -> None:
+        source = cls(catalog)
+        result = source.run(start_date=start_date, end_date=end_date)
+        async with AsyncTaskQueue[None](maxsize=storage_config.max_concurrency) as queue:
+            if isinstance(result, AsyncIterator):
+                async for asset in result:
+                    logger.info("Processing asset retrieved from source.", asset=asset.asset)
+                    await queue.add_task(source.catalog.store_asset_async(asset.asset, asset.data, notify=notify))
+            else:
+                for asset in await result:
+                    logger.info("Processing asset retrieved from source.", asset=asset.asset)
+                    await queue.add_task(source.catalog.store_asset_async(asset.asset, asset.data, notify=notify))
+
+    @classmethod
+    def handle_aws_lambda_event(
+        cls,
+        event: SourceEventType | None = None,
+        context: Context | None = None,
+        *,
+        loop: asyncio.AbstractEventLoop | None = None,
+    ) -> None:
+        logger.info("Starting Hyperion source.", source=cls.__name__, event=str(event), context=str(context))
+        catalog = Catalog.from_config()
+        loop = loop or get_loop()
+        queue = SQSQueue.from_config()
+        if isinstance(event, dict) and "Records" in event:
+            # We may presume this is an SQS Event
+            event = cast(SQSEvent, event)
+            for message in iter_messages_from_sqs_event(event):
+                if not isinstance(message, SourceBackfillMessage):
+                    logger.warning(
+                        "Only SourceBackfillMessage is supported, this message is probably not for us.", message=message
+                    )
+                    continue
+                if message.source != cls.source:
+                    logger.info("Message is not intended for us.", source=cls.source, message_source=message.source)
+                    continue
+                logger.info("Source triggered by an SQS Message.", source=cls.source, message=message)
+                loop.run_until_complete(
+                    cls._run(catalog, start_date=message.start_date, end_date=message.end_date, notify=message.notify)
+                )
+                if message.receipt_handle:
+                    queue.delete(message.receipt_handle)
+            return
+        if isinstance(event, dict) and "detail" in event:
+            # We may presume this is an EventBridgeEvent
+            event = cast(EventBridgeEvent, event)
+            logger.warning("EventBridge events can carry no config for now.")
+            loop.run_until_complete(cls._run(catalog, start_date=None, end_date=None))
+            return
+        logger.warning("No event was provided, assuming a no-config run.")
+        loop.run_until_complete(cls._run(catalog, start_date=None, end_date=None))

hyperion/typeutils.py

@@ -0,0 +1,52 @@
+import datetime
+from collections.abc import Sequence
+from dataclasses import asdict
+from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar
+
+if TYPE_CHECKING:
+    from _typeshed import DataclassInstance
+
+T = TypeVar("T")
+
+DateOrDelta: TypeAlias = datetime.datetime | datetime.timedelta
+
+
+def assert_type(variable: Any, assertion: type[T]) -> T:
+    if not isinstance(variable, assertion):
+        raise TypeError(f"Expected {assertion}, got {type(variable)}.")
+    return variable
+
+
+def dataclass_asdict(
+    dataclass: "DataclassInstance",
+    *,
+    exclude: Sequence[str] | None = None,
+    include: Sequence[str] | None = None,
+) -> dict[str, Any]:
+    """Convert a dataclass instance to a dictionary.
+
+    Args:
+        dataclass (DataclassInstance): Dataclass instance to convert.
+        exclude (Sequence[str], optional): Fields to exclude. Defaults to None.
+        include (Sequence[str], optional): Fields to include. Defaults to None.
+
+    Returns:
+        dict[str, Any]: Converted dictionary.
+
+    Raises:
+        ValueError: If include and exclude overlap.
+        ValueError: If include field is not found in dataclass.
+    """
+    exclude = exclude or []
+    include = include or []
+    dct = asdict(dataclass)
+    fields = list(dct.keys())
+    if set(include) & set(exclude):
+        raise ValueError("Include and exclude cannot overlap.")
+    if set(include) - set(fields):
+        raise ValueError("Include field not found in dataclass.")
+    if exclude or include:
+        for field in fields:
+            if (include and field not in include) or field in exclude:
+                del dct[field]
+    return dct