Make MSSQL merge exists implementation opt-in

Author: treysp

sqlmesh/core/engine_adapter/base.py

@@ -1939,6 +1939,7 @@ def merge(
         unique_key: t.Sequence[exp.Expression],
         when_matched: t.Optional[exp.Whens] = None,
         merge_filter: t.Optional[exp.Expression] = None,
+        **kwargs: t.Any,
     ) -> None:
         source_queries, columns_to_types = self._get_source_queries_and_columns_to_types(
             source_table, columns_to_types, target_table=target_table

sqlmesh/core/engine_adapter/mixins.py

@@ -32,6 +32,7 @@ def merge(
         unique_key: t.Sequence[exp.Expression],
         when_matched: t.Optional[exp.Whens] = None,
         merge_filter: t.Optional[exp.Expression] = None,
+        **kwargs: t.Any,
     ) -> None:
         logical_merge(
             self,

sqlmesh/core/engine_adapter/mssql.py

@@ -198,80 +198,97 @@ def merge(
         unique_key: t.Sequence[exp.Expression],
         when_matched: t.Optional[exp.Whens] = None,
         merge_filter: t.Optional[exp.Expression] = None,
+        **kwargs: t.Any,
     ) -> None:
-        source_queries, columns_to_types = self._get_source_queries_and_columns_to_types(
-            source_table, columns_to_types, target_table=target_table
-        )
-        columns_to_types = columns_to_types or self.columns(target_table)
-        on = exp.and_(
-            *(
-                add_table(part, MERGE_TARGET_ALIAS).eq(add_table(part, MERGE_SOURCE_ALIAS))
-                for part in unique_key
+        mssql_merge_exists = kwargs.get("physical_properties", {}).get("mssql_merge_exists")
+
+        if mssql_merge_exists:
+            source_queries, columns_to_types = self._get_source_queries_and_columns_to_types(
+                source_table, columns_to_types, target_table=target_table
             )
-        )
-        if merge_filter:
-            on = exp.and_(merge_filter, on)
-
-        match_expressions = []
-        if not when_matched:
-            match_condition = None
-            unique_key_names = [y.name for y in unique_key]
-            columns_to_types_no_keys = [c for c in columns_to_types if c not in unique_key_names]
-
-            target_columns_no_keys = [
-                exp.column(c, MERGE_TARGET_ALIAS) for c in columns_to_types_no_keys
-            ]
-            source_columns_no_keys = [
-                exp.column(c, MERGE_SOURCE_ALIAS) for c in columns_to_types_no_keys
-            ]
-
-            match_condition = exp.Exists(
-                this=exp.select(*target_columns_no_keys).except_(
-                    exp.select(*source_columns_no_keys)
+            columns_to_types = columns_to_types or self.columns(target_table)
+            on = exp.and_(
+                *(
+                    add_table(part, MERGE_TARGET_ALIAS).eq(add_table(part, MERGE_SOURCE_ALIAS))
+                    for part in unique_key
                 )
             )
+            if merge_filter:
+                on = exp.and_(merge_filter, on)
+
+            match_expressions = []
+            if not when_matched:
+                match_condition = None
+                unique_key_names = [y.name for y in unique_key]
+                columns_to_types_no_keys = [
+                    c for c in columns_to_types if c not in unique_key_names
+                ]
+
+                target_columns_no_keys = [
+                    exp.column(c, MERGE_TARGET_ALIAS) for c in columns_to_types_no_keys
+                ]
+                source_columns_no_keys = [
+                    exp.column(c, MERGE_SOURCE_ALIAS) for c in columns_to_types_no_keys
+                ]
+
+                match_condition = exp.Exists(
+                    this=exp.select(*target_columns_no_keys).except_(
+                        exp.select(*source_columns_no_keys)
+                    )
+                )
 
-            if target_columns_no_keys:
-                match_expressions.append(
-                    exp.When(
-                        matched=True,
-                        source=False,
-                        condition=match_condition,
-                        then=exp.Update(
+                if target_columns_no_keys:
+                    match_expressions.append(
+                        exp.When(
+                            matched=True,
+                            source=False,
+                            condition=match_condition,
+                            then=exp.Update(
+                                expressions=[
+                                    exp.column(col, MERGE_TARGET_ALIAS).eq(
+                                        exp.column(col, MERGE_SOURCE_ALIAS)
+                                    )
+                                    for col in columns_to_types_no_keys
+                                ],
+                            ),
+                        )
+                    )
+            else:
+                match_expressions.extend(when_matched.copy().expressions)
+
+            match_expressions.append(
+                exp.When(
+                    matched=False,
+                    source=False,
+                    then=exp.Insert(
+                        this=exp.Tuple(expressions=[exp.column(col) for col in columns_to_types]),
+                        expression=exp.Tuple(
                             expressions=[
-                                exp.column(col, MERGE_TARGET_ALIAS).eq(
-                                    exp.column(col, MERGE_SOURCE_ALIAS)
-                                )
-                                for col in columns_to_types_no_keys
-                            ],
+                                exp.column(col, MERGE_SOURCE_ALIAS) for col in columns_to_types
+                            ]
                         ),
-                    )
-                )
-        else:
-            match_expressions.extend(when_matched.copy().expressions)
-
-        match_expressions.append(
-            exp.When(
-                matched=False,
-                source=False,
-                then=exp.Insert(
-                    this=exp.Tuple(expressions=[exp.column(col) for col in columns_to_types]),
-                    expression=exp.Tuple(
-                        expressions=[
-                            exp.column(col, MERGE_SOURCE_ALIAS) for col in columns_to_types
-                        ]
                     ),
-                ),
+                )
             )
+            for source_query in source_queries:
+                with source_query as query:
+                    self._merge(
+                        target_table=target_table,
+                        query=query,
+                        on=on,
+                        whens=exp.Whens(expressions=match_expressions),
+                    )
+            return
+
+        super().merge(
+            target_table,
+            source_table,
+            columns_to_types,
+            unique_key,
+            when_matched,
+            merge_filter,
+            **kwargs,
         )
-        for source_query in source_queries:
-            with source_query as query:
-                self._merge(
-                    target_table=target_table,
-                    query=query,
-                    on=on,
-                    whens=exp.Whens(expressions=match_expressions),
-                )
 
     def _convert_df_datetime(self, df: DF, columns_to_types: t.Dict[str, exp.DataType]) -> None:
         import pandas as pd

sqlmesh/core/engine_adapter/postgres.py

@@ -109,6 +109,7 @@ def merge(
         unique_key: t.Sequence[exp.Expression],
         when_matched: t.Optional[exp.Whens] = None,
         merge_filter: t.Optional[exp.Expression] = None,
+        **kwargs: t.Any,
     ) -> None:
         # Merge isn't supported until Postgres 15
         merge_impl = (

sqlmesh/core/engine_adapter/redshift.py

@@ -353,6 +353,7 @@ def merge(
         unique_key: t.Sequence[exp.Expression],
         when_matched: t.Optional[exp.Whens] = None,
         merge_filter: t.Optional[exp.Expression] = None,
+        **kwargs: t.Any,
     ) -> None:
         if self.enable_merge:
             # By default we use the logical merge unless the user has opted in

sqlmesh/core/snapshot/evaluator.py

@@ -1614,6 +1614,7 @@ def insert(
                     end=kwargs.get("end"),
                     execution_time=kwargs.get("execution_time"),
                 ),
+                physical_properties=kwargs.get("physical_properties", model.physical_properties),
             )
 
     def append(
@@ -1634,6 +1635,7 @@ def append(
                 end=kwargs.get("end"),
                 execution_time=kwargs.get("execution_time"),
             ),
+            physical_properties=kwargs.get("physical_properties", model.physical_properties),
         )
 
 

tests/core/engine_adapter/test_mssql.py

@@ -474,7 +474,7 @@ def test_merge_pandas(
 
     assert to_sql_calls(adapter) == [
         f"""IF NOT EXISTS (SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = '__temp_target_{temp_table_id}') EXEC('CREATE TABLE [__temp_target_{temp_table_id}] ([id] INTEGER, [ts] DATETIME2, [val] INTEGER)');""",
-        f"MERGE INTO [target] AS [__MERGE_TARGET__] USING (SELECT CAST([id] AS INTEGER) AS [id], CAST([ts] AS DATETIME2) AS [ts], CAST([val] AS INTEGER) AS [val] FROM [__temp_target_{temp_table_id}]) AS [__MERGE_SOURCE__] ON [__MERGE_TARGET__].[id] = [__MERGE_SOURCE__].[id] WHEN MATCHED AND EXISTS(SELECT [__MERGE_TARGET__].[ts], [__MERGE_TARGET__].[val] EXCEPT SELECT [__MERGE_SOURCE__].[ts], [__MERGE_SOURCE__].[val]) THEN UPDATE SET [__MERGE_TARGET__].[ts] = [__MERGE_SOURCE__].[ts], [__MERGE_TARGET__].[val] = [__MERGE_SOURCE__].[val] WHEN NOT MATCHED THEN INSERT ([id], [ts], [val]) VALUES ([__MERGE_SOURCE__].[id], [__MERGE_SOURCE__].[ts], [__MERGE_SOURCE__].[val]);",
+        f"MERGE INTO [target] AS [__MERGE_TARGET__] USING (SELECT CAST([id] AS INTEGER) AS [id], CAST([ts] AS DATETIME2) AS [ts], CAST([val] AS INTEGER) AS [val] FROM [__temp_target_{temp_table_id}]) AS [__MERGE_SOURCE__] ON [__MERGE_TARGET__].[id] = [__MERGE_SOURCE__].[id] WHEN MATCHED THEN UPDATE SET [__MERGE_TARGET__].[id] = [__MERGE_SOURCE__].[id], [__MERGE_TARGET__].[ts] = [__MERGE_SOURCE__].[ts], [__MERGE_TARGET__].[val] = [__MERGE_SOURCE__].[val] WHEN NOT MATCHED THEN INSERT ([id], [ts], [val]) VALUES ([__MERGE_SOURCE__].[id], [__MERGE_SOURCE__].[ts], [__MERGE_SOURCE__].[val]);",
         f"DROP TABLE IF EXISTS [__temp_target_{temp_table_id}];",
     ]
 
@@ -498,11 +498,47 @@ def test_merge_pandas(
 
     assert to_sql_calls(adapter) == [
         f"""IF NOT EXISTS (SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = '__temp_target_{temp_table_id}') EXEC('CREATE TABLE [__temp_target_{temp_table_id}] ([id] INTEGER, [ts] DATETIME2, [val] INTEGER)');""",
-        f"MERGE INTO [target] AS [__MERGE_TARGET__] USING (SELECT CAST([id] AS INTEGER) AS [id], CAST([ts] AS DATETIME2) AS [ts], CAST([val] AS INTEGER) AS [val] FROM [__temp_target_{temp_table_id}]) AS [__MERGE_SOURCE__] ON [__MERGE_TARGET__].[id] = [__MERGE_SOURCE__].[id] AND [__MERGE_TARGET__].[ts] = [__MERGE_SOURCE__].[ts] WHEN MATCHED AND EXISTS(SELECT [__MERGE_TARGET__].[val] EXCEPT SELECT [__MERGE_SOURCE__].[val]) THEN UPDATE SET [__MERGE_TARGET__].[val] = [__MERGE_SOURCE__].[val] WHEN NOT MATCHED THEN INSERT ([id], [ts], [val]) VALUES ([__MERGE_SOURCE__].[id], [__MERGE_SOURCE__].[ts], [__MERGE_SOURCE__].[val]);",
+        f"MERGE INTO [target] AS [__MERGE_TARGET__] USING (SELECT CAST([id] AS INTEGER) AS [id], CAST([ts] AS DATETIME2) AS [ts], CAST([val] AS INTEGER) AS [val] FROM [__temp_target_{temp_table_id}]) AS [__MERGE_SOURCE__] ON [__MERGE_TARGET__].[id] = [__MERGE_SOURCE__].[id] AND [__MERGE_TARGET__].[ts] = [__MERGE_SOURCE__].[ts] WHEN MATCHED THEN UPDATE SET [__MERGE_TARGET__].[id] = [__MERGE_SOURCE__].[id], [__MERGE_TARGET__].[ts] = [__MERGE_SOURCE__].[ts], [__MERGE_TARGET__].[val] = [__MERGE_SOURCE__].[val] WHEN NOT MATCHED THEN INSERT ([id], [ts], [val]) VALUES ([__MERGE_SOURCE__].[id], [__MERGE_SOURCE__].[ts], [__MERGE_SOURCE__].[val]);",
+        f"DROP TABLE IF EXISTS [__temp_target_{temp_table_id}];",
+    ]
+
+
+def test_merge_exists(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name: t.Callable
+):
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.mssql.MSSQLEngineAdapter.table_exists",
+        return_value=False,
+    )
+
+    adapter = make_mocked_engine_adapter(MSSQLEngineAdapter)
+
+    temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
+    table_name = "target"
+    temp_table_id = "abcdefgh"
+    temp_table_mock.return_value = make_temp_table_name(table_name, temp_table_id)
+
+    df = pd.DataFrame({"id": [1, 2, 3], "ts": [1, 2, 3], "val": [4, 5, 6]})
+
+    # regular implementation
+    adapter.merge(
+        target_table=table_name,
+        source_table=df,
+        columns_to_types={
+            "id": exp.DataType.build("int"),
+            "ts": exp.DataType.build("TIMESTAMP"),
+            "val": exp.DataType.build("int"),
+        },
+        unique_key=[exp.to_identifier("id")],
+    )
+
+    assert to_sql_calls(adapter) == [
+        f"""IF NOT EXISTS (SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = '__temp_target_{temp_table_id}') EXEC('CREATE TABLE [__temp_target_{temp_table_id}] ([id] INTEGER, [ts] DATETIME2, [val] INTEGER)');""",
+        f"MERGE INTO [target] AS [__MERGE_TARGET__] USING (SELECT CAST([id] AS INTEGER) AS [id], CAST([ts] AS DATETIME2) AS [ts], CAST([val] AS INTEGER) AS [val] FROM [__temp_target_{temp_table_id}]) AS [__MERGE_SOURCE__] ON [__MERGE_TARGET__].[id] = [__MERGE_SOURCE__].[id] WHEN MATCHED THEN UPDATE SET [__MERGE_TARGET__].[id] = [__MERGE_SOURCE__].[id], [__MERGE_TARGET__].[ts] = [__MERGE_SOURCE__].[ts], [__MERGE_TARGET__].[val] = [__MERGE_SOURCE__].[val] WHEN NOT MATCHED THEN INSERT ([id], [ts], [val]) VALUES ([__MERGE_SOURCE__].[id], [__MERGE_SOURCE__].[ts], [__MERGE_SOURCE__].[val]);",
         f"DROP TABLE IF EXISTS [__temp_target_{temp_table_id}];",
     ]
 
-    # all model columns are keys
+    # merge exists implementation
     adapter.cursor.reset_mock()
     adapter._connection_pool.get().reset_mock()
     temp_table_mock.return_value = make_temp_table_name(table_name, temp_table_id)
@@ -512,11 +548,31 @@ def test_merge_pandas(
         columns_to_types={
             "id": exp.DataType.build("int"),
             "ts": exp.DataType.build("TIMESTAMP"),
+            "val": exp.DataType.build("int"),
         },
-        unique_key=[exp.to_identifier("id"), exp.to_column("ts")],
+        unique_key=[exp.to_identifier("id")],
+        physical_properties={"mssql_merge_exists": True},
     )
-    adapter._connection_pool.get().bulk_copy.assert_called_with(
-        f"__temp_target_{temp_table_id}", [(1, 1), (2, 2), (3, 3)]
+
+    assert to_sql_calls(adapter) == [
+        f"""IF NOT EXISTS (SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = '__temp_target_{temp_table_id}') EXEC('CREATE TABLE [__temp_target_{temp_table_id}] ([id] INTEGER, [ts] DATETIME2, [val] INTEGER)');""",
+        f"MERGE INTO [target] AS [__MERGE_TARGET__] USING (SELECT CAST([id] AS INTEGER) AS [id], CAST([ts] AS DATETIME2) AS [ts], CAST([val] AS INTEGER) AS [val] FROM [__temp_target_{temp_table_id}]) AS [__MERGE_SOURCE__] ON [__MERGE_TARGET__].[id] = [__MERGE_SOURCE__].[id] WHEN MATCHED AND EXISTS(SELECT [__MERGE_TARGET__].[ts], [__MERGE_TARGET__].[val] EXCEPT SELECT [__MERGE_SOURCE__].[ts], [__MERGE_SOURCE__].[val]) THEN UPDATE SET [__MERGE_TARGET__].[ts] = [__MERGE_SOURCE__].[ts], [__MERGE_TARGET__].[val] = [__MERGE_SOURCE__].[val] WHEN NOT MATCHED THEN INSERT ([id], [ts], [val]) VALUES ([__MERGE_SOURCE__].[id], [__MERGE_SOURCE__].[ts], [__MERGE_SOURCE__].[val]);",
+        f"DROP TABLE IF EXISTS [__temp_target_{temp_table_id}];",
+    ]
+
+    # merge exists and all model columns are keys
+    adapter.cursor.reset_mock()
+    adapter._connection_pool.get().reset_mock()
+    temp_table_mock.return_value = make_temp_table_name(table_name, temp_table_id)
+    adapter.merge(
+        target_table=table_name,
+        source_table=df,
+        columns_to_types={
+            "id": exp.DataType.build("int"),
+            "ts": exp.DataType.build("TIMESTAMP"),
+        },
+        unique_key=[exp.to_identifier("id"), exp.to_column("ts")],
+        physical_properties={"mssql_merge_exists": True},
     )
 
     assert to_sql_calls(adapter) == [

tests/core/test_snapshot_evaluator.py

@@ -2225,6 +2225,7 @@ def test_create_incremental_by_unique_key_updated_at_exp(adapter_mock, make_snap
                 )
             ]
         ),
+        physical_properties={},
     )
 
 
@@ -2307,6 +2308,7 @@ def test_create_incremental_by_unique_key_multiple_updated_at_exp(adapter_mock,
                 ),
             ],
         ),
+        physical_properties={},
     )
 
 
@@ -2452,6 +2454,7 @@ def test_create_incremental_by_unique_key_merge_filter(adapter_mock, make_snapsh
                 expression=exp.Literal(this="2020-01-01", is_string=True),
             ),
         ),
+        physical_properties={},
     )