Skip to content

Commit 2ebc7a0

Browse files
author
0xhanh
committed
Fix ClickHouse cluster CTAS execution
1 parent e6ac0e5 commit 2ebc7a0

2 files changed

Lines changed: 34 additions & 4 deletions

File tree

sqlmesh/core/engine_adapter/clickhouse.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -536,7 +536,7 @@ def _create_table(
536536
) -> None:
537537
"""Creates a table in the database.
538538
539-
Clickhouse Cloud requires doing CTAS in two steps.
539+
ClickHouse Cloud and cluster modes require doing CTAS in two steps.
540540
541541
First, we add the `EMPTY` property to the CTAS call to create a table with the proper
542542
schema, then insert the data with the CTAS query.
@@ -567,16 +567,19 @@ def _create_table(
567567
table_description,
568568
column_descriptions,
569569
table_kind,
570-
empty_ctas=(self.engine_run_mode.is_cloud and expression is not None),
570+
empty_ctas=(
571+
(self.engine_run_mode.is_cloud or self.engine_run_mode.is_cluster)
572+
and expression is not None
573+
),
571574
track_rows_processed=track_rows_processed,
572575
**kwargs,
573576
)
574577

575-
# execute the second INSERT step if on cloud and creating a table
578+
# execute the second INSERT step if on cloud or cluster and creating a table
576579
# - Additional clause is to avoid clickhouse-connect HTTP client bug where CTAS LIMIT 0
577580
# returns a success code but malformed response
578581
if (
579-
self.engine_run_mode.is_cloud
582+
(self.engine_run_mode.is_cloud or self.engine_run_mode.is_cluster)
580583
and table_kind != "VIEW"
581584
and expression
582585
and not (

tests/core/engine_adapter/test_clickhouse.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,33 @@ def test_create_table(adapter: ClickhouseEngineAdapter, mocker):
101101
]
102102

103103

104+
def test_ctas_in_cluster_mode_uses_empty_ctas_then_single_insert(
105+
adapter: ClickhouseEngineAdapter, mocker: MockerFixture
106+
):
107+
"""A distributed CREATE must not execute its SELECT once per cluster node."""
108+
mocker.patch.object(
109+
ClickhouseEngineAdapter,
110+
"cluster",
111+
new_callable=mocker.PropertyMock(return_value="default"),
112+
)
113+
mocker.patch.object(
114+
ClickhouseEngineAdapter,
115+
"engine_run_mode",
116+
new_callable=mocker.PropertyMock(return_value=EngineRunMode.CLUSTER),
117+
)
118+
119+
adapter.ctas(
120+
"foo",
121+
parse_one("SELECT 1 AS a", dialect=adapter.dialect),
122+
{"a": exp.DataType.build("Int8", dialect=adapter.dialect)},
123+
)
124+
125+
create_sql, insert_sql = to_sql_calls(adapter)
126+
assert 'ON CLUSTER "default"' in create_sql
127+
assert " EMPTY AS SELECT" in create_sql
128+
assert insert_sql.startswith('INSERT INTO "foo" ("a") SELECT')
129+
130+
104131
def test_rename_table(adapter: ClickhouseEngineAdapter, mocker):
105132
mocker.patch.object(
106133
ClickhouseEngineAdapter,

0 commit comments

Comments
 (0)