Repository: TobikoData/sqlmesh
Pull request #2250 (2024-03-12 03:57:05)

Fix: Adding of LIMIT 0 to union queries when generating a model CTAS query


Patch:
diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py
index 5819fef0..3e0d3e03 100644
--- a/sqlmesh/core/config/connection.py
+++ b/sqlmesh/core/config/connection.py
@@ -122,10 +122,6 @@ class BaseDuckDBConnectionConfig(ConnectionConfig):
     concurrent_tasks: Literal[1] = 1
     register_comments: bool = True
 
-    @property
-    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
-        return {"config": self.connector_config} if self.connector_config else {}
-
     @property
     def _engine_adapter(self) -> t.Type[EngineAdapter]:
         return engine_adapter.DuckDBEngineAdapter
@@ -150,6 +146,12 @@ class BaseDuckDBConnectionConfig(ConnectionConfig):
                 except Exception as e:
                     raise ConfigError(f"Failed to load extension {extension}: {e}")
 
+            for field, setting in self.connector_config.items():
+                try:
+                    cursor.execute(f"SET {field} = '{setting}'")
+                except Exception as e:
+                    raise ConfigError(f"Failed to set connector config {field} to {setting}: {e}")
+
             for i, (alias, path) in enumerate((getattr(self, "catalogs", None) or {}).items()):
                 # we parse_identifier and generate to ensure that `alias` has exactly one set of quotes
                 # regardless of whether it comes in quoted or not
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index 90bc570c..c24b8700 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -422,13 +422,14 @@ class _Model(ModelMeta, frozen=True):
         Return:
             The mocked out ctas query.
         """
-        # the query is expanded so it's been copied, it's safe to mutate.
-        query = self.render_query_or_raise(**render_kwarg)
+        query = self.render_query_or_raise(**render_kwarg).copy()
 
         for select in query.find_all(exp.Select):
             if select.args.get("from"):
                 select.where(exp.false(), copy=False)
-                if not isinstance(select.parent, exp.Union):
+                if not isinstance(select.parent, exp.Union) or (
+                    select.parent.parent is None and select.arg_key == "expression"
+                ):
                     select.limit(0, copy=False)
 
         if self.managed_columns:


Test patch:
diff --git a/tests/core/test_connection_config.py b/tests/core/test_connection_config.py
index 6b320a67..2e7bda2b 100644
--- a/tests/core/test_connection_config.py
+++ b/tests/core/test_connection_config.py
@@ -376,7 +376,6 @@ def test_duckdb(make_config):
         connector_config={"foo": "bar"},
     )
     assert isinstance(config, DuckDBConnectionConfig)
-    assert config._static_connection_kwargs == {"config": {"foo": "bar"}}
 
 
 @pytest.mark.parametrize(
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index d72afc72..9be879a0 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -1767,6 +1767,30 @@ def test_model_ctas_query():
         == 'SELECT 1 AS "a" FROM "b" AS "b" WHERE FALSE LIMIT 0'
     )
 
+    expressions = d.parse(
+        """
+        MODEL (name `a-b-c.table`, kind FULL, dialect bigquery);
+        SELECT 1 AS a FROM t UNION ALL SELECT 2 AS a FROM t UNION ALL SELECT 3 AS a FROM t UNION ALL SELECT 4 AS a FROM t
+    """
+    )
+
+    assert (
+        load_sql_based_model(expressions, dialect="bigquery").ctas_query().sql()
+        == 'SELECT 1 AS "a" FROM "t" AS "t" WHERE FALSE UNION ALL SELECT 2 AS "a" FROM "t" AS "t" WHERE FALSE UNION ALL SELECT 3 AS "a" FROM "t" AS "t" WHERE FALSE UNION ALL SELECT 4 AS "a" FROM "t" AS "t" WHERE FALSE LIMIT 0'
+    )
+
+    expressions = d.parse(
+        """
+        MODEL (name `a-b-c.table`, kind FULL, dialect bigquery);
+        SELECT 1 AS a FROM t UNION ALL SELECT 2 AS a FROM t
+    """
+    )
+
+    assert (
+        load_sql_based_model(expressions, dialect="bigquery").ctas_query().sql()
+        == 'SELECT 1 AS "a" FROM "t" AS "t" WHERE FALSE UNION ALL SELECT 2 AS "a" FROM "t" AS "t" WHERE FALSE LIMIT 0'
+    )
+
 
 def test_is_breaking_change():
     model = create_external_model("a", columns={"a": "int", "limit": "int"})