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

moose_lib/data_models.py

@@ -6,8 +6,19 @@ from inspect import isclass
 from uuid import UUID
 from datetime import datetime, date
 
-from typing import Literal, Tuple, Union, Any, get_origin, get_args, TypeAliasType, Annotated, Type, _BaseGenericAlias, \
-    GenericAlias
+from typing import (
+    Literal,
+    Tuple,
+    Union,
+    Any,
+    get_origin,
+    get_args,
+    TypeAliasType,
+    Annotated,
+    Type,
+    _BaseGenericAlias,
+    GenericAlias,
+)
 from pydantic import BaseModel, Field, PlainSerializer, GetCoreSchemaHandler, ConfigDict
 from pydantic_core import CoreSchema, core_schema
 import ipaddress
@@ -15,18 +26,39 @@ import ipaddress
 type Key[T: (str, int)] = T
 type JWT[T] = T
 
+# Integer type aliases for ClickHouse integer types
+type Int8 = Annotated[int, "int8"]
+type Int16 = Annotated[int, "int16"]
+type Int32 = Annotated[int, "int32"]
+type Int64 = Annotated[int, "int64"]
+type UInt8 = Annotated[int, "uint8"]
+type UInt16 = Annotated[int, "uint16"]
+type UInt32 = Annotated[int, "uint32"]
+type UInt64 = Annotated[int, "uint64"]
 
-@dataclasses.dataclass  # a BaseModel in the annotations will confuse pydantic
+# Float type aliases for ClickHouse float types
+type Float32 = Annotated[float, "float32"]
+type Float64 = Annotated[float, "float64"]
+
+
+@dataclasses.dataclass(
+    frozen=True
+)  # a BaseModel in the annotations will confuse pydantic
 class ClickhousePrecision:
     precision: int
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(frozen=True)
 class ClickhouseSize:
     size: int
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(frozen=True)
+class ClickhouseFixedStringSize:
+    size: int
+
+
+@dataclasses.dataclass(frozen=True)
 class ClickhouseDefault:
     expression: str
 
@@ -35,6 +67,60 @@ def clickhouse_default(expression: str) -> ClickhouseDefault:
     return ClickhouseDefault(expression=expression)
 
 
+@dataclasses.dataclass(frozen=True)
+class ClickHouseTTL:
+    expression: str
+
+
+@dataclasses.dataclass(frozen=True)
+class ClickHouseCodec:
+    expression: str
+
+
+@dataclasses.dataclass(frozen=True)
+class ClickHouseMaterialized:
+    """
+    ClickHouse MATERIALIZED column annotation.
+    The column value is computed at INSERT time and physically stored.
+    Cannot be explicitly inserted by users.
+
+    Args:
+        expression: ClickHouse SQL expression using column names (snake_case)
+
+    Examples:
+        # Extract date component
+        event_date: Annotated[date, ClickHouseMaterialized("toDate(event_time)")]
+
+        # Precompute hash
+        user_hash: Annotated[int, ClickHouseMaterialized("cityHash64(user_id)")]
+
+        # Complex expression with JSON
+        combination_hash: Annotated[
+            list[int],
+            ClickHouseMaterialized(
+                "arrayMap(kv -> cityHash64(kv.1, kv.2), "
+                "JSONExtractKeysAndValuesRaw(toString(log_blob)))"
+            )
+        ]
+
+    Notes:
+        - Expression uses ClickHouse column names, not Python field names
+        - MATERIALIZED and DEFAULT are mutually exclusive
+        - Can be combined with ClickHouseCodec for compression
+        - Changing the expression modifies the column in-place (existing values preserved)
+    """
+
+    expression: str
+
+
+@dataclasses.dataclass(frozen=True)
+class ClickHouseJson:
+    max_dynamic_paths: int | None = None
+    max_dynamic_types: int | None = None
+    skip_paths: tuple[str, ...] = ()
+    skip_regexps: tuple[str, ...] = ()
+
+
 def clickhouse_decimal(precision: int, scale: int) -> Type[Decimal]:
     return Annotated[Decimal, Field(max_digits=precision, decimal_places=scale)]
 
