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

moose_lib/data_models.py

@@ -87,7 +87,12 @@ class NamedTupleType(BaseModel):
     fields: list[tuple[str, "DataType"]]
 
 
-type DataType = str | DataEnum | ArrayType | Nested | NamedTupleType
+class MapType(BaseModel):
+    key_type: "DataType"
+    value_type: "DataType"
+
+
+type DataType = str | DataEnum | ArrayType | Nested | NamedTupleType | MapType
 
 
 def handle_jwt(field_type: type) -> Tuple[bool, type]:
@@ -185,6 +190,15 @@ def py_type_to_column_type(t: type, mds: list[Any]) -> Tuple[bool, list[Any], Da
     elif get_origin(t) is list:
         inner_optional, _, inner_type = py_type_to_column_type(get_args(t)[0], [])
         data_type = ArrayType(element_type=inner_type, element_nullable=inner_optional)
+    elif get_origin(t) is dict:
+        args = get_args(t)
+        if len(args) == 2:
+            key_optional, _, key_type = py_type_to_column_type(args[0], [])
+            value_optional, _, value_type = py_type_to_column_type(args[1], [])
+            # For dict types, we assume keys are required and values match their type
+            data_type = MapType(key_type=key_type, value_type=value_type)
+        else:
+            raise ValueError(f"Dict type must have exactly 2 type arguments, got {len(args)}")
     elif t is UUID:
         data_type = "UUID"
     elif t is Any:

moose_lib/dmv2/__init__.py

@@ -0,0 +1,142 @@
+"""
+Moose Data Model v2 (dmv2)
+
+This package provides the Python classes for defining Moose v2 data model resources.
+"""
+
+from .types import (
+    BaseTypedResource,
+    TypedMooseResource,
+    Columns,
+    T,
+    U,
+    T_none,
+    U_none,
+    ZeroOrMany,
+)
+
+from .olap_table import (
+    OlapConfig,
+    OlapTable,
+)
+
+from .stream import (
+    StreamConfig,
+    TransformConfig,
+    ConsumerConfig,
+    Stream,
+    DeadLetterModel,
+    DeadLetterQueue,
+)
+
+from .ingest_api import (
+    IngestConfig,
+    IngestConfigWithDestination,
+    IngestApi,
+)
+
+from .ingest_pipeline import (
+    IngestPipelineConfig,
+    IngestPipeline,
+)
+
+from .consumption import (
+    EgressConfig,
+    ConsumptionApi,
+)
+
+from .sql_resource import (
+    SqlResource,
+)
+
+from .view import (
+    View,
+)
+
+from .materialized_view import (
+    MaterializedViewOptions,
+    MaterializedView,
+)
+
+from .workflow import (
+    TaskConfig,
+    Task,
+    WorkflowConfig,
+    Workflow,
+)
+
+from .registry import (
+    get_tables,
+    get_table,
+    get_streams,
+    get_stream,
+    get_ingest_apis,
+    get_ingest_api,
+    get_consumption_apis,
+    get_consumption_api,
+    get_sql_resources,
+    get_sql_resource,
+    get_workflows,
+    get_workflow,
+)
+
+__all__ = [
+    # Types
+    'BaseTypedResource',
+    'TypedMooseResource',
+    'Columns',
+    'T',
+    'U',
+    'T_none',
+    'U_none',
+    'ZeroOrMany',
+
+    # OLAP Tables
+    'OlapConfig',
+    'OlapTable',
+
+    # Streams
+    'StreamConfig',
+    'TransformConfig',
+    'ConsumerConfig',
+    'Stream',
+    'DeadLetterModel',
+    'DeadLetterQueue',
+
+    # Ingestion
+    'IngestConfig',
+    'IngestConfigWithDestination',
+    'IngestPipelineConfig',
+    'IngestApi',
+    'IngestPipeline',
+
+    # Consumption
+    'EgressConfig',
+    'ConsumptionApi',
+
+    # SQL
+    'SqlResource',
+    'View',
+    'MaterializedViewOptions',
+    'MaterializedView',
+
+    # Workflow
+    'TaskConfig',
+    'Task',
+    'WorkflowConfig',
+    'Workflow',
+
+    # Registry
+    'get_tables',
+    'get_table',
+    'get_streams',
+    'get_stream',
+    'get_ingest_apis',
+    'get_ingest_api',
+    'get_consumption_apis',
+    'get_consumption_api',
+    'get_sql_resources',
+    'get_sql_resource',
+    'get_workflows',
+    'get_workflow',
+]

moose_lib/dmv2/_registry.py

@@ -0,0 +1,15 @@
+"""
+Internal registry dictionaries for Moose Data Model v2 (dmv2) resources.
+
+This module maintains the raw dictionaries that store all registered resources.
+It has no imports from other dmv2 modules to avoid circular dependencies.
+"""
+from typing import Dict, Any
+
+# Global registries for all resource types
+_tables: Dict[str, Any] = {}
+_streams: Dict[str, Any] = {}
+_ingest_apis: Dict[str, Any] = {}
+_egress_apis: Dict[str, Any] = {}
+_sql_resources: Dict[str, Any] = {}
+_workflows: Dict[str, Any] = {} 

moose_lib/dmv2/consumption.py

@@ -0,0 +1,101 @@
+"""
+Consumption (Egress) API definitions for Moose Data Model v2 (dmv2).
+
+This module provides classes for defining and configuring consumption APIs
+that allow querying data through user-defined functions.
+"""
+from typing import Any, Callable, Optional, Tuple
+from pydantic import BaseModel
+from pydantic.json_schema import JsonSchemaValue
+
+from .types import BaseTypedResource, T, U
+from ._registry import _egress_apis
+
+class EgressConfig(BaseModel):
+    """Configuration for Consumption (Egress) APIs.
+
+    Attributes:
+        version: Optional version string.
+        metadata: Optional metadata for the consumption API.
+    """
+    version: Optional[str] = None
+    metadata: Optional[dict] = None
+
+class ConsumptionApi(BaseTypedResource):
+    """Represents a Consumption (Egress) API endpoint.
+
+    Allows querying data, typically powered by a user-defined function.
+    Requires two Pydantic models: `T` for query parameters and `U` for the response body.
+
+    Args:
+        name: The name of the consumption API endpoint.
+        query_function: The callable that executes the query logic.
+                      It receives parameters matching model `T` (and potentially
+                      other runtime utilities) and should return data matching model `U`.
+        config: Optional configuration (currently only `version`).
+        t: A tuple containing the input (`T`) and output (`U`) Pydantic models
+           (passed via `ConsumptionApi[InputModel, OutputModel](...)`).
+
+    Attributes:
+        config (EgressConfig): Configuration for the API.
+        query_function (Callable[..., U]): The handler function for the API.
+        name (str): The name of the API.
+        model_type (type[T]): The Pydantic model for the input/query parameters.
+        return_type (type[U]): The Pydantic model for the response body.
+    """
+    config: EgressConfig
+    query_function: Callable[..., U]
+    _u: type[U]
+
+    def __class_getitem__(cls, items):
+        # Handle two type parameters
+        if not isinstance(items, tuple) or len(items) != 2:
+            raise ValueError(f"Use `{cls.__name__}[T, U](name='...')` to supply both input and output types")
+        input_type, output_type = items
+
+        def curried_constructor(*args, **kwargs):
+            return cls(t=(input_type, output_type), *args, **kwargs)
+
+        return curried_constructor
+
+    def __init__(self, name: str, query_function: Callable[..., U], config: EgressConfig = EgressConfig(), **kwargs):
+        super().__init__()
+        self._set_type(name, self._get_type(kwargs))
+        self.config = config
+        self.query_function = query_function
+        self.metadata = config.metadata
+        _egress_apis[name] = self
+
+    @classmethod
+    def _get_type(cls, keyword_args: dict):
+        t = keyword_args.get('t')
+        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 not isinstance(input_type, type) or not issubclass(input_type, BaseModel):
+            raise ValueError(f"Input type {input_type} is not a Pydantic model")
+        if not isinstance(output_type, type) or not issubclass(output_type, BaseModel):
+            raise ValueError(f"Output type {output_type} is not a Pydantic model")
+        return t
+
+    def _set_type(self, name: str, t: Tuple[type[T], type[U]]):
+        input_type, output_type = t
+        self._t = input_type
+        self._u = output_type
+        self.name = name
+
+    def return_type(self) -> type[U]:
+        """Get the Pydantic model type for the API's response body."""
+        return self._u
+
+    def get_response_schema(self) -> JsonSchemaValue:
+        """Generates the JSON schema for the API's response body model (`U`).
+
+        Returns:
+            A dictionary representing the JSON schema.
+        """
+        from pydantic.type_adapter import TypeAdapter
+        return TypeAdapter(self.return_type).json_schema(
+            ref_template='#/components/schemas/{model}'
+        )

moose_lib/dmv2/ingest_api.py

@@ -0,0 +1,64 @@
+"""
+Ingestion API definitions for Moose Data Model v2 (dmv2).
+
+This module provides classes for defining and configuring ingestion APIs
+that receive data and send it to streams.
+"""
+import dataclasses
+from typing import Any, Optional, Generic
+from pydantic import BaseModel
+
+from .types import TypedMooseResource, T
+from .stream import Stream, DeadLetterQueue
+from ._registry import _ingest_apis
+
+class IngestConfig(BaseModel):
+    """Basic configuration for an ingestion point.
+
+    Attributes:
+        version: Optional version string.
+        metadata: Optional metadata for the ingestion point.
+    """
+    version: Optional[str] = None
+    metadata: Optional[dict] = None
+
+@dataclasses.dataclass
+class IngestConfigWithDestination[T: BaseModel]:
+    """Ingestion configuration that includes the mandatory destination stream.
+
+    Attributes:
+        destination: The `Stream` where ingested data will be sent.
+        version: Optional version string.
+        metadata: Optional metadata for the ingestion configuration.
+    """
+    destination: Stream[T]
+    dead_letter_queue: Optional[DeadLetterQueue[T]] = None
+    version: Optional[str] = None
+    metadata: Optional[dict] = None
+
+class IngestApi(TypedMooseResource, Generic[T]):
+    """Represents an Ingestion API endpoint typed with a Pydantic model.
+
+    This endpoint receives data (matching schema `T`) and sends it to a configured
+    destination stream.
+
+    Args:
+        name: The name of the ingestion API endpoint.
+        config: Configuration specifying the destination stream and data format.
+        t: The Pydantic model defining the expected input data schema
+           (passed via `IngestApi[MyModel](...)`).
+
+    Attributes:
+        config (IngestConfigWithDestination[T]): The configuration for this API.
+        columns (Columns[T]): Helper for accessing input field names safely.
+        name (str): The name of the API.
+        model_type (type[T]): The Pydantic model associated with this API's input.
+    """
+    config: IngestConfigWithDestination[T]
+
+    def __init__(self, name: str, config: IngestConfigWithDestination[T], **kwargs):
+        super().__init__()
+        self._set_type(name, self._get_type(kwargs))
+        self.config = config
+        self.metadata = getattr(config, 'metadata', None)
+        _ingest_apis[name] = self

moose_lib/dmv2/ingest_pipeline.py

@@ -0,0 +1,156 @@
+"""
+Ingestion Pipeline definitions for Moose Data Model v2 (dmv2).
+
+This module provides classes for defining and configuring complete ingestion pipelines,
+which combine tables, streams, and ingestion APIs into a single cohesive unit.
+"""
+from typing import Any, Optional, Generic, TypeVar
+from pydantic import BaseModel
+
+from .types import TypedMooseResource, T
+from .olap_table import OlapTable, OlapConfig
+from .stream import Stream, StreamConfig, DeadLetterQueue
+from .ingest_api import IngestApi, IngestConfig, IngestConfigWithDestination
+
+class IngestPipelineConfig(BaseModel):
+    """Configuration for creating a complete ingestion pipeline.
+
+    Defines which components (table, stream, ingest API) should be created.
+    Set a component to `True` for default settings, `False` to disable, or provide
+    a specific config object (`OlapConfig`, `StreamConfig`, `IngestConfig`).
+
+    Attributes:
+        table: Configuration for the OLAP table component.
+        stream: Configuration for the stream component.
+        ingest: Configuration for the ingest API component.
+        version: Optional version string applied to all created components.
+        metadata: Optional metadata for the ingestion pipeline.
+    """
+    table: bool | OlapConfig = True
+    stream: bool | StreamConfig = True
+    ingest: bool | IngestConfig = True
+    dead_letter_queue: bool | StreamConfig = True
+    version: Optional[str] = None
+    metadata: Optional[dict] = None
+
+class IngestPipeline(TypedMooseResource, Generic[T]):
+    """Creates and configures a linked Table, Stream, and Ingest API pipeline.
+
+    Simplifies the common pattern of ingesting data through an API, processing it
+    in a stream, and storing it in a table.
+
+    Args:
+        name: The base name used for all created components (table, stream, API).
+        config: Specifies which components to create and their configurations.
+        t: The Pydantic model defining the data schema for all components
+           (passed via `IngestPipeline[MyModel](...)`).
+
+    Attributes:
+        table: The created `OlapTable` instance, if configured.
+        stream: The created `Stream` instance, if configured.
+        ingest_api: The created `IngestApi` instance, if configured.
+        dead_letter_queue: The created `DeadLetterQueue` instance, if configured.
+        columns (Columns[T]): Helper for accessing data field names safely.
+        name (str): The base name of the pipeline.
+        model_type (type[T]): The Pydantic model associated with this pipeline.
+    """
+    table: Optional[OlapTable[T]] = None
+    stream: Optional[Stream[T]] = None
+    ingest_api: Optional[IngestApi[T]] = None
+    dead_letter_queue: Optional[DeadLetterQueue[T]] = None
+    metadata: Optional[dict] = None
+
+    def get_table(self) -> OlapTable[T]:
+        """Retrieves the pipeline's OLAP table component.
+
+        Raises:
+            ValueError: If the table was not configured for this pipeline.
+
+        Returns:
+            The `OlapTable` instance.
+        """
+        if self.table is None:
+            raise ValueError("Table was not configured for this pipeline")
+        return self.table
+
+    def get_stream(self) -> Stream[T]:
+        """Retrieves the pipeline's stream component.
+
+        Raises:
+            ValueError: If the stream was not configured for this pipeline.
+
+        Returns:
+            The `Stream` instance.
+        """
+        if self.stream is None:
+            raise ValueError("Stream was not configured for this pipeline")
+        return self.stream
+
+    def get_dead_letter_queue(self) -> Stream[T]:
+        """Retrieves the pipeline's dead letter queue.
+
+        Raises:
+            ValueError: If the dead letter queue was not configured for this pipeline.
+
+        Returns:
+            The `Stream` instance.
+        """
+        if self.dead_letter_queue is None:
+            raise ValueError("DLQ was not configured for this pipeline")
+        return self.dead_letter_queue
+
+    def get_ingest_api(self) -> IngestApi[T]:
+        """Retrieves the pipeline's Ingestion API component.
+
+        Raises:
+            ValueError: If the Ingest API was not configured for this pipeline.
+
+        Returns:
+            The `IngestApi` instance.
+        """
+        if self.ingest_api is None:
+            raise ValueError("Ingest API was not configured for this pipeline")
+        return self.ingest_api
+
+    def __init__(self, name: str, config: IngestPipelineConfig, **kwargs):
+        super().__init__()
+        self._set_type(name, self._get_type(kwargs))
+        self.metadata = config.metadata
+        table_metadata = config.metadata
+        stream_metadata = config.metadata
+        ingest_metadata = config.metadata
+        if config.table:
+            table_config = OlapConfig() if config.table is True else config.table
+            if config.version:
+                table_config.version = config.version
+            table_config.metadata = table_metadata
+            self.table = OlapTable(name, table_config, t=self._t)
+        if config.stream:
+            stream_config = StreamConfig() if config.stream is True else config.stream
+            if config.table and stream_config.destination is not None:
+                raise ValueError("The destination of the stream should be the table created in the IngestPipeline")
+            stream_config.destination = self.table
+            if config.version:
+                stream_config.version = config.version
+            stream_config.metadata = stream_metadata
+            self.stream = Stream(name, stream_config, t=self._t)
+        if config.dead_letter_queue:
+            stream_config = StreamConfig() if config.dead_letter_queue is True else config.dead_letter_queue
+            if config.version:
+                stream_config.version = config.version
+            stream_config.metadata = stream_metadata
+            self.dead_letter_queue = DeadLetterQueue(f"{name}DeadLetterQueue", stream_config, t=self._t)
+        if config.ingest:
+            if self.stream is None:
+                raise ValueError("Ingest API needs a stream to write to.")
+            ingest_config_dict = (
+                IngestConfig() if config.ingest is True else config.ingest
+            ).model_dump()
+            ingest_config_dict["destination"] = self.stream
+            if config.version:
+                ingest_config_dict["version"] = config.version
+            if self.dead_letter_queue:
+                ingest_config_dict["dead_letter_queue"] = self.dead_letter_queue
+            ingest_config_dict["metadata"] = ingest_metadata
+            ingest_config = IngestConfigWithDestination(**ingest_config_dict)
+            self.ingest_api = IngestApi(name, ingest_config, t=self._t)

moose_lib/dmv2/materialized_view.py

@@ -0,0 +1,94 @@
+"""
+Materialized View definitions for Moose Data Model v2 (dmv2).
+
+This module provides classes for defining Materialized Views,
+including their SQL statements, target tables, and dependencies.
+"""
+from typing import Any, Optional, Union, Generic, TypeVar
+from pydantic import BaseModel, ConfigDict
+
+from moose_lib import ClickHouseEngines
+from .types import BaseTypedResource, T
+from .olap_table import OlapTable, OlapConfig
+from .sql_resource import SqlResource
+
+class MaterializedViewOptions(BaseModel):
+    """Configuration options for creating a Materialized View.
+
+    Attributes:
+        select_statement: The SQL SELECT statement defining the view's data.
+        select_tables: List of source tables/views the select statement reads from.
+        table_name: The name of the underlying target table storing the materialized data.
+        materialized_view_name: The name of the MATERIALIZED VIEW object itself.
+        engine: Optional ClickHouse engine for the target table.
+        order_by_fields: Optional ordering key for the target table (required for
+                         engines like ReplacingMergeTree).
+        model_config: ConfigDict for Pydantic validation
+    """
+    select_statement: str
+    select_tables: list[Union[OlapTable, SqlResource]]
+    table_name: str
+    materialized_view_name: str
+    engine: Optional[ClickHouseEngines] = None
+    order_by_fields: Optional[list[str]] = None
+    metadata: Optional[dict] = None
+    # Ensure arbitrary types are allowed for Pydantic validation
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+class MaterializedView(SqlResource, BaseTypedResource, Generic[T]):
+    """Represents a ClickHouse Materialized View.
+
+    Encapsulates the MATERIALIZED VIEW definition and the underlying target `OlapTable`
+    that stores the data.
+
+    Args:
+        options: Configuration defining the select statement, names, and dependencies.
+        t: The Pydantic model defining the schema of the target table
+           (passed via `MaterializedView[MyModel](...)`).
+
+    Attributes:
+        target_table (OlapTable[T]): The `OlapTable` instance storing the materialized data.
+        config (MaterializedViewOptions): The configuration options used to create the view.
+        name (str): The name of the MATERIALIZED VIEW object.
+        model_type (type[T]): The Pydantic model associated with the target table.
+        setup (list[str]): SQL commands to create the view and populate the target table.
+        teardown (list[str]): SQL command to drop the view.
+        pulls_data_from (list[SqlObject]): Source tables/views.
+        pushes_data_to (list[SqlObject]): The target table.
+    """
+    target_table: OlapTable[T]
+    config: MaterializedViewOptions
+
+    def __init__(
+            self,
+            options: MaterializedViewOptions,
+            **kwargs
+    ):
+        self._set_type(options.materialized_view_name, self._get_type(kwargs))
+
+        setup = [
+            f"CREATE MATERIALIZED VIEW IF NOT EXISTS {options.materialized_view_name} TO {options.table_name} AS {options.select_statement}",
+            f"INSERT INTO {options.table_name} {options.select_statement}"
+        ]
+        teardown = [f"DROP VIEW IF EXISTS {options.materialized_view_name}"]
+
+        target_table = OlapTable(
+            name=options.table_name,
+            config=OlapConfig(
+                order_by_fields=options.order_by_fields or [],
+                engine=options.engine
+            ),
+            t=self._t
+        )
+
+        super().__init__(
+            options.materialized_view_name,
+            setup,
+            teardown,
+            pulls_data_from=options.select_tables,
+            pushes_data_to=[target_table],
+            metadata=options.metadata
+        )
+
+        self.target_table = target_table
+        self.config = options

moose_lib/dmv2/olap_table.py

@@ -0,0 +1,57 @@
+"""
+OLAP table definitions for Moose Data Model v2 (dmv2).
+
+This module provides classes for defining and configuring OLAP tables,
+particularly for ClickHouse.
+"""
+from typing import Optional, Dict, Any, Generic
+from pydantic import BaseModel
+
+from moose_lib import ClickHouseEngines
+from .types import TypedMooseResource, T
+from ._registry import _tables
+
+class OlapConfig(BaseModel):
+    """Configuration for OLAP tables (e.g., ClickHouse tables).
+
+    Attributes:
+        order_by_fields: List of column names to use for the ORDER BY clause.
+                       Crucial for `ReplacingMergeTree` and performance.
+        deduplicate: If True, uses the ReplacingMergeTree engine for automatic
+                     deduplication based on `order_by_fields`. Equivalent to
+                     setting `engine=ClickHouseEngines.ReplacingMergeTree`.
+        engine: The ClickHouse table engine to use (e.g., MergeTree, ReplacingMergeTree).
+        version: Optional version string for tracking configuration changes.
+        metadata: Optional metadata for the table.
+    """
+    order_by_fields: list[str] = []
+    # equivalent to setting `engine=ClickHouseEngines.ReplacingMergeTree`
+    deduplicate: bool = False
+    engine: Optional[ClickHouseEngines] = None
+    version: Optional[str] = None
+    metadata: Optional[dict] = None
+
+class OlapTable(TypedMooseResource, Generic[T]):
+    """Represents an OLAP table (e.g., a ClickHouse table) typed with a Pydantic model.
+
+    Args:
+        name: The name of the OLAP table.
+        config: Configuration options for the table engine, ordering, etc.
+        t: The Pydantic model defining the table schema (passed via `OlapTable[MyModel](...)`).
+
+    Attributes:
+        config (OlapConfig): The configuration settings for this table.
+        columns (Columns[T]): Helper for accessing column names safely.
+        name (str): The name of the table.
+        model_type (type[T]): The Pydantic model associated with this table.
+        kind: The kind of the table (e.g., "OlapTable").
+    """
+    config: OlapConfig
+    kind: str = "OlapTable"
+
+    def __init__(self, name: str, config: OlapConfig = OlapConfig(), **kwargs):
+        super().__init__()
+        self._set_type(name, self._get_type(kwargs))
+        self.config = config
+        self.metadata = config.metadata
+        _tables[name] = self