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

moose_lib/clients/redis_client.py

@@ -120,13 +120,13 @@ class MooseCache:
 
             # Use provided TTL or default to 1 hour
             ttl = ttl_seconds if ttl_seconds is not None else 3600
-            
+
             # Store the value and its type metadata
             pipe = self._client.pipeline()
             pipe.setex(prefixed_key, ttl, string_value)
             pipe.setex(metadata_key, ttl, value_type)
             pipe.execute()
-            
+
         except Exception as e:
             print(f"Error setting cache key {key}: {e}")
             raise
@@ -174,13 +174,13 @@ class MooseCache:
             self._ensure_connected()
             prefixed_key = self._get_prefixed_key(key)
             metadata_key = f"{prefixed_key}:__type__"
-            
+
             # Get both the value and metadata in a single pipeline call
             pipe = self._client.pipeline()
             pipe.get(prefixed_key)
             pipe.get(metadata_key)
             results = pipe.execute()
-            
+
             value, stored_type = results[0], results[1]
 
             if value is None:
@@ -193,33 +193,45 @@ class MooseCache:
                         return value
                     elif type_hint is list:
                         # Type mismatch: stored as string but requested as list
-                        raise ValueError(f"Value was stored as string but requested as list")
+                        raise ValueError(
+                            f"Value was stored as string but requested as list"
+                        )
                     else:
-                        raise ValueError(f"Value was stored as string but requested as {type_hint.__name__}")
-                
+                        raise ValueError(
+                            f"Value was stored as string but requested as {type_hint.__name__}"
+                        )
+
                 elif stored_type == "list":
                     parsed_value = json.loads(value)
                     if type_hint is list:
                         return parsed_value
                     elif type_hint is str:
                         # Type mismatch: stored as list but requested as string
-                        raise ValueError(f"Value was stored as list but requested as string")
+                        raise ValueError(
+                            f"Value was stored as list but requested as string"
+                        )
                     else:
-                        raise ValueError(f"Value was stored as list but requested as {type_hint.__name__}")
-                
+                        raise ValueError(
+                            f"Value was stored as list but requested as {type_hint.__name__}"
+                        )
+
                 elif stored_type.startswith("pydantic:"):
                     parsed_value = json.loads(value)
                     if isinstance(type_hint, type) and issubclass(type_hint, BaseModel):
                         return type_hint.model_validate(parsed_value)
                     elif type_hint is str:
                         # Type mismatch: stored as Pydantic but requested as string
-                        raise ValueError(f"Value was stored as Pydantic model but requested as string")
+                        raise ValueError(
+                            f"Value was stored as Pydantic model but requested as string"
+                        )
                     elif type_hint is list:
                         # Type mismatch: stored as Pydantic but requested as list
-                        raise ValueError(f"Value was stored as Pydantic model but requested as list")
+                        raise ValueError(
+                            f"Value was stored as Pydantic model but requested as list"
+                        )
                     else:
                         return type_hint.model_validate(parsed_value)
-            
+
             # Backwards compatibility: no metadata found, use legacy behavior
             # But remove the problematic auto-detection for strings
             if type_hint is str:
@@ -256,7 +268,7 @@ class MooseCache:
             self._ensure_connected()
             prefixed_key = self._get_prefixed_key(key)
             metadata_key = f"{prefixed_key}:__type__"
-            
+
             # Delete both the value and its metadata
             pipe = self._client.pipeline()
             pipe.delete(prefixed_key)

moose_lib/commons.py

@@ -4,8 +4,9 @@ from datetime import datetime, timezone
 
 import requests
 import json
-from typing import Optional, Literal
+from typing import Optional, Literal, Any, Union, Callable
 import os
+from kafka import KafkaConsumer, KafkaProducer
 
 
 class CliLogData:
@@ -14,8 +15,12 @@ class CliLogData:
     ERROR = "Error"
     HIGHLIGHT = "Highlight"
 
-    def __init__(self, action: str, message: str,
-                 message_type: Optional[Literal[INFO, SUCCESS, ERROR, HIGHLIGHT]] = INFO):
+    def __init__(
+        self,
+        action: str,
+        message: str,
+        message_type: Optional[Literal[INFO, SUCCESS, ERROR, HIGHLIGHT]] = INFO,
+    ):
         self.message_type = message_type
         self.action = action
         self.message = message
@@ -32,7 +37,7 @@ def cli_log(log: CliLogData) -> None:
         # tries to send logs when moose hasn't fully started, the requests will fail.
         # The try catch is to ignore those errors.
         url = f"http://localhost:{moose_management_port}/logs"
-        headers = {'Content-Type': 'application/json'}
+        headers = {"Content-Type": "application/json"}
         requests.post(url, data=json.dumps(log.__dict__), headers=headers)
     except:
         pass
