moose-lib 0.6.90__py3-none-any.whl → 0.6.283__py3-none-any.whl

moose_lib/dmv2/registry.py

@@ -4,6 +4,7 @@ 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
@@ -11,6 +12,7 @@ from .ingest_api import IngestApi
 from .consumption import Api
 from .sql_resource import SqlResource
 from .workflow import Workflow
+from .web_app import WebApp
 from ._registry import (
     _tables,
     _streams,
@@ -20,36 +22,45 @@ from ._registry import (
     _workflows,
     _api_name_aliases,
     _api_path_map,
+    _web_apps,
 )
 
+
 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_apis() -> Dict[str, Api]:
     """Get all registered APIs."""
     return _apis
 
+
 def get_api(name: str) -> Optional[Api]:
     """Get a registered API by name or path.
 
@@ -62,35 +73,49 @@ def get_api(name: str) -> Optional[Api]:
     api = _apis.get(name)
     if api:
         return api
-    
+
     # Try alias lookup
     api = _api_name_aliases.get(name)
     if api:
         return api
-    
+
     # Try path-based lookup
     return _api_path_map.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)
 
 
+def get_web_apps() -> Dict[str, WebApp]:
+    """Get all registered WebApps."""
+    return _web_apps
+
+
+def get_web_app(name: str) -> Optional[WebApp]:
+    """Get a registered WebApp by name."""
+    return _web_apps.get(name)
+
+
 # Backward compatibility aliases (deprecated)
 get_consumption_apis = get_apis
 """@deprecated: Use get_apis instead of get_consumption_apis"""
 
 get_consumption_api = get_api
-"""@deprecated: Use get_api instead of get_consumption_api"""
+"""@deprecated: Use get_api instead of get_consumption_api"""

moose_lib/dmv2/sql_resource.py

@@ -4,11 +4,14 @@ 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
+from ._source_capture import get_source_file_from_stack
+
 
 class SqlResource:
     """Base class for SQL resources like Views and Materialized Views.
@@ -23,22 +26,25 @@ class SqlResource:
         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").
+        source_file: Optional path to the source file where this resource was defined.
     """
+
     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"]]
+    source_file: Optional[str]
 
     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: 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
@@ -46,4 +52,6 @@ class SqlResource:
         self.pulls_data_from = pulls_data_from or []
         self.pushes_data_to = pushes_data_to or []
         self.metadata = metadata
-        _sql_resources[name] = self
+        # Capture source file from call stack
+        self.source_file = get_source_file_from_stack()
+        _sql_resources[name] = self

moose_lib/dmv2/stream.py

@@ -4,16 +4,39 @@ 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
+import json
 from typing import Any, Optional, Callable, Union, Literal, Generic
-from pydantic import BaseModel, ConfigDict, AliasGenerator
+from pydantic import BaseModel, ConfigDict, AliasGenerator, Field
 from pydantic.alias_generators import to_camel
+from kafka import KafkaProducer
 
 from .types import TypedMooseResource, ZeroOrMany, T, U
 from .olap_table import OlapTable
 from ._registry import _streams
 from .life_cycle import LifeCycle
+from ..config.runtime import config_registry, RuntimeKafkaConfig
+from ..commons import get_kafka_producer
+
+
+class SubjectLatest(BaseModel):
+    name: str = Field(serialization_alias="subjectLatest")
+
+
+class SubjectVersion(BaseModel):
+    subject: str
+    version: int
+
+
+class SchemaById(BaseModel):
+    id: int
+
+
+class KafkaSchemaConfig(BaseModel):
+    kind: Literal["JSON", "AVRO", "PROTOBUF"]
+    reference: Union[SubjectLatest, SubjectVersion, SchemaById]
 
 
 class StreamConfig(BaseModel):
@@ -28,6 +51,7 @@ class StreamConfig(BaseModel):
         life_cycle: Determines how changes in code will propagate to the resources.
         default_dead_letter_queue: default dead letter queue used by transforms/consumers
     """
+
     parallelism: int = 1
     retention_period: int = 60 * 60 * 24 * 7  # 7 days
     destination: Optional[OlapTable] = None
@@ -37,6 +61,7 @@ class StreamConfig(BaseModel):
     default_dead_letter_queue: "Optional[DeadLetterQueue]" = None
     # allow DeadLetterQueue
     model_config = ConfigDict(arbitrary_types_allowed=True)
+    schema_config: Optional[KafkaSchemaConfig] = None
 
 
 class TransformConfig(BaseModel):
@@ -46,6 +71,7 @@ class TransformConfig(BaseModel):
         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)
