Skip to content

Commit 7cfb218

Browse files
mattiasthalenmattiasthalen
authored
Feat: use proper MERGE statement for Fabric instead of the logical equivalent (#5586)
1 parent 922ec89 commit 7cfb218

File tree

2 files changed

+144
-2
lines changed

2 files changed

+144
-2
lines changed

‎sqlmesh/core/engine_adapter/fabric.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
from functools import cached_property
88
from sqlglot import exp
99
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_result
10-
from sqlmesh.core.engine_adapter.mixins import LogicalMergeMixin
1110
from sqlmesh.core.engine_adapter.mssql import MSSQLEngineAdapter
1211
from sqlmesh.core.engine_adapter.shared import (
1312
InsertOverwriteStrategy,
@@ -19,7 +18,7 @@
1918
logger = logging.getLogger(__name__)
2019

2120

22-
class FabricEngineAdapter(LogicalMergeMixin, MSSQLEngineAdapter):
21+
class FabricEngineAdapter(MSSQLEngineAdapter):
2322
"""
2423
Adapter for Microsoft Fabric.
2524
"""

‎tests/core/engine_adapter/test_fabric.py‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import typing as t
44

5+
import pandas as pd # noqa: TID253
56
import pytest
67
from pytest_mock import MockerFixture
78
from sqlglot import exp, parse_one
@@ -88,3 +89,145 @@ def test_replace_query(adapter: FabricEngineAdapter, mocker: MockerFixture):
8889
"TRUNCATE TABLE [test_table];",
8990
"INSERT INTO [test_table] ([a]) SELECT [a] FROM [tbl];",
9091
]
92+
93+
94+
def test_merge_pandas(
95+
make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name: t.Callable
96+
):
97+
mocker.patch(
98+
"sqlmesh.core.engine_adapter.fabric.FabricEngineAdapter.table_exists",
99+
return_value=False,
100+
)
101+
102+
adapter = make_mocked_engine_adapter(FabricEngineAdapter)
103+
104+
temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
105+
table_name = "target"
106+
temp_table_id = "abcdefgh"
107+
temp_table_mock.return_value = make_temp_table_name(table_name, temp_table_id)
108+
109+
df = pd.DataFrame({"id": [1, 2, 3], "ts": [1, 2, 3], "val": [4, 5, 6]})
110+
111+
# 1 key
112+
adapter.merge(
113+
target_table=table_name,
114+
source_table=df,
115+
target_columns_to_types={
116+
"id": exp.DataType.build("int"),
117+
"ts": exp.DataType.build("TIMESTAMP"),
118+
"val": exp.DataType.build("int"),
119+
},
120+
unique_key=[exp.to_identifier("id")],
121+
)
122+
adapter._connection_pool.get().bulk_copy.assert_called_with(
123+
f"__temp_target_{temp_table_id}", [(1, 1, 4), (2, 2, 5), (3, 3, 6)]
124+
)
125+
126+
assert to_sql_calls(adapter) == [
127+
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] INT, [ts] DATETIME2(6), [val] INT)');""",
128+
f"MERGE INTO [target] AS [__MERGE_TARGET__] USING (SELECT CAST([id] AS INT) AS [id], CAST([ts] AS DATETIME2(6)) AS [ts], CAST([val] AS INT) 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__].[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]);",
129+
f"DROP TABLE IF EXISTS [__temp_target_{temp_table_id}];",
130+
]
131+
132+
# 2 keys
133+
adapter.cursor.reset_mock()
134+
adapter._connection_pool.get().reset_mock()
135+
temp_table_mock.return_value = make_temp_table_name(table_name, temp_table_id)
136+
adapter.merge(
137+
target_table=table_name,
138+
source_table=df,
139+
target_columns_to_types={
140+
"id": exp.DataType.build("int"),
141+
"ts": exp.DataType.build("TIMESTAMP"),
142+
"val": exp.DataType.build("int"),
143+
},
144+
unique_key=[exp.to_identifier("id"), exp.to_column("ts")],
145+
)
146+
adapter._connection_pool.get().bulk_copy.assert_called_with(
147+
f"__temp_target_{temp_table_id}", [(1, 1, 4), (2, 2, 5), (3, 3, 6)]
148+
)
149+
150+
assert to_sql_calls(adapter) == [
151+
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] INT, [ts] DATETIME2(6), [val] INT)');""",
152+
f"MERGE INTO [target] AS [__MERGE_TARGET__] USING (SELECT CAST([id] AS INT) AS [id], CAST([ts] AS DATETIME2(6)) AS [ts], CAST([val] AS INT) 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__].[val] = [__MERGE_SOURCE__].[val] WHEN NOT MATCHED THEN INSERT ([id], [ts], [val]) VALUES ([__MERGE_SOURCE__].[id], [__MERGE_SOURCE__].[ts], [__MERGE_SOURCE__].[val]);",
153+
f"DROP TABLE IF EXISTS [__temp_target_{temp_table_id}];",
154+
]
155+
156+
157+
def test_merge_exists(
158+
make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name: t.Callable
159+
):
160+
mocker.patch(
161+
"sqlmesh.core.engine_adapter.fabric.FabricEngineAdapter.table_exists",
162+
return_value=False,
163+
)
164+
165+
adapter = make_mocked_engine_adapter(FabricEngineAdapter)
166+
167+
temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
168+
table_name = "target"
169+
temp_table_id = "abcdefgh"
170+
temp_table_mock.return_value = make_temp_table_name(table_name, temp_table_id)
171+
172+
df = pd.DataFrame({"id": [1, 2, 3], "ts": [1, 2, 3], "val": [4, 5, 6]})
173+
174+
# regular implementation
175+
adapter.merge(
176+
target_table=table_name,
177+
source_table=df,
178+
target_columns_to_types={
179+
"id": exp.DataType.build("int"),
180+
"ts": exp.DataType.build("TIMESTAMP"),
181+
"val": exp.DataType.build("int"),
182+
},
183+
unique_key=[exp.to_identifier("id")],
184+
)
185+
186+
assert to_sql_calls(adapter) == [
187+
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] INT, [ts] DATETIME2(6), [val] INT)');""",
188+
f"MERGE INTO [target] AS [__MERGE_TARGET__] USING (SELECT CAST([id] AS INT) AS [id], CAST([ts] AS DATETIME2(6)) AS [ts], CAST([val] AS INT) 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__].[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]);",
189+
f"DROP TABLE IF EXISTS [__temp_target_{temp_table_id}];",
190+
]
191+
192+
# merge exists implementation
193+
adapter.cursor.reset_mock()
194+
adapter._connection_pool.get().reset_mock()
195+
temp_table_mock.return_value = make_temp_table_name(table_name, temp_table_id)
196+
adapter.merge(
197+
target_table=table_name,
198+
source_table=df,
199+
target_columns_to_types={
200+
"id": exp.DataType.build("int"),
201+
"ts": exp.DataType.build("TIMESTAMP"),
202+
"val": exp.DataType.build("int"),
203+
},
204+
unique_key=[exp.to_identifier("id")],
205+
physical_properties={"mssql_merge_exists": True},
206+
)
207+
208+
assert to_sql_calls(adapter) == [
209+
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] INT, [ts] DATETIME2(6), [val] INT)');""",
210+
f"MERGE INTO [target] AS [__MERGE_TARGET__] USING (SELECT CAST([id] AS INT) AS [id], CAST([ts] AS DATETIME2(6)) AS [ts], CAST([val] AS INT) 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]);",
211+
f"DROP TABLE IF EXISTS [__temp_target_{temp_table_id}];",
212+
]
213+
214+
# merge exists and all model columns are keys
215+
adapter.cursor.reset_mock()
216+
adapter._connection_pool.get().reset_mock()
217+
temp_table_mock.return_value = make_temp_table_name(table_name, temp_table_id)
218+
adapter.merge(
219+
target_table=table_name,
220+
source_table=df,
221+
target_columns_to_types={
222+
"id": exp.DataType.build("int"),
223+
"ts": exp.DataType.build("TIMESTAMP"),
224+
},
225+
unique_key=[exp.to_identifier("id"), exp.to_column("ts")],
226+
physical_properties={"mssql_merge_exists": True},
227+
)
228+
229+
assert to_sql_calls(adapter) == [
230+
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] INT, [ts] DATETIME2(6))');""",
231+
f"MERGE INTO [target] AS [__MERGE_TARGET__] USING (SELECT CAST([id] AS INT) AS [id], CAST([ts] AS DATETIME2(6)) AS [ts] FROM [__temp_target_{temp_table_id}]) AS [__MERGE_SOURCE__] ON [__MERGE_TARGET__].[id] = [__MERGE_SOURCE__].[id] AND [__MERGE_TARGET__].[ts] = [__MERGE_SOURCE__].[ts] WHEN NOT MATCHED THEN INSERT ([id], [ts]) VALUES ([__MERGE_SOURCE__].[id], [__MERGE_SOURCE__].[ts]);",
232+
f"DROP TABLE IF EXISTS [__temp_target_{temp_table_id}];",
233+
]

0 commit comments

Comments
 (0)