refactor make jinja registry in a method; ensure correct execution order

Author: themisvaltinos

sqlmesh/dbt/loader.py

@@ -30,6 +30,7 @@
 from sqlmesh.utils.jinja import (
     JinjaMacroRegistry,
     extract_macro_references_and_variables,
+    make_jinja_registry,
 )
 
 if sys.version_info >= (3, 12):
@@ -244,22 +245,27 @@ def _load_environment_statements(self, macros: MacroRegistry) -> t.List[Environm
         dialect = self.config.dialect
         for project in self._load_projects():
             context = project.context
+            hooks_by_package_name: t.Dict[str, EnvironmentStatements] = {}
             for package_name, package in project.packages.items():
                 context.set_and_render_variables(package.variables, package_name)
-                on_run_start: t.List[str] = []
-                on_run_end: t.List[str] = []
-                for hook in package.on_run_start.values():
-                    on_run_start.append(hook.sql)
-                for hook in package.on_run_end.values():
-                    on_run_end.append(hook.sql)
+                on_run_start: t.List[str] = [
+                    on_run_hook.sql
+                    for on_run_hook in sorted(package.on_run_start.values(), key=lambda h: h.index)
+                ]
+                on_run_end: t.List[str] = [
+                    on_run_hook.sql
+                    for on_run_hook in sorted(package.on_run_end.values(), key=lambda h: h.index)
+                ]
 
                 if statements := on_run_start + on_run_end:
                     jinja_references, used_variables = extract_macro_references_and_variables(
                         *(gen(stmt) for stmt in statements)
                     )
-                    jinja_registry = context.jinja_macros.copy()
-                    jinja_registry.root_macros = jinja_registry.packages.get(package_name) or {}
-                    jinja_registry = jinja_registry.trim(jinja_references)
+
+                    jinja_registry = make_jinja_registry(
+                        context.jinja_macros, package_name, jinja_references
+                    )
+
                     python_env = make_python_env(
                         [s for stmt in statements for s in d.parse(stmt, default_dialect=dialect)],
                         jinja_macro_references=jinja_references,
@@ -270,20 +276,26 @@ def _load_environment_statements(self, macros: MacroRegistry) -> t.List[Environm
                         path=self.config_path,
                     )
 
-                    environment_statements.append(
-                        EnvironmentStatements(
-                            before_all=[
-                                d.jinja_statement(stmt).sql(dialect=dialect)
-                                for stmt in on_run_start or []
-                            ],
-                            after_all=[
-                                d.jinja_statement(stmt).sql(dialect=dialect)
-                                for stmt in on_run_end or []
-                            ],
-                            python_env=python_env,
-                            jinja_macros=jinja_registry,
-                        )
+                    hooks_by_package_name[package_name] = EnvironmentStatements(
+                        before_all=[
+                            d.jinja_statement(stmt).sql(dialect=dialect)
+                            for stmt in on_run_start or []
+                        ],
+                        after_all=[
+                            d.jinja_statement(stmt).sql(dialect=dialect)
+                            for stmt in on_run_end or []
+                        ],
+                        python_env=python_env,
+                        jinja_macros=jinja_registry,
                     )
+                # Project hooks should be executed first and then rest of the packages
+                environment_statements = [
+                    statements
+                    for _, statements in sorted(
+                        hooks_by_package_name.items(),
+                        key=lambda item: 0 if item[0] == context.project_name else 1,
+                    )
+                ]
         return environment_statements
 
     def _compute_yaml_max_mtime_per_subfolder(self, root: Path) -> t.Dict[Path, float]:

sqlmesh/dbt/manifest.py

@@ -294,11 +294,11 @@ def _load_on_run_start_end(self) -> None:
                 node_path = Path(node.original_file_path)
                 if "on-run-start" in node.tags:
                     self._on_run_start_per_package[node.package_name][node_name] = HookConfig(
-                        sql=sql, path=node_path
+                        sql=sql, index=node.index or 0, path=node_path
                     )
                 else:
                     self._on_run_end_per_package[node.package_name][node_name] = HookConfig(
-                        sql=sql, path=node_path
+                        sql=sql, index=node.index or 0, path=node_path
                     )
 
     @property

sqlmesh/dbt/package.py

@@ -32,6 +32,7 @@ class HookConfig(PydanticModel):
     """Class to contain on run start / on run end hooks."""
 
     sql: str
+    index: int
     path: Path
 
 

sqlmesh/utils/jinja.py

@@ -608,3 +608,31 @@ def create_builtin_globals(
         c.GATEWAY: lambda: variables.get(c.GATEWAY, None),
         **global_vars,
     }
+
+
+def make_jinja_registry(
+    jinja_macros: JinjaMacroRegistry, package_name: str, jinja_references: t.Set[MacroReference]
+) -> JinjaMacroRegistry:
+    """
+    Creates a Jinja macro registry for a specific package.
+
+    This function takes an existing Jinja macro registry and returns a new
+    registry that includes only the macros associated with the specified
+    package and trims the registry to include only the macros referenced
+    in the provided set of macro references.
+
+    Args:
+        jinja_macros: The original Jinja macro registry containing all macros.
+        package_name: The name of the package for which to create the registry.
+        jinja_references: A set of macro references to retain in the new registry.
+
+    Returns:
+        A new JinjaMacroRegistry containing only the macros for the specified
+        package and the referenced macros.
+    """
+
+    jinja_registry = jinja_macros.copy()
+    jinja_registry.root_macros = jinja_registry.packages.get(package_name) or {}
+    jinja_registry = jinja_registry.trim(jinja_references)
+
+    return jinja_registry

tests/dbt/test_adapter.py

@@ -280,16 +280,20 @@ def test_on_run_start_end(copy_to_temp_path):
     sushi_context = Context(paths=copy_to_temp_path(project_root))
     assert len(sushi_context._environment_statements) == 2
 
-    # Root project on run start / on run end
+    # Root project's on run start / on run end should be first by checking the macros
     root_environment_statements = sushi_context._environment_statements[0]
+    assert "create_tables" in root_environment_statements.jinja_macros.root_macros
+
+    # Validate order of execution to be correct
     assert root_environment_statements.before_all == [
-        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS analytic_stats (physical_table VARCHAR, evaluation_time VARCHAR);\nJINJA_END;"
+        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS analytic_stats (physical_table VARCHAR, evaluation_time VARCHAR);\nJINJA_END;",
+        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS to_be_executed_last (col VARCHAR);\nJINJA_END;",
     ]
     assert root_environment_statements.after_all == [
-        "JINJA_STATEMENT_BEGIN;\n{{ create_tables(schemas) }}\nJINJA_END;"
+        "JINJA_STATEMENT_BEGIN;\n{{ create_tables(schemas) }}\nJINJA_END;",
+        "JINJA_STATEMENT_BEGIN;\nDROP TABLE to_be_executed_last;\nJINJA_END;",
     ]
 
-    assert "create_tables" in root_environment_statements.jinja_macros.root_macros
     assert root_environment_statements.jinja_macros.root_package_name == "sushi"
 
     rendered_before_all = render_statements(
@@ -311,25 +315,30 @@ def test_on_run_start_end(copy_to_temp_path):
     )
 
     assert rendered_before_all == [
-        "CREATE TABLE IF NOT EXISTS analytic_stats (physical_table TEXT, evaluation_time TEXT)"
+        "CREATE TABLE IF NOT EXISTS analytic_stats (physical_table TEXT, evaluation_time TEXT)",
+        "CREATE TABLE IF NOT EXISTS to_be_executed_last (col TEXT)",
     ]
 
     # The jinja macro should have resolved the schemas for this environment and generated corresponding statements
     assert sorted(rendered_after_all) == sorted(
         [
             "CREATE OR REPLACE TABLE schema_table_snapshots__dev AS SELECT 'snapshots__dev' AS schema",
             "CREATE OR REPLACE TABLE schema_table_sushi__dev AS SELECT 'sushi__dev' AS schema",
+            "DROP TABLE to_be_executed_last",
         ]
     )
 
     # Nested dbt_packages on run start / on run end
     packaged_environment_statements = sushi_context._environment_statements[1]
 
+    # Validate order of execution to be correct
     assert packaged_environment_statements.before_all == [
-        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS analytic_stats_packaged_project (physical_table VARCHAR, evaluation_time VARCHAR);\nJINJA_END;"
+        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS to_be_executed_first (col VARCHAR);\nJINJA_END;",
+        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS analytic_stats_packaged_project (physical_table VARCHAR, evaluation_time VARCHAR);\nJINJA_END;",
     ]
     assert packaged_environment_statements.after_all == [
-        "JINJA_STATEMENT_BEGIN;\n{{ packaged_tables(schemas) }}\nJINJA_END;"
+        "JINJA_STATEMENT_BEGIN;\nDROP TABLE to_be_executed_first\nJINJA_END;",
+        "JINJA_STATEMENT_BEGIN;\n{{ packaged_tables(schemas) }}\nJINJA_END;",
     ]
 
     assert "packaged_tables" in packaged_environment_statements.jinja_macros.root_macros
@@ -353,13 +362,19 @@ def test_on_run_start_end(copy_to_temp_path):
         environment_naming_info=EnvironmentNamingInfo(name="dev"),
     )
 
+    # Validate order of execution to match dbt's
     assert rendered_before_all == [
-        "CREATE TABLE IF NOT EXISTS analytic_stats_packaged_project (physical_table TEXT, evaluation_time TEXT)"
+        "CREATE TABLE IF NOT EXISTS to_be_executed_first (col TEXT)",
+        "CREATE TABLE IF NOT EXISTS analytic_stats_packaged_project (physical_table TEXT, evaluation_time TEXT)",
     ]
 
+    # This on run end statement should be executed first
+    assert rendered_after_all[0] == "DROP TABLE to_be_executed_first"
+
     # The table names is an indication of the rendering of the dbt_packages statements
     assert sorted(rendered_after_all) == sorted(
         [
+            "DROP TABLE to_be_executed_first",
             "CREATE OR REPLACE TABLE schema_table_snapshots__dev_nested_package AS SELECT 'snapshots__dev' AS schema",
             "CREATE OR REPLACE TABLE schema_table_sushi__dev_nested_package AS SELECT 'sushi__dev' AS schema",
         ]

tests/dbt/test_transformation.py

@@ -1003,15 +1003,33 @@ def test_dbt_version(sushi_test_project: Project):
 
 @pytest.mark.xdist_group("dbt_manifest")
 def test_dbt_on_run_start_end(sushi_test_project: Project):
+    # Validate perservation of dbt's order of execution
+    assert sushi_test_project.packages["sushi"].on_run_start["sushi-on-run-start-0"].index == 0
+    assert sushi_test_project.packages["sushi"].on_run_start["sushi-on-run-start-1"].index == 1
+    assert sushi_test_project.packages["sushi"].on_run_end["sushi-on-run-end-0"].index == 0
+    assert sushi_test_project.packages["sushi"].on_run_end["sushi-on-run-end-1"].index == 1
+    assert (
+        sushi_test_project.packages["customers"].on_run_start["customers-on-run-start-0"].index == 0
+    )
+    assert (
+        sushi_test_project.packages["customers"].on_run_start["customers-on-run-start-1"].index == 1
+    )
+    assert sushi_test_project.packages["customers"].on_run_end["customers-on-run-end-0"].index == 0
+    assert sushi_test_project.packages["customers"].on_run_end["customers-on-run-end-1"].index == 1
+
     assert (
         sushi_test_project.packages["customers"].on_run_start["customers-on-run-start-0"].sql
+        == "CREATE TABLE IF NOT EXISTS to_be_executed_first (col VARCHAR);"
+    )
+    assert (
+        sushi_test_project.packages["customers"].on_run_start["customers-on-run-start-1"].sql
         == "CREATE TABLE IF NOT EXISTS analytic_stats_packaged_project (physical_table VARCHAR, evaluation_time VARCHAR);"
     )
     assert (
-        sushi_test_project.packages["customers"].on_run_end["customers-on-run-end-0"].sql
+        sushi_test_project.packages["customers"].on_run_end["customers-on-run-end-1"].sql
         == "{{ packaged_tables(schemas) }}"
     )
-    assert sushi_test_project.packages["sushi"].on_run_end
+
     assert (
         sushi_test_project.packages["sushi"].on_run_start["sushi-on-run-start-0"].sql
         == "CREATE TABLE IF NOT EXISTS analytic_stats (physical_table VARCHAR, evaluation_time VARCHAR);"

tests/fixtures/dbt/sushi_test/dbt_project.yml

@@ -61,5 +61,7 @@ vars:
 
 on-run-start:
   - 'CREATE TABLE IF NOT EXISTS analytic_stats (physical_table VARCHAR, evaluation_time VARCHAR);'
+  - 'CREATE TABLE IF NOT EXISTS to_be_executed_last (col VARCHAR);'
 on-run-end:
-  - '{{ create_tables(schemas) }}'
+  - '{{ create_tables(schemas) }}'
+  - 'DROP TABLE to_be_executed_last;'

tests/fixtures/dbt/sushi_test/packages/customers/dbt_project.yml

@@ -33,6 +33,8 @@ vars:
 
 
 on-run-start:
+  - 'CREATE TABLE IF NOT EXISTS to_be_executed_first (col VARCHAR);'
   - 'CREATE TABLE IF NOT EXISTS analytic_stats_packaged_project (physical_table VARCHAR, evaluation_time VARCHAR);'
 on-run-end:
+  - 'DROP TABLE to_be_executed_first'
   - '{{ packaged_tables(schemas) }}'