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

tests/test_int_types.py

@@ -0,0 +1,211 @@
+from datetime import datetime
+from typing import Annotated
+from pydantic import BaseModel
+from moose_lib import (
+    Key,
+    Int8,
+    Int16,
+    Int32,
+    Int64,
+    UInt8,
+    UInt16,
+    UInt32,
+    UInt64,
+    Float32,
+    Float64,
+)
+from moose_lib.data_models import _to_columns
+
+
+def test_integer_type_aliases():
+    """Test that integer type aliases convert to correct ClickHouse types."""
+
+    class IntegerTypesTest(BaseModel):
+        id: Key[str]
+        created_at: datetime
+        tiny_int: Int8
+        small_int: Int16
+        medium_int: Int32
+        big_int: Int64
+        tiny_uint: UInt8
+        small_uint: UInt16
+        medium_uint: UInt32
+        big_uint: UInt64
+
+    columns = _to_columns(IntegerTypesTest)
+    by_name = {col.name: col for col in columns}
+
+    # Verify signed integer types
+    assert by_name["tiny_int"].data_type == "Int8"
+    assert by_name["small_int"].data_type == "Int16"
+    assert by_name["medium_int"].data_type == "Int32"
+    assert by_name["big_int"].data_type == "Int64"
+
+    # Verify unsigned integer types
+    assert by_name["tiny_uint"].data_type == "UInt8"
+    assert by_name["small_uint"].data_type == "UInt16"
+    assert by_name["medium_uint"].data_type == "UInt32"
+    assert by_name["big_uint"].data_type == "UInt64"
+
+    # Verify other fields still work
+    assert by_name["id"].data_type == "String"
+    assert by_name["created_at"].data_type == "DateTime"
+
+
+def test_float_type_aliases():
+    """Test that float type aliases convert to correct ClickHouse types."""
+
+    class FloatTypesTest(BaseModel):
+        id: Key[str]
+        precision_float: Float32
+        double_precision_float: Float64
+
+    columns = _to_columns(FloatTypesTest)
+    by_name = {col.name: col for col in columns}
+
+    assert by_name["precision_float"].data_type == "Float32"
+    assert by_name["double_precision_float"].data_type == "Float64"
+    assert by_name["id"].data_type == "String"
+
+
+def test_combined_numeric_types():
+    """Test combining integer and float types in a single model."""
+
+    class NumericTypesTest(BaseModel):
+        user_id: UInt64
+        age: UInt8
+        score: Int32
+        latitude: Float64
+        longitude: Float64
+        precision_value: Float32
+
+    columns = _to_columns(NumericTypesTest)
+    by_name = {col.name: col for col in columns}
+
+    assert by_name["user_id"].data_type == "UInt64"
+    assert by_name["age"].data_type == "UInt8"
+    assert by_name["score"].data_type == "Int32"
+    assert by_name["latitude"].data_type == "Float64"
+    assert by_name["longitude"].data_type == "Float64"
+    assert by_name["precision_value"].data_type == "Float32"
+
+
+def test_integer_types_as_keys():
+    """Test that integer types can be used as primary keys."""
+
+    class IntegerKeyTest(BaseModel):
+        user_id: Key[UInt64]
+        event_id: Key[Int64]
+        name: str
+
+    columns = _to_columns(IntegerKeyTest)
+    by_name = {col.name: col for col in columns}
+
+    assert by_name["user_id"].data_type == "UInt64"
+    assert by_name["user_id"].primary_key is True
+    assert by_name["event_id"].data_type == "Int64"
+    assert by_name["event_id"].primary_key is True
+    assert by_name["name"].data_type == "String"
+
+
+def test_optional_integer_types():
+    """Test that optional integer types work correctly."""
+    from typing import Optional
+
+    class OptionalIntTest(BaseModel):
+        required_count: UInt32
+        optional_count: Optional[UInt32]
+        optional_score: Optional[Int16]
+
+    columns = _to_columns(OptionalIntTest)
+    by_name = {col.name: col for col in columns}
+
+    assert by_name["required_count"].data_type == "UInt32"
+    assert by_name["required_count"].required is True
+
+    assert by_name["optional_count"].data_type == "UInt32"
+    assert by_name["optional_count"].required is False
+
+    assert by_name["optional_score"].data_type == "Int16"
+    assert by_name["optional_score"].required is False
+
+
+def test_uint_common_use_cases():
+    """Test common use cases for unsigned integers."""
+
+    class CommonUIntUseCases(BaseModel):
+        # User/entity IDs (always positive)
+        user_id: UInt64
+        # Counters (always positive or zero)
+        page_views: UInt32
+        click_count: UInt32
+        # Small enums/flags (0-255)
+        status_code: UInt8
+        # Port numbers (0-65535)
+        port: UInt16
+        # Timestamps as unix epoch
+        timestamp: UInt64
+
+    columns = _to_columns(CommonUIntUseCases)
+    by_name = {col.name: col for col in columns}
+
+    assert by_name["user_id"].data_type == "UInt64"
+    assert by_name["page_views"].data_type == "UInt32"
+    assert by_name["click_count"].data_type == "UInt32"
+    assert by_name["status_code"].data_type == "UInt8"
+    assert by_name["port"].data_type == "UInt16"
+    assert by_name["timestamp"].data_type == "UInt64"
+
+
+def test_int_common_use_cases():
+    """Test common use cases for signed integers."""
+
+    class CommonIntUseCases(BaseModel):
+        # Temperature (can be negative)
+        temperature: Int16
+        # Financial amounts (can be negative for debits)
+        balance: Int64
+        # Deltas/differences
+        delta: Int32
+        # Small range values
+        offset: Int8
+
+    columns = _to_columns(CommonIntUseCases)
+    by_name = {col.name: col for col in columns}
+
+    assert by_name["temperature"].data_type == "Int16"
+    assert by_name["balance"].data_type == "Int64"
+    assert by_name["delta"].data_type == "Int32"
+    assert by_name["offset"].data_type == "Int8"
+
+
+def test_default_int_still_works():
+    """Test that plain int without type annotation still works as before."""
+
+    class PlainIntTest(BaseModel):
+        plain_int: int
+        typed_int: UInt32
+
+    columns = _to_columns(PlainIntTest)
+    by_name = {col.name: col for col in columns}
+
+    # Plain int should still map to "Int" (default behavior)
+    assert by_name["plain_int"].data_type == "Int64"
+    # Typed int should map to specific type
+    assert by_name["typed_int"].data_type == "UInt32"
+
+
+def test_default_float_still_works():
+    """Test that plain float without type annotation still works as before."""
+
+    class PlainFloatTest(BaseModel):
+        plain_float: float
+        typed_float: Float32
+
+    columns = _to_columns(PlainFloatTest)
+    by_name = {col.name: col for col in columns}
+
+    # Plain float should still map to "Float64" (default behavior)
+    assert by_name["plain_float"].data_type == "Float64"
+    # Typed float should map to specific type
+    assert by_name["typed_float"].data_type == "Float32"

