moose-lib 0.4.218__py3-none-any.whl → 0.4.220__py3-none-any.whl

moose_lib/dmv2/registry.py

@@ -0,0 +1,62 @@
+"""
+Global registries for Moose Data Model v2 (dmv2) resources.
+
+This module provides functions to access the registered resources.
+The actual registry dictionaries are maintained in _registry.py to avoid circular dependencies.
+"""
+from typing import Optional, Dict
+from .olap_table import OlapTable
+from .stream import Stream
+from .ingest_api import IngestApi
+from .consumption import ConsumptionApi
+from .sql_resource import SqlResource
+from .workflow import Workflow
+from ._registry import _tables, _streams, _ingest_apis, _egress_apis, _sql_resources, _workflows
+
+def get_tables() -> Dict[str, OlapTable]:
+    """Get all registered OLAP tables."""
+    return _tables
+
+def get_table(name: str) -> Optional[OlapTable]:
+    """Get a registered OLAP table by name."""
+    return _tables.get(name)
+
+def get_streams() -> Dict[str, Stream]:
+    """Get all registered streams."""
+    return _streams
+
+def get_stream(name: str) -> Optional[Stream]:
+    """Get a registered stream by name."""
+    return _streams.get(name)
+
+def get_ingest_apis() -> Dict[str, IngestApi]:
+    """Get all registered ingestion APIs."""
+    return _ingest_apis
+
+def get_ingest_api(name: str) -> Optional[IngestApi]:
+    """Get a registered ingestion API by name."""
+    return _ingest_apis.get(name)
+
+def get_consumption_apis() -> Dict[str, ConsumptionApi]:
+    """Get all registered consumption APIs."""
+    return _egress_apis
+
+def get_consumption_api(name: str) -> Optional[ConsumptionApi]:
+    """Get a registered consumption API by name."""
+    return _egress_apis.get(name)
+
+def get_sql_resources() -> Dict[str, SqlResource]:
+    """Get all registered SQL resources."""
+    return _sql_resources
+
+def get_sql_resource(name: str) -> Optional[SqlResource]:
+    """Get a registered SQL resource by name."""
+    return _sql_resources.get(name)
+
+def get_workflows() -> Dict[str, Workflow]:
+    """Get all registered workflows."""
+    return _workflows
+
+def get_workflow(name: str) -> Optional[Workflow]:
+    """Get a registered workflow by name."""
+    return _workflows.get(name)

moose_lib/dmv2/sql_resource.py

@@ -0,0 +1,49 @@
+"""
+Base SQL resource definitions for Moose Data Model v2 (dmv2).
+
+This module provides the base class for SQL resources like Views and Materialized Views,
+handling common functionality like setup/teardown SQL commands and dependency tracking.
+"""
+from typing import Any, Optional, Union, List
+from pydantic import BaseModel
+
+from .olap_table import OlapTable
+from ._registry import _sql_resources
+
+class SqlResource:
+    """Base class for SQL resources like Views and Materialized Views.
+
+    Handles the definition of setup (CREATE) and teardown (DROP) SQL commands
+    and tracks data dependencies.
+
+    Attributes:
+        name (str): The name of the SQL resource (e.g., view name).
+        setup (list[str]): SQL commands to create the resource.
+        teardown (list[str]): SQL commands to drop the resource.
+        pulls_data_from (list[SqlObject]): List of tables/views this resource reads from.
+        pushes_data_to (list[SqlObject]): List of tables/views this resource writes to.
+        kind: The kind of the SQL resource (e.g., "SqlResource").
+    """
+    setup: list[str]
+    teardown: list[str]
+    name: str
+    kind: str = "SqlResource"
+    pulls_data_from: list[Union[OlapTable, "SqlResource"]]
+    pushes_data_to: list[Union[OlapTable, "SqlResource"]]
+
+    def __init__(
+            self,
+            name: str,
+            setup: list[str],
+            teardown: list[str],
+            pulls_data_from: Optional[list[Union[OlapTable, "SqlResource"]]] = None,
+            pushes_data_to: Optional[list[Union[OlapTable, "SqlResource"]]] = None,
+            metadata: dict = None
+    ):
+        self.name = name
+        self.setup = setup
+        self.teardown = teardown
+        self.pulls_data_from = pulls_data_from or []
+        self.pushes_data_to = pushes_data_to or []
+        self.metadata = metadata
+        _sql_resources[name] = self