@@ -48,25 +134,92 @@ def clickhouse_datetime64(precision: int) -> Type[datetime]:
     return Annotated[datetime, ClickhousePrecision(precision=precision)]
 
 
+def FixedString(size: int) -> ClickhouseFixedStringSize:
+    """
+    Creates a FixedString(N) annotation for fixed-length strings.
+
+    ClickHouse stores exactly N bytes, padding shorter values with null bytes.
+    Values exceeding N bytes will raise an exception.
+
+    Use for fixed-length data like hashes, IPs, UUIDs, MAC addresses.
+
+    Example:
+        md5_hash: Annotated[str, FixedString(16)]  # 16-byte MD5
+        sha256: Annotated[str, FixedString(32)]    # 32-byte SHA256
+        ipv6: Annotated[str, FixedString(16)]      # 16-byte IPv6
+    """
+    return ClickhouseFixedStringSize(size=size)
+
+
+type Point = Annotated[tuple[float, float], "Point"]
+type Ring = Annotated[list[tuple[float, float]], "Ring"]
+type LineString = Annotated[list[tuple[float, float]], "LineString"]
+type MultiLineString = Annotated[list[list[tuple[float, float]]], "MultiLineString"]
+type Polygon = Annotated[list[list[tuple[float, float]]], "Polygon"]
+type MultiPolygon = Annotated[list[list[list[tuple[float, float]]]], "MultiPolygon"]
+
+
 def aggregated[T](
-        result_type: Type[T],
-        agg_func: str,
-        param_types: list[type | GenericAlias | _BaseGenericAlias]
+    result_type: Type[T],
+    agg_func: str,
+    param_types: list[type | GenericAlias | _BaseGenericAlias],
 ) -> Type[T]:
-    return Annotated[result_type, AggregateFunction(agg_func=agg_func, param_types=param_types)]
+    return Annotated[
+        result_type,
+        AggregateFunction(agg_func=agg_func, param_types=tuple(param_types)),
+    ]
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(frozen=True)
 class AggregateFunction:
     agg_func: str
-    param_types: list[type | GenericAlias | _BaseGenericAlias]
+    param_types: tuple[type | GenericAlias | _BaseGenericAlias, ...]
 
     def to_dict(self):
         return {
             "functionName": self.agg_func,
             "argumentTypes": [
                 py_type_to_column_type(t, [])[2] for t in self.param_types
-            ]
+            ],
+        }
+
+
+def simple_aggregated[T](agg_func: str, arg_type: Type[T]) -> Type[T]:
+    """Helper to create a SimpleAggregateFunction type annotation.
+
+    SimpleAggregateFunction is a ClickHouse type for storing aggregated values directly
+    instead of intermediate states. It's more efficient for functions like sum, max, min, etc.
+
+    Args:
+        agg_func: The aggregation function name (e.g., "sum", "max", "anyLast")
+        arg_type: The argument type for the function (also the result type)
+
+    Returns:
+        An Annotated type with SimpleAggregateFunction metadata
+
+    Example:
+        ```python
+        from moose_lib import simple_aggregated
+
+        row_count: simple_aggregated("sum", int)
+        max_value: simple_aggregated("max", float)
+        last_status: simple_aggregated("anyLast", str)
+        ```
+    """
+    return Annotated[
+        arg_type, SimpleAggregateFunction(agg_func=agg_func, arg_type=arg_type)
+    ]
+
+
+@dataclasses.dataclass(frozen=True)
+class SimpleAggregateFunction:
+    agg_func: str
+    arg_type: type | GenericAlias | _BaseGenericAlias
+
+    def to_dict(self):
+        return {
+            "functionName": self.agg_func,
+            "argumentType": py_type_to_column_type(self.arg_type, [])[2],
         }
 
 
