diff --git a/docs/concepts/models/model_kinds.md b/docs/concepts/models/model_kinds.md index cde104790a..15f56ce351 100644 --- a/docs/concepts/models/model_kinds.md +++ b/docs/concepts/models/model_kinds.md @@ -618,11 +618,14 @@ MODEL ( * BigQuery * Databricks +* DuckDB (version 1.4 or later, see below) * Postgres * Redshift * Snowflake * Spark +In DuckDB's case, the native `MERGE` statement is used on version 1.4 or later when the target table is in a DuckDB, DuckLake or MotherDuck catalog. Older DuckDB versions, and target tables in other attached catalogs (such as PostgreSQL, SQLite or MySQL), use `DELETE` + `INSERT` instead, and `when_matched` and `merge_filter` raise an error there. Source rows are de-duplicated by `unique_key` before the merge so there is one row per key. Which of several duplicate rows is kept is unspecified. + In Redshift's case, to enable the use of the native `MERGE` statement, you need to pass the `enable_merge` flag in the connection and set it to `true`. It is disabled by default. ```yaml linenums="1" @@ -668,7 +671,7 @@ Depending on the target engine, models of the `INCREMENTAL_BY_UNIQUE_KEY` kind a | BigQuery | MERGE ON unique key | | Redshift | MERGE ON unique key | | Postgres | MERGE ON unique key | -| DuckDB | DELETE ON matched + INSERT new rows | +| DuckDB | MERGE ON unique key (version 1.4 or later, DuckDB, DuckLake and MotherDuck catalogs); otherwise DELETE ON matched + INSERT new rows | ## FULL Models of the `FULL` kind cause the dataset associated with a model to be fully refreshed (rewritten) upon each model evaluation. diff --git a/docs/integrations/engines/duckdb.md b/docs/integrations/engines/duckdb.md index aca58615b3..36b4427118 100644 --- a/docs/integrations/engines/duckdb.md +++ b/docs/integrations/engines/duckdb.md @@ -196,6 +196,10 @@ Some connections, like SQLite, do not support schema names and therefore objects Example: mounting a SQLite database with the name `sqlite` that has a table `example_table` will be accessible as `sqlite.main.example_table`. +##### Incremental by unique key models + +On DuckDB version 1.4 or later, [`INCREMENTAL_BY_UNIQUE_KEY`](../../concepts/models/model_kinds.md#incremental_by_unique_key) models targeting a DuckDB, DuckLake or MotherDuck catalog use a single, atomic `MERGE` statement and support `when_matched` and `merge_filter`. Older DuckDB versions and tables in other attached catalogs (such as PostgreSQL or SQLite) use `DELETE` + `INSERT`, which runs as two statements because DuckDB adapter transactions are disabled, and do not support `when_matched` or `merge_filter`. + ##### Sensitive fields in paths If a connector, like Postgres, requires sensitive information in the path, it might support defining environment variables instead. diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py index 2dd234f70a..a39d8b4a27 100644 --- a/sqlmesh/core/engine_adapter/base.py +++ b/sqlmesh/core/engine_adapter/base.py @@ -1771,6 +1771,10 @@ def update_table( ) -> None: self.execute(exp.update(table_name, properties, where=where)) + def _prepare_merge_source_query(self, query: Query, unique_key: t.Sequence[exp.Expr]) -> Query: + """Hook that lets an engine adjust the MERGE source query. By default it is used as is.""" + return query + def _merge( self, target_table: TableName, @@ -2316,7 +2320,7 @@ def merge( with source_query as query: self._merge( target_table=target_table, - query=query, + query=self._prepare_merge_source_query(query, unique_key), on=on, whens=exp.Whens(expressions=match_expressions), ) diff --git a/sqlmesh/core/engine_adapter/duckdb.py b/sqlmesh/core/engine_adapter/duckdb.py index 3666a801dd..982356ba83 100644 --- a/sqlmesh/core/engine_adapter/duckdb.py +++ b/sqlmesh/core/engine_adapter/duckdb.py @@ -1,13 +1,16 @@ from __future__ import annotations import typing as t +from functools import cached_property, partial from sqlglot import exp from pathlib import Path +from packaging import version +from sqlmesh.core.engine_adapter.base import MERGE_TARGET_ALIAS from sqlmesh.core.engine_adapter.mixins import ( GetCurrentCatalogFromFunctionMixin, - LogicalMergeMixin, RowDiffMixin, + logical_merge, ) from sqlmesh.core.engine_adapter.shared import ( CatalogSupport, @@ -21,11 +24,11 @@ if t.TYPE_CHECKING: from sqlmesh.core._typing import SchemaName, TableName - from sqlmesh.core.engine_adapter._typing import DF + from sqlmesh.core.engine_adapter._typing import DF, Query, QueryOrDF @set_catalog(override_mapping={"_get_data_objects": CatalogSupport.REQUIRES_SET_CATALOG}) -class DuckDBEngineAdapter(LogicalMergeMixin, GetCurrentCatalogFromFunctionMixin, RowDiffMixin): +class DuckDBEngineAdapter(GetCurrentCatalogFromFunctionMixin, RowDiffMixin): DIALECT = "duckdb" SUPPORTS_TRANSACTIONS = False SCHEMA_DIFFER_KWARGS = { @@ -37,6 +40,11 @@ class DuckDBEngineAdapter(LogicalMergeMixin, GetCurrentCatalogFromFunctionMixin, COMMENT_CREATION_VIEW = CommentCreationView.COMMENT_COMMAND_ONLY SUPPORTS_CREATE_DROP_CATALOG = True SUPPORTED_DROP_CASCADE_OBJECT_KINDS = ["SCHEMA", "TABLE", "VIEW"] + # MERGE INTO was added in DuckDB 1.4.0 + MIN_NATIVE_MERGE_VERSION = version.parse("1.4.0") + # Catalog types whose storage extension supports MERGE INTO. Other attached catalogs + # (postgres, sqlite, mysql, ...) fall back to DELETE + INSERT. + NATIVE_MERGE_CATALOG_TYPES = {"duckdb", "ducklake"} @property def catalog_support(self) -> CatalogSupport: @@ -72,6 +80,76 @@ def _drop_catalog(self, catalog_name: exp.Identifier) -> None: ) ) + @cached_property + def _duckdb_supports_merge(self) -> bool: + import duckdb + + return version.parse(duckdb.__version__) >= self.MIN_NATIVE_MERGE_VERSION + + def _supports_native_merge(self, target_table: TableName) -> bool: + if not self._duckdb_supports_merge: + return False + if self._is_motherduck: + return True + catalog = exp.to_table(target_table).catalog or self.get_current_catalog() + # Errors from the catalog lookup propagate on purpose; only a None/unknown type falls back. + return self._get_catalog_type(catalog) in self.NATIVE_MERGE_CATALOG_TYPES + + def merge( + self, + target_table: TableName, + source_table: QueryOrDF, + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]], + unique_key: t.Sequence[exp.Expr], + when_matched: t.Optional[exp.Whens] = None, + merge_filter: t.Optional[exp.Expr] = None, + source_columns: t.Optional[t.List[str]] = None, + **kwargs: t.Any, + ) -> None: + merge_impl = ( + super().merge + if self._supports_native_merge(target_table) + else partial(logical_merge, self) + ) + merge_impl( # type: ignore + target_table, + source_table, + target_columns_to_types, + unique_key, + when_matched=when_matched, + merge_filter=merge_filter, + source_columns=source_columns, + ) + + def _prepare_merge_source_query(self, query: Query, unique_key: t.Sequence[exp.Expr]) -> Query: + # A native MERGE inserts every unmatched source row, so duplicate keys in the source would + # produce duplicate keys in the target. Keep one row per key, like the DELETE + INSERT path. + return exp.select("*").distinct(*unique_key).from_(query.subquery("_src")) # type: ignore + + def _merge( + self, + target_table: TableName, + query: Query, + on: exp.Expr, + whens: exp.Whens, + ) -> None: + # DuckDB rejects a qualified column on the left-hand side of UPDATE SET and in the INSERT + # column list. Conditions and right-hand sides may keep the qualifier. + # TODO: sqlglot's DuckDB generator doesn't apply `merge_without_target_sql` yet. Remove this + # once the pinned sqlglot version does. + whens = whens.copy() + for when in whens.expressions: + then = when.args.get("then") + targets: t.List[exp.Expr] = [] + if isinstance(then, exp.Update): + targets = [e.this for e in then.expressions if isinstance(e, exp.EQ)] + elif isinstance(then, exp.Insert) and isinstance(then.this, exp.Tuple): + targets = then.this.expressions + for column in targets: + if isinstance(column, exp.Column) and column.table.upper() == MERGE_TARGET_ALIAS: + column.set("table", None) + super()._merge(target_table, query, on, whens) + def _df_to_source_queries( self, df: DF, diff --git a/tests/core/engine_adapter/integration/test_integration.py b/tests/core/engine_adapter/integration/test_integration.py index 6a2905d64d..16fad5dc03 100644 --- a/tests/core/engine_adapter/integration/test_integration.py +++ b/tests/core/engine_adapter/integration/test_integration.py @@ -31,6 +31,7 @@ from sqlmesh.core.environment import EnvironmentSuffixTarget from sqlmesh.core.dialect import select_from_values from sqlmesh.core.model import Model, load_sql_based_model +from sqlmesh.core.engine_adapter.duckdb import DuckDBEngineAdapter from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType from sqlmesh.core.engine_adapter.mixins import RowDiffMixin, LogicalMergeMixin from sqlmesh.core.model.definition import create_sql_model @@ -2857,11 +2858,16 @@ def test_incremental_by_unique_key_model_when_matched(ctx: TestContext): if not ctx.supports_merge: pytest.skip(f"{ctx.dialect} on {ctx.gateway} doesnt support merge") - # DuckDB and some other engines use logical_merge which doesn't support when_matched + # Engines that use logical_merge (e.g. MySQL) don't support when_matched if isinstance(ctx.engine_adapter, LogicalMergeMixin): pytest.skip( f"{ctx.dialect} on {ctx.gateway} uses logical merge which doesn't support when_matched" ) + # DuckDB only merges natively from 1.4.0; older versions fall back to logical_merge + if isinstance(ctx.engine_adapter, DuckDBEngineAdapter) and not ( + ctx.engine_adapter._duckdb_supports_merge + ): + pytest.skip("The installed duckdb doesn't support MERGE (requires >= 1.4.0)") def _mutate_config(current_gateway_name: str, config: Config): connection = config.gateways[current_gateway_name].connection diff --git a/tests/core/engine_adapter/test_duckdb.py b/tests/core/engine_adapter/test_duckdb.py index 6e8097d7eb..7c7268795d 100644 --- a/tests/core/engine_adapter/test_duckdb.py +++ b/tests/core/engine_adapter/test_duckdb.py @@ -1,11 +1,14 @@ import typing as t +from unittest.mock import MagicMock + import pandas as pd # noqa: TID253 import pytest from pytest_mock.plugin import MockerFixture from sqlglot import expressions as exp from sqlglot import parse_one from sqlmesh.core.engine_adapter import DuckDBEngineAdapter, EngineAdapter +from sqlmesh.utils.errors import SQLMeshError from tests.core.engine_adapter import to_sql_calls pytestmark = [pytest.mark.duckdb, pytest.mark.engine] @@ -232,3 +235,386 @@ def test_drop_object_cascade_by_catalog_type(make_mocked_engine_adapter: t.Calla 'DROP SCHEMA IF EXISTS "lake"."virt" CASCADE', 'DROP TABLE IF EXISTS "native"."phys"."t" CASCADE', ] + + +# --------------------------------------------------------------------------- +# Native MERGE +# --------------------------------------------------------------------------- + +MERGE_COLUMNS = { + "id": exp.DataType.build("int"), + "name": exp.DataType.build("text"), + "value": exp.DataType.build("int"), +} + + +def _installed_duckdb_supports_merge() -> bool: + import duckdb + from packaging import version + + return version.parse(duckdb.__version__) >= version.parse("1.4.0") + + +@pytest.fixture +def native_merge_adapter(duck_conn, mocker: MockerFixture) -> DuckDBEngineAdapter: + if not _installed_duckdb_supports_merge(): + pytest.skip("The installed duckdb doesn't support MERGE (requires >= 1.4.0)") + mocker.patch("duckdb.__version__", "1.4.0") + return DuckDBEngineAdapter(lambda: duck_conn) + + +def _set_up_merge_tables(duck_conn, target_rows: str = "", source_rows: str = "") -> None: + duck_conn.execute("CREATE TABLE tgt (id INT, name TEXT, value INT)") + duck_conn.execute("CREATE TABLE src (id INT, name TEXT, value INT)") + if target_rows: + duck_conn.execute(f"INSERT INTO tgt VALUES {target_rows}") + if source_rows: + duck_conn.execute(f"INSERT INTO src VALUES {source_rows}") + + +def _merge(adapter: DuckDBEngineAdapter, **kwargs: t.Any) -> None: + kwargs.setdefault("target_table", "tgt") + kwargs.setdefault("source_table", parse_one("SELECT id, name, value FROM src")) + kwargs.setdefault("target_columns_to_types", MERGE_COLUMNS) + kwargs.setdefault("unique_key", [exp.to_column("id")]) + adapter.merge(**kwargs) + + +def _rows(duck_conn, table: str = "tgt") -> t.List[t.Tuple]: + return duck_conn.execute(f"SELECT * FROM {table} ORDER BY id NULLS LAST, name").fetchall() + + +def _when_matched(sql: str) -> exp.Whens: + return t.cast( + exp.Whens, + parse_one( + f"MERGE INTO x USING y ON 1 = 1 {sql}", + dialect="duckdb", + ).args["whens"], + ) + + +def _mocked_adapter( + make_mocked_engine_adapter: t.Callable, + mocker: MockerFixture, + duckdb_version: str, + catalog_type: t.Optional[str] = "duckdb", + **kwargs: t.Any, +) -> DuckDBEngineAdapter: + mocker.patch("duckdb.__version__", duckdb_version) + adapter = make_mocked_engine_adapter(DuckDBEngineAdapter, **kwargs) + mocker.patch.object(adapter, "get_current_catalog", return_value="memory") + mocker.patch.object(adapter, "_get_catalog_type", return_value=catalog_type) + return adapter + + +def test_merge_native_sql(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture): + adapter = _mocked_adapter(make_mocked_engine_adapter, mocker, "1.4.0") + + adapter.merge( + target_table="tgt", + source_table=parse_one("SELECT id, name, value FROM src"), + target_columns_to_types=MERGE_COLUMNS, + unique_key=[exp.to_column("id")], + ) + + assert to_sql_calls(adapter) == [ + 'MERGE INTO "tgt" AS "__MERGE_TARGET__" USING ' + '(SELECT DISTINCT ON ("id") * FROM (SELECT "id", "name", "value" FROM "src") AS "_src") AS "__MERGE_SOURCE__" ' + 'ON "__MERGE_TARGET__"."id" = "__MERGE_SOURCE__"."id" ' + 'WHEN MATCHED THEN UPDATE SET "id" = "__MERGE_SOURCE__"."id", "name" = "__MERGE_SOURCE__"."name", "value" = "__MERGE_SOURCE__"."value" ' + 'WHEN NOT MATCHED THEN INSERT ("id", "name", "value") ' + 'VALUES ("__MERGE_SOURCE__"."id", "__MERGE_SOURCE__"."name", "__MERGE_SOURCE__"."value")' + ] + + +def test_merge_native_upsert(native_merge_adapter, duck_conn): + _set_up_merge_tables(duck_conn, "(1, 'a', 10), (2, 'b', 20)", "(2, 'b2', 21), (3, 'c', 30)") + _merge(native_merge_adapter) + assert _rows(duck_conn) == [(1, "a", 10), (2, "b2", 21), (3, "c", 30)] + + +def test_merge_native_when_matched_condition(native_merge_adapter, duck_conn): + _set_up_merge_tables(duck_conn, "(1, 'a', 10), (2, 'b', 20)", "(1, 'a2', 5), (2, 'b2', 25)") + _merge( + native_merge_adapter, + when_matched=_when_matched( + "WHEN MATCHED AND __MERGE_SOURCE__.value > __MERGE_TARGET__.value " + "THEN UPDATE SET __MERGE_TARGET__.value = __MERGE_SOURCE__.value" + ), + ) + # Only the second row satisfies the condition; the name column is not in the SET list. + assert _rows(duck_conn) == [(1, "a", 10), (2, "b", 25)] + + +def test_merge_native_when_matched_delete_and_multiple_whens(native_merge_adapter, duck_conn): + _set_up_merge_tables( + duck_conn, + "(1, 'a', 10), (2, 'b', 20), (3, 'c', 30)", + "(1, 'a2', 0), (2, 'b2', 21), (4, 'd', 40)", + ) + _merge( + native_merge_adapter, + when_matched=_when_matched( + "WHEN MATCHED AND __MERGE_SOURCE__.value = 0 THEN DELETE " + "WHEN MATCHED THEN UPDATE SET __MERGE_TARGET__.value = __MERGE_SOURCE__.value" + ), + ) + assert _rows(duck_conn) == [(2, "b", 21), (3, "c", 30), (4, "d", 40)] + + +def test_merge_native_when_not_matched_by_source(native_merge_adapter, duck_conn): + _set_up_merge_tables(duck_conn, "(1, 'a', 10), (2, 'b', 20)", "(2, 'b2', 21)") + _merge( + native_merge_adapter, + when_matched=_when_matched( + "WHEN MATCHED THEN UPDATE SET __MERGE_TARGET__.value = __MERGE_SOURCE__.value " + "WHEN NOT MATCHED BY SOURCE THEN DELETE" + ), + ) + assert _rows(duck_conn) == [(2, "b", 21)] + + +def test_merge_native_merge_filter( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, native_merge_adapter, duck_conn +): + merge_filter = parse_one("__MERGE_SOURCE__.value > __MERGE_TARGET__.value", dialect="duckdb") + + # The filter is AND-ed into the ON clause + mocked = _mocked_adapter(make_mocked_engine_adapter, mocker, "1.4.0") + mocked.merge( + target_table="tgt", + source_table=parse_one("SELECT id, name, value FROM src"), + target_columns_to_types=MERGE_COLUMNS, + unique_key=[exp.to_column("id")], + merge_filter=merge_filter.copy(), + ) + assert ( + 'ON "__MERGE_SOURCE__"."value" > "__MERGE_TARGET__"."value" ' + 'AND "__MERGE_TARGET__"."id" = "__MERGE_SOURCE__"."id"' + ) in to_sql_calls(mocked)[0] + + # Row 1 passes the filter and is updated, row 2 does not and is inserted again as a new row + _set_up_merge_tables(duck_conn, "(1, 'a', 10), (2, 'b', 20)", "(1, 'a2', 11), (2, 'b2', 5)") + _merge(native_merge_adapter, merge_filter=merge_filter) + assert _rows(duck_conn) == [(1, "a2", 11), (2, "b", 20), (2, "b2", 5)] + + +def test_merge_native_deduplicates_source_keys(native_merge_adapter, duck_conn): + _set_up_merge_tables( + duck_conn, + "(1, 'a', 10)", + "(1, 'a2', 11), (1, 'a3', 12), (5, 'e', 50), (5, 'e2', 51)", + ) + _merge(native_merge_adapter) + rows = _rows(duck_conn) + assert [r[0] for r in rows] == [1, 5] + assert rows[0] in {(1, "a2", 11), (1, "a3", 12)} + assert rows[1] in {(5, "e", 50), (5, "e2", 51)} + + +def test_merge_native_null_key_and_empty_source(native_merge_adapter, duck_conn): + _set_up_merge_tables(duck_conn, "(1, 'a', 10), (NULL, 'n', 0)", "(NULL, 'n2', 1)") + _merge(native_merge_adapter) + # NULL never matches, so the NULL-keyed source row is inserted as a new row + assert _rows(duck_conn) == [(1, "a", 10), (None, "n", 0), (None, "n2", 1)] + + duck_conn.execute("DELETE FROM src") + _merge(native_merge_adapter) + assert len(_rows(duck_conn)) == 3 + + +def test_merge_native_dataframe_source(native_merge_adapter, duck_conn): + _set_up_merge_tables(duck_conn, "(1, 'a', 10)") + df = pd.DataFrame({"id": [1, 2], "name": ["a2", "b"], "value": [11, 20]}) + _merge(native_merge_adapter, source_table=df) + assert _rows(duck_conn) == [(1, "a2", 11), (2, "b", 20)] + # The temp table that backs the DataFrame is dropped afterwards + assert duck_conn.execute("SELECT COUNT(*) FROM duckdb_tables()").fetchone() == (2,) + + +def test_merge_native_dataframe_source_cleanup_on_failure(native_merge_adapter, duck_conn): + _set_up_merge_tables(duck_conn, "(1, 'a', 10)") + df = pd.DataFrame({"id": [1], "name": ["a2"], "value": [11]}) + with pytest.raises(Exception): + _merge( + native_merge_adapter, + source_table=df, + when_matched=_when_matched("WHEN MATCHED THEN UPDATE SET __MERGE_TARGET__.nope = 1"), + ) + assert duck_conn.execute("SELECT COUNT(*) FROM duckdb_tables()").fetchone() == (2,) + + +def test_merge_native_quoted_case_sensitive_identifiers(native_merge_adapter, duck_conn): + duck_conn.execute('CREATE SCHEMA "MySchema"') + duck_conn.execute('CREATE TABLE "MySchema"."MyTarget" ("ID" INT, "Name" TEXT)') + duck_conn.execute('CREATE TABLE "MySchema"."MySource" ("ID" INT, "Name" TEXT)') + duck_conn.execute('INSERT INTO "MySchema"."MyTarget" VALUES (1, \'a\')') + duck_conn.execute("INSERT INTO \"MySchema\".\"MySource\" VALUES (1, 'a2'), (2, 'b')") + native_merge_adapter.merge( + target_table=exp.to_table('"MySchema"."MyTarget"'), + source_table=parse_one('SELECT "ID", "Name" FROM "MySchema"."MySource"'), + target_columns_to_types={ + "ID": exp.DataType.build("int"), + "Name": exp.DataType.build("text"), + }, + unique_key=[exp.to_column('"ID"')], + ) + assert duck_conn.execute('SELECT * FROM "MySchema"."MyTarget" ORDER BY 1').fetchall() == [ + (1, "a2"), + (2, "b"), + ] + + +def test_merge_native_unique_key_expression(native_merge_adapter, duck_conn): + _set_up_merge_tables(duck_conn, "(1, 'a', 10)", "(1, 'a', 11), (1, 'a', 12), (1, 'b', 13)") + _merge( + native_merge_adapter, + unique_key=[parse_one("COALESCE(id, 0) || '|' || name", dialect="duckdb")], + ) + # The existing 'a' row is matched through the key expression and updated (one of the two + # duplicate source rows wins), 'b' is inserted. A mis-qualified key in ON would insert instead. + rows = _rows(duck_conn) + assert len(rows) == 2 + assert rows[0] in {(1, "a", 11), (1, "a", 12)} + assert rows[1] == (1, "b", 13) + + +def test_merge_version_gate_falls_back_to_delete_insert( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +): + adapter = _mocked_adapter(make_mocked_engine_adapter, mocker, "1.3.2") + + adapter.merge( + target_table="tgt", + source_table=parse_one("SELECT id, name, value FROM src"), + target_columns_to_types=MERGE_COLUMNS, + unique_key=[exp.to_column("id")], + ) + sql_calls = to_sql_calls(adapter) + assert not any(call.startswith("MERGE") for call in sql_calls) + assert any(call.startswith('DELETE FROM "tgt"') for call in sql_calls) + assert any(call.startswith('INSERT INTO "tgt"') for call in sql_calls) + + with pytest.raises(SQLMeshError, match="`when_matched` is not supported"): + adapter.merge( + target_table="tgt", + source_table=parse_one("SELECT id, name, value FROM src"), + target_columns_to_types=MERGE_COLUMNS, + unique_key=[exp.to_column("id")], + when_matched=_when_matched("WHEN MATCHED THEN DELETE"), + ) + + +@pytest.mark.parametrize("catalog_type", ["postgres", "sqlite", "mysql", None]) +def test_merge_catalog_gate_falls_back_to_delete_insert( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, catalog_type +): + adapter = _mocked_adapter(make_mocked_engine_adapter, mocker, "1.5.0", catalog_type) + mocker.patch.object(adapter, "get_current_catalog", return_value="pg") + + adapter.merge( + target_table="pg.s.tgt", + source_table=parse_one("SELECT id, name, value FROM src"), + target_columns_to_types=MERGE_COLUMNS, + unique_key=[exp.to_column("id")], + ) + sql_calls = to_sql_calls(adapter) + assert not any(call.startswith("MERGE") for call in sql_calls) + assert any(call.startswith('DELETE FROM "pg"."s"."tgt"') for call in sql_calls) + # The catalog type is looked up for the target's own catalog + adapter._get_catalog_type.assert_called_with("pg") # type: ignore + + +def test_merge_catalog_gate_uses_target_catalog_when_qualified( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +): + adapter = _mocked_adapter(make_mocked_engine_adapter, mocker, "1.5.0", "ducklake") + get_type = t.cast(MagicMock, adapter._get_catalog_type) + + adapter.merge( + target_table="lake.s.tgt", + source_table=parse_one("SELECT id, name, value FROM src"), + target_columns_to_types=MERGE_COLUMNS, + unique_key=[exp.to_column("id")], + ) + get_type.assert_called_once_with("lake") + assert to_sql_calls(adapter)[0].startswith('MERGE INTO "lake"."s"."tgt"') + + +def test_merge_motherduck_uses_native_without_catalog_lookup( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +): + adapter = _mocked_adapter( + make_mocked_engine_adapter, mocker, "1.4.1", "motherduck", is_motherduck=True + ) + get_type = t.cast(MagicMock, adapter._get_catalog_type) + + adapter.merge( + target_table="md.s.tgt", + source_table=parse_one("SELECT id, name, value FROM src"), + target_columns_to_types=MERGE_COLUMNS, + unique_key=[exp.to_column("id")], + ) + get_type.assert_not_called() + assert to_sql_calls(adapter)[0].startswith('MERGE INTO "md"."s"."tgt"') + + +def test_merge_motherduck_old_client_falls_back( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +): + adapter = _mocked_adapter(make_mocked_engine_adapter, mocker, "1.3.2", is_motherduck=True) + + adapter.merge( + target_table="md.s.tgt", + source_table=parse_one("SELECT id, name, value FROM src"), + target_columns_to_types=MERGE_COLUMNS, + unique_key=[exp.to_column("id")], + ) + assert not any(call.startswith("MERGE") for call in to_sql_calls(adapter)) + + +def test_merge_does_not_strip_qualifiers_from_right_hand_side(native_merge_adapter, duck_conn): + _set_up_merge_tables(duck_conn, "(1, 'a', 10)", "(1, 'a', 1)") + _merge( + native_merge_adapter, + when_matched=_when_matched( + "WHEN MATCHED THEN UPDATE SET __MERGE_TARGET__.value = __MERGE_TARGET__.value + __MERGE_SOURCE__.value" + ), + ) + assert _rows(duck_conn) == [(1, "a", 11)] + + +@pytest.mark.isolated +@pytest.mark.xdist_group("ducklake_extension") +@pytest.mark.skipif( + not _installed_duckdb_supports_merge(), reason="Installed duckdb doesn't support MERGE" +) +def test_merge_native_ducklake(ducklake_conn, tmp_path, mocker: MockerFixture): + mocker.patch("duckdb.__version__", "1.4.0") + adapter = DuckDBEngineAdapter(lambda: ducklake_conn) + ducklake_conn.execute( + f"ATTACH 'ducklake:{tmp_path}/lake.ducklake' AS lake (DATA_PATH '{tmp_path}')" + ) + ducklake_conn.execute("CREATE SCHEMA lake.s") + ducklake_conn.execute("CREATE TABLE lake.s.tgt (id INT, name TEXT, value INT)") + ducklake_conn.execute("CREATE TABLE lake.s.src (id INT, name TEXT, value INT)") + ducklake_conn.execute("INSERT INTO lake.s.tgt VALUES (1, 'a', 10), (2, 'b', 20)") + ducklake_conn.execute( + "INSERT INTO lake.s.src VALUES (2, 'b2', 21), (3, 'c', 30), (3, 'c2', 31)" + ) + + spy = mocker.spy(adapter, "_merge") + adapter.merge( + target_table="lake.s.tgt", + source_table=parse_one("SELECT id, name, value FROM lake.s.src"), + target_columns_to_types=MERGE_COLUMNS, + unique_key=[exp.to_column("id")], + when_matched=_when_matched( + "WHEN MATCHED THEN UPDATE SET __MERGE_TARGET__.value = __MERGE_SOURCE__.value" + ), + ) + spy.assert_called_once() + rows = ducklake_conn.execute("SELECT * FROM lake.s.tgt ORDER BY id").fetchall() + assert [r[0] for r in rows] == [1, 2, 3] + assert rows[1] == (2, "b", 21)