Skip to content

Commit cf156d0

Browse files
committed
fix(clickhouse): suppress ON CLUSTER for object DDL in a Replicated database
Inside a `Replicated` database Keeper already propagates object DDL, so adding `ON CLUSTER` asks for a second, redundant fan-out and ClickHouse rejects the statement outright: Code: 80. DB::Exception: It's not initial query. ON CLUSTER is not allowed for Replicated database. (INCORRECT_QUERY) The ClickHouse adapter had no way to avoid this: `_on_cluster_sql()` took no argument and keyed only on `engine_run_mode.is_cluster`. The practical effect is total -- with a cluster configured, no table or view can be created in a `Replicated` database at all. Reproduced end to end against a 1x3 cluster on 26.6.3.62: `create_table` into such a database fails with the error above before this change and succeeds after it, propagating to all three replicas under one shared UUID. This cannot be a connection-level flag. One connection can hold both an Atomic and a `Replicated` database, and objects in the Atomic one still need `ON CLUSTER`, so the decision is made per target database by `_should_use_on_cluster()` and threaded through all 14 emission sites. The policy fails open throughout. An unknown target, an unresolvable database, or any introspection failure keeps the previous behaviour, and a single cached probe short-circuits the whole path when the server hosts no `Replicated` database at all -- so this is a no-op, and costs no extra round-trips, for any deployment that does not have one. Two carve-outs are deliberate. `CREATE DATABASE` is how a `Replicated` database comes into existence, so there is no engine to consult, and a database's own DDL log cannot propagate its own removal; both stay cluster-wide unconditionally. `alter_table` decides per expression rather than per call, because one call can carry alters against both kinds of database, and `RENAME` / `EXCHANGE` across a Replicated and a non-Replicated database are refused rather than rendered silently half-correct. `_build_table_properties_exp` and `_build_view_properties_exp` receive no reference to the object being created, so the base class now passes the target table down. The parameter is additive and every override across the nine adapter files already accepts `**kwargs`, so no other adapter changes; Athena already declares and receives the same `table` parameter via its own override. This is a correctness fix independent of managed models: any SQLMesh user running against a `Replicated` database is affected today for ordinary tables and views. Signed-off-by: mday-io <mdaytn@gmail.com>
1 parent 7c1e4e4 commit cf156d0

3 files changed

Lines changed: 453 additions & 17 deletions

File tree

‎sqlmesh/core/engine_adapter/base.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1043,12 +1043,15 @@ def _build_create_table_exp(
10431043
) -> exp.Create:
10441044
exists = False if replace else exists
10451045
catalog_name = None
1046+
target_table: t.Optional[exp.Table] = None
10461047
if not isinstance(table_name_or_schema, exp.Schema):
10471048
table_name_or_schema = exp.to_table(table_name_or_schema)
10481049
catalog_name = table_name_or_schema.catalog
1050+
target_table = table_name_or_schema
10491051
else:
10501052
if isinstance(table_name_or_schema.this, exp.Table):
10511053
catalog_name = table_name_or_schema.this.catalog
1054+
target_table = table_name_or_schema.this
10521055