moose_lib/dmv2/stream.py

@@ -0,0 +1,258 @@
+"""
+Stream definitions for Moose Data Model v2 (dmv2).
+
+This module provides classes for defining and configuring data streams,
+including stream transformations, consumers, and dead letter queues.
+"""
+import dataclasses
+import datetime
+from typing import Any, Optional, Callable, Union, Literal, Generic
+from pydantic import BaseModel, ConfigDict, AliasGenerator
+from pydantic.alias_generators import to_camel
+
+from .types import TypedMooseResource, ZeroOrMany, T, U
+from .olap_table import OlapTable
+from ._registry import _streams
+
+class StreamConfig(BaseModel):
+    """Configuration for data streams (e.g., Redpanda topics).
+
+    Attributes:
+        parallelism: Number of partitions for the stream.
+        retention_period: Data retention period in seconds (default: 7 days).
+        destination: Optional `OlapTable` where stream messages should be automatically ingested.
+        version: Optional version string for tracking configuration changes.
+        metadata: Optional metadata for the stream.
+    """
+    parallelism: int = 1
+    retention_period: int = 60 * 60 * 24 * 7  # 7 days
+    destination: Optional[OlapTable[Any]] = None
+    version: Optional[str] = None
+    metadata: Optional[dict] = None
+
+class TransformConfig(BaseModel):
+    """Configuration for stream transformations.
+
+    Attributes:
+        version: Optional version string to identify a specific transformation.
+                 Allows multiple transformations to the same destination if versions differ.
+    """
+    version: Optional[str] = None
+    dead_letter_queue: "Optional[DeadLetterQueue]" = None
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    metadata: Optional[dict] = None
+
+class ConsumerConfig(BaseModel):
+    """Configuration for stream consumers.
+
+    Attributes:
+        version: Optional version string to identify a specific consumer.
+                 Allows multiple consumers if versions differ.
+    """
+    version: Optional[str] = None
+    dead_letter_queue: "Optional[DeadLetterQueue]" = None
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+@dataclasses.dataclass
+class _RoutedMessage:
+    """Internal class representing a message routed to a specific stream."""
+    destination: "Stream[Any]"
+    values: ZeroOrMany[Any]
+
+@dataclasses.dataclass
+class ConsumerEntry(Generic[T]):
+    """Internal class representing a consumer with its configuration."""
+    consumer: Callable[[T], None]
+    config: ConsumerConfig
+
+@dataclasses.dataclass
+class TransformEntry(Generic[T]):
+    """Internal class representing a transformation with its configuration."""
+    destination: "Stream[Any]"
+    transformation: Callable[[T], ZeroOrMany[Any]]
+    config: TransformConfig
+
+class Stream(TypedMooseResource, Generic[T]):
+    """Represents a data stream (e.g., a Redpanda topic) typed with a Pydantic model.
+
+    Allows defining transformations to other streams and adding consumers.
+
+    Args:
+        name: The name of the stream.
+        config: Configuration options for the stream (parallelism, retention, destination).
+        t: The Pydantic model defining the stream message schema (passed via `Stream[MyModel](...)`).
+
+    Attributes:
+        config (StreamConfig): Configuration settings for this stream.
+        transformations (dict[str, list[TransformEntry[T]]]): Dictionary mapping destination stream names
+                                                            to lists of transformation functions.
+        consumers (list[ConsumerEntry[T]]): List of consumers attached to this stream.
+        columns (Columns[T]): Helper for accessing message field names safely.
+        name (str): The name of the stream.
+        model_type (type[T]): The Pydantic model associated with this stream.
+    """
+    config: StreamConfig
+    transformations: dict[str, list[TransformEntry[T]]]
+    consumers: list[ConsumerEntry[T]]
+    _multipleTransformations: Optional[Callable[[T], list[_RoutedMessage]]] = None
+
+    def __init__(self, name: str, config: StreamConfig = StreamConfig(), **kwargs):
+        super().__init__()
+        self._set_type(name, self._get_type(kwargs))
+        self.config = config
+        self.metadata = config.metadata
+        self.consumers = []
+        self.transformations = {}
+        _streams[name] = self
+
+    def add_transform(self, destination: "Stream[U]", transformation: Callable[[T], ZeroOrMany[U]],
+                      config: TransformConfig = None):
+        """Adds a transformation step from this stream to a destination stream.
+
+        The transformation function receives a record of type `T` and should return
+        a record of type `U`, a list of `U` records, or `None` to filter.
+
+        Args:
+            destination: The target `Stream` for the transformed records.
+            transformation: A callable that performs the transformation.
+            config: Optional configuration, primarily for setting a version.
+        """
+        config = config or TransformConfig()
+        if destination.name in self.transformations:
+            existing_transforms = self.transformations[destination.name]
+            # Check if a transform with this version already exists
+            has_version = any(t.config.version == config.version for t in existing_transforms)
+            if not has_version:
+                existing_transforms.append(
+                    TransformEntry(destination=destination, transformation=transformation, config=config))
+        else:
+            self.transformations[destination.name] = [
+                TransformEntry(destination=destination, transformation=transformation, config=config)]
+
+    def add_consumer(self, consumer: Callable[[T], None], config: ConsumerConfig = None):
+        """Adds a consumer function to be executed for each record in the stream.
+
+        Consumers are typically used for side effects like logging or triggering external actions.
+
+        Args:
+            consumer: A callable that accepts a record of type `T`.
+            config: Optional configuration, primarily for setting a version.
+        """
+        config = config or ConsumerConfig()
+        has_version = any(c.config.version == config.version for c in self.consumers)
+        if not has_version:
+            self.consumers.append(ConsumerEntry(consumer=consumer, config=config))
+
+    def has_consumers(self) -> bool:
+        """Checks if any consumers have been added to this stream.
+
+        Returns:
+            True if the stream has one or more consumers, False otherwise.
+        """
+        return len(self.consumers) > 0
+
+    def routed(self, values: ZeroOrMany[T]) -> _RoutedMessage:
+        """Creates a `_RoutedMessage` for use in multi-transform functions.
+
+        Wraps the value(s) to be sent with this stream as the destination.
+
+        Args:
+            values: A single record, a list of records, or None.
+
+        Returns:
+            A `_RoutedMessage` object.
+        """
+        return _RoutedMessage(destination=self, values=values)
+
+    def set_multi_transform(self, transformation: Callable[[T], list[_RoutedMessage]]):
+        """Sets a transformation function capable of routing records to multiple streams.
+
+        The provided function takes a single input record (`T`) and must return a list
+        of `_RoutedMessage` objects, created using the `.routed()` method of the
+        target streams.
+
+        Example:
+            def my_multi_transform(record: InputModel) -> list[_RoutedMessage]:
+                output1 = transform_for_stream1(record)
+                output2 = transform_for_stream2(record)
+                return [
+                    stream1.routed(output1),
+                    stream2.routed(output2)
+                ]
+            input_stream.set_multi_transform(my_multi_transform)
+
+        Note: Only one multi-transform function can be set per stream.
+
+        Args:
+            transformation: The multi-routing transformation function.
+        """
+        self._multipleTransformations = transformation
+
+class DeadLetterModel(BaseModel, Generic[T]):
+    """Model for dead letter queue messages.
+
+    Attributes:
+        original_record: The original record that failed processing.
+        error_message: Description of the error that occurred.
+        error_type: Type of error (e.g., "ValidationError").
+        failed_at: Timestamp when the error occurred.
+        source: Source of the error ("api", "transform", or "table").
+    """
+    model_config = ConfigDict(alias_generator=AliasGenerator(
+        serialization_alias=to_camel,
+    ))
+    original_record: Any
+    error_message: str
+    error_type: str
+    failed_at: datetime.datetime
+    source: Literal["api", "transform", "table"]
+
+    def as_typed(self) -> T:
+        return self._t.model_validate(self.original_record)
+
+class DeadLetterQueue(Stream, Generic[T]):
+    """A specialized Stream for handling failed records.
+
+    Dead letter queues store records that failed during processing, along with
+    error information to help diagnose and potentially recover from failures.
+
+    Attributes:
+        All attributes inherited from Stream.
+    """
+
+    _model_type: type[T]
+
+    def __init__(self, name: str, config: StreamConfig = StreamConfig(), **kwargs):
+        """Initialize a new DeadLetterQueue.
+
+        Args:
+            name: The name of the dead letter queue stream.
+            config: Configuration for the stream.
+        """
+        self._model_type = self._get_type(kwargs)
+        kwargs["t"] = DeadLetterModel[self._model_type]
+        super().__init__(name, config, **kwargs)
+
+    def add_transform(self, destination: Stream[U], transformation: Callable[[DeadLetterModel[T]], ZeroOrMany[U]],
+                      config: TransformConfig = None):
+        def wrapped_transform(record: DeadLetterModel[T]):
+            record._t = self._model_type
+            return transformation(record)
+
+        config = config or TransformConfig()
+        super().add_transform(destination, wrapped_transform, config)
+
+    def add_consumer(self, consumer: Callable[[DeadLetterModel[T]], None], config: ConsumerConfig = None):
+        def wrapped_consumer(record: DeadLetterModel[T]):
+            record._t = self._model_type
+            return consumer(record)
+
+        config = config or ConsumerConfig()
+        super().add_consumer(wrapped_consumer, config)
+
+    def set_multi_transform(self, transformation: Callable[[DeadLetterModel[T]], list[_RoutedMessage]]):
+        def wrapped_transform(record: DeadLetterModel[T]):
+            record._t = self._model_type
+            return transformation(record)
+
+        super().set_multi_transform(wrapped_transform)

