Repository: TobikoData/sqlmesh
Pull request #3208 (2024-10-01 17:27:42)

Feat: add support for parameterized python model names
We don't currently support parameterized python models, i.e. whose names contain variables. This PR adds support for them.


Patch:
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index de895b69..9f44199a 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -1166,7 +1166,7 @@ def _unquote_schema(schema: t.Dict) -> t.Dict:
 
 @contextmanager
 def normalize_and_quote(
-    query: E, dialect: str, default_catalog: t.Optional[str], quote: bool = True
+    query: E, dialect: DialectType, default_catalog: t.Optional[str], quote: bool = True
 ) -> t.Iterator[E]:
     qualify_tables(query, catalog=default_catalog, dialect=dialect)
     normalize_identifiers(query, dialect=dialect)
diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py
index ea3b5557..975b81c6 100644
--- a/sqlmesh/core/macros.py
+++ b/sqlmesh/core/macros.py
@@ -44,6 +44,7 @@ from sqlmesh.utils.jinja import JinjaMacroRegistry, has_jinja
 from sqlmesh.utils.metaprogramming import Executable, prepare_env, print_exception
 
 if t.TYPE_CHECKING:
+    from sqlglot.dialects.dialect import DialectType
     from sqlmesh.core._typing import TableName
     from sqlmesh.core.engine_adapter import EngineAdapter
     from sqlmesh.core.snapshot import Snapshot
@@ -144,7 +145,7 @@ class MacroEvaluator:
 
     def __init__(
         self,
-        dialect: str = "",
+        dialect: DialectType = "",
         python_env: t.Optional[t.Dict[str, Executable]] = None,
         jinja_env: t.Optional[Environment] = None,
         schema: t.Optional[MappingSchema] = None,
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index 592bb63f..46ae495f 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -46,6 +46,7 @@ from sqlmesh.utils.metaprogramming import (
 )
 
 if t.TYPE_CHECKING:
+    from sqlglot.dialects.dialect import DialectType
     from sqlmesh.core._typing import TableName
     from sqlmesh.core.audit import ModelAudit, Audit
     from sqlmesh.core.context import ExecutionContext
@@ -1600,24 +1601,15 @@ def load_sql_based_model(
         if prop.name.lower() == "audits":
             model_audits = prop.args.get("value")
 
-    meta_python_env = _python_env(
-        expressions=meta,
-        jinja_macro_references=None,
+    meta_renderer = _meta_renderer(
+        expression=meta,
         module_path=module_path,
-        macros=macros or macro.get_registry(),
+        macros=macros,
+        jinja_macros=jinja_macros,
         variables=variables,
         path=path,
-    )
-    meta_renderer = ExpressionRenderer(
-        meta,
-        dialect,
-        [],
-        path=path,
-        jinja_macro_registry=jinja_macros,
-        python_env=meta_python_env,
+        dialect=dialect,
         default_catalog=default_catalog,
-        quote_identifiers=False,
-        normalize_identifiers=False,
     )
 
     rendered_meta_exprs = meta_renderer.render()
@@ -1954,6 +1946,19 @@ def create_python_model(
             )
         )
 
+    dialect = kwargs.get("dialect")
+    name_renderer = _meta_renderer(
+        expression=d.parse_one(name, dialect=dialect),
+        module_path=module_path,
+        macros=macros,
+        jinja_macros=jinja_macros,
+        variables=variables,
+        path=path,
+        dialect=dialect,
+        default_catalog=kwargs.get("default_catalog"),
+    )
+    name = t.cast(t.List[exp.Expression], name_renderer.render())[0].sql(dialect=dialect)
+
     parsed_depends_on, referenced_variables = (
         _parse_dependencies(python_env, entrypoint) if python_env is not None else (set(), set())
     )
@@ -2385,6 +2390,37 @@ def _refs_to_sql(values: t.Any) -> exp.Expression:
     return exp.Tuple(expressions=values)
 
 
+def _meta_renderer(
+    expression: exp.Expression,
+    module_path: Path,
+    path: Path,
+    jinja_macros: t.Optional[JinjaMacroRegistry] = None,
+    macros: t.Optional[MacroRegistry] = None,
+    dialect: DialectType = None,
+    variables: t.Optional[t.Dict[str, t.Any]] = None,
+    default_catalog: t.Optional[str] = None,
+) -> ExpressionRenderer:
+    meta_python_env = _python_env(
+        expressions=expression,
+        jinja_macro_references=None,
+        module_path=module_path,
+        macros=macros or macro.get_registry(),
+        variables=variables,
+        path=path,
+    )
+    return ExpressionRenderer(
+        expression,
+        dialect,
+        [],
+        path=path,
+        jinja_macro_registry=jinja_macros,
+        python_env=meta_python_env,
+        default_catalog=default_catalog,
+        quote_identifiers=False,
+        normalize_identifiers=False,
+    )
+
+
 META_FIELD_CONVERTER: t.Dict[str, t.Callable] = {
     "start": lambda value: exp.Literal.string(value),
     "cron": lambda value: exp.Literal.string(value),
diff --git a/sqlmesh/core/renderer.py b/sqlmesh/core/renderer.py
index c466ed62..46ef207b 100644
--- a/sqlmesh/core/renderer.py
+++ b/sqlmesh/core/renderer.py
@@ -28,6 +28,7 @@ from sqlmesh.utils.metaprogramming import Executable, prepare_env
 
 if t.TYPE_CHECKING:
     from sqlglot._typing import E
+    from sqlglot.dialects.dialect import DialectType
 
     from sqlmesh.core.snapshot import DeployabilityIndex, Snapshot
 
@@ -39,7 +40,7 @@ class BaseExpressionRenderer:
     def __init__(
         self,
         expression: exp.Expression,
-        dialect: str,
+        dialect: DialectType,
         macro_definitions: t.List[d.MacroDef],
         path: Path = Path(),
         jinja_macro_registry: t.Optional[JinjaMacroRegistry] = None,


Test patch:
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 8dd74f06..b62abf4e 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -4565,7 +4565,7 @@ def test_variables_jinja():
 
 def test_variables_python_model(mocker: MockerFixture) -> None:
     @model(
-        "test_variables_python_model",
+        "foo_@{bar}",
         kind="full",
         columns={"a": "string", "b": "string", "c": "string"},
     )
@@ -4580,12 +4580,13 @@ def test_variables_python_model(mocker: MockerFixture) -> None:
             ]
         )
 
-    python_model = model.get_registry()["test_variables_python_model"].model(
+    python_model = model.get_registry()["foo_@{bar}"].model(
         module_path=Path("."),
         path=Path("."),
-        variables={"test_var_a": "test_value", "test_var_unused": 2},
+        variables={"test_var_a": "test_value", "test_var_unused": 2, "bar": "suffix"},
     )
 
+    assert python_model.name == "foo_suffix"
     assert python_model.python_env[c.SQLMESH_VARS] == Executable.value({"test_var_a": "test_value"})
 
     context = ExecutionContext(mocker.Mock(), {}, None, None)