@@ -59,7 +64,11 @@ class Logger:
             elif message_type == CliLogData.HIGHLIGHT:
                 moose_scripts_logger.warning(message)
         else:
-            cli_log(CliLogData(action=self.action, message=message, message_type=message_type))
+            cli_log(
+                CliLogData(
+                    action=self.action, message=message, message_type=message_type
+                )
+            )
 
     def info(self, message: str) -> None:
         self._log(message, CliLogData.INFO)
@@ -100,3 +109,83 @@ class EnhancedJSONEncoder(json.JSONEncoder):
         if dataclasses.is_dataclass(o):
             return dataclasses.asdict(o)
         return super().default(o)
+
+
+def _build_kafka_kwargs(
+    broker: Union[str, list[str]],
+    sasl_username: Optional[str] = None,
+    sasl_password: Optional[str] = None,
+    sasl_mechanism: Optional[str] = None,
+    security_protocol: Optional[str] = None,
+) -> dict[str, Any]:
+    """Builds common Kafka client kwargs from provided parameters."""
+    kwargs: dict[str, Any] = {
+        "bootstrap_servers": broker,
+    }
+    if sasl_mechanism:
+        kwargs["sasl_mechanism"] = sasl_mechanism
+    if sasl_username is not None:
+        kwargs["sasl_plain_username"] = sasl_username
+    if sasl_password is not None:
+        kwargs["sasl_plain_password"] = sasl_password
+    if security_protocol is not None:
+        kwargs["security_protocol"] = security_protocol
+    return kwargs
+
+
+def get_kafka_consumer(
+    *,
+    broker: Union[str, list[str]],
+    client_id: str,
+    group_id: str,
+    sasl_username: Optional[str] = None,
+    sasl_password: Optional[str] = None,
+    sasl_mechanism: Optional[str] = None,
+    security_protocol: Optional[str] = None,
+    value_deserializer=lambda m: json.loads(m.decode("utf-8")),
+    **extra_kwargs: Any,
+) -> KafkaConsumer:
+    """Creates a configured KafkaConsumer with optional SASL/security settings."""
+    kwargs = _build_kafka_kwargs(
+        broker,
+        sasl_username=sasl_username,
+        sasl_password=sasl_password,
+        sasl_mechanism=sasl_mechanism,
+        security_protocol=security_protocol,
+    )
+    return KafkaConsumer(
+        client_id=client_id,
+        group_id=group_id,
+        value_deserializer=value_deserializer,
+        **kwargs,
+        **extra_kwargs,
+    )
+
+
+def get_kafka_producer(
+    *,
+    broker: Union[str, list[str]],
+    sasl_username: Optional[str] = None,
+    sasl_password: Optional[str] = None,
+    sasl_mechanism: Optional[str] = None,
+    security_protocol: Optional[str] = None,
+    max_request_size: Optional[int] = None,
+    value_serializer: Optional[Callable[[Any], bytes]] = None,
+    **extra_kwargs: Any,
+) -> KafkaProducer:
+    """Creates a configured KafkaProducer with optional SASL/security settings."""
+    kwargs = _build_kafka_kwargs(
+        broker,
+        sasl_username=sasl_username,
+        sasl_password=sasl_password,
+        sasl_mechanism=sasl_mechanism,
+        security_protocol=security_protocol,
+    )
+    if max_request_size is not None:
+        kwargs["max_request_size"] = max_request_size
+    kwargs["max_in_flight_requests_per_connection"] = 1
+    if value_serializer is not None:
+        kwargs["value_serializer"] = value_serializer
+    # Allow callers to pass through additional Kafka configs like linger_ms, acks, retries, etc.
+    kwargs.update(extra_kwargs)
+    return KafkaProducer(**kwargs)

moose_lib/config/config_file.py

@@ -4,14 +4,17 @@ Configuration file handling for Moose.
 This module provides functionality for reading and parsing the moose.config.toml file,
 which contains project-wide configuration settings.
 """
+
 import os
 import tomllib
 from dataclasses import dataclass
 from typing import Optional
 
+
 @dataclass
 class ClickHouseConfig:
     """ClickHouse configuration settings from moose.config.toml."""
+
     host: str
     host_port: int
     user: str
@@ -20,11 +23,29 @@ class ClickHouseConfig:
     use_ssl: bool = False
     native_port: Optional[int] = None
 
+
+@dataclass
+class KafkaConfig:
+    """Redpanda/Kafka configuration settings from moose.config.toml."""
+
+    broker: str
+    message_timeout_ms: int
+    sasl_username: Optional[str] = None
+    sasl_password: Optional[str] = None
+    sasl_mechanism: Optional[str] = None
+    security_protocol: Optional[str] = None
+    namespace: Optional[str] = None
+    schema_registry_url: Optional[str] = None
+
+
 @dataclass
 class ProjectConfig:
     """Project configuration from moose.config.toml."""
+
     language: str
     clickhouse_config: ClickHouseConfig
+    kafka_config: Optional[KafkaConfig] = None
+
 
 def find_config_file(start_dir: str = os.getcwd()) -> Optional[str]:
     """Find moose.config.toml by walking up directory tree.