moose_lib/dmv2/types.py

@@ -0,0 +1,95 @@
+"""
+Shared types and base classes for Moose Data Model v2 (dmv2).
+
+This module provides the core type definitions and base classes used across
+the dmv2 package, including generic type parameters, type aliases, and base
+resource classes.
+"""
+from typing import Any, Generic, TypeVar, Union
+from pydantic import BaseModel
+from pydantic.fields import FieldInfo
+
+T = TypeVar('T', bound=BaseModel)
+U = TypeVar('U', bound=BaseModel)
+T_none = TypeVar('T_none', bound=Union[BaseModel, None])
+U_none = TypeVar('U_none', bound=Union[BaseModel, None])
+type ZeroOrMany[T] = Union[T, list[T], None]
+
+class Columns(Generic[T]):
+    """Provides runtime checked column name access for Moose resources.
+
+    Instead of using string literals for column names, you can use attribute access
+    on this object, which will verify the name against the Pydantic model's fields.
+
+    Example:
+        >>> class MyModel(BaseModel):
+        ...     user_id: int
+        ...     event_name: str
+        >>> cols = Columns(MyModel)
+        >>> print(cols.user_id)  # Output: user_id
+        >>> print(cols.non_existent) # Raises AttributeError
+
+    Args:
+        model: The Pydantic model type whose fields represent the columns.
+    """
+    _fields: dict[str, FieldInfo]
+
+    def __init__(self, model: type[T]):
+        self._fields = model.model_fields
+
+    def __getattr__(self, item: str) -> str:
+        if item in self._fields:
+            return item  # or some Column representation
+        raise AttributeError(f"{item} is not a valid column name")
+
+class BaseTypedResource(Generic[T]):
+    """Base class for Moose resources that are typed with a Pydantic model.
+
+    Handles the association of a Pydantic model `T` with a Moose resource,
+    providing type validation and access to the model type.
+
+    Attributes:
+        name (str): The name of the Moose resource.
+    """
+    _t: type[T]
+    name: str
+
+    @classmethod
+    def _get_type(cls, keyword_args: dict):
+        t = keyword_args.get('t')
+        if t is None:
+            raise ValueError(f"Use `{cls.__name__}[T](name='...')` to supply the Pydantic model type`")
+        if not isinstance(t, type) or not issubclass(t, BaseModel):
+            raise ValueError(f"{t} is not a Pydantic model")
+        return t
+
+    @property
+    def model_type(self) -> type[T]:
+        """Get the Pydantic model type associated with this resource."""
+        return self._t
+
+    def _set_type(self, name: str, t: type[T]):
+        """Internal method to set the resource name and associated Pydantic type."""
+        self._t = t
+        self.name = name
+
+    def __class_getitem__(cls, item: type[BaseModel]):
+        def curried_constructor(*args, **kwargs):
+            return cls(t=item, *args, **kwargs)
+
+        return curried_constructor
+
+class TypedMooseResource(BaseTypedResource, Generic[T]):
+    """Base class for Moose resources that have columns derived from a Pydantic model.
+
+    Extends `BaseTypedResource` by adding a `Columns` helper for type-safe
+    column name access.
+
+    Attributes:
+        columns (Columns[T]): An object providing attribute access to column names.
+    """
+    columns: Columns[T]
+
+    def _set_type(self, name: str, t: type[T]):
+        super()._set_type(name, t)
+        self.columns = Columns[T](self._t)