10531056
properties = (
10541057
self._build_table_properties_exp(
@@ -1057,6 +1060,9 @@ def _build_create_table_exp(
10571060
target_columns_to_types=target_columns_to_types,
10581061
table_description=table_description,
10591062
table_kind=table_kind,
1063+
# Passed so an adapter can vary properties by target object, not just by
1064+
# connection. Additive: every override accepts **kwargs.
1065+
table=target_table,
10601066
)
10611067
if kwargs or table_description
10621068
else None
@@ -1364,6 +1370,8 @@ def create_view(
13641370
else None
13651371
),
13661372
physical_cluster=create_kwargs.pop("physical_cluster", None),
1373+
# See the note on the table-properties call site above.
1374+
table=exp.to_table(view_name),
13671375
)
13681376
if create_view_properties:
13691377
for view_property in create_view_properties.expressions:

‎sqlmesh/core/engine_adapter/clickhouse.py‎

Lines changed: 209 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import typing as t
44
import logging
55
import re
6+
from functools import cached_property
67
from sqlglot import exp, maybe_parse
78
from sqlmesh.core.dialect import to_schema
89
from sqlmesh.core.engine_adapter.mixins import LogicalMergeMixin
@@ -31,6 +32,13 @@
3132

3233
logger = logging.getLogger(__name__)
3334

35+
# `system.databases.engine` value for a Keeper-coordinated database. Object DDL inside one
36+
# must not carry `ON CLUSTER`: the database's own DDL log already propagates it to every
37+
# replica, so ClickHouse rejects the redundant second fan-out outright with code 80,
38+
# `It's not initial query. ON CLUSTER is not allowed for Replicated database.`
39+
# Matched as a prefix so engine variants are covered without another release.
40+
REPLICATED_DATABASE_ENGINE_PREFIX = "Replicated"
41+
3442

3543
class ClickhouseEngineAdapter(EngineAdapterWithIndexSupport, LogicalMergeMixin):
3644
DIALECT = "clickhouse"
@@ -209,9 +217,15 @@ def create_schema(
209217
from sqlmesh.utils.errors import SQLMeshError
210218

211219
properties_copy = properties.copy()
220+
# Always cluster-wide, deliberately. This is the statement that creates the
221+
# database, including a `Replicated` one, so there is no database engine to
222+
# consult yet and nothing for Keeper to propagate through.
212223
if self.engine_run_mode.is_cluster:
213224
properties_copy.append(exp.OnCluster(this=exp.to_identifier(self.cluster)))
214225

226+
# A database that previously resolved as absent may now exist, with an engine.
227+
self._clear_database_engine_cache()
228+
215229
# ClickHouse does not support catalogs. When a virtual catalog has been injected
216230
# (self._default_catalog is set), strip it from the schema name. This mirrors the
217231
# SINGLE_CATALOG_ONLY branch in the set_catalog decorator, which does not apply here
@@ -524,7 +538,8 @@ def _create_table_like(
524538
dialect=self.dialect, identify=True
525539
)
526540
self.execute(
527-
f"CREATE TABLE {target_table_sql}{self._on_cluster_sql()} AS {source_table_sql}"
541+
f"CREATE TABLE {target_table_sql}"
542+
f"{self._on_cluster_sql(target_table_name)} AS {source_table_sql}"
528543
)
529544

530545
def _get_partition_ids(
@@ -688,10 +703,14 @@ def _exchange_tables(
688703
dialect=self.dialect, identify=True
689704
)
690705

706+
on_cluster_sql = (
707+
self._on_cluster_sql(old_table_name)
708+
if self._assert_same_on_cluster_scope("EXCHANGE TABLES", old_table_name, new_table_name)
709+
else ""
710+
)
711+
691712
try:
692-
self.execute(
693-
f"EXCHANGE TABLES {old_table_sql} AND {new_table_sql}{self._on_cluster_sql()}"
694-
)
713+
self.execute(f"EXCHANGE TABLES {old_table_sql} AND {new_table_sql}{on_cluster_sql}")
695714
except DatabaseError as e:
696715
if "NOT_IMPLEMENTED" in str(e):
697716
# If someone is using an old Clickhouse version, an OS that doesn't support atomic exchanges,
@@ -717,11 +736,18 @@ def _rename_table(
717736
dialect=self.dialect, identify=True
718737
)
719738

720-
self.execute(f"RENAME TABLE {old_table_sql} TO {new_table_sql}{self._on_cluster_sql()}")
739+
on_cluster_sql = (
740+
self._on_cluster_sql(old_table_name)
741+
if self._assert_same_on_cluster_scope("RENAME TABLE", old_table_name, new_table_name)
742+
else ""
743+
)
744+
745+
self.execute(f"RENAME TABLE {old_table_sql} TO {new_table_sql}{on_cluster_sql}")
721746

722747
def delete_from(self, table_name: TableName, where: t.Union[str, exp.Expr]) -> None:
723-
delete_expr = exp.delete(self._strip_virtual_catalog(table_name), where)
724-
if self.engine_run_mode.is_cluster:
748+
target_table = self._strip_virtual_catalog(table_name)
749+
delete_expr = exp.delete(target_table, where)
750+
if self._should_use_on_cluster(target_table):
725751
delete_expr.set("cluster", exp.OnCluster(this=exp.to_identifier(self.cluster)))
726752
self.execute(delete_expr)
727753

@@ -740,7 +766,12 @@ def alter_table(
740766
if self._default_catalog and isinstance(alter_expression.this, exp.Table):
741767
if alter_expression.this.catalog == self._default_catalog:
742768
alter_expression.this.set("catalog", None)
743-
if self.engine_run_mode.is_cluster:
769+
# Decided per expression, not hoisted: one call can legitimately carry
770+
# alters against both a Replicated and a non-Replicated database.
771+
altered_table = (
772+
alter_expression.this if isinstance(alter_expression.this, exp.Table) else None
773+
)
774+
if self._should_use_on_cluster(altered_table):
744775
alter_expression.set(
745776
"cluster", exp.OnCluster(this=exp.to_identifier(self.cluster))
746777
)
@@ -764,14 +795,24 @@ def _drop_object(
764795
kind: What kind of object to drop. Defaults to TABLE
765796
**drop_args: Any extra arguments to set on the Drop expression
766797
"""
798+
# Dropping the database itself is always cluster-wide: the Keeper-backed DDL log
799+
# being dropped cannot propagate its own removal, and `name` here is a database
800+
# rather than an object inside one.
801+
is_database = kind.upper() in ("SCHEMA", "DATABASE")
802+
use_on_cluster = (
803+
self.engine_run_mode.is_cluster if is_database else self._should_use_on_cluster(name)
804+
)
805+
806+
if is_database:
807+
# Cheap to rebuild and easy to get wrong for one entry; drop the lot.
808+
self._clear_database_engine_cache()
809+
767810
super()._drop_object(
768811
name=name,
769812
exists=exists,
770813
kind=kind,
771814
cascade=cascade,
772-
cluster=exp.OnCluster(this=exp.to_identifier(self.cluster))
773-
if self.engine_run_mode.is_cluster
774-
else None,
815+
cluster=exp.OnCluster(this=exp.to_identifier(self.cluster)) if use_on_cluster else None,
775816
**drop_args,
776817
)
777818

@@ -873,6 +914,7 @@ def _build_table_properties_exp(
873914
table_description: t.Optional[str] = None,
874915
table_kind: t.Optional[str] = None,
875916
empty_ctas: bool = False,
917+
table: t.Optional[exp.Table] = None,
876918
**kwargs: t.Any,
877919
) -> t.Optional[exp.Properties]:
878920
properties: t.List[exp.Expr] = []
@@ -951,7 +993,7 @@ def _build_table_properties_exp(
951993
):
952994
properties.append(partitioned_by_prop)
953995

954-
if self.engine_run_mode.is_cluster:
996+
if self._should_use_on_cluster(table):
955997
properties.append(exp.OnCluster(this=exp.to_identifier(self.cluster)))
956998

957999
if empty_ctas:
@@ -976,14 +1018,15 @@ def _build_view_properties_exp(
9761018
self,
9771019
view_properties: t.Optional[t.Dict[str, exp.Expr]] = None,
9781020
table_description: t.Optional[str] = None,
1021+
table: t.Optional[exp.Table] = None,
9791022
**kwargs: t.Any,
9801023
) -> t.Optional[exp.Properties]:
9811024
"""Creates a SQLGlot table properties expression for view"""
9821025
properties: t.List[exp.Expr] = []
9831026

9841027
view_properties_copy = view_properties.copy() if view_properties else {}
9851028

986-
if self.engine_run_mode.is_cluster:
1029+
if self._should_use_on_cluster(table):
9871030
properties.append(exp.OnCluster(this=exp.to_identifier(self.cluster)))
9881031

9891032
if view_properties_copy:
@@ -1008,7 +1051,7 @@ def _build_create_comment_table_exp(
10081051
truncated_comment = self._truncate_table_comment(table_comment)
10091052
comment_sql = exp.Literal.string(truncated_comment).sql(dialect=self.dialect)
10101053

1011-
return f"ALTER TABLE {table_sql}{self._on_cluster_sql()} MODIFY COMMENT {comment_sql}"
1054+
return f"ALTER TABLE {table_sql}{self._on_cluster_sql(table)} MODIFY COMMENT {comment_sql}"
10121055

10131056
def _build_create_comment_column_exp(
10141057
self,
@@ -1024,10 +1067,159 @@ def _build_create_comment_column_exp(
10241067
truncated_comment = self._truncate_table_comment(column_comment)
10251068
comment_sql = exp.Literal.string(truncated_comment).sql(dialect=self.dialect)
10261069

1027-
return f"ALTER TABLE {table_sql}{self._on_cluster_sql()} COMMENT COLUMN {column_sql} {comment_sql}"
1070+
return (
1071+
f"ALTER TABLE {table_sql}{self._on_cluster_sql(table)}"
1072+
f" COMMENT COLUMN {column_sql} {comment_sql}"
1073+
)
10281074

1029-
def _on_cluster_sql(self) -> str:
1030-
if self.engine_run_mode.is_cluster:
1075+
@cached_property
1076+
def _database_engine_cache(self) -> t.Dict[str, t.Optional[str]]:
1077+
return {}
1078+
1079+
@cached_property
1080+
def _has_replicated_database(self) -> bool:
1081+
"""Whether this server hosts any Keeper-coordinated database at all.
1082+
1083+
One probe per connection, and the reason this change costs nothing on a
1084+
deployment that has no `Replicated` database: when the answer is no, every
1085+
`ON CLUSTER` decision short-circuits without resolving a target or querying
1086+
`system.databases` again.
1087+
"""
1088+
try:
1089+
row = self.fetchone(
1090+
exp.select(exp.func("count"))
1091+
.from_("system.databases")
1092+
.where(
1093+
exp.column("engine").like(
1094+
exp.Literal.string(f"{REPLICATED_DATABASE_ENGINE_PREFIX}%")
1095+
)
1096+
)
1097+
)
1098+
except Exception:
1099+
return False
1100+
return bool(row and row[0])
1101+
1102+
@cached_property
1103+
def _connection_database(self) -> t.Optional[str]:
1104+
"""The database an unqualified object resolves to on this connection."""
1105+
try:
1106+
row = self.fetchone("SELECT currentDatabase()")
1107+
except Exception:
1108+
return None
1109+
return str(row[0]) if row and row[0] else None
1110+
1111+
def _clear_database_engine_cache(self, database: t.Optional[str] = None) -> None:
1112+
if database is None:
1113+
self._database_engine_cache.clear()
1114+
else:
1115+
self._database_engine_cache.pop(database, None)
1116+
# Creating or dropping a database can also change whether any Replicated one
1117+
# exists, which is what the short-circuit above depends on.
1118+
self.__dict__.pop("_has_replicated_database", None)
1119+
1120+
def _database_engine(self, database: str) -> t.Optional[str]:
1121+
"""The engine of a ClickHouse database, or None when it cannot be resolved.
1122+
1123+
Deliberately connection-local rather than replica-wide: a `Replicated` database's
1124+
engine is uniform by construction, and `clusterAllReplicas` fails outright when any
1125+
host is down, which would turn an unrelated outage into a DDL failure.
1126+
"""
1127+
cache = self._database_engine_cache
1128+
if database in cache:
1129+
return cache[database]
1130+
1131+
engine: t.Optional[str] = None
1132+
try:
1133+
row = self.fetchone(
1134+
exp.select("engine")
1135+
.from_("system.databases")
1136+
.where(exp.column("name").eq(exp.Literal.string(database)))
1137+
)
1138+
if row and row[0]:
1139+
engine = str(row[0])
1140+
except Exception:
1141+
# Unresolvable for any reason - absent, permission-denied, introspection
1142+
# failure. Cache the miss but let the caller fall back to emitting.
1143+
engine = None
1144+
1145+
cache[database] = engine
1146+
return engine
1147+
1148+
def _on_cluster_target_database(self, target: t.Optional[TableName]) -> t.Optional[str]:
1149+
"""The database an object DDL statement targets, or None when unknown."""
1150+
if target is None:
1151+
return None
1152+
1153+
table = exp.to_table(target, dialect=self.dialect) if isinstance(target, str) else target
1154+
if not isinstance(table, exp.Table):
1155+
return None
1156+
1157+
table = self._strip_virtual_catalog(table)
1158+
database = table.db
1159+
# An unqualified name resolves to the connection's database, not to "unknown".
1160+
return database or self._connection_database
1161+
1162+
def _should_use_on_cluster(self, target: t.Optional[TableName] = None) -> bool:
1163+
"""Whether object DDL for `target` should carry `ON CLUSTER`.
1164+
1165+
Inside a `Replicated` database Keeper already propagates object DDL, so adding
1166+
`ON CLUSTER` asks for a second, redundant fan-out and ClickHouse refuses the
1167+
statement with code 80 `INCORRECT_QUERY`. Without this, no object can be created
1168+
in such a database at all while the adapter is in cluster mode.
1169+
Suppression cannot be a connection-level flag: one connection can hold both an
1170+
Atomic and a `Replicated` database, and objects in the Atomic one still need
1171+
`ON CLUSTER`.
1172+
1173+
Fails open. An unknown target, an unresolvable database, or any introspection
1174+
failure keeps today's behaviour, so this is a no-op for every deployment that has
1175+
no `Replicated` database.
1176+
"""
1177+
if not self.engine_run_mode.is_cluster:
1178+
return False
1179+
1180+
if not self._has_replicated_database:
1181+
return True
1182+
1183+
database = self._on_cluster_target_database(target)
1184+
if not database:
1185+
return True
1186+
1187+
engine = self._database_engine(database)
1188+
if engine is None:
1189+
return True
1190+
1191+
return not engine.startswith(REPLICATED_DATABASE_ENGINE_PREFIX)
1192+
1193+
def _on_cluster_sql(self, target: t.Optional[TableName] = None) -> str:
1194+
"""Render the `ON CLUSTER` clause for object DDL against `target`.
1195+
1196+
Omitting `target` means "target unknown" and preserves the pre-existing
1197+
behaviour of always emitting in cluster mode.
1198+
"""
1199+
if self._should_use_on_cluster(target):
10311200
cluster_name = exp.to_identifier(self.cluster, quoted=True).sql(dialect=self.dialect) # type: ignore
10321201
return f" ON CLUSTER {cluster_name} "
10331202
return ""
1203+
1204+
def _assert_same_on_cluster_scope(
1205+
self, operation: str, first: TableName, second: TableName
1206+
) -> bool:
1207+
"""Resolve one `ON CLUSTER` decision for a two-table statement.
1208+
1209+
`RENAME` and `EXCHANGE` can span databases, and a single statement cannot be both
1210+
cluster-wide and Keeper-propagated. Refuse rather than pick one and be silently
1211+
half-correct on the other.
1212+
"""
1213+
from sqlmesh.utils.errors import SQLMeshError
1214+
1215+
first_scope = self._should_use_on_cluster(first)
1216+
second_scope = self._should_use_on_cluster(second)
1217+
if first_scope != second_scope:
1218+
first_sql = exp.to_table(first, dialect=self.dialect).sql(dialect=self.dialect)
1219+
second_sql = exp.to_table(second, dialect=self.dialect).sql(dialect=self.dialect)
1220+
raise SQLMeshError(
1221+
f"Cannot {operation} between a Replicated database and a non-Replicated one: "
1222+
f"{first_sql} and {second_sql} disagree on whether object DDL carries "
1223+
"ON CLUSTER. Move both objects into databases with the same engine."
1224+
)
1225+
return first_scope

0 commit comments

Comments
 (0)