tests/test_kafka_config.py

@@ -0,0 +1,141 @@
+"""Tests for Kafka engine configuration."""
+
+import pytest
+from moose_lib import OlapTable, OlapConfig
+from moose_lib.blocks import ClickHouseEngines, KafkaEngine
+from pydantic import BaseModel
+
+
+class SampleEvent(BaseModel):
+    event_id: str
+    user_id: str
+    timestamp: str
+
+
+def test_kafka_engine_required_fields():
+    engine = KafkaEngine(
+        broker_list="kafka:9092",
+        topic_list="events",
+        group_name="moose_consumer",
+        format="JSONEachRow",
+    )
+    assert engine.broker_list == "kafka:9092"
+    assert engine.topic_list == "events"
+    assert engine.group_name == "moose_consumer"
+    assert engine.format == "JSONEachRow"
+
+
+def test_kafka_engine_missing_broker_list():
+    with pytest.raises(ValueError, match="Kafka engine requires 'broker_list'"):
+        KafkaEngine(
+            broker_list="", topic_list="events", group_name="c", format="JSONEachRow"
+        )
+
+
+def test_kafka_engine_missing_topic_list():
+    with pytest.raises(ValueError, match="Kafka engine requires 'topic_list'"):
+        KafkaEngine(
+            broker_list="kafka:9092",
+            topic_list="",
+            group_name="c",
+            format="JSONEachRow",
+        )
+
+
+def test_kafka_engine_missing_group_name():
+    with pytest.raises(ValueError, match="Kafka engine requires 'group_name'"):
+        KafkaEngine(
+            broker_list="kafka:9092",
+            topic_list="events",
+            group_name="",
+            format="JSONEachRow",
+        )
+
+
+def test_kafka_engine_missing_format():
+    with pytest.raises(ValueError, match="Kafka engine requires 'format'"):
+        KafkaEngine(
+            broker_list="kafka:9092", topic_list="events", group_name="c", format=""
+        )
+
+
+def test_kafka_engine_rejects_order_by():
+    with pytest.raises(ValueError, match="KafkaEngine does not support ORDER BY"):
+        OlapTable[SampleEvent](
+            "kafka_table",
+            OlapConfig(
+                engine=KafkaEngine(
+                    broker_list="kafka:9092",
+                    topic_list="events",
+                    group_name="c",
+                    format="JSONEachRow",
+                ),
+                order_by_fields=["event_id"],
+            ),
+        )
+
+
+def test_kafka_engine_rejects_partition_by():
+    with pytest.raises(ValueError, match="KafkaEngine does not support PARTITION BY"):
+        OlapTable[SampleEvent](
+            "kafka_table",
+            OlapConfig(
+                engine=KafkaEngine(
+                    broker_list="kafka:9092",
+                    topic_list="events",
+                    group_name="c",
+                    format="JSONEachRow",
+                ),
+                partition_by="toYYYYMM(timestamp)",
+            ),
+        )
+
+
+def test_kafka_engine_rejects_sample_by():
+    with pytest.raises(ValueError, match="KafkaEngine does not support SAMPLE BY"):
+        OlapTable[SampleEvent](
+            "kafka_table",
+            OlapConfig(
+                engine=KafkaEngine(
+                    broker_list="kafka:9092",
+                    topic_list="events",
+                    group_name="c",
+                    format="JSONEachRow",
+                ),
+                sample_by_expression="event_id",
+            ),
+        )
+
+
+def test_kafka_engine_accepts_valid_config():
+    table = OlapTable[SampleEvent](
+        "kafka_table",
+        OlapConfig(
+            engine=KafkaEngine(
+                broker_list="kafka:9092",
+                topic_list="events",
+                group_name="c",
+                format="JSONEachRow",
+            ),
+            settings={"kafka_num_consumers": "2"},
+        ),
+    )
+    assert table.name == "kafka_table"
+    assert isinstance(table.config.engine, KafkaEngine)
+    assert table.config.settings["kafka_num_consumers"] == "2"
+
+
+def test_kafka_engine_serialization():
+    from moose_lib.internal import _convert_engine_instance_to_config_dict
+
+    engine = KafkaEngine(
+        broker_list="kafka-1:9092,kafka-2:9092",
+        topic_list="events,logs",
+        group_name="moose_group",
+        format="JSONEachRow",
+    )
+    config_dict = _convert_engine_instance_to_config_dict(engine)
+
+    assert config_dict.engine == "Kafka"
+    assert config_dict.broker_list == "kafka-1:9092,kafka-2:9092"
+    assert config_dict.topic_list == "events,logs"

