Fix: Allow SELECT * in models that reference external tables (#1007)

Author: izeigerman

sqlmesh/core/dialect.py

@@ -10,6 +10,7 @@
 from sqlglot import Dialect, Generator, Parser, Tokenizer, TokenType, exp
 from sqlglot.dialects.dialect import DialectType
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
+from sqlglot.optimizer.scope import traverse_scope
 from sqlglot.tokens import Token
 
 from sqlmesh.core.constants import MAX_MODEL_DEFINITION_SIZE
@@ -645,3 +646,21 @@ def extract_columns_to_types(query: exp.Subqueryable) -> t.Dict[str, exp.DataTyp
         expression.output_name: expression.type or exp.DataType.build("unknown")
         for expression in query.selects
     }
+
+
+def find_tables(expression: exp.Expression, dialect: DialectType = None) -> t.Set[str]:
+    """Find all tables referenced in a query.
+
+    Args:
+        expressions: The query to find the tables in.
+        dialect: The dialect to use for normalization of table names.
+
+    Returns:
+        A Set of all the table names.
+    """
+    return {
+        normalize_model_name(table, dialect=dialect)
+        for scope in traverse_scope(expression)
+        for table in scope.tables
+        if isinstance(table.this, exp.Identifier) and exp.table_name(table) not in scope.cte_sources
+    }

sqlmesh/core/loader.py

@@ -54,18 +54,14 @@ def update_model_schemas(
             continue
 
         model.update_schema(schema)
+
         cache_hit = optimized_query_cache.with_optimized_query(model)
-        schema.add_table(name, model.columns_to_types, dialect=model.dialect)
-
-        if any(dep not in models for dep in model.depends_on):
-            if "*" in model.columns_to_types:
-                raise ConfigError(
-                    f"Can't expand SELECT * expression for model '{name}' at '{model._path}'."
-                    " Either specify external source projections expliticly or"
-                    ' add source tables as "external models" using the command'
-                    " 'sqlmesh create_external_models'."
-                )
-        elif isinstance(model, SqlModel) and model.mapping_schema and not cache_hit:
+
+        columns_to_types = model.columns_to_types
+        if columns_to_types is not None:
+            schema.add_table(name, columns_to_types, dialect=model.dialect)
+
+        if isinstance(model, SqlModel) and model.mapping_schema and not cache_hit:
             query = model.render_query()
             if query is not None:
                 try:

sqlmesh/core/model/definition.py

@@ -16,7 +16,6 @@
 from sqlglot import diff, exp
 from sqlglot.diff import Insert, Keep
 from sqlglot.helper import ensure_list
-from sqlglot.optimizer.scope import traverse_scope
 from sqlglot.schema import MappingSchema, nested_set
 from sqlglot.time import format_time
 
@@ -460,6 +459,18 @@ def update_schema(self, schema: MappingSchema) -> None:
                     tuple(str(part) for part in table.parts),
                     {k: str(v) for k, v in mapping_schema.items()},
                 )
+            else:
+                # Reset the entire mapping if at least one upstream dependency is missing from the mapping
+                # to prevent partial mappings from being used.
+                logger.warning(
+                    "Missing schema for model '%s' referenced in model '%s'. Run `sqlmesh create_external_models` "
+                    "and / or make sure that the model '%s' can be rendered at parse time",
+                    dep,
+                    self.name,
+                    dep,
+                )
+                self.mapping_schema.clear()
+                return
 
     @property
     def depends_on(self) -> t.Set[str]:
@@ -468,16 +479,7 @@ def depends_on(self) -> t.Set[str]:
         Returns:
             A list of all the upstream table names.
         """
-        if self.depends_on_ is not None:
-            return self.depends_on_ - {self.name}
-
-        if self._depends_on is None:
-            query = self.render_query(optimize=False)
-            if query is None:
-                self._depends_on = set()
-            else:
-                self._depends_on = _find_tables(query) - {self.name}
-        return self._depends_on
+        return self.depends_on_ or set()
 
     @property
     def columns_to_types(self) -> t.Optional[t.Dict[str, exp.DataType]]:
@@ -537,7 +539,7 @@ def depends_on_past(self) -> bool:
             query = self.render_query(optimize=False)
             if query is None:
                 return False
-            return self.name in _find_tables(query)
+            return self.name in d.find_tables(query, dialect=self.dialect)
 
         return self._depends_on_past
 
@@ -559,7 +561,7 @@ def validate_definition(self) -> None:
                 )
 
             columns_to_types = self.columns_to_types
-            if columns_to_types is not None and "*" not in columns_to_types:
+            if columns_to_types is not None:
                 column_names = {c.lower() for c in columns_to_types}
                 missing_keys = unique_partition_keys - column_names
                 if missing_keys:
@@ -742,6 +744,18 @@ def render_definition(self, include_python: bool = True) -> t.List[exp.Expressio
     def is_sql(self) -> bool:
         return True
 
+    @property
+    def depends_on(self) -> t.Set[str]:
+        if self._depends_on is None:
+            self._depends_on = self.depends_on_ or set()
+
+            query = self.render_query(optimize=False)
+            if query is not None:
+                self._depends_on |= d.find_tables(query, dialect=self.dialect)
+
+            self._depends_on -= {self.name}
+        return self._depends_on
+
     @property
     def columns_to_types(self) -> t.Optional[t.Dict[str, exp.DataType]]:
         if self.columns_to_types_ is not None:
@@ -753,6 +767,9 @@ def columns_to_types(self) -> t.Optional[t.Dict[str, exp.DataType]]:
                 return None
             self._columns_to_types = d.extract_columns_to_types(query)
 
+        if "*" in self._columns_to_types:
+            return None
+
         return self._columns_to_types
 
     @property
@@ -1511,23 +1528,6 @@ def _validate_model_fields(klass: t.Type[_Model], provided_fields: t.Set[str], p
         raise_config_error(f"Invalid extra fields {extra_fields} in the model definition", path)
 
 
-def _find_tables(expression: exp.Expression) -> t.Set[str]:
-    """Find all tables referenced in a query.
-
-    Args:
-        expressions: The list of expressions to find tables for.
-
-    Returns:
-        A Set of all the table names.
-    """
-    return {
-        exp.table_name(table)
-        for scope in traverse_scope(expression)
-        for table in scope.tables
-        if isinstance(table.this, exp.Identifier) and exp.table_name(table) not in scope.cte_sources
-    }
-
-
 def _python_env(
     expressions: t.Union[exp.Expression, t.List[exp.Expression]],
     jinja_macro_references: t.Optional[t.Set[MacroReference]],

sqlmesh/core/renderer.py

@@ -14,7 +14,7 @@
 from sqlglot.optimizer.qualify_columns import quote_identifiers
 from sqlglot.optimizer.qualify_tables import qualify_tables
 from sqlglot.optimizer.simplify import simplify
-from sqlglot.schema import ensure_schema
+from sqlglot.schema import MappingSchema
 
 from sqlmesh.core import constants as c
 from sqlmesh.core import dialect as d
@@ -263,7 +263,6 @@ def render(
                     ),
                 )
             except ParsetimeAdapterCallError:
-                logger.debug("Failed to render query at parse time:\n%s", self._expression)
                 return None
 
             if not query:
@@ -325,10 +324,21 @@ def time_column_filter(self, start: TimeLike, end: TimeLike) -> exp.Between:
 
     def _optimize_query(self, query: exp.Subqueryable) -> exp.Subqueryable:
         # We don't want to normalize names in the schema because that's handled by the optimizer
-        schema = ensure_schema(self.schema, dialect=self._dialect, normalize=False)
+        schema = MappingSchema(self.schema, dialect=self._dialect, normalize=False)
         original = query
         failure = False
 
+        if not schema.empty:
+            for dependency in d.find_tables(query, dialect=self._dialect):
+                if schema.find(exp.to_table(dependency, dialect=self._dialect)) is None:
+                    logger.warning(
+                        "Query cannot be optimized due to missing schema for model '%s'. "
+                        "Make sure that the model query can be rendered at parse time",
+                        dependency,
+                    )
+                    schema = MappingSchema(None, dialect=self._dialect, normalize=False)
+                    break
+
         try:
             if not schema.empty:
                 query = query.copy()

tests/core/test_model.py

@@ -1,9 +1,11 @@
 from datetime import date
 from pathlib import Path
+from unittest.mock import patch
 
 import pytest
 from pytest_mock.plugin import MockerFixture
 from sqlglot import exp, parse, parse_one
+from sqlglot.schema import MappingSchema
 
 import sqlmesh.core.dialect as d
 from sqlmesh.core.config import Config
@@ -1227,7 +1229,7 @@ def test_case_sensitivity(assert_exp_eq):
             """
             MODEL (name example.source, kind EMBEDDED);
 
-            SELECT "id", "name", "payload" FROM db.schema."table"
+            SELECT 'id' AS "id", 'name' AS "name", 'payload' AS "payload"
             """
         ),
         dialect="snowflake",
@@ -1258,10 +1260,9 @@ def test_case_sensitivity(assert_exp_eq):
           "SOURCE"."payload" AS "payload"
         FROM (
           SELECT
-            "id" AS "id",
-            "name" AS "name",
-            "payload" AS "payload"
-          FROM "DB"."SCHEMA"."table" AS "table"
+            'id' AS "id",
+            'name' AS "name",
+            'payload' AS "payload"
         ) AS "SOURCE"
         """,
     )