@@ -59,6 +85,7 @@ class ConsumerConfig(BaseModel):
         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)
@@ -67,6 +94,7 @@ class ConsumerConfig(BaseModel):
 @dataclasses.dataclass
 class _RoutedMessage:
     """Internal class representing a message routed to a specific stream."""
+
     destination: "Stream[Any]"
     values: ZeroOrMany[Any]
 
@@ -74,6 +102,7 @@ class _RoutedMessage:
 @dataclasses.dataclass
 class ConsumerEntry(Generic[T]):
     """Internal class representing a consumer with its configuration."""
+
     consumer: Callable[[T], None]
     config: ConsumerConfig
 
@@ -81,6 +110,7 @@ class ConsumerEntry(Generic[T]):
 @dataclasses.dataclass
 class TransformEntry(Generic[T]):
     """Internal class representing a transformation with its configuration."""
+
     destination: "Stream[Any]"
     transformation: Callable[[T], ZeroOrMany[Any]]
     config: TransformConfig
@@ -105,11 +135,14 @@ class Stream(TypedMooseResource, Generic[T]):
         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
     default_dead_letter_queue: "Optional[DeadLetterQueue[T]]" = None
+    _memoized_producer: Optional[KafkaProducer] = None
+    _kafka_config_hash: Optional[str] = None
 
     def __init__(self, name: str, config: "StreamConfig" = None, **kwargs):
         super().__init__()
@@ -121,8 +154,12 @@ class Stream(TypedMooseResource, Generic[T]):
         self.default_dead_letter_queue = self.config.default_dead_letter_queue
         _streams[name] = self
 
-    def add_transform(self, destination: "Stream[U]", transformation: Callable[[T], ZeroOrMany[U]],
-                      config: TransformConfig = None):
+    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
@@ -135,23 +172,37 @@ class Stream(TypedMooseResource, Generic[T]):
         """
         config = config or TransformConfig()
         if (
-                self.default_dead_letter_queue is not None
-                and config.dead_letter_queue is None
+            self.default_dead_letter_queue is not None
+            and config.dead_letter_queue is None
         ):
             config = config.model_copy()
             config.dead_letter_queue = self.default_dead_letter_queue
         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)
+            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))
+                    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):
+                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.
@@ -162,8 +213,8 @@ class Stream(TypedMooseResource, Generic[T]):
         """
         config = config or ConsumerConfig()
         if (
-                self.default_dead_letter_queue is not None
-                and config.dead_letter_queue is None
+            self.default_dead_letter_queue is not None
+            and config.dead_letter_queue is None
         ):
             config = config.model_copy()
             config.dead_letter_queue = self.default_dead_letter_queue