tests/test_materialized.py

@@ -0,0 +1,74 @@
+from datetime import datetime, date
+from typing import Annotated, Any
+from pydantic import BaseModel
+from moose_lib import Key, ClickHouseMaterialized, ClickHouseCodec, UInt64
+from moose_lib.data_models import _to_columns
+import pytest
+
+
+def test_materialized_basic():
+    """Test basic MATERIALIZED annotation converts to correct expression."""
+
+    class MaterializedTest(BaseModel):
+        timestamp: datetime
+        event_date: Annotated[date, ClickHouseMaterialized("toDate(timestamp)")]
+
+    columns = _to_columns(MaterializedTest)
+    by_name = {col.name: col for col in columns}
+
+    assert by_name["timestamp"].materialized is None
+    assert by_name["event_date"].materialized == "toDate(timestamp)"
+
+
+def test_materialized_hash():
+    """Test MATERIALIZED with hash function."""
+
+    class HashTest(BaseModel):
+        user_id: str
+        user_hash: Annotated[UInt64, ClickHouseMaterialized("cityHash64(user_id)")]
+
+    columns = _to_columns(HashTest)
+    by_name = {col.name: col for col in columns}
+
+    assert by_name["user_id"].materialized is None
+    assert by_name["user_hash"].materialized == "cityHash64(user_id)"
+
+
+def test_materialized_with_codec():
+    """Test MATERIALIZED combined with CODEC."""
+
+    class MaterializedCodecTest(BaseModel):
+        log_blob: Annotated[Any, ClickHouseCodec("ZSTD(3)")]
+        combination_hash: Annotated[
+            list[UInt64],
+            ClickHouseMaterialized(
+                "arrayMap(kv -> cityHash64(kv.1, kv.2), JSONExtractKeysAndValuesRaw(toString(log_blob)))"
+            ),
+            ClickHouseCodec("ZSTD(1)"),
+        ]
+
+    columns = _to_columns(MaterializedCodecTest)
+    by_name = {col.name: col for col in columns}
+
+    assert by_name["log_blob"].materialized is None
+    assert by_name["log_blob"].codec == "ZSTD(3)"
+    assert (
+        by_name["combination_hash"].materialized
+        == "arrayMap(kv -> cityHash64(kv.1, kv.2), JSONExtractKeysAndValuesRaw(toString(log_blob)))"
+    )
+    assert by_name["combination_hash"].codec == "ZSTD(1)"
+
+
+def test_materialized_mutually_exclusive_with_default():
+    """Test that MATERIALIZED and DEFAULT are mutually exclusive."""
+    from moose_lib import clickhouse_default
+
+    class BadModel(BaseModel):
+        bad_field: Annotated[
+            str,
+            clickhouse_default("'default_value'"),
+            ClickHouseMaterialized("'materialized_value'"),
+        ]
+
+    with pytest.raises(ValueError, match="cannot have both DEFAULT and MATERIALIZED"):
+        _to_columns(BadModel)

