Merge branch 'themis/dbt_custom_materialisations' of https://github.com/TobikoData/sqlmesh into themis/dbt_custom_materialisations

Author: themisvaltinos

sqlmesh/core/snapshot/evaluator.py

@@ -67,7 +67,7 @@
     SnapshotTableCleanupTask,
 )
 from sqlmesh.core.snapshot.execution_tracker import QueryExecutionTracker
-from sqlmesh.utils import random_id, CorrelationId
+from sqlmesh.utils import random_id, CorrelationId, AttributeDict
 from sqlmesh.utils.concurrency import (
     concurrent_apply_to_snapshots,
     concurrent_apply_to_values,
@@ -2731,12 +2731,12 @@ def _execute_materialization(
         **kwargs: t.Any,
     ) -> None:
         jinja_macros = model.jinja_macros
-        existing_globals = jinja_macros.global_objs.copy()
 
         # For vdes we need to use the table, since we don't know the schema/table at parse time
         parts = exp.to_table(table_name, dialect=self.adapter.dialect)
 
-        relation_info = existing_globals.pop("this")
+        existing_globals = jinja_macros.global_objs
+        relation_info = existing_globals.get("this")
         if isinstance(relation_info, dict):
             relation_info["database"] = parts.catalog
             relation_info["identifier"] = parts.name
@@ -2750,29 +2750,29 @@ def _execute_materialization(
             "identifier": parts.name,
             "target": existing_globals.get("target", {"type": self.adapter.dialect}),
             "execution_dt": kwargs.get("execution_time"),
+            "engine_adapter": self.adapter,
+            "sql": str(query_or_df),
+            "is_first_insert": is_first_insert,
+            "create_only": create_only,
+            # FIXME: Add support for transaction=False
+            "pre_hooks": [
+                AttributeDict({"sql": s.this.this, "transaction": True})
+                for s in model.pre_statements
+            ],
+            "post_hooks": [
+                AttributeDict({"sql": s.this.this, "transaction": True})
+                for s in model.post_statements
+            ],
+            "model_instance": model,
+            **kwargs,
         }
 
-        context = jinja_macros._create_builtin_globals(
-            {"engine_adapter": self.adapter, **jinja_globals}
-        )
-
-        context.update(
-            {
-                "sql": str(query_or_df),
-                "is_first_insert": is_first_insert,
-                "create_only": create_only,
-                "pre_hooks": model.render_pre_statements(**render_kwargs),
-                "post_hooks": model.render_post_statements(**render_kwargs),
-                **kwargs,
-            }
-        )
-
         try:
-            jinja_env = jinja_macros.build_environment(**context)
+            jinja_env = jinja_macros.build_environment(**jinja_globals)
             template = jinja_env.from_string(self.materialization_template)
 
             try:
-                template.render(**context)
+                template.render()
             except MacroReturnVal as ret:
                 # this is a successful return from a macro call (dbt uses this list of Relations to update their relation cache)
                 returned_relations = ret.value.get("relations", [])

sqlmesh/dbt/builtin.py

@@ -547,6 +547,7 @@ def create_builtin_globals(
             "run_hooks": adapter.run_hooks,
             "graph": adapter.graph,
             "selected_resources": list(jinja_globals.get("selected_models") or []),
+            "write": lambda input: None,  # We don't support writing yet
         }
     )
 

sqlmesh/dbt/manifest.py

@@ -395,6 +395,12 @@ def _load_models_and_seeds(self) -> None:
                 dependencies = dependencies.union(
                     self._extra_dependencies(sql, node.package_name, track_all_model_attrs=True)
                 )
+                for hook in [*node_config.get("pre-hook", []), *node_config.get("post-hook", [])]:
+                    dependencies = dependencies.union(
+                        self._extra_dependencies(
+                            hook["sql"], node.package_name, track_all_model_attrs=True
+                        )
+                    )
                 dependencies = dependencies.union(
                     self._flatten_dependencies_from_macros(dependencies.macros, node.package_name)
                 )

sqlmesh/utils/jinja.py

@@ -369,6 +369,7 @@ def build_environment(self, **kwargs: t.Any) -> Environment:
         context.update(builtin_globals)
         context.update(root_macros)
         context.update(package_macros)
+        context["render"] = lambda input: env.from_string(input).render()
 
         env.globals.update(context)
         env.filters.update(self._environment.filters)

tests/dbt/test_custom_materializations.py

@@ -37,7 +37,7 @@ def test_custom_materialization_manifest_loading():
     assert custom_incremental.name == "custom_incremental"
     assert custom_incremental.adapter == "default"
     assert "make_temp_relation(new_relation)" in custom_incremental.definition
-    assert "run_hooks(pre_hooks, inside_transaction=False)" in custom_incremental.definition
+    assert "run_hooks(pre_hooks)" in custom_incremental.definition
     assert " {{ return({'relations': [new_relation]}) }}" in custom_incremental.definition
 
 

tests/fixtures/dbt/sushi_test/macros/materializations/custom_incremental.sql

@@ -19,7 +19,7 @@
   {%- set time_column = config.get('time_column') -%}
   {%- set interval_config = config.get('interval') -%}
 
-  {{ run_hooks(pre_hooks, inside_transaction=False) }}
+  {{ run_hooks(pre_hooks) }}
 
   {%- if existing_relation is none -%}
     {# The first insert creates new table if it doesn't exist #}
@@ -55,7 +55,7 @@
     {%- endcall -%}
   {%- endif -%}
 
-  {{ run_hooks(post_hooks, inside_transaction=False) }}
+  {{ run_hooks(post_hooks) }}
 
   {{ return({'relations': [new_relation]}) }}
-{%- endmaterialization -%}
+{%- endmaterialization -%}