moose_lib/dmv2/view.py

@@ -0,0 +1,36 @@
+"""
+View definitions for Moose Data Model v2 (dmv2).
+
+This module provides classes for defining standard SQL Views,
+including their SQL statements and dependencies.
+"""
+from typing import Union, List, Optional
+from pydantic import BaseModel
+
+from .sql_resource import SqlResource
+from .olap_table import OlapTable
+
+class View(SqlResource):
+    """Represents a standard SQL database View.
+
+    Args:
+        name: The name of the view to be created.
+        select_statement: The SQL SELECT statement defining the view.
+        base_tables: A list of `OlapTable`, `View`, or `MaterializedView` objects
+                     that this view depends on.
+        metadata: Optional metadata for the view.
+
+    Attributes:
+        name (str): The name of the view.
+        setup (list[str]): SQL command to create the view.
+        teardown (list[str]): SQL command to drop the view.
+        pulls_data_from (list[SqlObject]): Source tables/views.
+    """
+
+    def __init__(self, name: str, select_statement: str, base_tables: list[Union[OlapTable, SqlResource]],
+                 metadata: dict = None):
+        setup = [
+            f"CREATE VIEW IF NOT EXISTS {name} AS {select_statement}".strip()
+        ]
+        teardown = [f"DROP VIEW IF EXISTS {name}"]
+        super().__init__(name, setup, teardown, pulls_data_from=base_tables, metadata=metadata)

