Add merge query to engine adapter

sqlmesh/core/engine_adapter.py

@@ -273,6 +273,53 @@ def update_table(
     ) -> None:
         self.execute(exp.update(table_name, properties, where=where))
 
+    def merge(
+        self,
+        target_table: str,
+        source_table: str,
+        columns: t.Iterable[str],
+        unique_keys: t.Iterable[str],
+    ):
+        on = exp.and_(
+            *(
+                exp.EQ(
+                    this=exp.column(key, target_table),
+                    expression=exp.column(key, source_table),
+                )
+                for key in unique_keys
+            )
+        )
+        when_matched = exp.When(
+            this="MATCHED",
+            then=exp.update(
+                None,
+                properties={
+                    exp.column(col, target_table): exp.column(col, source_table)
+                    for col in columns
+                },
+            ),
+        )
+        when_not_matched = exp.When(
+            this=exp.Not(this="MATCHED"),
+            then=exp.Insert(
+                this=exp.Tuple(expressions=[exp.column(col) for col in columns]),
+                expression=exp.Tuple(
+                    expressions=[exp.column(col, source_table) for col in columns]
+                ),
+            ),
+        )
+        self.execute(
+            exp.Merge(
+                this=target_table,
+                using=source_table,
+                on=on,
+                expressions=[
+                    when_matched,
+                    when_not_matched,
+                ],
+            )
+        )
+
     def fetchone(self, query: t.Union[exp.Expression, str]) -> t.Tuple:
         self.execute(query)
         return self.cursor.fetchone()

tests/core/test_engine_adapter.py

@@ -273,3 +273,35 @@ def test_transaction_duckdb(adapter: EngineAdapter, duck_conn):
     except Exception:
         pass
     assert duck_conn.execute("SELECT * FROM test_table").fetchall() == [(1,)]
+
+
+def test_merge(mocker: MockerFixture):
+    connection_mock = mocker.NonCallableMock()
+    cursor_mock = mocker.Mock()
+    connection_mock.cursor.return_value = cursor_mock
+
+    adapter = EngineAdapter(lambda: connection_mock, "spark")  # type: ignore
+    adapter.merge(
+        target_table="target",
+        source_table="source",
+        columns=["id", "ts", "val"],
+        unique_keys=["id"],
+    )
+    cursor_mock.execute.assert_called_once_with(
+        "MERGE INTO target USING source ON `target`.`id` = `source`.`id` "
+        "WHEN MATCHED THEN UPDATE SET `target`.`id` = `source`.`id`, `target`.`ts` = `source`.`ts`, `target`.`val` = `source`.`val` "
+        "WHEN NOT MATCHED THEN INSERT (`id`, `ts`, `val`) VALUES (`source`.`id`, `source`.`ts`, `source`.`val`)"
+    )
+
+    cursor_mock.reset_mock()
+    adapter.merge(
+        target_table="target",
+        source_table="source",
+        columns=["id", "ts", "val"],
+        unique_keys=["id", "ts"],
+    )
+    cursor_mock.execute.assert_called_once_with(
+        "MERGE INTO target USING source ON `target`.`id` = `source`.`id` AND `target`.`ts` = `source`.`ts` "
+        "WHEN MATCHED THEN UPDATE SET `target`.`id` = `source`.`id`, `target`.`ts` = `source`.`ts`, `target`.`val` = `source`.`val` "
+        "WHEN NOT MATCHED THEN INSERT (`id`, `ts`, `val`) VALUES (`source`.`id`, `source`.`ts`, `source`.`val`)"
+    )