Feat: Expose the gateway name as a macro (#1509)

Author: izeigerman

examples/sushi/config.py

@@ -31,7 +31,8 @@
 
 # A configuration used for SQLMesh tests.
 test_config = Config(
-    default_connection=DuckDBConnectionConfig(),
+    gateways={"in_memory": GatewayConfig(connection=DuckDBConnectionConfig())},
+    default_gateway="in_memory",
     auto_categorize_changes=CategorizerConfig(sql=AutoCategorizationMode.SEMI),
     model_defaults=ModelDefaultsConfig(dialect="duckdb"),
 )

sqlmesh/core/config/root.py

@@ -5,6 +5,7 @@
 import zlib
 
 from pydantic import Field
+from sqlglot.helper import first
 
 from sqlmesh.core import constants as c
 from sqlmesh.core.config import EnvironmentSuffixTarget
@@ -122,7 +123,7 @@ def get_gateway(self, name: t.Optional[str] = None) -> GatewayConfig:
                 if "" in self.gateways:
                     return self.gateways[""]
 
-                return next(iter(self.gateways.values()))
+                return first(self.gateways.values())
 
             if name not in self.gateways:
                 raise ConfigError(f"Missing gateway with name '{name}'.")
@@ -152,6 +153,14 @@ def get_scheduler(self, gateway_name: t.Optional[str] = None) -> SchedulerConfig
     def get_state_schema(self, gateway_name: t.Optional[str] = None) -> t.Optional[str]:
         return self.get_gateway(gateway_name).state_schema
 
+    @property
+    def default_gateway_name(self) -> str:
+        if self.default_gateway:
+            return self.default_gateway
+        if "" in self.gateways:
+            return ""
+        return first(self.gateways)
+
     @property
     def dialect(self) -> t.Optional[str]:
         return self.model_defaults.dialect

sqlmesh/core/context.py

@@ -388,6 +388,7 @@ def load(self, update_schemas: bool = True) -> Context:
             gc.disable()
             project = self._loader.load(self, update_schemas)
             self._macros = project.macros
+            self._jinja_macros = project.jinja_macros
             self._models = project.models
             self._metrics = project.metrics
             self._standalone_audits.clear()

sqlmesh/core/loader.py

@@ -33,6 +33,7 @@
 from sqlmesh.utils.dag import DAG
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.jinja import JinjaMacroRegistry, MacroExtractor
+from sqlmesh.utils.metaprogramming import Executable
 from sqlmesh.utils.yaml import YAML
 
 if t.TYPE_CHECKING:
@@ -256,6 +257,10 @@ def _load_scripts(self) -> t.Tuple[MacroRegistry, JinjaMacroRegistry]:
         macros = macro.get_registry()
         macro.set_registry(standard_macros)
 
+        gateway_name = self._context.gateway or self._context.config.default_gateway_name
+        macros["gateway"] = Executable.value(gateway_name)
+        jinja_macros.add_globals({"gateway": gateway_name})
+
         return macros, jinja_macros
 
     def _load_models(

sqlmesh/core/macros.py

@@ -114,6 +114,8 @@ def __init__(
                 self.macros[normalize_macro_name(k)] = self.env[v.name or k]
             elif v.is_import and getattr(self.env.get(k), "__sqlmesh_macro__", None):
                 self.macros[normalize_macro_name(k)] = self.env[k]
+            elif v.is_value:
+                self.locals[k] = self.env[k]
 
     def send(
         self, name: str, *args: t.Any

sqlmesh/core/model/definition.py

@@ -1823,10 +1823,14 @@ def _python_env(
     expressions = ensure_list(expressions)
     for expression in expressions:
         if not isinstance(expression, d.Jinja):
-            for macro_func in expression.find_all(d.MacroFunc):
-                if macro_func.__class__ is d.MacroFunc:
-                    name = macro_func.this.name.lower()
+            for macro_func_or_var in expression.find_all(d.MacroFunc, d.MacroVar):
+                if macro_func_or_var.__class__ is d.MacroFunc:
+                    name = macro_func_or_var.this.name.lower()
                     used_macros[name] = macros[name]
+                elif macro_func_or_var.__class__ is d.MacroVar:
+                    name = macro_func_or_var.name
+                    if name in macros:
+                        used_macros[name] = macros[name]
 
     for macro_ref in jinja_macro_references or set():
         if macro_ref.package is None and macro_ref.name in macros:

sqlmesh/utils/metaprogramming.py

@@ -304,7 +304,7 @@ def __lt__(self, other: t.Any) -> bool:
 
 
 class Executable(PydanticModel):
-    payload: t.Any
+    payload: str
     kind: ExecutableKind = ExecutableKind.DEFINITION
     name: t.Optional[str] = None
     path: t.Optional[str] = None
@@ -326,6 +326,10 @@ def is_statement(self) -> bool:
     def is_value(self) -> bool:
         return self.kind == ExecutableKind.VALUE
 
+    @classmethod
+    def value(cls, v: t.Any) -> Executable:
+        return Executable(payload=repr(v), kind=ExecutableKind.VALUE)
+
 
 def serialize_env(env: t.Dict[str, t.Any], path: Path) -> t.Dict[str, Executable]:
     """Serializes a python function into a self contained dictionary.
@@ -370,7 +374,7 @@ def serialize_env(env: t.Dict[str, t.Any], path: Path) -> t.Dict[str, Executable
                 kind=ExecutableKind.IMPORT,
             )
         else:
-            serialized[k] = Executable(payload=repr(v), kind=ExecutableKind.VALUE)
+            serialized[k] = Executable.value(v)
 
     return serialized
 

tests/core/test_context.py

@@ -287,7 +287,7 @@ def test():
     context = Context(paths=str(tmp_path), config=config)
 
     assert ["db.actual_test"] == list(context.models)
-    assert "test" == list(context._macros)[-1]
+    assert "test" in context._macros
 
 
 def test_plan_apply(sushi_context) -> None:
@@ -477,3 +477,41 @@ def test_default_schema_and_config(sushi_context_pre_scheduling) -> None:
         MappingSchema({"a": {"col": "int"}}), default_schema="schema", default_catalog="catalog"
     )
     assert c.mapping_schema == {"catalog": {"schema": {"a": {"col": "int"}}}}
+
+
+def test_gateway_macro(sushi_context: Context) -> None:
+    sushi_context.upsert_model(
+        load_sql_based_model(
+            parse(
+                """
+            MODEL(name sushi.test_gateway_macro);
+            SELECT @gateway AS gateway;
+            """
+            ),
+            macros=sushi_context._macros,
+        )
+    )
+
+    assert (
+        sushi_context.render("sushi.test_gateway_macro").sql()
+        == "SELECT 'in_memory' AS \"gateway\""
+    )
+
+    sushi_context.upsert_model(
+        load_sql_based_model(
+            parse(
+                """
+            MODEL(name sushi.test_gateway_macro_jinja);
+            JINJA_QUERY_BEGIN;
+            SELECT '{{ gateway }}' AS gateway_jinja;
+            JINJA_END;
+            """
+            ),
+            jinja_macros=sushi_context._jinja_macros,
+        )
+    )
+
+    assert (
+        sushi_context.render("sushi.test_gateway_macro_jinja").sql()
+        == "SELECT 'in_memory' AS \"gateway_jinja\""
+    )