moose_lib/dmv2/workflow.py

@@ -0,0 +1,156 @@
+"""
+Workflow definitions for Moose Data Model v2 (dmv2).
+
+This module provides classes for defining and configuring workflows composed of tasks,
+including task dependencies, configurations, and execution functions.
+"""
+import dataclasses
+from typing import Any, Optional, Dict, List, Callable, Union, Awaitable, Generic
+from pydantic import BaseModel
+
+from .types import TypedMooseResource, T_none, U_none
+from ._registry import _workflows
+
+type TaskRunFunc[T_none, U_none] = Union[
+    # Case 1: No input, no output
+    Callable[[], None],
+    # Case 2: No input, with output
+    Callable[[], Union[U_none, Awaitable[U_none]]],
+    # Case 3: With input, no output
+    Callable[[T_none], None],
+    # Case 4: With input, with output
+    Callable[[T_none], Union[U_none, Awaitable[U_none]]]
+]
+
+@dataclasses.dataclass
+class TaskConfig(Generic[T_none, U_none]):
+    """Configuration for a Task.
+
+    Attributes:
+        run: The handler function that executes the task logic.
+        on_complete: Optional list of tasks to run after this task completes.
+        timeout: Optional timeout string (e.g. "5m", "1h").
+        retries: Optional number of retry attempts.
+    """
+    run: TaskRunFunc[T_none, U_none]
+    on_complete: Optional[list["Task[U_none, Any]"]] = None
+    timeout: Optional[str] = None
+    retries: Optional[int] = None
+
+class Task(TypedMooseResource, Generic[T_none, U_none]):
+    """Represents a task that can be executed as part of a workflow.
+
+    Tasks are the basic unit of work in a workflow, with typed input and output.
+    They can be chained together using the on_complete configuration.
+
+    Args:
+        name: The name of the task.
+        config: Configuration specifying the task's behavior.
+        t: The Pydantic model defining the task's input schema
+           (passed via `Task[InputModel, OutputModel](...)`).
+           OutputModel can be None for tasks that don't return a value.
+
+    Attributes:
+        config (TaskConfig[T, U]): The configuration for this task.
+        columns (Columns[T]): Helper for accessing input field names safely.
+        name (str): The name of the task.
+        model_type (type[T]): The Pydantic model associated with this task's input.
+    """
+    config: TaskConfig[T_none, U_none]
+
+    def __init__(self, name: str, config: TaskConfig[T_none, U_none], **kwargs):
+        super().__init__()
+        self._set_type(name, self._get_type(kwargs))
+        self.config = config
+
+    @classmethod
+    def _get_type(cls, keyword_args: dict):
+        t = keyword_args.get('t')
+        if t is None:
+            raise ValueError(f"Use `{cls.__name__}[T, U](name='...')` to supply both input and output types")
+        if not isinstance(t, tuple) or len(t) != 2:
+            raise ValueError(f"Use `{cls.__name__}[T, U](name='...')` to supply both input and output types")
+
+        input_type, output_type = t
+        if input_type is not None and (not isinstance(input_type, type) or not issubclass(input_type, BaseModel)):
+            raise ValueError(f"Input type {input_type} is not a Pydantic model or None")
+        if output_type is not None and (not isinstance(output_type, type) or not issubclass(output_type, BaseModel)):
+            raise ValueError(f"Output type {output_type} is not a Pydantic model or None")
+        return t
+
+    def _set_type(self, name: str, t: tuple[type[T_none], type[U_none]]):
+        input_type, output_type = t
+        self._t = input_type
+        self._u = output_type
+        self.name = name
+
+@dataclasses.dataclass
+class WorkflowConfig:
+    """Configuration for a workflow.
+
+    Attributes:
+        starting_task: The first task to execute in the workflow.
+        retries: Optional number of retry attempts for the entire workflow.
+        timeout: Optional timeout string for the entire workflow.
+        schedule: Optional cron-like schedule string for recurring execution.
+    """
+    starting_task: Task[Any, Any]
+    retries: Optional[int] = None
+    timeout: Optional[str] = None
+    schedule: Optional[str] = None
+
+class Workflow:
+    """Represents a workflow composed of one or more tasks.
+
+    Workflows define a sequence of tasks to be executed, with optional
+    scheduling, retries, and timeouts at the workflow level.
+
+    Args:
+        name: The name of the workflow.
+        config: Configuration specifying the workflow's behavior.
+
+    Attributes:
+        name (str): The name of the workflow.
+        config (WorkflowConfig): The configuration for this workflow.
+    """
+    def __init__(self, name: str, config: WorkflowConfig):
+        self.name = name
+        self.config = config
+        # Register the workflow in the internal registry
+        _workflows[name] = self
+
+    def get_task_names(self) -> list[str]:
+        """Get a list of all task names in this workflow.
+
+        Returns:
+            list[str]: List of task names in the workflow, including all child tasks
+        """
+        def collect_task_names(task: Task) -> list[str]:
+            names = [task.name]
+            if task.config.on_complete:
+                for child in task.config.on_complete:
+                    names.extend(collect_task_names(child))
+            return names
+
+        return collect_task_names(self.config.starting_task)
+
+    def get_task(self, task_name: str) -> Optional[Task]:
+        """Find a task in this workflow by name.
+
+        Args:
+            task_name: The name of the task to find
+
+        Returns:
+            Optional[Task]: The task if found, None otherwise
+        """
+        def find_task(task: Task) -> Optional[Task]:
+            if task.name == task_name:
+                return task
+            if task.config.on_complete:
+                for child in task.config.on_complete:
+                    found = find_task(child)
+                    if found:
+                        return found
+            return None
+
+        return find_task(self.config.starting_task)