@@ -48,6 +69,7 @@ def find_config_file(start_dir: str = os.getcwd()) -> Optional[str]:
         current_dir = parent_dir
     return None
 
+
 def read_project_config() -> ProjectConfig:
     """Read and parse moose.config.toml.
 
@@ -71,12 +93,32 @@ def read_project_config() -> ProjectConfig:
             password=config_data["clickhouse_config"]["password"],
             db_name=config_data["clickhouse_config"]["db_name"],
             use_ssl=config_data["clickhouse_config"].get("use_ssl", False),
-            native_port=config_data["clickhouse_config"].get("native_port")
+            native_port=config_data["clickhouse_config"].get("native_port"),
         )
 
+        def _parse_kafka(section_name: str) -> Optional[KafkaConfig]:
+            sec = config_data.get(section_name)
+            if sec is None:
+                return None
+            return KafkaConfig(
+                broker=sec["broker"],
+                message_timeout_ms=sec.get("message_timeout_ms", 1000),
+                sasl_username=sec.get("sasl_username"),
+                sasl_password=sec.get("sasl_password"),
+                sasl_mechanism=sec.get("sasl_mechanism"),
+                security_protocol=sec.get("security_protocol"),
+                namespace=sec.get("namespace"),
+                schema_registry_url=sec.get("schema_registry_url"),
+            )
+
+        kafka_cfg = _parse_kafka("kafka_config")
+        if kafka_cfg is None:
+            kafka_cfg = _parse_kafka("redpanda_config")
+
         return ProjectConfig(
             language=config_data["language"],
-            clickhouse_config=clickhouse_config
+            clickhouse_config=clickhouse_config,
+            kafka_config=kafka_cfg,
         )
     except Exception as e:
         raise RuntimeError(f"Failed to parse moose.config.toml: {e}")

moose_lib/config/runtime.py

@@ -4,13 +4,16 @@ Runtime configuration management for Moose.
 This module provides a singleton registry for managing runtime configuration settings,
 particularly for ClickHouse connections.
 """
+
 import os
 from dataclasses import dataclass
 from typing import Optional
 
+
 @dataclass
 class RuntimeClickHouseConfig:
     """Runtime ClickHouse configuration settings."""
+
     host: str
     port: str
     username: str
@@ -18,17 +21,34 @@ class RuntimeClickHouseConfig:
     database: str
     use_ssl: bool
 
+
+@dataclass
+class RuntimeKafkaConfig:
+    """Runtime Kafka configuration settings."""
+
+    broker: str
+    message_timeout_ms: int
+    sasl_username: Optional[str]
+    sasl_password: Optional[str]
+    sasl_mechanism: Optional[str]
+    security_protocol: Optional[str]
+    namespace: Optional[str]
+    schema_registry_url: Optional[str]
+
+
 class ConfigurationRegistry:
     """Singleton registry for managing runtime configuration.
 
     This class provides a centralized way to manage and access runtime configuration
     settings, with fallback to file-based configuration when runtime settings are not set.
     """
-    _instance: Optional['ConfigurationRegistry'] = None
+
+    _instance: Optional["ConfigurationRegistry"] = None
     _clickhouse_config: Optional[RuntimeClickHouseConfig] = None
+    _kafka_config: Optional[RuntimeKafkaConfig] = None
 
     @classmethod
-    def get_instance(cls) -> 'ConfigurationRegistry':
+    def get_instance(cls) -> "ConfigurationRegistry":
         """Get the singleton instance of ConfigurationRegistry.
 
         Returns:
@@ -46,6 +66,14 @@ class ConfigurationRegistry:
         """
         self._clickhouse_config = config
 
