Load variables from gateway specified in YAML

Author: VaggelisD

sqlmesh/core/context.py

@@ -1780,17 +1780,10 @@ def test(
         if verbosity >= Verbosity.VERBOSE:
             pd.set_option("display.max_columns", None)
 
-        # Merge the root variables with the gateway's variables
-        variables = {
-            **self.config.variables,
-            **self.config.get_gateway(self.selected_gateway).variables,
-        }
-
         test_meta = load_model_tests(
-            configs=self.configs,
+            loaders=self._loaders,
             tests=tests,
             patterns=match_patterns,
-            variables=variables,
         )
 
         return run_tests(

sqlmesh/core/test/__init__.py

@@ -3,9 +3,6 @@
 from sqlmesh.core.test.definition import ModelTest as ModelTest, generate_test as generate_test
 from sqlmesh.core.test.discovery import (
     ModelTestMetadata as ModelTestMetadata,
-    filter_tests_by_patterns as filter_tests_by_patterns,
-    get_all_model_tests as get_all_model_tests,
-    load_model_test_file as load_model_test_file,
     load_model_tests as load_model_tests,
 )
 from sqlmesh.core.test.result import ModelTextTestResult as ModelTextTestResult

sqlmesh/core/test/discovery.py

@@ -14,7 +14,7 @@
 from sqlmesh.utils.yaml import load as yaml_load
 
 if t.TYPE_CHECKING:
-    from sqlmesh.core.config.loader import C
+    from sqlmesh.core.loader import Loader
 
 
 class ModelTestMetadata(PydanticModel):
@@ -31,7 +31,8 @@ def __hash__(self) -> int:
 
 
 def load_model_test_file(
-    path: pathlib.Path, variables: dict[str, t.Any] | None = None
+    path: pathlib.Path,
+    get_variables: t.Callable[[t.Optional[str]], t.Dict[str, str]],
 ) -> dict[str, ModelTestMetadata]:
     """Load a single model test file.
 
@@ -42,7 +43,7 @@ def load_model_test_file(
         A list of ModelTestMetadata named tuples.
     """
     model_test_metadata = {}
-    contents = yaml_load(path, variables=variables)
+    contents = yaml_load(path, get_variables=get_variables)
 
     for test_name, value in contents.items():
         model_test_metadata[test_name] = ModelTestMetadata(
@@ -53,8 +54,8 @@ def load_model_test_file(
 
 def discover_model_tests(
     path: pathlib.Path,
+    get_variables: t.Callable[[t.Optional[str]], t.Dict[str, str]],
     ignore_patterns: list[str] | None = None,
-    variables: dict[str, t.Any] | None = None,
 ) -> Iterator[ModelTestMetadata]:
     """Discover model tests.
 
@@ -78,7 +79,7 @@ def discover_model_tests(
                 break
         else:
             for model_test_metadata in load_model_test_file(
-                yaml_file, variables=variables
+                yaml_file, get_variables=get_variables
             ).values():
                 yield model_test_metadata
 
@@ -103,34 +104,16 @@ def filter_tests_by_patterns(
     )
 
 
-def get_all_model_tests(
-    *paths: pathlib.Path,
-    patterns: list[str] | None = None,
-    ignore_patterns: list[str] | None = None,
-    variables: dict[str, t.Any] | None = None,
-) -> list[ModelTestMetadata]:
-    model_test_metadatas = [
-        meta
-        for path in paths
-        for meta in discover_model_tests(pathlib.Path(path), ignore_patterns, variables=variables)
-    ]
-    if patterns:
-        model_test_metadatas = filter_tests_by_patterns(model_test_metadatas, patterns)
-    return model_test_metadatas
-
-
 def load_model_tests(
-    configs: dict[pathlib.Path, C],
+    loaders: list[Loader],
     tests: t.Optional[t.List[str]] = None,
     patterns: list[str] | None = None,
-    variables: dict[str, t.Any] | None = None,
 ) -> list[ModelTestMetadata]:
     """Load model tests into a list of ModelTestMetadata which will be propagated to the test runner.
 
     Args:
         tests: A list of tests to load; If not specified, all tests are loaded
-        patterns: A list of patterns to match against.
-        variables: A dictionary of variables to use when loading the tests.
+        patterns: A list of patterns that'll be used to filter tests by file name.
         configs: A dictionary of configs to use when loading all the tests.
     """
     test_meta = []
@@ -139,24 +122,27 @@ def load_model_tests(
         for test in tests:
             filename, test_name = test.split("::", maxsplit=1) if "::" in test else (test, "")
 
-            test_file = load_model_test_file(pathlib.Path(filename), variables=variables)
+            test_file = load_model_test_file(
+                pathlib.Path(filename), get_variables=loaders[0]._get_variables
+            )
             if test_name:
                 test_meta.append(test_file[test_name])
             else:
                 test_meta.extend(test_file.values())
-
-        if patterns:
-            test_meta = filter_tests_by_patterns(test_meta, patterns)
-
     else:
-        for path, config in configs.items():
+        for loader in loaders:
             test_meta.extend(
-                get_all_model_tests(
-                    path / c.TESTS,
-                    patterns=patterns,
-                    ignore_patterns=config.ignore_patterns,
-                    variables=variables,
-                )
+                [
+                    meta
+                    for meta in discover_model_tests(
+                        pathlib.Path(loader.config_path / c.TESTS),
+                        ignore_patterns=loader.config.ignore_patterns,  # type: ignore
+                        get_variables=loader._get_variables,
+                    )
+                ]
             )
 
+    if patterns:
+        test_meta = filter_tests_by_patterns(test_meta, patterns)
+
     return test_meta

sqlmesh/core/test/runner.py

@@ -14,9 +14,6 @@
 from sqlmesh.core.test.definition import ModelTest as ModelTest, generate_test as generate_test
 from sqlmesh.core.test.discovery import (
     ModelTestMetadata as ModelTestMetadata,
-    filter_tests_by_patterns as filter_tests_by_patterns,
-    get_all_model_tests as get_all_model_tests,
-    load_model_test_file as load_model_test_file,
 )
 from sqlmesh.core.config.connection import BaseDuckDBConnectionConfig
 

sqlmesh/magics.py

@@ -26,13 +26,12 @@
 from rich.jupyter import JupyterRenderable
 from sqlmesh.cli.example_project import ProjectTemplate, init_example_project
 from sqlmesh.core import analytics
-from sqlmesh.core import constants as c
 from sqlmesh.core.config import load_configs
 from sqlmesh.core.console import create_console, set_console, configure_console
 from sqlmesh.core.context import Context
 from sqlmesh.core.dialect import format_model_expressions, parse
 from sqlmesh.core.model import load_sql_based_model
-from sqlmesh.core.test import ModelTestMetadata, get_all_model_tests
+from sqlmesh.core.test import ModelTestMetadata, load_model_tests
 from sqlmesh.utils import sqlglot_dialects, yaml, Verbosity
 from sqlmesh.utils.errors import MagicError, MissingContextException, SQLMeshError
 
@@ -273,15 +272,7 @@ def test(self, context: Context, line: str, test_def_raw: t.Optional[str] = None
         if not args.test_name and not args.ls:
             raise MagicError("Must provide either test name or `--ls` to list tests")
 
-        test_meta = []
-
-        for path, config in context.configs.items():
-            test_meta.extend(
-                get_all_model_tests(
-                    path / c.TESTS,
-                    ignore_patterns=config.ignore_patterns,
-                )
-            )
+        test_meta = load_model_tests(loaders=context._loaders)
 
         tests: t.Dict[str, t.Dict[str, ModelTestMetadata]] = defaultdict(dict)
         for model_test_metadata in test_meta:

sqlmesh/utils/yaml.py

@@ -5,6 +5,7 @@
 from decimal import Decimal
 from os import getenv
 from pathlib import Path
+import re
 
 from ruamel import yaml
 
@@ -17,6 +18,9 @@
 }
 
 
+gateway_pattern = re.compile(r"gateway:\s*([^\s]+)")
+
+
 def YAML(typ: t.Optional[str] = "safe") -> yaml.YAML:
     yaml_obj = yaml.YAML(typ=typ)
 
@@ -36,6 +40,7 @@ def load(
     render_jinja: bool = True,
     allow_duplicate_keys: bool = False,
     variables: t.Optional[t.Dict[str, t.Any]] = None,
+    get_variables: t.Callable[[t.Optional[str]], t.Dict[str, str]] | None = None,
 ) -> t.Dict:
     """Loads a YAML object from either a raw string or a file."""
     path: t.Optional[Path] = None
@@ -45,6 +50,10 @@ def load(
         with open(source, "r", encoding="utf-8") as file:
             source = file.read()
 
+    if get_variables:
+        gateway = gateway_pattern.search(source)
+        variables = get_variables(gateway.group(1) if gateway else None)
+
     if render_jinja:
         source = ENVIRONMENT.from_string(source).render(
             {

tests/core/test_test.py

@@ -1563,26 +1563,32 @@ def test_variable_usage(tmp_path: Path) -> None:
     init_example_project(tmp_path, dialect="duckdb")
 
     variables = {"gold": "gold_db", "silver": "silver_db"}
-
-    # Case 1: Test root variables
-    config = Config(
-        default_connection=DuckDBConnectionConfig(),
-        model_defaults=ModelDefaultsConfig(dialect="duckdb"),
-        variables=variables,
-    )
-    context = Context(paths=tmp_path, config=config)
+    incorrect_variables = {"gold": "foo", "silver": "bar"}
 
     parent = _create_model(
         "SELECT 1 AS id, '2022-01-02'::DATE AS ds, @start_ts AS start_ts",
         meta="MODEL (name silver_db.sch.b, kind INCREMENTAL_BY_TIME_RANGE(time_column ds))",
     )
-    parent = t.cast(SqlModel, context.upsert_model(parent))
 
     child = _create_model(
         "SELECT ds, @IF(@VAR('myvar'), id, id + 1) AS id FROM silver_db.sch.b WHERE ds BETWEEN @start_ds and @end_ds",
         meta="MODEL (name gold_db.sch.a, kind INCREMENTAL_BY_TIME_RANGE(time_column ds))",
     )
-    child = t.cast(SqlModel, context.upsert_model(child))
+
+    def init_context(config: Config, **kwargs):
+        context = Context(paths=tmp_path, config=config, **kwargs)
+        context.upsert_model(parent)
+        context.upsert_model(child)
+        return context
+
+    # Case 1: Test root variables
+    config = Config(
+        default_connection=DuckDBConnectionConfig(),
+        model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+        variables=variables,
+    )
+
+    context = init_context(config)
 
     test_file = tmp_path / "tests" / "test_parameterized_model_names.yaml"
     test_file.write_text(
@@ -1617,28 +1623,76 @@ def test_variable_usage(tmp_path: Path) -> None:
     assert len(results.successes) == 2
 
     # Case 2: Test gateway variables
-    context.config = Config(
+    config = Config(
         gateways={
-            "": GatewayConfig(connection=DuckDBConnectionConfig(), variables=variables),
+            "main": GatewayConfig(connection=DuckDBConnectionConfig(), variables=variables),
         },
         model_defaults=ModelDefaultsConfig(dialect="duckdb"),
     )
 
+    context = init_context(config, gateway="main")
+
     results = context.test()
 
     assert not results.failures
     assert not results.errors
     assert len(results.successes) == 2
 
     # Case 3: Test gateway variables overriding root variables
-    context.config = Config(
+    config = Config(
         gateways={
-            "": GatewayConfig(connection=DuckDBConnectionConfig(), variables=variables),
+            "main": GatewayConfig(connection=DuckDBConnectionConfig(), variables=variables),
         },
         model_defaults=ModelDefaultsConfig(dialect="duckdb"),
-        variables={"gold": "foo", "silver": "bar"},
+        variables=incorrect_variables,
     )
 
+    context = init_context(config, gateway="main")
+
+    results = context.test()
+
+    assert not results.failures
+    assert not results.errors
+    assert len(results.successes) == 2
+
+    # Case 4: Use variable from the defined gateway
+    test_file = tmp_path / "tests" / "test_parameterized_model_names.yaml"
+    test_file.write_text(
+        """
+test_parameterized_model_names:
+  model: {{ var('gold') }}.sch.a
+  gateway: secondary
+  vars:
+    myvar: True
+    start_ds: 2022-01-01
+    end_ds: 2022-01-03
+  inputs:
+    {{ var('silver') }}.sch.b:
+      - ds: 2022-01-01
+        id: 1
+      - ds: 2022-01-01
+        id: 2
+  outputs:
+    query:
+      - ds: 2022-01-01
+        id: 1
+      - ds: 2022-01-01
+        id: 2
+        """
+    )
+
+    config = Config(
+        gateways={
+            "main": GatewayConfig(
+                connection=DuckDBConnectionConfig(), variables=incorrect_variables
+            ),
+            "secondary": GatewayConfig(connection=DuckDBConnectionConfig(), variables=variables),
+        },
+        model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+    )
+
+    context = init_context(config, gateway="main")
+
     results = context.test()
 
     assert not results.failures