@@ -216,6 +267,160 @@ class Stream(TypedMooseResource, Generic[T]):
         """
         self._multipleTransformations = transformation
 
+    def _build_full_topic_name(self, namespace: Optional[str]) -> str:
+        """Build full topic name with optional namespace and version suffix."""
+        version_suffix = (
+            f"_{self.config.version.replace('.', '_')}" if self.config.version else ""
+        )
+        base = f"{self.name}{version_suffix}"
+        return f"{namespace}.{base}" if namespace else base
+
+    def _create_kafka_config_hash(self, cfg: RuntimeKafkaConfig) -> str:
+        import hashlib
+
+        config_string = ":".join(
+            str(x)
+            for x in (
+                cfg.broker,
+                cfg.message_timeout_ms,
+                cfg.sasl_username,
+                cfg.sasl_password,
+                cfg.sasl_mechanism,
+                cfg.security_protocol,
+                cfg.namespace,
+            )
+        )
+        return hashlib.sha256(config_string.encode()).hexdigest()[:16]
+
+    def _parse_brokers(self, broker_string: str) -> list[str]:
+        if not broker_string:
+            return []
+        return [b.strip() for b in broker_string.split(",") if b.strip()]
+
+    def _get_memoized_producer(self) -> tuple[KafkaProducer, RuntimeKafkaConfig]:
+        """Create or reuse a KafkaProducer using runtime configuration."""
+        cfg: RuntimeKafkaConfig = config_registry.get_kafka_config()
+        current_hash = self._create_kafka_config_hash(cfg)
+
+        if (
+            self._memoized_producer is not None
+            and self._kafka_config_hash == current_hash
+        ):
+            return self._memoized_producer, cfg
+
+        # Close previous producer if config changed
+        if (
+            self._memoized_producer is not None
+            and self._kafka_config_hash != current_hash
+        ):
+            try:
+                self._memoized_producer.flush()
+                self._memoized_producer.close()
+            except Exception:
+                pass
+            finally:
+                self._memoized_producer = None
+
+        brokers = self._parse_brokers(cfg.broker)
+        if not brokers:
+            raise RuntimeError(f"No valid broker addresses found in: '{cfg.broker}'")
+
+        producer = get_kafka_producer(
+            broker=brokers,
+            sasl_username=cfg.sasl_username,
+            sasl_password=cfg.sasl_password,
+            sasl_mechanism=cfg.sasl_mechanism,
+            security_protocol=cfg.security_protocol,
+            value_serializer=lambda v: v.model_dump_json().encode("utf-8"),
+            acks="all",
+        )
+
+        self._memoized_producer = producer
+        self._kafka_config_hash = current_hash
+        return producer, cfg
+
+    def close_producer(self) -> None:
+        """Closes the memoized Kafka producer if it exists."""
+        if self._memoized_producer is not None:
+            try:
+                self._memoized_producer.flush()
+                self._memoized_producer.close()
+            except Exception:
+                pass
+            finally:
+                self._memoized_producer = None
+                self._kafka_config_hash = None
+
+    def send(self, values: ZeroOrMany[T]) -> None:
+        """Send one or more records to this stream's Kafka topic.
+
+        If `schema_registry` (JSON) is configured, resolve schema id and
+        send using Confluent wire format (0x00 + 4-byte schema id + JSON bytes).
+        Otherwise, values are JSON-serialized.
+        """
+        # Normalize inputs to a flat list of records
+        filtered: list[T] = []
+        if isinstance(values, list):
+            for v in values:
+                if v is None:
+                    continue
+                else:
+                    filtered.append(v)
+        elif values is not None:
+            filtered.append(values)  # type: ignore[arg-type]
+
+        if len(filtered) == 0:
+            return
+
+        # ensure all records are instances of the stream's model type
+        model_type = self._t
+        for rec in filtered:
+            if not isinstance(rec, model_type):
+                raise TypeError(
+                    f"Stream '{self.name}' expects instances of {model_type.__name__}, "
+                    f"got {type(rec).__name__}"
+                )
+
+        producer, cfg = self._get_memoized_producer()
+        topic = self._build_full_topic_name(cfg.namespace)
+
+        sr = self.config.schema_config
+        if sr is not None:
+            if sr.kind != "JSON":
+                raise NotImplementedError("Currently JSON Schema is supported.")
+            try:
+                from confluent_kafka.schema_registry import SchemaRegistryClient
+                from confluent_kafka.schema_registry.json_schema import JSONSerializer
+            except Exception as e:
+                raise RuntimeError(
+                    "confluent-kafka[json,schemaregistry] is required for Schema Registry JSON"
+                ) from e
+
+            sr_url = cfg.schema_registry_url
+            if not sr_url:
+                raise RuntimeError("Schema Registry URL not configured")
+            client = SchemaRegistryClient({"url": sr_url})
+
+            if isinstance(sr.reference, SchemaById):
+                schema = client.get_schema(sr.reference.id)
+            elif isinstance(sr.reference, SubjectLatest):
+                schema = client.get_latest_version(sr.reference.name).schema
+            else:
+                schema = client.get_version(
+                    sr.reference.subject, sr.reference.version
+                ).schema
+
+            serializer = JSONSerializer(schema, client)
+
+            for rec in filtered:
+                value_bytes = serializer(rec.model_dump())
+                producer.send(topic, value=value_bytes)
+            producer.flush()
+        else:
+            for rec in filtered:
+                producer.send(topic, value=rec)
+            producer.flush()
+
 
 class DeadLetterModel(BaseModel, Generic[T]):
     """Model for dead letter queue messages.
