Repository: TobikoData/sqlmesh
Pull request #2711 (2024-05-31 14:59:11)

Fix: Avoid using concat_ws when there's only one key is present
For example, MSSQL requires 3 arguments to the CONCAT_WS function: https://learn.microsoft.com/en-us/sql/t-sql/functions/concat-ws-transact-sql?view=sql-server-ver15


Patch:
diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py
index 091a1799..37923bea 100644
--- a/sqlmesh/core/config/connection.py
+++ b/sqlmesh/core/config/connection.py
@@ -223,10 +223,15 @@ class MotherDuckConnectionConfig(BaseDuckDBConnectionConfig):
     @property
     def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
         """kwargs that are for execution config only"""
+        from sqlmesh import __version__
+
         connection_str = f"md:{self.database}"
         if self.token:
             connection_str += f"?motherduck_token={self.token}"
-        return {"database": connection_str}
+        return {
+            "database": connection_str,
+            "config": {"custom_user_agent": f"SQLMesh/{__version__}"},
+        }
 
 
 class DuckDBAttachOptions(BaseConfig):
diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py
index 92c35579..6e2d3e87 100644
--- a/sqlmesh/core/engine_adapter/base.py
+++ b/sqlmesh/core/engine_adapter/base.py
@@ -2027,7 +2027,7 @@ class EngineAdapter:
             columns_to_types = self.columns(target_table)
 
         temp_table = self._get_temp_table(target_table)
-        key_exp = exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", *key)
+        key_exp = exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", *key) if len(key) > 1 else key[0]
         column_names = list(columns_to_types or [])
 
         with self.transaction():


Test patch:
diff --git a/tests/core/engine_adapter/test_base.py b/tests/core/engine_adapter/test_base.py
index d5a4bbd7..c9d3bb95 100644
--- a/tests/core/engine_adapter/test_base.py
+++ b/tests/core/engine_adapter/test_base.py
@@ -2677,8 +2677,18 @@ def test_pre_ping(mocker: MockerFixture, make_mocked_engine_adapter: t.Callable)
     adapter._connection_pool.get().close.assert_called_once()
 
 
+@pytest.mark.parametrize(
+    "partitioned_by",
+    [
+        ["DATETIME_TRUNC(ds, MONTH)", "b"],
+        ["b"],
+    ],
+)
 def test_insert_overwrite_by_partition_query(
-    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name: t.Callable
+    make_mocked_engine_adapter: t.Callable,
+    mocker: MockerFixture,
+    make_temp_table_name: t.Callable,
+    partitioned_by: t.List[str],
 ):
     adapter = make_mocked_engine_adapter(EngineAdapter)
 
@@ -2690,10 +2700,7 @@ def test_insert_overwrite_by_partition_query(
     adapter.insert_overwrite_by_partition(
         table_name,
         parse_one("SELECT a, ds, b FROM tbl"),
-        partitioned_by=[
-            d.parse_one("DATETIME_TRUNC(ds, MONTH)"),
-            d.parse_one("b"),
-        ],
+        partitioned_by=[d.parse_one(k) for k in partitioned_by],
         columns_to_types={
             "a": exp.DataType.build("int"),
             "ds": exp.DataType.build("DATETIME"),
@@ -2701,10 +2708,16 @@ def test_insert_overwrite_by_partition_query(
         },
     )
 
+    expected_delete_stmt = (
+        'DELETE FROM "test_schema"."test_table" WHERE CONCAT_WS(\'__SQLMESH_DELIM__\', DATETIME_TRUNC("ds", MONTH), "b") IN (SELECT DISTINCT CONCAT_WS(\'__SQLMESH_DELIM__\', DATETIME_TRUNC("ds", MONTH), "b") FROM "test_schema"."__temp_test_table_abcdefgh")'
+        if len(partitioned_by) > 1
+        else 'DELETE FROM "test_schema"."test_table" WHERE "b" IN (SELECT DISTINCT "b" FROM "test_schema"."__temp_test_table_abcdefgh")'
+    )
+
     sql_calls = to_sql_calls(adapter)
     assert sql_calls == [
         'CREATE TABLE "test_schema"."__temp_test_table_abcdefgh" AS SELECT "a", "ds", "b" FROM "tbl"',
-        'DELETE FROM "test_schema"."test_table" WHERE CONCAT_WS(\'__SQLMESH_DELIM__\', DATETIME_TRUNC("ds", MONTH), "b") IN (SELECT DISTINCT CONCAT_WS(\'__SQLMESH_DELIM__\', DATETIME_TRUNC("ds", MONTH), "b") FROM "test_schema"."__temp_test_table_abcdefgh")',
+        expected_delete_stmt,
         'INSERT INTO "test_schema"."test_table" ("a", "ds", "b") SELECT "a", "ds", "b" FROM "test_schema"."__temp_test_table_abcdefgh"',
         'DROP TABLE IF EXISTS "test_schema"."__temp_test_table_abcdefgh"',
     ]
diff --git a/tests/core/engine_adapter/test_mixins.py b/tests/core/engine_adapter/test_mixins.py
index 805744b8..2f8b1c74 100644
--- a/tests/core/engine_adapter/test_mixins.py
+++ b/tests/core/engine_adapter/test_mixins.py
@@ -28,15 +28,13 @@ def test_logical_merge(make_mocked_engine_adapter: t.Callable, mocker: MockerFix
             "ts": exp.DataType(this=exp.DataType.Type.TIMESTAMP),
             "val": exp.DataType(this=exp.DataType.Type.INT),
         },
-        unique_key=["id"],
+        unique_key=[parse_one("id")],
     )
 
     adapter.cursor.execute.assert_has_calls(
         [
             call('''CREATE TABLE "temporary" AS SELECT "id", "ts", "val" FROM "source"'''),
-            call(
-                """DELETE FROM "target" WHERE CONCAT_WS('__SQLMESH_DELIM__', "id") IN (SELECT CONCAT_WS('__SQLMESH_DELIM__', "id") FROM "temporary")"""
-            ),
+            call("""DELETE FROM "target" WHERE "id" IN (SELECT "id" FROM "temporary")"""),
             call(
                 'INSERT INTO "target" ("id", "ts", "val") SELECT DISTINCT ON ("id") "id", "ts", "val" FROM "temporary"'
             ),
@@ -53,7 +51,7 @@ def test_logical_merge(make_mocked_engine_adapter: t.Callable, mocker: MockerFix
             "ts": exp.DataType(this=exp.DataType.Type.TIMESTAMP),
             "val": exp.DataType(this=exp.DataType.Type.INT),
         },
-        unique_key=["id", "ts"],
+        unique_key=[parse_one("id"), parse_one("ts")],
     )
 
     adapter.cursor.execute.assert_has_calls(