@@ -79,7 +232,9 @@ def enum_value_serializer(value: int | str):
 
 class EnumValue(BaseModel):
     name: str
-    value: Annotated[int | str, PlainSerializer(enum_value_serializer, return_type=dict)]
+    value: Annotated[
+        int | str, PlainSerializer(enum_value_serializer, return_type=dict)
+    ]
 
 
 class DataEnum(BaseModel):
@@ -107,7 +262,17 @@ class MapType(BaseModel):
     value_type: "DataType"
 
 
-type DataType = str | DataEnum | ArrayType | Nested | NamedTupleType | MapType
+class JsonOptions(BaseModel):
+    max_dynamic_paths: int | None = None
+    max_dynamic_types: int | None = None
+    typed_paths: list[tuple[str, "DataType"]] = []
+    skip_paths: list[str] = []
+    skip_regexps: list[str] = []
+
+
+type DataType = (
+    str | DataEnum | ArrayType | Nested | NamedTupleType | MapType | JsonOptions
+)
 
 
 def handle_jwt(field_type: type) -> Tuple[bool, type]:
@@ -146,12 +311,105 @@ class Column(BaseModel):
     primary_key: bool
     default: str | None = None
     annotations: list[Tuple[str, Any]] = []
+    ttl: str | None = None
+    codec: str | None = None
+    materialized: str | None = None
 
     def to_expr(self):
         # Lazy import to avoid circular dependency at import time
         from .query_builder import ColumnRef
+
         return ColumnRef(self)
 
+    def __str__(self) -> str:
+        """Return properly quoted identifier for SQL interpolation.
+
+        This enables Column objects to be used directly in f-strings and
+        string concatenation for SQL query construction.
+
+        Returns:
+            Backtick-quoted identifier safe for ClickHouse SQL.
+
+        Example:
+            >>> col = Column(name="user_id", ...)
+            >>> f"SELECT {col} FROM users"
+            "SELECT `user_id` FROM users"
+        """
+        from .utilities.sql import quote_identifier
+
+        return quote_identifier(self.name)
+
+    def __format__(self, format_spec: str) -> str:
+        """Format Column for f-string interpolation with format specifiers.
+
+        Supports format specs:
+        - 'col', 'c', 'column': Returns quoted identifier
+        - '' (empty): Returns quoted identifier (default)
+
+        Args:
+            format_spec: Format specification string
+
+        Returns:
+            Backtick-quoted identifier
+
+        Example:
+            >>> col = Column(name="email", ...)
+            >>> f"SELECT {col:col} FROM users"
+            "SELECT `email` FROM users"
+        """
+        # All format specs return quoted identifier
+        # This provides flexibility for user preference
+        from .utilities.sql import quote_identifier
+
+        return quote_identifier(self.name)
+
+
+def _is_point_type(t: type) -> bool:
+    origin = get_origin(t)
+    if origin is tuple:
+        args = get_args(t)
+        return len(args) == 2 and all(arg is float for arg in args)
+    return False
+
+
+def _is_list_of(inner_check: Any, t: type) -> bool:
+    origin = get_origin(t)
+    if origin is list:
+        args = get_args(t)
+        return len(args) == 1 and inner_check(args[0])
+    return False
+
+
+def _validate_geometry_type(requested: str, t: type) -> None:
+    """
+    Validates that the provided Python type matches the expected structure
+    for the requested geometry annotation.
+    """
+    match requested:
+        case "Point":
+            if not _is_point_type(t):
+                raise ValueError("Point must be typed as tuple[float, float]")
+        case "Ring" | "LineString":
+            if not _is_list_of(_is_point_type, t):
+                raise ValueError(
+                    f"{requested} must be typed as list[tuple[float, float]]"
+                )
+        case "MultiLineString" | "Polygon":
+            if not _is_list_of(lambda x: _is_list_of(_is_point_type, x), t):
+                raise ValueError(
+                    f"{requested} must be typed as list[list[tuple[float, float]]]"
+                )
+        case "MultiPolygon":
+            if not _is_list_of(
+                lambda x: _is_list_of(lambda y: _is_list_of(_is_point_type, y), x),
+                t,
+            ):
+                raise ValueError(
+                    "MultiPolygon must be typed as list[list[list[tuple[float, float]]]]"
+                )
+        case _:
+            raise ValueError(f"Unknown geometry type annotation: {requested}")
+
 
 def py_type_to_column_type(t: type, mds: list[Any]) -> Tuple[bool, list[Any], DataType]:
     # handle Annotated[Optional[Annotated[...], ...]
