Repository: TobikoData/sqlmesh
Pull request #3272 (2024-10-21 15:53:49)

Fix: Support merging for Postgres version < 15
Support for Merge statement has only been introduced in Postgres 15.


Patch:
diff --git a/sqlmesh/core/engine_adapter/mixins.py b/sqlmesh/core/engine_adapter/mixins.py
index 0aa35145..a55b4c00 100644
--- a/sqlmesh/core/engine_adapter/mixins.py
+++ b/sqlmesh/core/engine_adapter/mixins.py
@@ -28,22 +28,13 @@ class LogicalMergeMixin(EngineAdapter):
         unique_key: t.Sequence[exp.Expression],
         when_matched: t.Optional[t.Union[exp.When, t.List[exp.When]]] = None,
     ) -> None:
-        """
-        Merge implementation for engine adapters that do not support merge natively.
-
-        The merge is executed as follows:
-        1. Create a temporary table containing the new data to merge.
-        2. Delete rows from target table where unique_key cols match a row in the temporary table.
-        3. Insert the temporary table contents into the target table. Any duplicate, non-unique rows
-           within the temporary table are ommitted.
-        4. Drop the temporary table.
-        """
-        if when_matched:
-            raise SQLMeshError(
-                "This engine does not support MERGE expressions and therefore `when_matched` is not supported."
-            )
-        self._replace_by_key(
-            target_table, source_table, columns_to_types, unique_key, is_unique_key=True
+        logical_merge(
+            self,
+            target_table,
+            source_table,
+            columns_to_types,
+            unique_key,
+            when_matched=when_matched,
         )
 
 
@@ -396,3 +387,30 @@ class ClusteredByMixin(EngineAdapter):
             kind="TABLE",
             actions=[exp.Command(this="DROP", expression="CLUSTERING KEY")],
         )
+
+
+def logical_merge(
+    engine_adapter: EngineAdapter,
+    target_table: TableName,
+    source_table: QueryOrDF,
+    columns_to_types: t.Optional[t.Dict[str, exp.DataType]],
+    unique_key: t.Sequence[exp.Expression],
+    when_matched: t.Optional[t.Union[exp.When, t.List[exp.When]]] = None,
+) -> None:
+    """
+    Merge implementation for engine adapters that do not support merge natively.
+
+    The merge is executed as follows:
+    1. Create a temporary table containing the new data to merge.
+    2. Delete rows from target table where unique_key cols match a row in the temporary table.
+    3. Insert the temporary table contents into the target table. Any duplicate, non-unique rows
+       within the temporary table are ommitted.
+    4. Drop the temporary table.
+    """
+    if when_matched:
+        raise SQLMeshError(
+            "This engine does not support MERGE expressions and therefore `when_matched` is not supported."
+        )
+    engine_adapter._replace_by_key(
+        target_table, source_table, columns_to_types, unique_key, is_unique_key=True
+    )
diff --git a/sqlmesh/core/engine_adapter/postgres.py b/sqlmesh/core/engine_adapter/postgres.py
index 8fdc4ce2..bd90857e 100644
--- a/sqlmesh/core/engine_adapter/postgres.py
+++ b/sqlmesh/core/engine_adapter/postgres.py
@@ -2,19 +2,21 @@ from __future__ import annotations
 
 import logging
 import typing as t
+from functools import partial
 from sqlglot import exp
 
 from sqlmesh.core.engine_adapter.base_postgres import BasePostgresEngineAdapter
 from sqlmesh.core.engine_adapter.mixins import (
     GetCurrentCatalogFromFunctionMixin,
     PandasNativeFetchDFSupportMixin,
+    logical_merge,
 )
 from sqlmesh.core.engine_adapter.shared import set_catalog
 from sqlmesh.core.schema_diff import SchemaDiffer
 
 if t.TYPE_CHECKING:
     from sqlmesh.core._typing import TableName
-    from sqlmesh.core.engine_adapter._typing import DF
+    from sqlmesh.core.engine_adapter._typing import DF, QueryOrDF
 
 logger = logging.getLogger(__name__)
 
@@ -94,3 +96,25 @@ class PostgresEngineAdapter(
                 exists=exists,
             )
         )
