moose-lib 0.6.148.dev3442438466__py3-none-any.whl → 0.6.283__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (59) hide show
  1. moose_lib/__init__.py +34 -3
  2. moose_lib/blocks.py +416 -52
  3. moose_lib/clients/redis_client.py +26 -14
  4. moose_lib/commons.py +37 -30
  5. moose_lib/config/config_file.py +5 -1
  6. moose_lib/config/runtime.py +73 -34
  7. moose_lib/data_models.py +331 -61
  8. moose_lib/dmv2/__init__.py +69 -73
  9. moose_lib/dmv2/_registry.py +2 -1
  10. moose_lib/dmv2/_source_capture.py +37 -0
  11. moose_lib/dmv2/consumption.py +55 -32
  12. moose_lib/dmv2/ingest_api.py +9 -2
  13. moose_lib/dmv2/ingest_pipeline.py +35 -16
  14. moose_lib/dmv2/life_cycle.py +3 -1
  15. moose_lib/dmv2/materialized_view.py +24 -14
  16. moose_lib/dmv2/moose_model.py +165 -0
  17. moose_lib/dmv2/olap_table.py +299 -151
  18. moose_lib/dmv2/registry.py +18 -3
  19. moose_lib/dmv2/sql_resource.py +16 -8
  20. moose_lib/dmv2/stream.py +75 -23
  21. moose_lib/dmv2/types.py +14 -8
  22. moose_lib/dmv2/view.py +13 -6
  23. moose_lib/dmv2/web_app.py +11 -6
  24. moose_lib/dmv2/web_app_helpers.py +5 -1
  25. moose_lib/dmv2/workflow.py +37 -9
  26. moose_lib/internal.py +340 -56
  27. moose_lib/main.py +87 -56
  28. moose_lib/query_builder.py +18 -5
  29. moose_lib/query_param.py +54 -20
  30. moose_lib/secrets.py +122 -0
  31. moose_lib/streaming/streaming_function_runner.py +233 -117
  32. moose_lib/utilities/sql.py +0 -1
  33. {moose_lib-0.6.148.dev3442438466.dist-info → moose_lib-0.6.283.dist-info}/METADATA +18 -1
  34. moose_lib-0.6.283.dist-info/RECORD +63 -0
  35. tests/__init__.py +1 -1
  36. tests/conftest.py +6 -5
  37. tests/test_backward_compatibility.py +85 -0
  38. tests/test_cluster_validation.py +85 -0
  39. tests/test_codec.py +75 -0
  40. tests/test_column_formatting.py +80 -0
  41. tests/test_fixedstring.py +43 -0
  42. tests/test_iceberg_config.py +105 -0
  43. tests/test_int_types.py +211 -0
  44. tests/test_kafka_config.py +141 -0
  45. tests/test_materialized.py +74 -0
  46. tests/test_metadata.py +37 -0
  47. tests/test_moose.py +21 -30
  48. tests/test_moose_model.py +153 -0
  49. tests/test_olap_table_moosemodel.py +89 -0
  50. tests/test_olap_table_versioning.py +52 -58
  51. tests/test_query_builder.py +97 -9
  52. tests/test_redis_client.py +10 -3
  53. tests/test_s3queue_config.py +211 -110
  54. tests/test_secrets.py +239 -0
  55. tests/test_simple_aggregate.py +42 -40
  56. tests/test_web_app.py +11 -5
  57. moose_lib-0.6.148.dev3442438466.dist-info/RECORD +0 -47
  58. {moose_lib-0.6.148.dev3442438466.dist-info → moose_lib-0.6.283.dist-info}/WHEEL +0 -0
  59. {moose_lib-0.6.148.dev3442438466.dist-info → moose_lib-0.6.283.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,153 @@
1
+ """Tests for MooseModel base class with column descriptors"""
2
+
3
+ from pydantic import BaseModel
4
+ from moose_lib.dmv2.moose_model import MooseModel
5
+ from moose_lib.data_models import Column
6
+
7
+
8
+ def test_moosemodel_inherits_from_basemodel():
9
+ """MooseModel should be a valid Pydantic BaseModel"""
10
+
11
+ class User(MooseModel):
12
+ user_id: int
13
+ email: str
14
+
15
+ # Should work as normal Pydantic model
16
+ instance = User(user_id=123, email="test@example.com")
17
+ assert instance.user_id == 123
18
+ assert instance.email == "test@example.com"
19
+
20
+
21
+ def test_moosemodel_adds_column_descriptors():
22
+ """MooseModel metaclass should add Column descriptors for each field"""
23
+
24
+ class User(MooseModel):
25
+ user_id: int
26
+ email: str
27
+ age: int
28
+
29
+ # Check Column descriptors exist at class level
30
+ assert hasattr(User, "user_id")
31
+ assert isinstance(User.user_id, Column)
32
+ assert User.user_id.name == "user_id"
33
+
34
+ assert hasattr(User, "email")
35
+ assert isinstance(User.email, Column)
36
+ assert User.email.name == "email"
37
+
38
+ assert hasattr(User, "age")
39
+ assert isinstance(User.age, Column)
40
+ assert User.age.name == "age"
41
+
42
+
43
+ def test_moosemodel_column_format_spec():
44
+ """Column descriptors should support format specs"""
45
+
46
+ class Product(MooseModel):
47
+ product_id: int
48
+ product_name: str
49
+
50
+ # Test format spec
51
+ result = f"{Product.product_id:col}"
52
+ assert result == "`product_id`"
53
+
54
+ result = f"{Product.product_name:c}"
55
+ assert result == "`product_name`"
56
+
57
+
58
+ def test_moosemodel_adds_cols_property():
59
+ """MooseModel should add .cols property for backward compatibility"""
60
+
61
+ class Order(MooseModel):
62
+ order_id: int
63
+ total: float
64
+
65
+ # Check .cols property exists
66
+ assert hasattr(Order, "cols")
67
+ assert hasattr(Order.cols, "order_id")
68
+ assert hasattr(Order.cols, "total")
69
+
70
+ # Verify .cols.field returns Column
71
+ assert isinstance(Order.cols.order_id, Column)
72
+ assert Order.cols.order_id.name == "order_id"
73
+
74
+
75
+ def test_moosemodel_instance_attributes_separate():
76
+ """Instance attributes should be separate from class Column descriptors"""
77
+
78
+ class User(MooseModel):
79
+ user_id: int
80
+ email: str
81
+
82
+ # Class level: Column objects
83
+ assert isinstance(User.user_id, Column)
84
+
85
+ # Instance level: actual values
86
+ instance = User(user_id=456, email="user@test.com")
87
+ assert instance.user_id == 456
88
+ assert isinstance(instance.user_id, int)
89
+ assert instance.email == "user@test.com"
90
+
91
+
92
+ def test_moosemodel_backward_compatible_with_basemodel():
93
+ """MooseModel should be usable wherever BaseModel is expected"""
94
+
95
+ class User(MooseModel):
96
+ user_id: int
97
+ email: str
98
+
99
+ # Check it's a BaseModel subclass
100
+ assert issubclass(User, BaseModel)
101
+
102
+ # Check Pydantic features work
103
+ assert hasattr(User, "model_fields")
104
+ assert hasattr(User, "model_validate")
105
+ assert hasattr(User, "model_dump")
106
+
107
+ instance = User(user_id=789, email="another@test.com")
108
+ dumped = instance.model_dump()
109
+ assert dumped == {"user_id": 789, "email": "another@test.com"}
110
+
111
+
112
+ def test_moosemodel_empty_model():
113
+ """MooseModel should handle models with no fields"""
114
+
115
+ class EmptyModel(MooseModel):
116
+ pass
117
+
118
+ # Should not crash
119
+ instance = EmptyModel()
120
+ assert instance is not None
121
+
122
+
123
+ def test_moosemodel_cols_bracket_access():
124
+ """MooseModel.cols should support bracket notation"""
125
+
126
+ class User(MooseModel):
127
+ user_id: int
128
+ email: str
129
+
130
+ # Bracket access
131
+ col = User.cols["user_id"]
132
+ assert isinstance(col, Column)
133
+ assert col.name == "user_id"
134
+
135
+ col2 = User.cols["email"]
136
+ assert col2.name == "email"
137
+
138
+
139
+ def test_moosemodel_in_sql_fstring():
140
+ """MooseModel columns should work in SQL f-strings"""
141
+
142
+ class Analytics(MooseModel):
143
+ event_id: int
144
+ timestamp: str
145
+ value: float
146
+
147
+ # Test complete SQL construction
148
+ query = f"SELECT {Analytics.event_id:col}, {Analytics.timestamp:col}, {Analytics.value:col} FROM analytics WHERE {Analytics.event_id:col} > 100"
149
+
150
+ expected = (
151
+ "SELECT `event_id`, `timestamp`, `value` FROM analytics WHERE `event_id` > 100"
152
+ )
153
+ assert query == expected
@@ -0,0 +1,89 @@
1
+ """Tests for OlapTable with MooseModel integration"""
2
+
3
+ from moose_lib.dmv2 import OlapTable, OlapConfig, MooseModel
4
+ from moose_lib.data_models import Column
5
+
6
+
7
+ def test_olaptable_works_with_moosemodel():
8
+ """OlapTable should accept MooseModel types"""
9
+
10
+ class User(MooseModel):
11
+ user_id: int
12
+ email: str
13
+
14
+ table = OlapTable[User]("users", OlapConfig())
15
+
16
+ assert table.name == "users"
17
+ assert table.model_type == User
18
+
19
+
20
+ def test_olaptable_moosemodel_direct_column_access():
21
+ """OlapTable with MooseModel should enable direct column access via model"""
22
+
23
+ class Product(MooseModel):
24
+ product_id: int
25
+ name: str
26
+ price: float
27
+
28
+ table = OlapTable[Product]("products")
29
+
30
+ # Access columns through the model class
31
+ assert isinstance(Product.product_id, Column)
32
+ assert Product.product_id.name == "product_id"
33
+
34
+ # Should work in f-strings
35
+ query = f"SELECT {Product.product_id:col}, {Product.name:col} FROM {table.name}"
36
+ assert query == "SELECT `product_id`, `name` FROM products"
37
+
38
+
39
+ def test_olaptable_moosemodel_cols_backward_compat():
40
+ """OlapTable with MooseModel should maintain .cols backward compatibility"""
41
+
42
+ class Order(MooseModel):
43
+ order_id: int
44
+ total: float
45
+
46
+ table = OlapTable[Order]("orders")
47
+
48
+ # OLD pattern still works
49
+ assert hasattr(Order, "cols")
50
+ assert isinstance(Order.cols.order_id, Column)
51
+
52
+ # Can use in queries
53
+ query = f"SELECT {Order.cols.order_id} FROM orders"
54
+ assert "`order_id`" in query
55
+
56
+
57
+ def test_olaptable_with_basemodel_still_works():
58
+ """OlapTable should still work with regular BaseModel (backward compat)"""
59
+
60
+ from pydantic import BaseModel
61
+
62
+ class LegacyModel(BaseModel):
63
+ legacy_id: int
64
+ legacy_name: str
65
+
66
+ # Should not crash
67
+ table = OlapTable[LegacyModel]("legacy")
68
+
69
+ # Old .cols pattern should still work
70
+ assert hasattr(table, "cols")
71
+
72
+ # Note: LegacyModel.legacy_id won't be a Column (no metaclass)
73
+ # This is expected - only MooseModel gets the new feature
74
+
75
+
76
+ def test_olaptable_model_property():
77
+ """OlapTable should provide access to the model class"""
78
+
79
+ class Analytics(MooseModel):
80
+ event_id: int
81
+ timestamp: str
82
+
83
+ table = OlapTable[Analytics]("analytics")
84
+
85
+ # Should be able to access model type
86
+ assert table.model_type == Analytics
87
+
88
+ # Can use for column access
89
+ assert isinstance(table.model_type.event_id, Column)
@@ -6,7 +6,13 @@ can coexist and that the infrastructure map generation handles versioned keys co
6
6
  """
7
7
 
8
8
  import pytest
9
- from moose_lib import OlapTable, OlapConfig, ClickHouseEngines, MergeTreeEngine, ReplacingMergeTreeEngine
9
+ from moose_lib import (
10
+ OlapTable,
11
+ OlapConfig,
12
+ ClickHouseEngines,
13
+ MergeTreeEngine,
14
+ ReplacingMergeTreeEngine,
15
+ )
10
16
  from moose_lib.dmv2.registry import get_tables
11
17
  from moose_lib.internal import to_infra_map
12
18
  from pydantic import BaseModel
@@ -15,6 +21,7 @@ from typing import Optional
15
21
 
16
22
  class UserEvent(BaseModel):
17
23
  """Sample model for testing OlapTable versioning."""
24
+
18
25
  user_id: str
19
26
  event_type: str
20
27
  timestamp: float
@@ -23,6 +30,7 @@ class UserEvent(BaseModel):
23
30
 
24
31
  class UserEventV2(BaseModel):
25
32
  """Updated model with additional fields for version testing."""
33
+
26
34
  user_id: str
27
35
  event_type: str
28
36
  timestamp: float
@@ -39,29 +47,29 @@ def test_multiple_olap_table_versions_can_coexist():
39
47
  OlapConfig(
40
48
  version="1.0",
41
49
  engine=MergeTreeEngine(),
42
- order_by_fields=["user_id", "timestamp"]
43
- )
50
+ order_by_fields=["user_id", "timestamp"],
51
+ ),
44
52
  )
45
-
53
+
46
54
  # Create version 2.0 of the table with different configuration
47
55
  table_v2 = OlapTable[UserEventV2](
48
56
  "UserEvents",
49
57
  OlapConfig(
50
58
  version="2.0",
51
59
  engine=ReplacingMergeTreeEngine(),
52
- order_by_fields=["user_id", "timestamp", "session_id"]
53
- )
60
+ order_by_fields=["user_id", "timestamp", "session_id"],
61
+ ),
54
62
  )
55
-
63
+
56
64
  # Both tables should be registered successfully
57
65
  tables = get_tables()
58
66
  assert "UserEvents_1.0" in tables
59
67
  assert "UserEvents_2.0" in tables
60
-
68
+
61
69
  # Verify they are different instances
62
70
  assert tables["UserEvents_1.0"] is table_v1
63
71
  assert tables["UserEvents_2.0"] is table_v2
64
-
72
+
65
73
  # Verify configurations are different
66
74
  assert table_v1.config.version == "1.0"
67
75
  assert table_v2.config.version == "2.0"
@@ -73,24 +81,19 @@ def test_unversioned_and_versioned_tables_can_coexist():
73
81
  """Test that unversioned and versioned tables with the same name can coexist."""
74
82
  # Create unversioned table
75
83
  unversioned_table = OlapTable[UserEvent](
76
- "EventData",
77
- OlapConfig(engine=MergeTreeEngine())
84
+ "EventData", OlapConfig(engine=MergeTreeEngine())
78
85
  )
79
-
86
+
80
87
  # Create versioned table with same name
81
88
  versioned_table = OlapTable[UserEvent](
82
- "EventData",
83
- OlapConfig(
84
- version="1.5",
85
- engine=MergeTreeEngine()
86
- )
89
+ "EventData", OlapConfig(version="1.5", engine=MergeTreeEngine())
87
90
  )
88
-
91
+
89
92
  # Both should be registered
90
93
  tables = get_tables()
91
94
  assert "EventData" in tables # Unversioned
92
95
  assert "EventData_1.5" in tables # Versioned
93
-
96
+
94
97
  assert tables["EventData"] is unversioned_table
95
98
  assert tables["EventData_1.5"] is versioned_table
96
99
 
@@ -99,15 +102,16 @@ def test_duplicate_version_registration_fails():
99
102
  """Test that registering the same table name and version twice fails."""
100
103
  # Create first table
101
104
  OlapTable[UserEvent](
102
- "DuplicateTest",
103
- OlapConfig(version="1.0", engine=MergeTreeEngine())
105
+ "DuplicateTest", OlapConfig(version="1.0", engine=MergeTreeEngine())
104
106
  )
105
-
107
+
106
108
  # Attempting to create another table with same name and version should fail
107
- with pytest.raises(ValueError, match="OlapTable with name DuplicateTest and version 1.0 already exists"):
109
+ with pytest.raises(
110
+ ValueError,
111
+ match="OlapTable with name DuplicateTest and version 1.0 already exists",
112
+ ):
108
113
  OlapTable[UserEvent](
109
- "DuplicateTest",
110
- OlapConfig(version="1.0", engine=MergeTreeEngine())
114
+ "DuplicateTest", OlapConfig(version="1.0", engine=MergeTreeEngine())
111
115
  )
112
116
 
113
117
 
@@ -117,48 +121,45 @@ def test_infrastructure_map_uses_versioned_keys():
117
121
  table_v1 = OlapTable[UserEvent](
118
122
  "InfraMapTest",
119
123
  OlapConfig(
120
- version="1.0",
121
- engine=MergeTreeEngine(),
122
- order_by_fields=["user_id"]
123
- )
124
+ version="1.0", engine=MergeTreeEngine(), order_by_fields=["user_id"]
125
+ ),
124
126
  )
125
-
127
+
126
128
  table_v2 = OlapTable[UserEvent](
127
129
  "InfraMapTest",
128
130
  OlapConfig(
129
- version="2.0",
131
+ version="2.0",
130
132
  engine=ReplacingMergeTreeEngine(),
131
- order_by_fields=["user_id", "timestamp"]
132
- )
133
+ order_by_fields=["user_id", "timestamp"],
134
+ ),
133
135
  )
134
-
136
+
135
137
  unversioned_table = OlapTable[UserEvent](
136
- "UnversionedInfraTest",
137
- OlapConfig(engine=MergeTreeEngine())
138
+ "UnversionedInfraTest", OlapConfig(engine=MergeTreeEngine())
138
139
  )
139
-
140
+
140
141
  # Generate infrastructure map
141
142
  tables_registry = get_tables()
142
143
  infra_map = to_infra_map()
143
-
144
+
144
145
  # Verify versioned keys are used in infrastructure map
145
146
  assert "InfraMapTest_1.0" in infra_map["tables"]
146
147
  assert "InfraMapTest_2.0" in infra_map["tables"]
147
148
  assert "UnversionedInfraTest" in infra_map["tables"]
148
-
149
+
149
150
  # Verify table configurations in infra map
150
151
  v1_config = infra_map["tables"]["InfraMapTest_1.0"]
151
152
  v2_config = infra_map["tables"]["InfraMapTest_2.0"]
152
153
  unversioned_config = infra_map["tables"]["UnversionedInfraTest"]
153
-
154
+
154
155
  assert v1_config["name"] == "InfraMapTest"
155
156
  assert v1_config["version"] == "1.0"
156
157
  assert v1_config["engineConfig"]["engine"] == "MergeTree"
157
-
158
+
158
159
  assert v2_config["name"] == "InfraMapTest"
159
160
  assert v2_config["version"] == "2.0"
160
161
  assert v2_config["engineConfig"]["engine"] == "ReplacingMergeTree"
161
-
162
+
162
163
  assert unversioned_config["name"] == "UnversionedInfraTest"
163
164
  assert unversioned_config.get("version") is None
164
165
 
@@ -167,22 +168,18 @@ def test_version_with_dots_handled_correctly():
167
168
  """Test that versions with dots are handled correctly in keys."""
168
169
  # Create table with semantic version
169
170
  table = OlapTable[UserEvent](
170
- "SemanticVersionTest",
171
- OlapConfig(
172
- version="1.2.3",
173
- engine=MergeTreeEngine()
174
- )
171
+ "SemanticVersionTest", OlapConfig(version="1.2.3", engine=MergeTreeEngine())
175
172
  )
176
-
173
+
177
174
  # Should be registered with version in key
178
175
  tables = get_tables()
179
176
  assert "SemanticVersionTest_1.2.3" in tables
180
177
  assert tables["SemanticVersionTest_1.2.3"] is table
181
-
178
+
182
179
  # Verify in infrastructure map
183
180
  infra_map = to_infra_map()
184
181
  assert "SemanticVersionTest_1.2.3" in infra_map["tables"]
185
-
182
+
186
183
  table_config = infra_map["tables"]["SemanticVersionTest_1.2.3"]
187
184
  assert table_config["version"] == "1.2.3"
188
185
 
@@ -192,25 +189,22 @@ def test_backward_compatibility_with_legacy_engines():
192
189
  # Create table with legacy enum engine (should show deprecation warning)
193
190
  table = OlapTable[UserEvent](
194
191
  "LegacyEngineTest",
195
- OlapConfig(
196
- version="1.0",
197
- engine=ClickHouseEngines.ReplacingMergeTree
198
- )
192
+ OlapConfig(version="1.0", engine=ClickHouseEngines.ReplacingMergeTree),
199
193
  )
200
-
194
+
201
195
  # Should still be registered correctly
202
196
  tables = get_tables()
203
197
  assert "LegacyEngineTest_1.0" in tables
204
198
  assert tables["LegacyEngineTest_1.0"] is table
205
-
199
+
206
200
  # Should work in infrastructure map
207
201
  infra_map = to_infra_map()
208
202
  assert "LegacyEngineTest_1.0" in infra_map["tables"]
209
-
203
+
210
204
  table_config = infra_map["tables"]["LegacyEngineTest_1.0"]
211
205
  assert table_config["version"] == "1.0"
212
206
  assert table_config["engineConfig"]["engine"] == "ReplacingMergeTree"
213
207
 
214
208
 
215
209
  if __name__ == "__main__":
216
- pytest.main([__file__, "-v"])
210
+ pytest.main([__file__, "-v"])
@@ -1,7 +1,7 @@
1
1
  from datetime import datetime
2
2
 
3
3
  from moose_lib.query_builder import Query, col
4
- from moose_lib.dmv2 import IngestPipeline, IngestPipelineConfig
4
+ from moose_lib.dmv2 import IngestPipeline, IngestPipelineConfig, OlapTable, OlapConfig
5
5
  from pydantic import BaseModel
6
6
  from moose_lib.data_models import Key
7
7
 
@@ -14,15 +14,19 @@ class Bar(BaseModel):
14
14
 
15
15
 
16
16
  def test_simple_select_and_where():
17
- bar_model = IngestPipeline[Bar]("Bar", IngestPipelineConfig(
18
- ingest=False,
19
- stream=True,
20
- table=True,
21
- dead_letter_queue=True
22
- ))
17
+ bar_model = IngestPipeline[Bar](
18
+ "Bar",
19
+ IngestPipelineConfig(
20
+ ingest=False, stream=True, table=True, dead_letter_queue=True
21
+ ),
22
+ )
23
23
  bar_cols = bar_model.get_table().cols
24
24
 
25
- q1 = Query().from_(bar_model.get_table()).select(bar_cols.has_text, bar_cols.text_length)
25
+ q1 = (
26
+ Query()
27
+ .from_(bar_model.get_table())
28
+ .select(bar_cols.has_text, bar_cols.text_length)
29
+ )
26
30
  assert q1.to_sql() == 'SELECT "Bar"."has_text", "Bar"."text_length" FROM Bar'
27
31
 
28
32
  q2 = (
@@ -32,7 +36,91 @@ def test_simple_select_and_where():
32
36
  .where(col(bar_cols.has_text).eq(True))
33
37
  )
34
38
  sql, params = q2.to_sql_and_params()
35
- assert sql == 'SELECT "Bar"."has_text", "Bar"."text_length" FROM Bar WHERE "Bar"."has_text" = {p0: Bool}'
39
+ assert (
40
+ sql
41
+ == 'SELECT "Bar"."has_text", "Bar"."text_length" FROM Bar WHERE "Bar"."has_text" = {p0: Bool}'
42
+ )
36
43
  assert params == {"p0": True}
37
44
 
38
45
 
46
+ def test_table_with_database_config():
47
+ """Test that tables with database config generate correct SQL with two identifiers"""
48
+
49
+ class TestModel(BaseModel):
50
+ id: int
51
+ name: str
52
+
53
+ # Table without database
54
+ table_without_db = OlapTable[TestModel]("my_table_no_db", OlapConfig())
55
+
56
+ # Table with database
57
+ table_with_db = OlapTable[TestModel](
58
+ "my_table_with_db", OlapConfig(database="my_database")
59
+ )
60
+
61
+ # Test Query builder with table that has database
62
+ q1 = (
63
+ Query()
64
+ .from_(table_with_db)
65
+ .select(table_with_db.cols.id, table_with_db.cols.name)
66
+ )
67
+ sql1 = q1.to_sql()
68
+ # The Query builder should handle the database-qualified table reference
69
+ assert "my_database" in sql1 or "my_table" in sql1
70
+
71
+ # Test string interpolation format for QueryClient.execute()
72
+ # When a table with database is used, it should generate two separate Identifier parameters
73
+ from string import Formatter
74
+
75
+ # Simulate what happens in QueryClient.execute() with a table that has database
76
+ template = "SELECT * FROM {table}"
77
+ variables = {"table": table_with_db}
78
+
79
+ params = {}
80
+ values = {}
81
+ i = 0
82
+
83
+ for _, variable_name, _, _ in Formatter().parse(template):
84
+ if variable_name:
85
+ value = variables[variable_name]
86
+ if isinstance(value, OlapTable) and value.config.database:
87
+ # Should use two separate Identifier parameters
88
+ params[variable_name] = f"{{p{i}: Identifier}}.{{p{i + 1}: Identifier}}"
89
+ values[f"p{i}"] = value.config.database
90
+ values[f"p{i + 1}"] = value.name
91
+ i += 2
92
+ else:
93
+ params[variable_name] = f"{{p{i}: Identifier}}"
94
+ values[f"p{i}"] = value.name
95
+ i += 1
96
+
97
+ clickhouse_query = template.format_map(params)
98
+
99
+ assert clickhouse_query == "SELECT * FROM {p0: Identifier}.{p1: Identifier}"
100
+ assert values == {"p0": "my_database", "p1": "my_table_with_db"}
101
+
102
+ # Test with table without database
103
+ variables_no_db = {"table": table_without_db}
104
+ params_no_db = {}
105
+ values_no_db = {}
106
+ i = 0
107
+
108
+ for _, variable_name, _, _ in Formatter().parse(template):
109
+ if variable_name:
110
+ value = variables_no_db[variable_name]
111
+ if isinstance(value, OlapTable) and value.config.database:
112
+ params_no_db[variable_name] = (
113
+ f"{{p{i}: Identifier}}.{{p{i + 1}: Identifier}}"
114
+ )
115
+ values_no_db[f"p{i}"] = value.config.database
116
+ values_no_db[f"p{i + 1}"] = value.name
117
+ i += 2
118
+ else:
119
+ params_no_db[variable_name] = f"{{p{i}: Identifier}}"
120
+ values_no_db[f"p{i}"] = value.name
121
+ i += 1
122
+
123
+ clickhouse_query_no_db = template.format_map(params_no_db)
124
+
125
+ assert clickhouse_query_no_db == "SELECT * FROM {p0: Identifier}"
126
+ assert values_no_db == {"p0": "my_table_no_db"}
@@ -6,10 +6,12 @@ import pytest
6
6
  from pydantic import BaseModel
7
7
  from moose_lib import MooseCache
8
8
 
9
+
9
10
  class Config(BaseModel):
10
11
  baz: int
11
12
  qux: bool
12
13
 
14
+
13
15
  @pytest.mark.integration
14
16
  def test_cache_strings():
15
17
  cache = MooseCache()
@@ -26,6 +28,7 @@ def test_cache_strings():
26
28
  # Clean up
27
29
  cache.clear_keys("test")
28
30
 
31
+
29
32
  @pytest.mark.integration
30
33
  def test_cache_pydantic():
31
34
  cache = MooseCache()
@@ -47,6 +50,7 @@ def test_cache_pydantic():
47
50
  # Clean up
48
51
  cache.clear_keys("test")
49
52
 
53
+
50
54
  @pytest.mark.integration
51
55
  def test_cache_ttl():
52
56
  cache = MooseCache()
@@ -66,6 +70,7 @@ def test_cache_ttl():
66
70
  # Clean up
67
71
  cache.clear_keys("test")
68
72
 
73
+
69
74
  @pytest.mark.integration
70
75
  def test_cache_nonexistent():
71
76
  cache = MooseCache()
@@ -75,6 +80,7 @@ def test_cache_nonexistent():
75
80
  assert cache.get("nonexistent", str) is None
76
81
  assert cache.get("nonexistent", Config) is None
77
82
 
83
+
78
84
  @pytest.mark.integration
79
85
  def test_cache_invalid_type():
80
86
  cache = MooseCache()
@@ -86,6 +92,7 @@ def test_cache_invalid_type():
86
92
  with pytest.raises(TypeError):
87
93
  cache.get("test", dict)
88
94
 
95
+
89
96
  @pytest.mark.integration
90
97
  def test_atexit_cleanup():
91
98
  # Create a test script that will be run in a separate process
@@ -107,9 +114,9 @@ sys.exit(0)
107
114
 
108
115
  try:
109
116
  # Run the script and capture output
110
- result = subprocess.run([sys.executable, "test_atexit.py"],
111
- capture_output=True,
112
- text=True)
117
+ result = subprocess.run(
118
+ [sys.executable, "test_atexit.py"], capture_output=True, text=True
119
+ )
113
120
 
114
121
  # Check if we see both the connection and disconnection messages
115
122
  output = result.stdout + result.stderr