@@ -1386,3 +1387,72 @@ def runtime_macro(**kwargs) -> None:
         match=r"Dependencies must be provided explicitly for models that can be rendered only at runtime at.*",
     ):
         load_model(expressions)
+
+
+def test_update_schema():
+    expressions = d.parse(
+        """
+        MODEL (name db.table);
+
+        SELECT a, b FROM table_a JOIN table_b
+        """
+    )
+
+    model = load_model(expressions)
+
+    schema = MappingSchema(normalize=False)
+    schema.add_table("table_a", {"a": exp.DataType.build("int")})
+
+    # Make sure that the partial schema is not applied.
+    model.update_schema(schema)
+    assert not model.mapping_schema
+
+    schema.add_table("table_b", {"b": exp.DataType.build("int")})
+
+    model.update_schema(schema)
+    assert model.mapping_schema == {
+        "table_a": {"a": "INT"},
+        "table_b": {"b": "INT"},
+    }
+
+
+def test_user_provided_depends_on():
+    expressions = d.parse(
+        """
+        MODEL (name db.table, depends_on [table_b]);
+
+        SELECT a FROM table_a
+        """
+    )
+
+    model = load_model(expressions)
+
+    assert model.depends_on == {"table_a", "table_b"}
+
+
+def test_check_schema_mapping_when_rendering_at_runtime(assert_exp_eq):
+    expressions = d.parse(
+        """
+        MODEL (name db.table, depends_on [table_b]);
+
+        SELECT * FROM table_a JOIN table_b
+        """
+    )
+
+    model = load_model(expressions)
+
+    # Simulate a query that cannot be rendered at parse time.
+    with patch.object(SqlModel, "render_query", return_value=None) as render_query_mock:
+        schema = MappingSchema(normalize=False)
+        schema.add_table("table_b", {"b": exp.DataType.build("int")})
+        model.update_schema(schema)
+
+    assert "table_b" in model.mapping_schema
+    assert model.depends_on == {"table_b"}
+
+    render_query_mock.assert_called_once()
+
+    # Simulate rendering at runtime.
+    assert_exp_eq(
+        model.render_query(), """SELECT * FROM "table_a" AS "table_a", "table_b" AS "table_b" """
+    )