+
+    def merge(
+        self,
+        target_table: TableName,
+        source_table: QueryOrDF,
+        columns_to_types: t.Optional[t.Dict[str, exp.DataType]],
+        unique_key: t.Sequence[exp.Expression],
+        when_matched: t.Optional[t.Union[exp.When, t.List[exp.When]]] = None,
+    ) -> None:
+        # Merge isn't supported until Postgres 15
+        merge_impl = (
+            super().merge
+            if self._connection_pool.get().server_version >= 150000
+            else partial(logical_merge, self)
+        )
+        merge_impl(  # type: ignore
+            target_table,
+            source_table,
+            columns_to_types,
+            unique_key,
+            when_matched=when_matched,
+        )


Test patch:
diff --git a/tests/core/engine_adapter/test_postgres.py b/tests/core/engine_adapter/test_postgres.py
index 3e622ef8..5400c6e6 100644
--- a/tests/core/engine_adapter/test_postgres.py
+++ b/tests/core/engine_adapter/test_postgres.py
@@ -3,7 +3,7 @@ import typing as t
 import pytest
 from pytest_mock import MockFixture
 from pytest_mock.plugin import MockerFixture
-from sqlglot import exp
+from sqlglot import exp, parse_one
 from sqlglot.helper import ensure_list
 
 from sqlmesh.core.engine_adapter import PostgresEngineAdapter
@@ -89,3 +89,55 @@ def test_create_table_like(make_mocked_engine_adapter: t.Callable):
     adapter.cursor.execute.assert_called_once_with(
         'CREATE TABLE IF NOT EXISTS "target_table" (LIKE "source_table" INCLUDING ALL)'
     )
+
+
+def test_merge_version_gte_15(make_mocked_engine_adapter: t.Callable):
+    adapter = make_mocked_engine_adapter(PostgresEngineAdapter)
+    adapter._connection_pool.get().server_version = 150000
+
+    adapter.merge(
+        target_table="target",
+        source_table=t.cast(exp.Select, parse_one('SELECT "ID", ts, val FROM source')),
+        columns_to_types={
+            "ID": exp.DataType.build("int"),
+            "ts": exp.DataType.build("timestamp"),
+            "val": exp.DataType.build("int"),
+        },
+        unique_key=[exp.to_identifier("ID", quoted=True)],
+    )
+
+    sql_calls = to_sql_calls(adapter)
+    assert sql_calls == [
+        """MERGE INTO "target" AS "__MERGE_TARGET__" USING (SELECT "ID", "ts", "val" FROM "source") AS "__MERGE_SOURCE__" ON "__MERGE_TARGET__"."ID" = "__MERGE_SOURCE__"."ID" WHEN MATCHED THEN UPDATE SET "ID" = "__MERGE_SOURCE__"."ID", "ts" = "__MERGE_SOURCE__"."ts", "val" = "__MERGE_SOURCE__"."val" WHEN NOT MATCHED THEN INSERT ("ID", "ts", "val") VALUES ("__MERGE_SOURCE__"."ID", "__MERGE_SOURCE__"."ts", "__MERGE_SOURCE__"."val")"""
+    ]
+
+
+def test_merge_version_lt_15(
+    make_mocked_engine_adapter: t.Callable, make_temp_table_name: t.Callable, mocker: MockerFixture
+):
+    adapter = make_mocked_engine_adapter(PostgresEngineAdapter)
+    adapter._connection_pool.get().server_version = 140000
+
+    temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
+    table_name = "test"
+    temp_table_id = "abcdefgh"
+    temp_table_mock.return_value = make_temp_table_name(table_name, temp_table_id)
+
+    adapter.merge(
+        target_table="target",
+        source_table=t.cast(exp.Select, parse_one('SELECT "ID", ts, val FROM source')),
+        columns_to_types={
+            "ID": exp.DataType.build("int"),
+            "ts": exp.DataType.build("timestamp"),
+            "val": exp.DataType.build("int"),
+        },
+        unique_key=[exp.to_identifier("ID", quoted=True)],
+    )
+
+    sql_calls = to_sql_calls(adapter)
+    assert sql_calls == [
+        'CREATE TABLE "__temp_test_abcdefgh" AS SELECT CAST("ID" AS INT) AS "ID", CAST("ts" AS TIMESTAMP) AS "ts", CAST("val" AS INT) AS "val" FROM (SELECT "ID", "ts", "val" FROM "source") AS "_subquery"',
+        'DELETE FROM "target" WHERE "ID" IN (SELECT "ID" FROM "__temp_test_abcdefgh")',
+        'INSERT INTO "target" ("ID", "ts", "val") SELECT DISTINCT ON ("ID") "ID", "ts", "val" FROM "__temp_test_abcdefgh"',
+        'DROP TABLE IF EXISTS "__temp_test_abcdefgh"',
+    ]