Merge branch 'main' into main

Author: williebsweet

docs/guides/signals.md

@@ -24,6 +24,10 @@ It then divides those into _batches_ (configured with the model's [batch_size](.
 
 Signal checking functions examines a batch of time intervals. The function is always called with a batch of time intervals (DateTimeRanges). It can also optionally be called with key word arguments. It may return `True` if all intervals are ready for evaluation, `False` if no intervals are ready, or the time intervals themselves if only some are ready. A checking function is defined with the `@signal` decorator.
 
+!!! note "One model, multiple signals"
+
+    Multiple signals may be specified for a model. SQLMesh categorizes a candidate interval as ready for evaluation if **all** the signal checking functions determine it is ready.
+
 ## Defining a signal
 
 To define a signal, create a `signals` directory in your project folder. Define your signal in a file named `__init__.py` in that directory (you can have additional python file names as well).

sqlmesh/core/console.py

@@ -2163,13 +2163,12 @@ def log_test_results(self, result: ModelTextTestResult, target_dialect: str) ->
             self._print("-" * divider_length)
             self._print("Test Failure Summary", style="red")
             self._print("=" * divider_length)
-            failures = len(result.failures) + len(result.errors)
+            fail_and_error_tests = result.get_fail_and_error_tests()
             self._print(f"{message} \n")
 
-            self._print(f"Failed tests ({failures}):")
-            for test, _ in result.failures + result.errors:
-                if isinstance(test, ModelTest):
-                    self._print(f" • {test.path}::{test.test_name}")
+            self._print(f"Failed tests ({len(fail_and_error_tests)}):")
+            for test in fail_and_error_tests:
+                self._print(f" • {test.path}::{test.test_name}")
             self._print("=" * divider_length, end="\n\n")
 
     def _captured_unit_test_results(self, result: ModelTextTestResult) -> str:
@@ -2721,28 +2720,15 @@ def _log_test_details(
         Args:
             result: The unittest test result that contains metrics like num success, fails, ect.
         """
-
         if result.wasSuccessful():
             self._print("\n", end="")
             return
 
-        errors = result.errors
-        failures = result.failures
-        skipped = result.skipped
-
-        infos = []
-        if failures:
-            infos.append(f"failures={len(failures)}")
-        if errors:
-            infos.append(f"errors={len(errors)}")
-        if skipped:
-            infos.append(f"skipped={skipped}")
-
         if unittest_char_separator:
             self._print(f"\n{unittest.TextTestResult.separator1}\n\n", end="")
 
         for (test_case, failure), test_failure_tables in zip_longest(  # type: ignore
-            failures, result.failure_tables
+            result.failures, result.failure_tables
         ):
             self._print(unittest.TextTestResult.separator2)
             self._print(f"FAIL: {test_case}")
@@ -2758,7 +2744,7 @@ def _log_test_details(
                     self._print(failure_table)
                     self._print("\n", end="")
 
-        for test_case, error in errors:
+        for test_case, error in result.errors:
             self._print(unittest.TextTestResult.separator2)
             self._print(f"ERROR: {test_case}")
             self._print(f"{unittest.TextTestResult.separator2}")
@@ -3080,27 +3066,27 @@ def log_test_results(self, result: ModelTextTestResult, target_dialect: str) ->
             fail_shared_style = {**shared_style, **fail_color}
             header = str(h("span", {"style": fail_shared_style}, "-" * divider_length))
             message = str(h("span", {"style": fail_shared_style}, "Test Failure Summary"))
+            fail_and_error_tests = result.get_fail_and_error_tests()
             failed_tests = [
                 str(
                     h(
                         "span",
                         {"style": fail_shared_style},
-                        f"Failed tests ({len(result.failures) + len(result.errors)}):",
+                        f"Failed tests ({len(fail_and_error_tests)}):",
                     )
                 )
             ]
 
-            for test, _ in result.failures + result.errors:
-                if isinstance(test, ModelTest):
-                    failed_tests.append(
-                        str(
-                            h(
-                                "span",
-                                {"style": fail_shared_style},
-                                f" • {test.model.name}::{test.test_name}",
-                            )
+            for test in fail_and_error_tests:
+                failed_tests.append(
+                    str(
+                        h(
+                            "span",
+                            {"style": fail_shared_style},
+                            f" • {test.model.name}::{test.test_name}",
                         )
                     )
+                )
             failures = "<br>".join(failed_tests)
             footer = str(h("span", {"style": fail_shared_style}, "=" * divider_length))
             error_output = widgets.Textarea(output, layout={"height": "300px", "width": "100%"})
@@ -3508,10 +3494,10 @@ def log_test_results(self, result: ModelTextTestResult, target_dialect: str) ->
             self._log_test_details(result, unittest_char_separator=False)
             self._print("```\n\n")
 
-            failures = len(result.failures) + len(result.errors)
+            fail_and_error_tests = result.get_fail_and_error_tests()
             self._print(f"**{message}**\n")
-            self._print(f"**Failed tests ({failures}):**")
-            for test, _ in result.failures + result.errors:
+            self._print(f"**Failed tests ({len(fail_and_error_tests)}):**")
+            for test in fail_and_error_tests:
                 if isinstance(test, ModelTest):
                     self._print(f" • `{test.model.name}`::`{test.test_name}`\n\n")
 

sqlmesh/core/dialect.py

@@ -13,6 +13,7 @@
 from sqlglot import Dialect, Generator, ParseError, Parser, Tokenizer, TokenType, exp
 from sqlglot.dialects.dialect import DialectType
 from sqlglot.dialects import DuckDB, Snowflake
+import sqlglot.dialects.athena as athena
 from sqlglot.helper import seq_get
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 from sqlglot.optimizer.qualify_columns import quote_identifiers
@@ -1014,6 +1015,14 @@ def extend_sqlglot() -> None:
     generators = {Generator}
 
     for dialect in Dialect.classes.values():
+        # Athena picks a different Tokenizer / Parser / Generator depending on the query
+        # so this ensures that the extra ones it defines are also extended
+        if dialect == athena.Athena:
+            tokenizers.add(athena._TrinoTokenizer)
+            parsers.add(athena._TrinoParser)
+            generators.add(athena._TrinoGenerator)
+            generators.add(athena._HiveGenerator)
+
         if hasattr(dialect, "Tokenizer"):
             tokenizers.add(dialect.Tokenizer)
         if hasattr(dialect, "Parser"):

sqlmesh/core/engine_adapter/bigquery.py

@@ -324,14 +324,26 @@ def create_mapping_schema(
             bq_table = self._get_table(table)
             columns = create_mapping_schema(bq_table.schema)
 
-            if (
-                include_pseudo_columns
-                and bq_table.time_partitioning
-                and not bq_table.time_partitioning.field
-            ):
-                columns["_PARTITIONTIME"] = exp.DataType.build("TIMESTAMP", dialect="bigquery")
-                if bq_table.time_partitioning.type_ == "DAY":
-                    columns["_PARTITIONDATE"] = exp.DataType.build("DATE")
+            if include_pseudo_columns:
+                if bq_table.time_partitioning and not bq_table.time_partitioning.field:
+                    columns["_PARTITIONTIME"] = exp.DataType.build("TIMESTAMP", dialect="bigquery")
+                    if bq_table.time_partitioning.type_ == "DAY":
+                        columns["_PARTITIONDATE"] = exp.DataType.build("DATE")
+                if bq_table.table_id.endswith("*"):
+                    columns["_TABLE_SUFFIX"] = exp.DataType.build("STRING", dialect="bigquery")
+                if (
+                    bq_table.external_data_configuration is not None
+                    and bq_table.external_data_configuration.source_format
+                    in (
+                        "CSV",
+                        "NEWLINE_DELIMITED_JSON",
+                        "AVRO",
+                        "PARQUET",
+                        "ORC",
+                        "DATASTORE_BACKUP",
+                    )
+                ):
+                    columns["_FILE_NAME"] = exp.DataType.build("STRING", dialect="bigquery")
 
         return columns
 

sqlmesh/core/linter/rule.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import abc
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
 
 from sqlmesh.core.model import Model
@@ -49,12 +49,21 @@ class TextEdit:
     new_text: str
 
 
+@dataclass(frozen=True)
+class CreateFile:
+    """Create a new file with the provided text."""
+
+    path: Path
+    text: str
+
+
 @dataclass(frozen=True)
 class Fix:
     """A fix that can be applied to resolve a rule violation."""
 
     title: str
-    edits: t.List[TextEdit]
+    edits: t.List[TextEdit] = field(default_factory=list)
+    create_files: t.List[CreateFile] = field(default_factory=list)
 
 
 class _Rule(abc.ABCMeta):

sqlmesh/core/linter/rules/builtin.py

@@ -14,7 +14,15 @@
     get_range_of_model_block,
     read_range_from_string,
 )
-from sqlmesh.core.linter.rule import Rule, RuleViolation, Range, Fix, TextEdit, Position
+from sqlmesh.core.linter.rule import (
+    Rule,
+    RuleViolation,
+    Range,
+    Fix,
+    TextEdit,
+    Position,
+    CreateFile,
+)
 from sqlmesh.core.linter.definition import RuleSet
 from sqlmesh.core.model import Model, SqlModel, ExternalModel
 from sqlmesh.utils.lineage import extract_references_from_query, ExternalModelReference
@@ -227,7 +235,16 @@ def create_fix(self, model_name: str) -> t.Optional[Fix]:
 
         external_models_path = root / EXTERNAL_MODELS_YAML
         if not external_models_path.exists():
-            return None
+            return Fix(
+                title="Add external model file",
+                edits=[],
+                create_files=[
+                    CreateFile(
+                        path=external_models_path,
+                        text=f"- name: '{model_name}'\n",
+                    )
+                ],
+            )
 
         # Figure out the position to insert the new external model at the end of the file, whether
         # needs new line or not.

sqlmesh/core/loader.py

@@ -710,6 +710,8 @@ def _load_materializations(self) -> None:
     def _load_signals(self) -> UniqueKeyDict[str, signal]:
         """Loads signals for the built-in scheduler."""
 
+        base_signals = signal.get_registry()
+
         signals_max_mtime: t.Optional[float] = None
 
         for path in self._glob_paths(
@@ -729,7 +731,10 @@ def _load_signals(self) -> UniqueKeyDict[str, signal]:
 
         self._signals_max_mtime = signals_max_mtime
 
-        return signal.get_registry()
+        signals = signal.get_registry()
+        signal.set_registry(base_signals)
+
+        return signals
 
     def _load_audits(
         self, macros: MacroRegistry, jinja_macros: JinjaMacroRegistry

sqlmesh/core/model/common.py

@@ -157,7 +157,7 @@ def _add_variables_to_python_env(
 
     if blueprint_variables:
         blueprint_variables = {
-            k: SqlValue(sql=v.sql(dialect=dialect)) if isinstance(v, exp.Expression) else v
+            k.lower(): SqlValue(sql=v.sql(dialect=dialect)) if isinstance(v, exp.Expression) else v
             for k, v in blueprint_variables.items()
         }
         python_env[c.SQLMESH_BLUEPRINT_VARS] = Executable.value(

sqlmesh/core/test/definition.py

@@ -317,6 +317,13 @@ def _to_hashable(x: t.Any) -> t.Any:
             #
             # This is a bit of a hack, but it's a way to get the best of both worlds.
             args: t.List[t.Any] = []
+
+            failed_subtest = ""
+
+            if subtest := getattr(self, "_subtest", None):
+                if cte := subtest.params.get("cte"):
+                    failed_subtest = f" (CTE {cte})"
+
             if expected.shape != actual.shape:
                 _raise_if_unexpected_columns(expected.columns, actual.columns)
 
@@ -325,13 +332,13 @@ def _to_hashable(x: t.Any) -> t.Any:
                 missing_rows = _row_difference(expected, actual)
                 if not missing_rows.empty:
                     args[0] += f"\n\nMissing rows:\n\n{missing_rows}"
-                    args.append(df_to_table("Missing rows", missing_rows))
+                    args.append(df_to_table(f"Missing rows{failed_subtest}", missing_rows))
 
                 unexpected_rows = _row_difference(actual, expected)
 
                 if not unexpected_rows.empty:
                     args[0] += f"\n\nUnexpected rows:\n\n{unexpected_rows}"
-                    args.append(df_to_table("Unexpected rows", unexpected_rows))
+                    args.append(df_to_table(f"Unexpected rows{failed_subtest}", unexpected_rows))
 
             else:
                 diff = expected.compare(actual).rename(columns={"self": "exp", "other": "act"})
@@ -341,7 +348,8 @@ def _to_hashable(x: t.Any) -> t.Any:
                 diff.rename(columns={"exp": "Expected", "act": "Actual"}, inplace=True)
                 if self.verbosity == Verbosity.DEFAULT:
                     args.extend(
-                        df_to_table("Data mismatch", df) for df in _split_df_by_column_pairs(diff)
+                        df_to_table(f"Data mismatch{failed_subtest}", df)
+                        for df in _split_df_by_column_pairs(diff)
                     )
                 else:
                     from pandas import MultiIndex
@@ -351,7 +359,8 @@ def _to_hashable(x: t.Any) -> t.Any:
                         col_diff = diff[col]
                         if not col_diff.empty:
                             table = df_to_table(
-                                f"[bold red]Column '{col}' mismatch[/bold red]", col_diff
+                                f"[bold red]Column '{col}' mismatch{failed_subtest}[/bold red]",
+                                col_diff,
                             )
                             args.append(table)
 

sqlmesh/core/test/result.py

@@ -4,6 +4,8 @@
 import typing as t
 import unittest
 
+from sqlmesh.core.test.definition import ModelTest
+
 if t.TYPE_CHECKING:
     ErrorType = t.Union[
         t.Tuple[type[BaseException], BaseException, types.TracebackType],
@@ -42,7 +44,10 @@ def addSubTest(
             exctype, value, tb = err
             err = (exctype, value, None)  # type: ignore
 
-        super().addSubTest(test, subtest, err)
+            if err[0] and issubclass(err[0], test.failureException):
+                self.addFailure(test, err)
+            else:
+                self.addError(test, err)
 
     def _print_char(self, char: str) -> None:
         from sqlmesh.core.console import TerminalConsole
@@ -117,4 +122,14 @@ def merge(self, other: ModelTextTestResult) -> None:
             skipped_args = other.skipped[0]
             self.addSkip(skipped_args[0], skipped_args[1])
 
-        self.testsRun += 1
+        self.testsRun += other.testsRun
+
+    def get_fail_and_error_tests(self) -> t.List[ModelTest]:
+        # If tests contain failed subtests (e.g testing CTE outputs) we don't want
+        # to report it as different test failures
+        test_name_to_test = {
+            test.test_name: test
+            for test, _ in self.failures + self.errors
+            if isinstance(test, ModelTest)
+        }
+        return list(test_name_to_test.values())