@@ -162,18 +420,45 @@ def py_type_to_column_type(t: type, mds: list[Any]) -> Tuple[bool, list[Any], Da
     data_type: DataType
 
     if t is str:
-        data_type = "String"
+        # Check for FixedString annotation
+        fixed_string_size = next(
+            (md.size for md in mds if isinstance(md, ClickhouseFixedStringSize)), None
+        )
+        if fixed_string_size:
+            data_type = f"FixedString({fixed_string_size})"
+        else:
+            data_type = "String"
+    elif t is bytes:
+        # Check for FixedString annotation
+        fixed_string_size = next(
+            (md.size for md in mds if isinstance(md, ClickhouseFixedStringSize)), None
+        )
+        if fixed_string_size:
+            data_type = f"FixedString({fixed_string_size})"
+        else:
+            # Regular bytes without FixedString annotation
+            data_type = "String"
     elif t is int:
         # Check for int size annotations
-        int_size = next((md for md in mds if isinstance(md, str) and re.match(r'^u?int\d+$', md)), None)
+        int_size = next(
+            (md for md in mds if isinstance(md, str) and re.match(r"^u?int\d+$", md)),
+            None,
+        )
         if int_size:
             data_type = int_size.replace("u", "U").replace("i", "I")
         else:
-            data_type = "Int"
+            data_type = "Int64"
     elif t is float:
         size = next((md for md in mds if isinstance(md, ClickhouseSize)), None)
         if size is None:
-            bit_size = next((md for md in mds if isinstance(md, str) and re.match(r'^float\d+$', md)), None)
+            bit_size = next(
+                (
+                    md
+                    for md in mds
+                    if isinstance(md, str) and re.match(r"^float\d+$", md)
+                ),
+                None,
+            )
             if bit_size:
                 if bit_size == "float32":
                     data_type = "Float32"
@@ -191,12 +476,16 @@ def py_type_to_column_type(t: type, mds: list[Any]) -> Tuple[bool, list[Any], Da
             raise ValueError(f"Unsupported float size {size.size}")
     elif t is Decimal:
         precision = next((md.max_digits for md in mds if hasattr(md, "max_digits")), 10)
-        scale = next((md.decimal_places for md in mds if hasattr(md, "decimal_places")), 0)
+        scale = next(
+            (md.decimal_places for md in mds if hasattr(md, "decimal_places")), 0
+        )
         data_type = f"Decimal({precision}, {scale})"
     elif t is bool:
         data_type = "Boolean"
     elif t is datetime:
-        precision = next((md for md in mds if isinstance(md, ClickhousePrecision)), None)
+        precision = next(
+            (md for md in mds if isinstance(md, ClickhousePrecision)), None
+        )
         if precision is None:
             data_type = "DateTime"
         else:
@@ -213,39 +502,114 @@ def py_type_to_column_type(t: type, mds: list[Any]) -> Tuple[bool, list[Any], Da
         data_type = "IPv4"
     elif t is ipaddress.IPv6Address:
         data_type = "IPv6"
+    elif any(
+        md
+        in [  # this check has to happen before t is matched against tuple/list
+            "Point",
+            "Ring",
+            "LineString",
+            "MultiLineString",
+            "Polygon",
+            "MultiPolygon",
+        ]
+        for md in mds
+    ):
+        data_type = next(
+            md
+            for md in mds
+            if md
+            in [
+                "Point",
+                "Ring",
+                "LineString",
+                "MultiLineString",
+                "Polygon",
+                "MultiPolygon",
+            ]
+        )
+        _validate_geometry_type(data_type, t)
     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)
+            # Special case: dict[str, Any] should be JSON type (matches TypeScript's Record<string, any>)
+            # This is useful for storing arbitrary extra fields in a JSON column
+            if args[0] is str and args[1] is Any:
+                data_type = "Json"
+            else:
+                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)}")
+            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:
         data_type = "Json"
