Various sqlmesh and clickhouse updates for concurrency

.sqlfluff

@@ -4,7 +4,7 @@ templater = dbt
 runaway_limit = 10
 max_line_length = 80
 indent_unit = space
-exclude_rules = AL09
+exclude_rules = AL09, RF02
 
 [sqlfluff:indentation]
 tab_space_size = 4

pyproject.toml

@@ -54,7 +54,7 @@ sqlalchemy = "^2.0.25"
 textual = "^0.52.1"
 redis = "^5.0.7"
 githubkit = "^0.11.6"
-sqlmesh = { git = "https://github.com/opensource-observer/sqlmesh.git", extras = [
+sqlmesh = { git = "https://github.com/opensource-observer/sqlmesh.git", rev = "test-clickhouse-engine-updates", extras = [
     "gcppostgres",
 ] }
 dagster-duckdb = "^0.24.0"

warehouse/metrics_mesh/config.py

@@ -14,6 +14,16 @@
 
 dotenv.load_dotenv()
 
+
+def pool_manager_factory(config: ClickhouseConnectionConfig):
+    from clickhouse_connect.driver import httputil
+
+    return httputil.get_pool_manager(
+        num_pools=config.concurrent_tasks,
+        max_size=config.concurrent_tasks,
+    )
+
+
 config = Config(
     model_defaults=ModelDefaultsConfig(dialect="clickhouse", start="2024-08-01"),
     gateways={
@@ -32,8 +42,14 @@
                 password=os.environ.get("SQLMESH_CLICKHOUSE_PASSWORD", ""),
                 port=int(os.environ.get("SQLMESH_CLICKHOUSE_PORT", "443")),
                 concurrent_tasks=int(
-                    os.environ.get("SQLMESH_CLICKHOUSE_CONCURRENT_TASKS", "8")
+                    os.environ.get("SQLMESH_CLICKHOUSE_CONCURRENT_TASKS", "16")
                 ),
+                send_receive_timeout=1800,
+                connection_settings={"allow_nondeterministic_mutations": 1},
+                connection_pool_options={
+                    "maxsize": 24,
+                    "retries": 3,
+                },
             ),
             state_connection=GCPPostgresConnectionConfig(
                 instance_connection_string=os.environ.get(

warehouse/metrics_mesh/macros/time_aggregation_bucket.py

@@ -1,11 +1,11 @@
+from sqlglot import expressions as exp
 from sqlmesh import macro
 from sqlmesh.core.macros import MacroEvaluator
-from sqlglot import expressions as exp
 
 
 @macro()
 def time_aggregation_bucket(
-    evaluator: MacroEvaluator, timeExp: exp.Expression, rollup: str
+    evaluator: MacroEvaluator, time_exp: exp.Expression, rollup: str
 ):
     from sqlmesh.core.dialect import parse_one
 
@@ -25,7 +25,7 @@ def time_aggregation_bucket(
                     unit=exp.Var(this=rollup_to_interval[rollup]),
                 ),
                 exp.Cast(
-                    this=timeExp,
+                    this=time_exp,
                     to=exp.DataType(this=exp.DataType.Type.DATE, nested=False),
                 ),
             ],
@@ -38,6 +38,6 @@ def time_aggregation_bucket(
     return exp.Anonymous(
         this=rollup_to_clickhouse_function[rollup],
         expressions=[
-            timeExp,
+            time_exp,
         ],
     )

warehouse/metrics_mesh/macros/to_unix_timestamp.py

@@ -0,0 +1,28 @@
+from sqlglot import expressions as exp
+from sqlmesh import macro
+from sqlmesh.core.macros import MacroEvaluator
+
+
+@macro()
+def str_to_unix_timestamp(
+    evaluator: MacroEvaluator,
+    time_exp: exp.Expression,
+):
+    from sqlmesh.core.dialect import parse_one
+
+    if evaluator.runtime_stage in ["loading", "creating"]:
+        return parse_one("1::Uint32", dialect="clickhouse")
+
+    if evaluator.engine_adapter.dialect == "duckdb":
+        return exp.TimeToUnix(
+            this=exp.StrToTime(
+                this=time_exp,
+                format=exp.Array(
+                    expressions=[exp.Literal(this="%Y-%m-%d", is_string=True)]
+                ),
+            )
+        )
+    return exp.Anonymous(
+        this="toUnixTimestamp",
+        expressions=[time_exp],
+    )

warehouse/metrics_mesh/models/events_daily_to_artifact.sql

@@ -2,11 +2,12 @@ MODEL (
   name metrics.events_daily_to_artifact,
   kind INCREMENTAL_BY_TIME_RANGE (
     time_column bucket_day,
-    batch_size 30
+    batch_size 90
   ),
   start '2015-01-01',
   cron '@daily',
   dialect 'clickhouse',
+  partitioned_by (event_type, DATE_TRUNC('MONTH', bucket_day)),
   grain (
     bucket_day,
     event_type,
@@ -20,8 +21,20 @@ MODEL (
     event_type String,
     from_artifact_id String,
     to_artifact_id String,
+    _sign Int8,
+    _version UInt32,
     amount Float64
-  )
+  ),
+  physical_properties (
+    ORDER_BY = (
+      event_type,
+      event_source,
+      from_artifact_id,
+      to_artifact_id,
+      bucket_day
+    ),
+  ),
+  storage_format "VersionedCollapsingMergeTree(_sign, _version)",
 );
 WITH events AS (
   SELECT DISTINCT from_artifact_id,
@@ -38,6 +51,8 @@ SELECT from_artifact_id,
   event_source,
   event_type,
   DATE_TRUNC('DAY', time::DATE) AS bucket_day,
+  1 as _sign,
+  @str_to_unix_timestamp(@execution_ds) as _version,
   SUM(amount) AS amount
 FROM events
 GROUP BY from_artifact_id,

warehouse/metrics_tools/lib/factories/factory.py

@@ -125,6 +125,7 @@ def generate_models_from_query(
                     "name": ModelKindName.INCREMENTAL_BY_TIME_RANGE,
                     "time_column": "metrics_sample_date",
                     "batch_size": 1,
+                    "batch_concurrency": 1,
                 },
                 dialect="clickhouse",
                 columns=columns,
@@ -349,3 +350,4 @@ def timeseries_metrics(
 #             raise Exception("no queries generated from the evaluated queries")
 #         top_level_select = top_level_select.with_(f"all_{cte_suffix}", union_cte)
 #         return top_level_select
+#         return top_level_select