33import typing as t
44import logging
55import re
6+ from functools import cached_property
67from sqlglot import exp , maybe_parse
78from sqlmesh .core .dialect import to_schema
89from sqlmesh .core .engine_adapter .mixins import LogicalMergeMixin
3132
3233logger = 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
3543class 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