Feat!: Support pre-post statements in python models at creation time

sqlmesh/core/loader.py

@@ -329,7 +329,7 @@ def _load_models(
         """
         models = self._load_sql_models(macros, jinja_macros, audits)
         models.update(self._load_external_models(gateway))
-        models.update(self._load_python_models())
+        models.update(self._load_python_models(macros, jinja_macros))
 
         return models
 
@@ -392,7 +392,9 @@ def _load() -> Model:
 
         return models
 
-    def _load_python_models(self) -> UniqueKeyDict[str, Model]:
+    def _load_python_models(
+        self, macros: MacroRegistry, jinja_macros: JinjaMacroRegistry
+    ) -> UniqueKeyDict[str, Model]:
         """Loads the python models into a Dict"""
         models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
         registry = model_registry.registry()
@@ -418,6 +420,8 @@ def _load_python_models(self) -> UniqueKeyDict[str, Model]:
                             path=path,
                             module_path=context_path,
                             defaults=config.model_defaults.dict(),
+                            macros=macros,
+                            jinja_macros=jinja_macros,
                             dialect=config.model_defaults.dialect,
                             time_column_format=config.time_column_format,
                             physical_schema_override=config.physical_schema_override,

sqlmesh/core/model/decorator.py

@@ -8,8 +8,10 @@
 from sqlglot import exp
 from sqlglot.dialects.dialect import DialectType
 
+from sqlmesh.core.macros import MacroRegistry
+from sqlmesh.utils.jinja import JinjaMacroRegistry
 from sqlmesh.core import constants as c
-from sqlmesh.core.dialect import MacroFunc
+from sqlmesh.core.dialect import MacroFunc, parse_one
 from sqlmesh.core.model.definition import (
     Model,
     create_python_model,
@@ -75,6 +77,8 @@ def model(
         module_path: Path,
         path: Path,
         defaults: t.Optional[t.Dict[str, t.Any]] = None,
+        macros: t.Optional[MacroRegistry] = None,
+        jinja_macros: t.Optional[JinjaMacroRegistry] = None,
         dialect: t.Optional[str] = None,
         time_column_format: str = c.DEFAULT_TIME_COLUMN_FORMAT,
         physical_schema_override: t.Optional[t.Dict[str, str]] = None,
@@ -123,7 +127,9 @@ def model(
         for key in ("pre_statements", "post_statements"):
             statements = common_kwargs.get(key)
             if statements:
-                common_kwargs[key] = [exp.maybe_parse(s, dialect=dialect) for s in statements]
+                common_kwargs[key] = [
+                    parse_one(s, dialect=dialect) if isinstance(s, str) else s for s in statements
+                ]
 
         if self.is_sql:
             query = MacroFunc(this=exp.Anonymous(this=entrypoint))
@@ -132,5 +138,12 @@ def model(
             )
 
         return create_python_model(
-            self.name, entrypoint, columns=self.columns, dialect=dialect, **common_kwargs
+            self.name,
+            entrypoint,
+            module_path=module_path,
+            macros=macros,
+            jinja_macros=jinja_macros,
+            columns=self.columns,
+            dialect=dialect,
+            **common_kwargs,
         )