tests/test_metadata.py

@@ -0,0 +1,37 @@
+"""Tests for metadata handling in OlapTable."""
+
+from moose_lib.dmv2 import OlapTable, OlapConfig
+from pydantic import BaseModel
+
+
+class SampleModel(BaseModel):
+    """Sample model for testing metadata."""
+
+    id: str
+    name: str
+
+
+def test_respect_user_provided_source():
+    """Test that user-provided source file path is not overwritten."""
+
+    user_provided_path = "custom/path/to/model.py"
+
+    config = OlapConfig(metadata={"source": {"file": user_provided_path}})
+
+    table = OlapTable[SampleModel]("test_user_provided", config=config)
+
+    assert table.metadata is not None
+    assert table.metadata["source"]["file"] == user_provided_path
+
+
+def test_preserve_metadata_with_auto_capture():
+    """Test that user metadata is preserved while auto-capturing source."""
+
+    config = OlapConfig(metadata={"description": "A test table"})
+
+    table = OlapTable[SampleModel]("test_preserve_metadata", config=config)
+
+    assert table.metadata is not None
+    assert isinstance(table.metadata, dict)
+    assert table.metadata["description"] == "A test table"
+    assert "test_metadata.py" in table.metadata["source"]["file"]

tests/test_moose.py

@@ -3,8 +3,13 @@ from dataclasses import dataclass
 from typing import Optional
 from pydantic import BaseModel
 
-from moose_lib.query_param import convert_pydantic_definition, QueryField, ArrayType, map_params_to_class, \
-    convert_dataclass_definition
+from moose_lib.query_param import (
+    convert_pydantic_definition,
+    QueryField,
+    ArrayType,
+    map_params_to_class,
+    convert_dataclass_definition,
+)
 
 
 @dataclass
@@ -24,28 +29,28 @@ class QueryParamPydantic(BaseModel):
 
 query_fields = [
     QueryField(
-        name='optional_field',
-        data_type='Int',
+        name="optional_field",
+        data_type="Int64",
         has_default=False,
         required=False,
     ),
     QueryField(
-        name='date_field',
-        data_type='DateTime',
+        name="date_field",
+        data_type="DateTime",
         has_default=False,
         required=True,
     ),
     QueryField(
-        name='list_field',
+        name="list_field",
         data_type=ArrayType(
-            element_type='String',
+            element_type="String",
         ),
         has_default=False,
         required=True,
     ),
     QueryField(
-        name='int_field',
-        data_type='Int',
+        name="int_field",
+        data_type="Int64",
         has_default=True,
         required=False,
     ),
@@ -58,34 +63,20 @@ parsed_date = datetime.datetime.fromisoformat(datestr)
 def test_pydantic():
     assert convert_pydantic_definition(QueryParamPydantic) == query_fields
     assert map_params_to_class(
-        {
-            "date_field": [datestr],
-            "list_field": ["123"],
-            "int_field": ["1"]
-        },
+        {"date_field": [datestr], "list_field": ["123"], "int_field": ["1"]},
         query_fields,
-        QueryParamPydantic
+        QueryParamPydantic,
     ) == QueryParamPydantic(
-        optional_field=None,
-        date_field=parsed_date,
-        list_field=["123"],
-        int_field=1
+        optional_field=None, date_field=parsed_date, list_field=["123"], int_field=1
     )
 
 
 def test_dataclass():
     assert convert_dataclass_definition(QueryParamDataClass) == query_fields
     assert map_params_to_class(
-        {
-            "date_field": [datestr],
-            "list_field": ["123"],
-            "int_field": ["1"]
-        },
+        {"date_field": [datestr], "list_field": ["123"], "int_field": ["1"]},
         query_fields,
-        QueryParamDataClass
+        QueryParamDataClass,
     ) == QueryParamDataClass(
-        optional_field=None,
-        date_field=parsed_date,
-        list_field=["123"],
-        int_field=1
+        optional_field=None, date_field=parsed_date, list_field=["123"], int_field=1
     )