@@ -227,9 +432,12 @@ class DeadLetterModel(BaseModel, Generic[T]):
         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,
-    ))
+
+    model_config = ConfigDict(
+        alias_generator=AliasGenerator(
+            serialization_alias=to_camel,
+        )
+    )
     original_record: Any
     error_message: str
     error_type: str
@@ -261,10 +469,16 @@ class DeadLetterQueue(Stream, Generic[T]):
         """
         self._model_type = self._get_type(kwargs)
         kwargs["t"] = DeadLetterModel[self._model_type]
-        super().__init__(name, config if config is not None else StreamConfig(), **kwargs)
-
-    def add_transform(self, destination: Stream[U], transformation: Callable[[DeadLetterModel[T]], ZeroOrMany[U]],
-                      config: TransformConfig = None):
+        super().__init__(
+            name, config if config is not None else StreamConfig(), **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)
@@ -272,7 +486,11 @@ class DeadLetterQueue(Stream, Generic[T]):
         config = config or TransformConfig()
         super().add_transform(destination, wrapped_transform, config)
 
-    def add_consumer(self, consumer: Callable[[DeadLetterModel[T]], None], config: ConsumerConfig = None):
+    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)
@@ -280,7 +498,9 @@ class DeadLetterQueue(Stream, Generic[T]):
         config = config or ConsumerConfig()
         super().add_consumer(wrapped_consumer, config)
 
-    def set_multi_transform(self, transformation: Callable[[DeadLetterModel[T]], list[_RoutedMessage]]):
+    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)

moose_lib/dmv2/types.py

@@ -5,6 +5,7 @@ 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
 
 import typing_extensions
@@ -12,10 +13,10 @@ from pydantic import BaseModel
 from pydantic.fields import FieldInfo
 from ..data_models import Column
 
-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])
+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]
 
 
@@ -33,6 +34,7 @@ class Cols:
         >>> print(table.cols.user_id)  # Output: a column object
         >>> print(table.cols.non_existent) # Raises AttributeError
     """
+
     _columns: dict[str, Column]
 
     def __init__(self, columns: list[Column]):
@@ -52,7 +54,7 @@ class Cols:
         return self.__getattr__(item)
 
 
-@typing_extensions.deprecated('use cols in OlapTable instead')
+@typing_extensions.deprecated("use cols in OlapTable instead")
 class Columns(Generic[T]):
     """Provides runtime checked column name access for Moose resources.
 
@@ -70,6 +72,7 @@ class Columns(Generic[T]):
     Args:
         model: The Pydantic model type whose fields represent the columns.
     """
+
     _fields: dict[str, FieldInfo]
 
     def __init__(self, model: type[T]):
@@ -90,14 +93,17 @@ class BaseTypedResource(Generic[T]):
     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')
+        t = keyword_args.get("t")
         if t is None:
-            raise ValueError(f"Use `{cls.__name__}[T](name='...')` to supply the Pydantic model type`")
+            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
@@ -130,7 +136,7 @@ class TypedMooseResource(BaseTypedResource, Generic[T]):
     """
 
     @property
-    @typing_extensions.deprecated('use cols in OlapTable instead', category=None)
+    @typing_extensions.deprecated("use cols in OlapTable instead", category=None)
     def columns(self):
         return Columns[T](self._t)
 

moose_lib/dmv2/view.py

@@ -4,12 +4,14 @@ 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.
 
@@ -27,10 +29,15 @@ class View(SqlResource):
         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()
-        ]
+    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)
+        super().__init__(
+            name, setup, teardown, pulls_data_from=base_tables, metadata=metadata
+        )