Repository: TobikoData/sqlmesh
Pull request #2800 (2024-06-20 23:04:20)

Fix: Dont set LIMIT 0 in nested queries that are part of a recursive CTE when generating a CTAS query


Patch:
diff --git a/setup.py b/setup.py
index 4f7f9449..3a21a965 100644
--- a/setup.py
+++ b/setup.py
@@ -49,7 +49,7 @@ setup(
         "requests",
         "rich[jupyter]",
         "ruamel.yaml",
-        "sqlglot[rs]~=25.2.0",
+        "sqlglot[rs]~=25.3.0",
     ],
     extras_require={
         "bigquery": [
diff --git a/sqlmesh/core/console.py b/sqlmesh/core/console.py
index cbecb9f6..e7d11749 100644
--- a/sqlmesh/core/console.py
+++ b/sqlmesh/core/console.py
@@ -61,7 +61,7 @@ class Console(abc.ABC):
     """Abstract base class for defining classes used for displaying information to the user and also interact
     with them when their input is needed."""
 
-    INDIRECTLY_MODIFIED_DISPLAY_THRESHOLD = 20
+    INDIRECTLY_MODIFIED_DISPLAY_THRESHOLD = 10
 
     @abc.abstractmethod
     def start_plan_evaluation(self, plan: Plan) -> None:
@@ -702,20 +702,20 @@ class TerminalConsole(Console):
         tree = Tree(f"[bold]{header}:")
         if added_snapshot_ids:
             added_tree = Tree("[bold][added]Added:")
-            for s_id in added_snapshot_ids:
+            for s_id in sorted(added_snapshot_ids):
                 snapshot = context_diff.snapshots[s_id]
                 added_tree.add(
                     f"[added]{snapshot.display_name(environment_naming_info, default_catalog, dialect=self.dialect)}"
                 )
-            tree.add(added_tree)
+            tree.add(self._limit_model_names(added_tree, self.verbose))
         if removed_snapshot_ids:
             removed_tree = Tree("[bold][removed]Removed:")
-            for s_id in removed_snapshot_ids:
+            for s_id in sorted(removed_snapshot_ids):
                 snapshot_table_info = context_diff.removed_snapshots[s_id]
                 removed_tree.add(
                     f"[removed]{snapshot_table_info.display_name(environment_naming_info, default_catalog, dialect=self.dialect)}"
                 )
-            tree.add(removed_tree)
+            tree.add(self._limit_model_names(removed_tree, self.verbose))
         if modified_snapshot_ids:
             direct = Tree("[bold][direct]Directly Modified:")
             indirect = Tree("[bold][indirect]Indirectly Modified:")
@@ -863,6 +863,8 @@ class TerminalConsole(Console):
             backfill.add(
                 f"{snapshot.display_name(plan.environment_naming_info, default_catalog, dialect=self.dialect)}: {missing.format_intervals(snapshot.node.interval_unit)}{preview_modifier}"
             )
+        if backfill:
+            backfill = self._limit_model_names(backfill, self.verbose)
         self._print(backfill)
 
     def _prompt_effective_from(
@@ -1510,10 +1512,17 @@ class MarkdownConsole(CaptureTerminalConsole):
         added_snapshot_models = {s for s in added_snapshots if s.is_model}
         if added_snapshot_models:
             self._print("\n**Added Models:**")
-            for snapshot in sorted(added_snapshot_models):
-                self._print(
-                    f"- `{snapshot.display_name(environment_naming_info, default_catalog, dialect=self.dialect)}`"
-                )
+            added_models = sorted(added_snapshot_models)
+            list_length = len(added_models)
+            if not self.verbose and list_length > self.INDIRECTLY_MODIFIED_DISPLAY_THRESHOLD:
+                self._print(added_models[0])
+                self._print(f"- `.... {list_length-2} more ....`\n")
+                self._print(added_models[-1])
+            else:
+                for snapshot in added_models:
+                    self._print(
+                        f"- `{snapshot.display_name(environment_naming_info, default_catalog, dialect=self.dialect)}`"
+                    )
 
         added_snapshot_audits = {s for s in added_snapshots if s.is_audit}
         if added_snapshot_audits:
@@ -1531,10 +1540,17 @@ class MarkdownConsole(CaptureTerminalConsole):
         removed_model_snapshot_table_infos = {s for s in removed_snapshot_table_infos if s.is_model}
         if removed_model_snapshot_table_infos:
             self._print("\n**Removed Models:**")
-            for snapshot_table_info in sorted(removed_model_snapshot_table_infos):
-                self._print(
-                    f"- `{snapshot_table_info.display_name(environment_naming_info, default_catalog, dialect=self.dialect)}`"
-                )
+            removed_models = sorted(removed_model_snapshot_table_infos)
+            list_length = len(removed_models)
+            if not self.verbose and list_length > self.INDIRECTLY_MODIFIED_DISPLAY_THRESHOLD:
+                self._print(removed_models[0])
+                self._print(f"- `.... {list_length-2} more ....`\n")
+                self._print(removed_models[-1])
+            else:
+                for snapshot_table_info in removed_models:
+                    self._print(
+                        f"- `{snapshot_table_info.display_name(environment_naming_info, default_catalog, dialect=self.dialect)}`"
+                    )
 
         removed_audit_snapshot_table_infos = {s for s in removed_snapshot_table_infos if s.is_audit}
         if removed_audit_snapshot_table_infos:
@@ -1606,6 +1622,7 @@ class MarkdownConsole(CaptureTerminalConsole):
         if not missing_intervals:
             return
         self._print("\n**Models needing backfill (missing dates):**")
+        snapshots = []
         for missing in missing_intervals:
             snapshot = plan.context_diff.snapshots[missing.snapshot_id]
             if not snapshot.is_model:
@@ -1615,10 +1632,19 @@ class MarkdownConsole(CaptureTerminalConsole):
             if not plan.deployability_index.is_deployable(snapshot):
                 preview_modifier = " (**preview**)"
 
-            self._print(
+            snapshots.append(
                 f"* `{snapshot.display_name(plan.environment_naming_info, default_catalog, dialect=self.dialect)}`: {missing.format_intervals(snapshot.node.interval_unit)}{preview_modifier}"
             )
 
+        length = len(snapshots)
+        if not self.verbose and length > self.INDIRECTLY_MODIFIED_DISPLAY_THRESHOLD:
+            self._print(snapshots[0])
+            self._print(f"- `.... {length-2} more ....`\n")
+            self._print(snapshots[-1])
+        else:
+            for snap in snapshots:
+                self._print(snap)
+
     def _show_categorized_snapshots(self, plan: Plan, default_catalog: t.Optional[str]) -> None:
         context_diff = plan.context_diff
         for snapshot in plan.categorized:
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index 2f121ccf..eea88685 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -412,15 +412,17 @@ def _parse_table_parts(
 ) -> exp.Table:
     index = self._index
     table = self.__parse_table_parts(schema=schema, is_db_reference=is_db_reference)  # type: ignore
+
     table_arg = table.this
+    name = table_arg.name
 
-    if isinstance(table_arg, exp.Var) and table_arg.name.startswith(SQLMESH_MACRO_PREFIX):
+    if isinstance(table_arg, exp.Var) and name.startswith(SQLMESH_MACRO_PREFIX):
         # Macro functions do not clash with the staged file syntax, so we can safely parse them
-        if self._prev.token_type == TokenType.STRING or "(" in table_arg.name:
+        if self._prev.token_type == TokenType.STRING or any(ch in name for ch in ("(", "{")):
             self._retreat(index)
             return Parser._parse_table_parts(self, schema=schema, is_db_reference=is_db_reference)
 
-        table_arg.replace(MacroVar(this=table_arg.name[1:]))
+        table_arg.replace(MacroVar(this=name[1:]))
         return StagedFilePath(**table.args)
 
     return table
diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py
index 2ef6b7ad..3075a9d5 100644
--- a/sqlmesh/core/engine_adapter/base.py
+++ b/sqlmesh/core/engine_adapter/base.py
@@ -187,7 +187,8 @@ class EngineAdapter:
             raise SQLMeshError(
                 "It is expected that if a DataFrame is passed in then columns_to_types is set"
             )
-        if query_or_df.empty:
+
+        if isinstance(query_or_df, pd.DataFrame) and query_or_df.empty:
             raise SQLMeshError(
                 "Cannot construct source query from an empty DataFrame. This error is commonly "
                 "related to Python models that produce no data. For such models, consider yielding "
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index f5c3e444..bbd656fb 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -423,8 +423,12 @@ class _Model(ModelMeta, frozen=True):
         query = self.render_query_or_raise(**render_kwarg).copy()
 
         for select_or_set_op in query.find_all(exp.Select, exp.SetOperation):
-            cte = select_or_set_op.find_ancestor(exp.With, exp.Select, exp.Subquery)
-            skip_limit = isinstance(cte, exp.With) and cte.recursive
+            skip_limit = False
+            ancestor = select_or_set_op.parent
+            while ancestor and not skip_limit:
+                if isinstance(ancestor, exp.With) and ancestor.recursive:
+                    skip_limit = True
+                ancestor = ancestor.parent
 
             if isinstance(select_or_set_op, exp.Select) and select_or_set_op.args.get("from"):
                 select_or_set_op.where(exp.false(), copy=False)
diff --git a/sqlmesh/core/model/meta.py b/sqlmesh/core/model/meta.py
index 6f635551..86a2e6e9 100644
--- a/sqlmesh/core/model/meta.py
+++ b/sqlmesh/core/model/meta.py
@@ -156,6 +156,9 @@ class ModelMeta(_Node):
         def _normalize(value: t.Any) -> t.Any:
             return normalize_identifiers(value, dialect=dialect) if normalize else value
 
+        if isinstance(v, exp.Paren):
+            v = [v.unnest()]
+
         if isinstance(v, (exp.Tuple, exp.Array)):
             return [_normalize(e).name for e in v.expressions]
         if isinstance(v, exp.Expression):


Test patch:
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 97f8aa97..687fbc22 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -939,7 +939,8 @@ def test_audits():
             audits (
                 audit_a,
                 audit_b(key='value')
-            )
+            ),
+            tags (foo)
         );
         SELECT 1, ds;
     """
@@ -950,6 +951,7 @@ def test_audits():
         ("audit_a", {}),
         ("audit_b", {"key": exp.Literal.string("value")}),
     ]
+    assert model.tags == ["foo"]
 
 
 def test_description(sushi_context):
@@ -2161,6 +2163,45 @@ def test_model_ctas_query():
         == 'WITH RECURSIVE "a" AS (SELECT * FROM "x" AS "x" WHERE FALSE), "b" AS (SELECT * FROM "a" AS "a" WHERE FALSE UNION ALL SELECT * FROM "a" AS "a" WHERE FALSE) SELECT * FROM "b" AS "b" WHERE FALSE LIMIT 0'
     )
 
+    expressions = d.parse(
+        """
+        MODEL (name `a-b-c.table`, kind FULL, dialect bigquery);
+        WITH RECURSIVE a AS (
+            SELECT * FROM (SELECT * FROM (SELECT * FROM x))
+        ), b AS (
+            SELECT * FROM a UNION ALL SELECT * FROM a
+        )
+        SELECT * FROM b
+
+    """
+    )
+
+    assert (
+        load_sql_based_model(expressions, dialect="bigquery").ctas_query().sql()
+        == 'WITH RECURSIVE "a" AS (SELECT * FROM (SELECT * FROM (SELECT * FROM "x" AS "x" WHERE FALSE) AS "_q_0" WHERE FALSE) AS "_q_1" WHERE FALSE), "b" AS (SELECT * FROM "a" AS "a" WHERE FALSE UNION ALL SELECT * FROM "a" AS "a" WHERE FALSE) SELECT * FROM "b" AS "b" WHERE FALSE LIMIT 0'
+    )
+
+    expressions = d.parse(
+        """
+        MODEL (name `a-b-c.table`, kind FULL, dialect bigquery);
+        WITH RECURSIVE a AS (
+            WITH nested_a AS (
+                SELECT * FROM (SELECT * FROM (SELECT * FROM x))
+            )
+            SELECT * FROM nested_a
+        ), b AS (
+            SELECT * FROM a UNION ALL SELECT * FROM a
+        )
+        SELECT * FROM b
+
+    """
+    )
+
+    assert (
+        load_sql_based_model(expressions, dialect="bigquery").ctas_query().sql()
+        == 'WITH RECURSIVE "a" AS (WITH "nested_a" AS (SELECT * FROM (SELECT * FROM (SELECT * FROM "x" AS "x" WHERE FALSE) AS "_q_0" WHERE FALSE) AS "_q_1" WHERE FALSE) SELECT * FROM "nested_a" AS "nested_a" WHERE FALSE), "b" AS (SELECT * FROM "a" AS "a" WHERE FALSE UNION ALL SELECT * FROM "a" AS "a" WHERE FALSE) SELECT * FROM "b" AS "b" WHERE FALSE LIMIT 0'
+    )
+
 
 def test_is_breaking_change():
     model = create_external_model("a", columns={"a": "int", "limit": "int"})
@@ -4043,6 +4084,24 @@ def test_variables_in_templates() -> None:
         == "SELECT 'combo' AS \"col_test_value_overridden_value_col_in_memory\""
     )
 
+    model = load_sql_based_model(
+        parse(
+            """
+        MODEL(
+          name @{some_var}.bar,
+          dialect snowflake
+        );
+
+        SELECT 1 AS c
+        """
+        ),
+        variables={
+            "some_var": "foo",
+        },
+    )
+
+    assert model.name == "foo.bar"
+
 
 def test_variables_jinja():
     expressions = parse(