+    elif any(isinstance(md, ClickHouseJson) for md in mds) and issubclass(t, BaseModel):
+        # Annotated[SomePydanticClass, ClickHouseJson(...)]
+        columns = _to_columns(t)
+        for c in columns:
+            if c.default is not None:
+                raise ValueError(
+                    "Default in inner field. Put ClickHouseDefault in top level field."
+                )
+        # Enforce extra='allow' for JSON-mapped models
+        if t.model_config.get("extra") != "allow":
+            raise ValueError(
+                f"Model {t.__name__} with ClickHouseJson must have model_config with extra='allow'. "
+                "Add: model_config = ConfigDict(extra='allow')"
+            )
+        opts = next(md for md in mds if isinstance(md, ClickHouseJson))
+
+        # Build typed_paths from fields as tuples of (name, type)
+        typed_paths: list[tuple[str, DataType]] = []
+        for c in columns:
+            typed_paths.append((c.name, c.data_type))
+
+        has_any_option = (
+            opts.max_dynamic_paths is not None
+            or opts.max_dynamic_types is not None
+            or len(typed_paths) > 0
+            or len(opts.skip_paths) > 0
+            or len(opts.skip_regexps) > 0
+        )
+
+        if not has_any_option:
+            data_type = "Json"
+        else:
+            data_type = JsonOptions(
+                max_dynamic_paths=opts.max_dynamic_paths,
+                max_dynamic_types=opts.max_dynamic_types,
+                typed_paths=typed_paths,
+                skip_paths=list(opts.skip_paths),
+                skip_regexps=list(opts.skip_regexps),
+            )
     elif get_origin(t) is Literal and all(isinstance(arg, str) for arg in get_args(t)):
         data_type = "String"
         mds.append("LowCardinality")
     elif not isclass(t):
         raise ValueError(f"Unknown type {t}")
     elif issubclass(t, BaseModel):
+        columns = _to_columns(t)
+        for c in columns:
+            if c.default is not None:
+                raise ValueError(
+                    "Default in inner field. Put ClickHouseDefault in top level field."
+                )
         if any(md == "ClickHouseNamedTuple" for md in mds):
             data_type = NamedTupleType(
-                fields=[(
-                    column.name,
-                    column.data_type
-                ) for column in _to_columns(t)],
+                fields=[(column.name, column.data_type) for column in columns],
             )
         else:
             data_type = Nested(
                 name=t.__name__,
-                columns=_to_columns(t),
+                columns=columns,
             )
     elif issubclass(t, Enum):
         values = [EnumValue(name=member.name, value=member.value) for member in t]
