From b6c3e72b46dc1361323cf6b55773076a263b0505 Mon Sep 17 00:00:00 2001 From: mday-io Date: Mon, 31 Aug 2026 11:03:39 -0400 Subject: [PATCH] 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. --- sqlmesh/core/engine_adapter/base.py | 8 + sqlmesh/core/engine_adapter/clickhouse.py | 226 ++++++++++++++++-- tests/core/engine_adapter/test_clickhouse.py | 236 +++++++++++++++++++ 3 files changed, 453 insertions(+), 17 deletions(-) diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py index bd435db76f..eebd17726e 100644 --- a/sqlmesh/core/engine_adapter/base.py +++ b/sqlmesh/core/engine_adapter/base.py @@ -1043,12 +1043,15 @@ def _build_create_table_exp( ) -> exp.Create: exists = False if replace else exists catalog_name = None + target_table: t.Optional[exp.Table] = None if not isinstance(table_name_or_schema, exp.Schema): table_name_or_schema = exp.to_table(table_name_or_schema) catalog_name = table_name_or_schema.catalog + target_table = table_name_or_schema else: if isinstance(table_name_or_schema.this, exp.Table): catalog_name = table_name_or_schema.this.catalog + target_table = table_name_or_schema.this properties = ( self._build_table_properties_exp( @@ -1057,6 +1060,9 @@ def _build_create_table_exp( target_columns_to_types=target_columns_to_types, table_description=table_description, table_kind=table_kind, + # Passed so an adapter can vary properties by target object, not just by + # connection. Additive: every override accepts **kwargs. + table=target_table, ) if kwargs or table_description else None @@ -1352,6 +1358,8 @@ def create_view( else None ), physical_cluster=create_kwargs.pop("physical_cluster", None), + # See the note on the table-properties call site above. + table=exp.to_table(view_name), ) if create_view_properties: for view_property in create_view_properties.expressions: diff --git a/sqlmesh/core/engine_adapter/clickhouse.py b/sqlmesh/core/engine_adapter/clickhouse.py index d1f67e0564..627f71c60c 100644 --- a/sqlmesh/core/engine_adapter/clickhouse.py +++ b/sqlmesh/core/engine_adapter/clickhouse.py @@ -3,6 +3,7 @@ import typing as t import logging import re +from functools import cached_property from sqlglot import exp, maybe_parse from sqlmesh.core.dialect import to_schema from sqlmesh.core.engine_adapter.mixins import LogicalMergeMixin @@ -30,6 +31,13 @@ logger = logging.getLogger(__name__) +# `system.databases.engine` value for a Keeper-coordinated database. Object DDL inside one +# must not carry `ON CLUSTER`: the database's own DDL log already propagates it to every +# replica, so ClickHouse rejects the redundant second fan-out outright with code 80, +# `It's not initial query. ON CLUSTER is not allowed for Replicated database.` +# Matched as a prefix so engine variants are covered without another release. +REPLICATED_DATABASE_ENGINE_PREFIX = "Replicated" + class ClickhouseEngineAdapter(EngineAdapterWithIndexSupport, LogicalMergeMixin): DIALECT = "clickhouse" @@ -192,9 +200,15 @@ def create_schema( from sqlmesh.utils.errors import SQLMeshError properties_copy = properties.copy() + # Always cluster-wide, deliberately. This is the statement that creates the + # database, including a `Replicated` one, so there is no database engine to + # consult yet and nothing for Keeper to propagate through. if self.engine_run_mode.is_cluster: properties_copy.append(exp.OnCluster(this=exp.to_identifier(self.cluster))) + # A database that previously resolved as absent may now exist, with an engine. + self._clear_database_engine_cache() + # ClickHouse does not support catalogs. When a virtual catalog has been injected # (self._default_catalog is set), strip it from the schema name. This mirrors the # SINGLE_CATALOG_ONLY branch in the set_catalog decorator, which does not apply here @@ -501,7 +515,8 @@ def _create_table_like( ) -> None: """Create table with identical structure as source table""" self.execute( - f"CREATE TABLE {target_table_name}{self._on_cluster_sql()} AS {source_table_name}" + f"CREATE TABLE {target_table_name}{self._on_cluster_sql(target_table_name)}" + f" AS {source_table_name}" ) def _get_partition_ids( @@ -661,10 +676,14 @@ def _exchange_tables( old_table_sql = exp.to_table(old_table_name).sql(dialect=self.dialect, identify=True) new_table_sql = exp.to_table(new_table_name).sql(dialect=self.dialect, identify=True) + on_cluster_sql = ( + self._on_cluster_sql(old_table_name) + if self._assert_same_on_cluster_scope("EXCHANGE TABLES", old_table_name, new_table_name) + else "" + ) + try: - self.execute( - f"EXCHANGE TABLES {old_table_sql} AND {new_table_sql}{self._on_cluster_sql()}" - ) + self.execute(f"EXCHANGE TABLES {old_table_sql} AND {new_table_sql}{on_cluster_sql}") except DatabaseError as e: if "NOT_IMPLEMENTED" in str(e): # If someone is using an old Clickhouse version, an OS that doesn't support atomic exchanges, @@ -686,11 +705,18 @@ def _rename_table( old_table_sql = exp.to_table(old_table_name).sql(dialect=self.dialect, identify=True) new_table_sql = exp.to_table(new_table_name).sql(dialect=self.dialect, identify=True) - self.execute(f"RENAME TABLE {old_table_sql} TO {new_table_sql}{self._on_cluster_sql()}") + on_cluster_sql = ( + self._on_cluster_sql(old_table_name) + if self._assert_same_on_cluster_scope("RENAME TABLE", old_table_name, new_table_name) + else "" + ) + + self.execute(f"RENAME TABLE {old_table_sql} TO {new_table_sql}{on_cluster_sql}") def delete_from(self, table_name: TableName, where: t.Union[str, exp.Expr]) -> None: - delete_expr = exp.delete(self._strip_virtual_catalog(table_name), where) - if self.engine_run_mode.is_cluster: + target_table = self._strip_virtual_catalog(table_name) + delete_expr = exp.delete(target_table, where) + if self._should_use_on_cluster(target_table): delete_expr.set("cluster", exp.OnCluster(this=exp.to_identifier(self.cluster))) self.execute(delete_expr) @@ -708,7 +734,12 @@ def alter_table( if self._default_catalog and isinstance(alter_expression.this, exp.Table): if alter_expression.this.catalog == self._default_catalog: alter_expression.this.set("catalog", None) - if self.engine_run_mode.is_cluster: + # Decided per expression, not hoisted: one call can legitimately carry + # alters against both a Replicated and a non-Replicated database. + altered_table = ( + alter_expression.this if isinstance(alter_expression.this, exp.Table) else None + ) + if self._should_use_on_cluster(altered_table): alter_expression.set( "cluster", exp.OnCluster(this=exp.to_identifier(self.cluster)) ) @@ -732,14 +763,24 @@ def _drop_object( kind: What kind of object to drop. Defaults to TABLE **drop_args: Any extra arguments to set on the Drop expression """ + # Dropping the database itself is always cluster-wide: the Keeper-backed DDL log + # being dropped cannot propagate its own removal, and `name` here is a database + # rather than an object inside one. + is_database = kind.upper() in ("SCHEMA", "DATABASE") + use_on_cluster = ( + self.engine_run_mode.is_cluster if is_database else self._should_use_on_cluster(name) + ) + + if is_database: + # Cheap to rebuild and easy to get wrong for one entry; drop the lot. + self._clear_database_engine_cache() + super()._drop_object( name=name, exists=exists, kind=kind, cascade=cascade, - cluster=exp.OnCluster(this=exp.to_identifier(self.cluster)) - if self.engine_run_mode.is_cluster - else None, + cluster=exp.OnCluster(this=exp.to_identifier(self.cluster)) if use_on_cluster else None, **drop_args, ) @@ -841,6 +882,7 @@ def _build_table_properties_exp( table_description: t.Optional[str] = None, table_kind: t.Optional[str] = None, empty_ctas: bool = False, + table: t.Optional[exp.Table] = None, **kwargs: t.Any, ) -> t.Optional[exp.Properties]: properties: t.List[exp.Expr] = [] @@ -919,7 +961,7 @@ def _build_table_properties_exp( ): properties.append(partitioned_by_prop) - if self.engine_run_mode.is_cluster: + if self._should_use_on_cluster(table): properties.append(exp.OnCluster(this=exp.to_identifier(self.cluster))) if empty_ctas: @@ -944,6 +986,7 @@ def _build_view_properties_exp( self, view_properties: t.Optional[t.Dict[str, exp.Expr]] = None, table_description: t.Optional[str] = None, + table: t.Optional[exp.Table] = None, **kwargs: t.Any, ) -> t.Optional[exp.Properties]: """Creates a SQLGlot table properties expression for view""" @@ -951,7 +994,7 @@ def _build_view_properties_exp( view_properties_copy = view_properties.copy() if view_properties else {} - if self.engine_run_mode.is_cluster: + if self._should_use_on_cluster(table): properties.append(exp.OnCluster(this=exp.to_identifier(self.cluster))) if view_properties_copy: @@ -976,7 +1019,7 @@ def _build_create_comment_table_exp( truncated_comment = self._truncate_table_comment(table_comment) comment_sql = exp.Literal.string(truncated_comment).sql(dialect=self.dialect) - return f"ALTER TABLE {table_sql}{self._on_cluster_sql()} MODIFY COMMENT {comment_sql}" + return f"ALTER TABLE {table_sql}{self._on_cluster_sql(table)} MODIFY COMMENT {comment_sql}" def _build_create_comment_column_exp( self, @@ -992,10 +1035,159 @@ def _build_create_comment_column_exp( truncated_comment = self._truncate_table_comment(column_comment) comment_sql = exp.Literal.string(truncated_comment).sql(dialect=self.dialect) - return f"ALTER TABLE {table_sql}{self._on_cluster_sql()} COMMENT COLUMN {column_sql} {comment_sql}" + return ( + f"ALTER TABLE {table_sql}{self._on_cluster_sql(table)}" + f" COMMENT COLUMN {column_sql} {comment_sql}" + ) - def _on_cluster_sql(self) -> str: - if self.engine_run_mode.is_cluster: + @cached_property + def _database_engine_cache(self) -> t.Dict[str, t.Optional[str]]: + return {} + + @cached_property + def _has_replicated_database(self) -> bool: + """Whether this server hosts any Keeper-coordinated database at all. + + One probe per connection, and the reason this change costs nothing on a + deployment that has no `Replicated` database: when the answer is no, every + `ON CLUSTER` decision short-circuits without resolving a target or querying + `system.databases` again. + """ + try: + row = self.fetchone( + exp.select(exp.func("count")) + .from_("system.databases") + .where( + exp.column("engine").like( + exp.Literal.string(f"{REPLICATED_DATABASE_ENGINE_PREFIX}%") + ) + ) + ) + except Exception: + return False + return bool(row and row[0]) + + @cached_property + def _connection_database(self) -> t.Optional[str]: + """The database an unqualified object resolves to on this connection.""" + try: + row = self.fetchone("SELECT currentDatabase()") + except Exception: + return None + return str(row[0]) if row and row[0] else None + + def _clear_database_engine_cache(self, database: t.Optional[str] = None) -> None: + if database is None: + self._database_engine_cache.clear() + else: + self._database_engine_cache.pop(database, None) + # Creating or dropping a database can also change whether any Replicated one + # exists, which is what the short-circuit above depends on. + self.__dict__.pop("_has_replicated_database", None) + + def _database_engine(self, database: str) -> t.Optional[str]: + """The engine of a ClickHouse database, or None when it cannot be resolved. + + Deliberately connection-local rather than replica-wide: a `Replicated` database's + engine is uniform by construction, and `clusterAllReplicas` fails outright when any + host is down, which would turn an unrelated outage into a DDL failure. + """ + cache = self._database_engine_cache + if database in cache: + return cache[database] + + engine: t.Optional[str] = None + try: + row = self.fetchone( + exp.select("engine") + .from_("system.databases") + .where(exp.column("name").eq(exp.Literal.string(database))) + ) + if row and row[0]: + engine = str(row[0]) + except Exception: + # Unresolvable for any reason - absent, permission-denied, introspection + # failure. Cache the miss but let the caller fall back to emitting. + engine = None + + cache[database] = engine + return engine + + def _on_cluster_target_database(self, target: t.Optional[TableName]) -> t.Optional[str]: + """The database an object DDL statement targets, or None when unknown.""" + if target is None: + return None + + table = exp.to_table(target, dialect=self.dialect) if isinstance(target, str) else target + if not isinstance(table, exp.Table): + return None + + table = self._strip_virtual_catalog(table) + database = table.db + # An unqualified name resolves to the connection's database, not to "unknown". + return database or self._connection_database + + def _should_use_on_cluster(self, target: t.Optional[TableName] = None) -> bool: + """Whether object DDL for `target` should carry `ON CLUSTER`. + + Inside a `Replicated` database Keeper already propagates object DDL, so adding + `ON CLUSTER` asks for a second, redundant fan-out and ClickHouse refuses the + statement with code 80 `INCORRECT_QUERY`. Without this, no object can be created + in such a database at all while the adapter is in cluster mode. + Suppression 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`. + + Fails open. An unknown target, an unresolvable database, or any introspection + failure keeps today's behaviour, so this is a no-op for every deployment that has + no `Replicated` database. + """ + if not self.engine_run_mode.is_cluster: + return False + + if not self._has_replicated_database: + return True + + database = self._on_cluster_target_database(target) + if not database: + return True + + engine = self._database_engine(database) + if engine is None: + return True + + return not engine.startswith(REPLICATED_DATABASE_ENGINE_PREFIX) + + def _on_cluster_sql(self, target: t.Optional[TableName] = None) -> str: + """Render the `ON CLUSTER` clause for object DDL against `target`. + + Omitting `target` means "target unknown" and preserves the pre-existing + behaviour of always emitting in cluster mode. + """ + if self._should_use_on_cluster(target): cluster_name = exp.to_identifier(self.cluster, quoted=True).sql(dialect=self.dialect) # type: ignore return f" ON CLUSTER {cluster_name} " return "" + + def _assert_same_on_cluster_scope( + self, operation: str, first: TableName, second: TableName + ) -> bool: + """Resolve one `ON CLUSTER` decision for a two-table statement. + + `RENAME` and `EXCHANGE` can span databases, and a single statement cannot be both + cluster-wide and Keeper-propagated. Refuse rather than pick one and be silently + half-correct on the other. + """ + from sqlmesh.utils.errors import SQLMeshError + + first_scope = self._should_use_on_cluster(first) + second_scope = self._should_use_on_cluster(second) + if first_scope != second_scope: + first_sql = exp.to_table(first, dialect=self.dialect).sql(dialect=self.dialect) + second_sql = exp.to_table(second, dialect=self.dialect).sql(dialect=self.dialect) + raise SQLMeshError( + f"Cannot {operation} between a Replicated database and a non-Replicated one: " + f"{first_sql} and {second_sql} disagree on whether object DDL carries " + "ON CLUSTER. Move both objects into databases with the same engine." + ) + return first_scope diff --git a/tests/core/engine_adapter/test_clickhouse.py b/tests/core/engine_adapter/test_clickhouse.py index a3dfe0fdda..99f90d29f9 100644 --- a/tests/core/engine_adapter/test_clickhouse.py +++ b/tests/core/engine_adapter/test_clickhouse.py @@ -11,10 +11,23 @@ from pytest_mock.plugin import MockerFixture from sqlmesh.core import dialect as d from sqlglot.optimizer.qualify_columns import quote_identifiers +from sqlmesh.utils.errors import SQLMeshError pytestmark = [pytest.mark.clickhouse, pytest.mark.engine] +@pytest.fixture(autouse=True) +def no_replicated_database_by_default(mocker) -> None: + """Default every adapter in this module to a server with no `Replicated` database. + + That is the world these tests describe, and in it `ON CLUSTER` behaviour is exactly + what it was before the Replicated-aware policy existed. Seeding the probe keeps the + introspection query out of the asserted SQL and makes the assumption explicit. + `test_on_cluster_*` re-patches this to cover the other side. + """ + mocker.patch.object(ClickhouseEngineAdapter, "_has_replicated_database", False) + + @pytest.fixture def adapter(make_mocked_engine_adapter, mocker) -> ClickhouseEngineAdapter: mocker.patch.object( @@ -1637,3 +1650,226 @@ def test_create_view_source_rejects_unexpected_virtual_catalog( "__clickhouse_gw__.my_db.connection_test__dev", parse_one("SELECT * FROM unexpected_catalog.my_db.physical_view"), ) + + +# Replicated-aware ON CLUSTER policy. +# +# Inside a `Replicated` database Keeper already propagates object DDL, so emitting +# `ON CLUSTER` fans the statement out a second time and double-applies it. The policy is +# per target database, because one connection can hold both an Atomic and a `Replicated` +# database and objects in the Atomic one still need `ON CLUSTER`. + + +def _on_cluster_adapter( + make_mocked_engine_adapter: t.Callable, + mocker, + engines: t.Dict[str, t.Optional[str]], + connection_database: t.Optional[str] = "default", +) -> ClickhouseEngineAdapter: + """An adapter on a cluster whose databases have the given engines.""" + adapter = make_mocked_engine_adapter(ClickhouseEngineAdapter, cluster="my_cluster") + mocker.patch.object(ClickhouseEngineAdapter, "_has_replicated_database", True) + mocker.patch.object(ClickhouseEngineAdapter, "_connection_database", connection_database) + mocker.patch.object( + ClickhouseEngineAdapter, + "_database_engine", + lambda self, database: engines.get(database), + ) + return adapter + + +def test_on_cluster_predicate(make_mocked_engine_adapter: t.Callable, mocker): + adapter = _on_cluster_adapter( + make_mocked_engine_adapter, + mocker, + {"atomic_db": "Atomic", "replicated_db": "Replicated", "memory_db": "Memory"}, + ) + + # A Replicated target is the only case that suppresses. + assert adapter._should_use_on_cluster(exp.to_table("atomic_db.t")) is True + assert adapter._should_use_on_cluster(exp.to_table("replicated_db.t")) is False + assert adapter._should_use_on_cluster(exp.to_table("memory_db.t")) is True + + # Fails open: an unknown target or an unresolvable database keeps today's behaviour. + assert adapter._should_use_on_cluster(None) is True + assert adapter._should_use_on_cluster(exp.to_table("absent_db.t")) is True + + # A bare name resolves to the connection's database rather than counting as unknown. + mocker.patch.object(ClickhouseEngineAdapter, "_connection_database", "replicated_db") + assert adapter._should_use_on_cluster(exp.to_table("t")) is False + + +def test_on_cluster_predicate_off_cluster_never_emits( + make_mocked_engine_adapter: t.Callable, mocker +): + adapter = make_mocked_engine_adapter(ClickhouseEngineAdapter) + assert adapter._should_use_on_cluster(exp.to_table("any_db.t")) is False + + +def test_on_cluster_predicate_short_circuits_without_a_replicated_database( + make_mocked_engine_adapter: t.Callable, +): + """The no-op guarantee: no Replicated database means no lookup and no behaviour change.""" + adapter = make_mocked_engine_adapter(ClickhouseEngineAdapter, cluster="my_cluster") + + assert adapter._should_use_on_cluster(exp.to_table("any_db.t")) is True + assert to_sql_calls(adapter) == [] + + +def test_on_cluster_predicate_strips_the_virtual_catalog( + make_mocked_engine_adapter: t.Callable, mocker +): + adapter = _on_cluster_adapter(make_mocked_engine_adapter, mocker, {"my_db": "Replicated"}) + adapter.inject_virtual_catalog("clickhouse_gw") + + assert adapter._should_use_on_cluster(exp.to_table("__clickhouse_gw__.my_db.t")) is False + + +def test_database_engine_lookup_is_cached_and_fails_open( + make_mocked_engine_adapter: t.Callable, mocker +): + adapter = make_mocked_engine_adapter(ClickhouseEngineAdapter, cluster="my_cluster") + + fetchone = mocker.patch.object(ClickhouseEngineAdapter, "fetchone", return_value=("Atomic",)) + assert adapter._database_engine("some_db") == "Atomic" + assert adapter._database_engine("some_db") == "Atomic" + assert fetchone.call_count == 1, "engine lookup should be cached per database" + + # An introspection failure resolves to None, which the predicate treats as "emit". + mocker.patch.object(ClickhouseEngineAdapter, "fetchone", side_effect=Exception("boom")) + assert adapter._database_engine("other_db") is None + mocker.patch.object(ClickhouseEngineAdapter, "_has_replicated_database", True) + mocker.patch.object(ClickhouseEngineAdapter, "_connection_database", "default") + assert adapter._should_use_on_cluster(exp.to_table("other_db.t")) is True + + +def test_on_cluster_suppressed_for_replicated_object_ddl( + make_mocked_engine_adapter: t.Callable, mocker +): + """Per-site rendering: the same statement against Atomic and Replicated targets.""" + adapter = _on_cluster_adapter( + make_mocked_engine_adapter, mocker, {"a_db": "Atomic", "r_db": "Replicated"} + ) + + adapter.delete_from(exp.to_table("a_db.t"), "x = 1") + adapter.delete_from(exp.to_table("r_db.t"), "x = 1") + adapter.drop_table("a_db.t") + adapter.drop_table("r_db.t") + adapter._create_table_like("a_db.copy", "a_db.t", exists=False) + adapter._create_table_like("r_db.copy", "r_db.t", exists=False) + + assert to_sql_calls(adapter) == [ + 'DELETE FROM "a_db"."t" ON CLUSTER "my_cluster" WHERE "x" = 1', + 'DELETE FROM "r_db"."t" WHERE "x" = 1', + 'DROP TABLE IF EXISTS "a_db"."t" ON CLUSTER "my_cluster"', + 'DROP TABLE IF EXISTS "r_db"."t"', + # The doubled space is pre-existing: _on_cluster_sql renders a padded clause. + 'CREATE TABLE a_db.copy ON CLUSTER "my_cluster" AS a_db.t', + "CREATE TABLE r_db.copy AS r_db.t", + ] + + +def test_on_cluster_suppressed_for_replicated_create_table_and_view( + make_mocked_engine_adapter: t.Callable, mocker +): + """Covers the two shared-code property builders, which get their target from base.""" + adapter = _on_cluster_adapter( + make_mocked_engine_adapter, mocker, {"a_db": "Atomic", "r_db": "Replicated"} + ) + columns_to_types = {"id": exp.DataType.build("UInt64", dialect="clickhouse")} + + adapter.create_table("a_db.t", columns_to_types) + adapter.create_table("r_db.t", columns_to_types) + adapter.create_view("a_db.v", parse_one("SELECT 1 AS id")) + adapter.create_view("r_db.v", parse_one("SELECT 1 AS id")) + + calls = to_sql_calls(adapter) + assert 'ON CLUSTER "my_cluster"' in calls[0] + assert "ON CLUSTER" not in calls[1] + assert 'ON CLUSTER "my_cluster"' in calls[2] + assert "ON CLUSTER" not in calls[3] + + +def test_on_cluster_suppressed_for_replicated_comments( + make_mocked_engine_adapter: t.Callable, mocker +): + adapter = _on_cluster_adapter( + make_mocked_engine_adapter, mocker, {"a_db": "Atomic", "r_db": "Replicated"} + ) + + assert 'ON CLUSTER "my_cluster"' in adapter._build_create_comment_table_exp( + exp.to_table("a_db.t"), "hi", "TABLE" + ) + assert "ON CLUSTER" not in adapter._build_create_comment_table_exp( + exp.to_table("r_db.t"), "hi", "TABLE" + ) + assert 'ON CLUSTER "my_cluster"' in adapter._build_create_comment_column_exp( + exp.to_table("a_db.t"), "id", "hi" + ) + assert "ON CLUSTER" not in adapter._build_create_comment_column_exp( + exp.to_table("r_db.t"), "id", "hi" + ) + + +def test_database_ddl_is_always_on_cluster(make_mocked_engine_adapter: t.Callable, mocker): + """Carve-out. Creating or dropping the database itself is never Keeper-propagated. + + `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. + """ + adapter = _on_cluster_adapter(make_mocked_engine_adapter, mocker, {"r_db": "Replicated"}) + + adapter.create_schema("r_db") + adapter.drop_schema("r_db") + + for call in to_sql_calls(adapter): + assert 'ON CLUSTER "my_cluster"' in call, call + + +def test_alter_table_decides_on_cluster_per_expression( + make_mocked_engine_adapter: t.Callable, mocker +): + """One call can carry alters against both kinds of database; it cannot be hoisted.""" + adapter = _on_cluster_adapter( + make_mocked_engine_adapter, mocker, {"a_db": "Atomic", "r_db": "Replicated"} + ) + + adapter.alter_table( + [ + parse_one("ALTER TABLE a_db.t ADD COLUMN x UInt64", dialect="clickhouse"), + parse_one("ALTER TABLE r_db.t ADD COLUMN x UInt64", dialect="clickhouse"), + ] + ) + + calls = to_sql_calls(adapter) + assert 'ON CLUSTER "my_cluster"' in calls[0] + assert "ON CLUSTER" not in calls[1] + + +def test_cross_engine_rename_and_exchange_are_refused( + make_mocked_engine_adapter: t.Callable, mocker +): + """A single statement cannot be both cluster-wide and Keeper-propagated.""" + adapter = _on_cluster_adapter( + make_mocked_engine_adapter, mocker, {"a_db": "Atomic", "r_db": "Replicated"} + ) + + with pytest.raises(SQLMeshError, match="RENAME TABLE between a Replicated database"): + adapter._rename_table("a_db.t", "r_db.t") + + with pytest.raises(SQLMeshError, match="EXCHANGE TABLES between a Replicated database"): + adapter._exchange_tables("a_db.t", "r_db.t") + + +def test_creating_a_database_invalidates_the_engine_cache( + make_mocked_engine_adapter: t.Callable, mocker +): + """A database that resolved as absent may exist, with an engine, after CREATE.""" + adapter = make_mocked_engine_adapter(ClickhouseEngineAdapter, cluster="my_cluster") + adapter.__dict__["_has_replicated_database"] = True + adapter._database_engine_cache["new_db"] = None + + adapter.create_schema("new_db") + + assert "new_db" not in adapter._database_engine_cache + assert "_has_replicated_database" not in adapter.__dict__