+    def set_kafka_config(self, config: "RuntimeKafkaConfig") -> None:
+        """Set the runtime Kafka configuration.
+
+        Args:
+            config: The Kafka configuration to use.
+        """
+        self._kafka_config = config
+
     def get_clickhouse_config(self) -> RuntimeClickHouseConfig:
         """Get the current ClickHouse configuration.
 
@@ -93,18 +121,122 @@ class ConfigurationRegistry:
                 username=env_user or config.clickhouse_config.user,
                 password=env_password or config.clickhouse_config.password,
                 database=env_db or config.clickhouse_config.db_name,
-                use_ssl=(env_use_ssl if env_use_ssl is not None else config.clickhouse_config.use_ssl),
+                use_ssl=(
+                    env_use_ssl
+                    if env_use_ssl is not None
+                    else config.clickhouse_config.use_ssl
+                ),
             )
         except Exception as e:
             raise RuntimeError(f"Failed to get ClickHouse configuration: {e}")
 
+    def get_kafka_config(self) -> "RuntimeKafkaConfig":
+        """Get the current Kafka configuration.
+
+        If runtime configuration is not set, falls back to reading from moose.config.toml
+        and environment variables (Redpanda- and Kafka-prefixed).
+
+        Returns:
+            The current Kafka configuration.
+        """
+        if self._kafka_config:
+            return self._kafka_config
+
+        from .config_file import read_project_config
+
+        def _env(name: str) -> Optional[str]:
+            val = os.environ.get(name)
+            if val is None:
+                return None
+            trimmed = val.strip()
+            return trimmed if trimmed else None
+
+        try:
+            config = read_project_config()
+
+            # Prefer Redpanda-prefixed env vars; fallback to Kafka-prefixed
+            broker = _env("MOOSE_REDPANDA_CONFIG__BROKER") or _env(
+                "MOOSE_KAFKA_CONFIG__BROKER"
+            )
+            message_timeout_ms = _env(
+                "MOOSE_REDPANDA_CONFIG__MESSAGE_TIMEOUT_MS"
+            ) or _env("MOOSE_KAFKA_CONFIG__MESSAGE_TIMEOUT_MS")
+            sasl_username = _env("MOOSE_REDPANDA_CONFIG__SASL_USERNAME") or _env(
+                "MOOSE_KAFKA_CONFIG__SASL_USERNAME"
+            )
+            sasl_password = _env("MOOSE_REDPANDA_CONFIG__SASL_PASSWORD") or _env(
+                "MOOSE_KAFKA_CONFIG__SASL_PASSWORD"
+            )
+            sasl_mechanism = _env("MOOSE_REDPANDA_CONFIG__SASL_MECHANISM") or _env(
+                "MOOSE_KAFKA_CONFIG__SASL_MECHANISM"
+            )
+            security_protocol = _env(
+                "MOOSE_REDPANDA_CONFIG__SECURITY_PROTOCOL"
+            ) or _env("MOOSE_KAFKA_CONFIG__SECURITY_PROTOCOL")
+            namespace = _env("MOOSE_REDPANDA_CONFIG__NAMESPACE") or _env(
+                "MOOSE_KAFKA_CONFIG__NAMESPACE"
+            )
+            schema_registry_url = _env(
+                "MOOSE_REDPANDA_CONFIG__SCHEMA_REGISTRY_URL"
+            ) or _env("MOOSE_KAFKA_CONFIG__SCHEMA_REGISTRY_URL")
+
+            file_kafka = config.kafka_config
+
+            def _to_int(value: Optional[str], fallback: int) -> int:
+                try:
+                    return int(value) if value is not None else fallback
+                except Exception:
+                    return fallback
+
+            return RuntimeKafkaConfig(
+                broker=broker
+                or (file_kafka.broker if file_kafka else "localhost:19092"),
+                message_timeout_ms=_to_int(
+                    message_timeout_ms,
+                    file_kafka.message_timeout_ms if file_kafka else 1000,
+                ),
+                sasl_username=(
+                    sasl_username
+                    if sasl_username is not None
+                    else (file_kafka.sasl_username if file_kafka else None)
+                ),
+                sasl_password=(
+                    sasl_password
+                    if sasl_password is not None
+                    else (file_kafka.sasl_password if file_kafka else None)
+                ),
+                sasl_mechanism=(
+                    sasl_mechanism
+                    if sasl_mechanism is not None
+                    else (file_kafka.sasl_mechanism if file_kafka else None)
+                ),
+                security_protocol=(
+                    security_protocol
+                    if security_protocol is not None
+                    else (file_kafka.security_protocol if file_kafka else None)
+                ),
+                namespace=(
+                    namespace
+                    if namespace is not None
+                    else (file_kafka.namespace if file_kafka else None)
+                ),
+                schema_registry_url=(
+                    schema_registry_url
+                    if schema_registry_url is not None
+                    else (file_kafka.schema_registry_url if file_kafka else None)
+                ),
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failed to get Kafka configuration: {e}")
+
     def has_runtime_config(self) -> bool:
         """Check if runtime configuration is set.
 
         Returns:
-            True if runtime configuration is set, False otherwise.
+            True if either runtime clickhouse or kafka configuration is set, False otherwise.
         """
-        return self._clickhouse_config is not None
+        return self._clickhouse_config is not None or self._kafka_config is not None
+
 
 # Create singleton instance
 config_registry = ConfigurationRegistry.get_instance()