@@ -258,26 +622,36 @@ def py_type_to_column_type(t: type, mds: list[Any]) -> Tuple[bool, list[Any], Da
 def _to_columns(model: type[BaseModel]) -> list[Column]:
     """Convert Pydantic model fields to Column definitions."""
     columns = []
+    # Get raw annotations from the model class to preserve type aliases
+    raw_annotations = getattr(model, "__annotations__", {})
+
     for field_name, field_info in model.model_fields.items():
-        # Get the field type annotation
-        field_type = field_info.annotation
+        # Use raw annotation if available (preserves type aliases and their metadata)
+        # Fall back to field_info.annotation if not found in __annotations__
+        field_type = raw_annotations.get(field_name, field_info.annotation)
         if field_type is None:
             raise ValueError(f"Missing type for {field_name}")
         primary_key, field_type = handle_key(field_type)
         is_jwt, field_type = handle_jwt(field_type)
 
-        optional, mds, data_type = py_type_to_column_type(field_type, field_info.metadata)
+        optional, mds, data_type = py_type_to_column_type(
+            field_type, field_info.metadata
+        )
 
         annotations = []
         for md in mds:
-            if isinstance(md, AggregateFunction):
-                annotations.append(
-                    ("aggregationFunction", md.to_dict())
-                )
-            if md == "LowCardinality":
-                annotations.append(
-                    ("LowCardinality", True)
-                )
+            if isinstance(md, AggregateFunction) and all(
+                key != "aggregationFunction" for (key, _) in annotations
+            ):
+                annotations.append(("aggregationFunction", md.to_dict()))
+            if isinstance(md, SimpleAggregateFunction) and all(
+                key != "simpleAggregationFunction" for (key, _) in annotations
+            ):
+                annotations.append(("simpleAggregationFunction", md.to_dict()))
+            if md == "LowCardinality" and all(
+                key != "LowCardinality" for (key, _) in annotations
+            ):
+                annotations.append(("LowCardinality", True))
 
         column_name = field_name if field_info.alias is None else field_info.alias
 
@@ -287,6 +661,31 @@ def _to_columns(model: type[BaseModel]) -> list[Column]:
             None,
         )
 
+        # Extract MATERIALIZED expression from metadata, if provided
+        materialized_expr = next(
+            (md.expression for md in mds if isinstance(md, ClickHouseMaterialized)),
+            None,
+        )
+
+        # Validate mutual exclusivity of DEFAULT and MATERIALIZED
+        if default_expr and materialized_expr:
+            raise ValueError(
+                f"Column '{column_name}' cannot have both DEFAULT and MATERIALIZED. "
+                f"Use one or the other."
+            )
+
+        # Extract TTL expression from metadata, if provided
+        ttl_expr = next(
+            (md.expression for md in mds if isinstance(md, ClickHouseTTL)),
+            None,
+        )
+
+        # Extract CODEC expression from metadata, if provided
+        codec_expr = next(
+            (md.expression for md in mds if isinstance(md, ClickHouseCodec)),
+            None,
+        )
+
         columns.append(
             Column(
                 name=column_name,
@@ -295,7 +694,10 @@ def _to_columns(model: type[BaseModel]) -> list[Column]:
                 unique=False,
                 primary_key=primary_key,
                 default=default_expr,
+                materialized=materialized_expr,
                 annotations=annotations,
+                ttl=ttl_expr,
+                codec=codec_expr,
             )
         )
     return columns
@@ -303,7 +705,9 @@ def _to_columns(model: type[BaseModel]) -> list[Column]:
 
 class StringToEnumMixin:
     @classmethod
-    def __get_pydantic_core_schema__(cls, _source_type: Any, _handler: GetCoreSchemaHandler) -> CoreSchema:
+    def __get_pydantic_core_schema__(
+        cls, _source_type: Any, _handler: GetCoreSchemaHandler
+    ) -> CoreSchema:
         def validate(value: Any, _: Any) -> Any:
             if isinstance(value, str):
                 try:
@@ -312,14 +716,15 @@ class StringToEnumMixin:
                     raise ValueError(f"Invalid enum name: {value}")
             return cls(value)  # fallback to default enum validation
 
-        return core_schema.with_info_before_validator_function(validate, core_schema.enum_schema(cls, list(cls)))
+        return core_schema.with_info_before_validator_function(
+            validate, core_schema.enum_schema(cls, list(cls))
+        )
 
 
 def is_array_nested_type(data_type: DataType) -> bool:
     """Type guard to check if a data type is Array(Nested(...))."""
-    return (
-            isinstance(data_type, ArrayType) and
-            isinstance(data_type.element_type, Nested)
+    return isinstance(data_type, ArrayType) and isinstance(
+        data_type.element_type, Nested
     )