Feat(dbt): Add dbt graph context variable support

Author: themisvaltinos

sqlmesh/dbt/adapter.py

@@ -11,6 +11,7 @@
 from sqlmesh.core.snapshot import DeployabilityIndex, Snapshot, to_table_mapping
 from sqlmesh.utils.errors import ConfigError, ParsetimeAdapterCallError
 from sqlmesh.utils.jinja import JinjaMacroRegistry
+from sqlmesh.utils import AttributeDict
 
 if t.TYPE_CHECKING:
     import agate
@@ -158,6 +159,20 @@ def compare_dbr_version(self, major: int, minor: int) -> int:
         # Always return -1 to fallback to Spark macro implementations.
         return -1
 
+    @property
+    def graph(self) -> t.Any:
+        return AttributeDict(
+            {
+                "exposures": {},
+                "groups": {},
+                "metrics": {},
+                "nodes": {},
+                "sources": {},
+                "semantic_models": {},
+                "saved_queries": {},
+            }
+        )
+
 
 class ParsetimeAdapter(BaseAdapter):
     def get_relation(self, database: str, schema: str, identifier: str) -> t.Optional[BaseRelation]:
@@ -246,6 +261,10 @@ def __init__(
             **table_mapping,
         }
 
+    @property
+    def graph(self) -> t.Any:
+        return self.jinja_globals.get("flat_graph", super().graph)
+
     def get_relation(
         self, database: t.Optional[str], schema: str, identifier: str
     ) -> t.Optional[BaseRelation]:

sqlmesh/dbt/builtin.py

@@ -452,6 +452,7 @@ def create_builtin_globals(
             "load_result": sql_execution.load_result,
             "run_query": sql_execution.run_query,
             "statement": sql_execution.statement,
+            "graph": adapter.graph,
         }
     )
 

sqlmesh/dbt/context.py

@@ -11,6 +11,7 @@
 from sqlmesh.dbt.manifest import ManifestHelper
 from sqlmesh.dbt.target import TargetConfig
 from sqlmesh.utils import AttributeDict
+from sqlmesh.utils.conversions import serializable
 from sqlmesh.utils.errors import ConfigError, SQLMeshError
 from sqlmesh.utils.jinja import (
     JinjaGlobalAttribute,
@@ -195,6 +196,49 @@ def refs(self) -> t.Dict[str, t.Union[ModelConfig, SeedConfig]]:
                     self._refs[f"{config_name}_v{model.version}"] = model
         return self._refs
 
+    @property
+    def flat_graph(self) -> t.Dict[str, t.Any]:
+        if self._manifest is None:
+            return {
+                "exposures": {},
+                "groups": {},
+                "metrics": {},
+                "nodes": {},
+                "sources": {},
+                "semantic_models": {},
+                "saved_queries": {},
+            }
+
+        manifest = self._manifest._manifest
+        return {
+            "exposures": {
+                k: serializable(v.to_dict(omit_none=False))
+                for k, v in getattr(manifest, "exposures", {}).items()
+            },
+            "groups": {
+                k: serializable(v.to_dict(omit_none=False))
+                for k, v in getattr(manifest, "groups", {}).items()
+            },
+            "metrics": {
+                k: serializable(v.to_dict(omit_none=False))
+                for k, v in getattr(manifest, "metrics", {}).items()
+            },
+            "nodes": {
+                k: serializable(v.to_dict(omit_none=False)) for k, v in manifest.nodes.items()
+            },
+            "sources": {
+                k: serializable(v.to_dict(omit_none=False)) for k, v in manifest.sources.items()
+            },
+            "semantic_models": {
+                k: serializable(v.to_dict(omit_none=False))
+                for k, v in getattr(manifest, "semantic_models", {}).items()
+            },
+            "saved_queries": {
+                k: serializable(v.to_dict(omit_none=False))
+                for k, v in getattr(manifest, "saved_queries", {}).items()
+            },
+        }
+
     @property
     def target(self) -> TargetConfig:
         if not self._target:
@@ -242,6 +286,9 @@ def jinja_globals(self) -> t.Dict[str, JinjaGlobalAttribute]:
         # pass user-specified default dialect if we have already loaded the config
         if self.sqlmesh_config.dialect:
             output["dialect"] = self.sqlmesh_config.dialect
+        # Pass flat graph structure like dbt
+        if self._manifest is not None:
+            output["flat_graph"] = AttributeDict(self.flat_graph)
         return output
 
     def context_for_dependencies(self, dependencies: Dependencies) -> DbtContext:

sqlmesh/utils/conversions.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import typing as t
+from datetime import date, datetime
 
 
 def ensure_bool(val: t.Any) -> bool:
@@ -19,3 +20,13 @@ def try_str_to_bool(val: str) -> t.Union[str, bool]:
         return maybe_bool == "true"
 
     return val
+
+
+def serializable(obj: t.Any) -> t.Any:
+    if isinstance(obj, (date, datetime)):
+        return obj.isoformat()
+    if isinstance(obj, dict):
+        return {k: serializable(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [serializable(item) for item in obj]
+    return obj

tests/dbt/test_transformation.py

@@ -1606,6 +1606,7 @@ def test_on_run_start_end():
     assert root_environment_statements.after_all == [
         "JINJA_STATEMENT_BEGIN;\n{{ create_tables(schemas) }}\nJINJA_END;",
         "JINJA_STATEMENT_BEGIN;\nDROP TABLE to_be_executed_last;\nJINJA_END;",
+        "JINJA_STATEMENT_BEGIN;\n{{ graph_usage() }}\nJINJA_END;",
     ]
 
     assert root_environment_statements.jinja_macros.root_package_name == "sushi"
@@ -1626,6 +1627,7 @@ def test_on_run_start_end():
         snapshots=sushi_context.snapshots,
         runtime_stage=RuntimeStage.AFTER_ALL,
         environment_naming_info=EnvironmentNamingInfo(name="dev"),
+        engine_adapter=sushi_context.engine_adapter,
     )
 
     assert rendered_before_all == [
@@ -1635,13 +1637,14 @@ def test_on_run_start_end():
     ]
 
     # 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",
-        ]
-    )
+    expected_statements = [
+        "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",
+        "CREATE OR REPLACE TABLE graph_table AS SELECT 'model.sushi.simple_model_a' AS unique_id, 'table' AS materialized UNION ALL SELECT 'model.sushi.waiters' AS unique_id, 'ephemeral' AS materialized UNION ALL SELECT 'model.sushi.simple_model_b' AS unique_id, 'table' AS materialized UNION ALL SELECT 'model.sushi.waiter_as_customer_by_day' AS unique_id, 'incremental' AS materialized UNION ALL SELECT 'model.sushi.top_waiters' AS unique_id, 'view' AS materialized UNION ALL SELECT 'model.customers.customers' AS unique_id, 'view' AS materialized UNION ALL SELECT 'model.customers.customer_revenue_by_day' AS unique_id, 'incremental' AS materialized UNION ALL SELECT 'model.sushi.waiter_revenue_by_day.v1' AS unique_id, 'incremental' AS materialized UNION ALL SELECT 'model.sushi.waiter_revenue_by_day.v2' AS unique_id, 'incremental' AS materialized",
+    ]
+
+    assert sorted(rendered_after_all) == sorted(expected_statements)
 
     # Nested dbt_packages on run start / on run end
     packaged_environment_statements = sushi_context._environment_statements[1]
@@ -1675,6 +1678,7 @@ def test_on_run_start_end():
         snapshots=sushi_context.snapshots,
         runtime_stage=RuntimeStage.AFTER_ALL,
         environment_naming_info=EnvironmentNamingInfo(name="dev"),
+        engine_adapter=sushi_context.engine_adapter,
     )
 
     # Validate order of execution to match dbt's

tests/fixtures/dbt/sushi_test/dbt_project.yml

@@ -70,4 +70,5 @@ on-run-start:
   - "{{ log_value('on-run-start') }}"
 on-run-end:
   - '{{ create_tables(schemas) }}'
-  - 'DROP TABLE to_be_executed_last;'
+  - 'DROP TABLE to_be_executed_last;'
+  - '{{ graph_usage() }}'

tests/fixtures/dbt/sushi_test/macros/graph_usage.sql

@@ -0,0 +1,18 @@
+{% macro graph_usage() %}
+{% if execute %}
+  {% set model_nodes = graph.nodes.values()
+     | selectattr("resource_type", "equalto", "model")
+     | list %}
+
+  {% set out = [] %}
+  {% for node in model_nodes %}
+    {% set line = "select '" ~ node.unique_id ~ "' as unique_id, '" ~ node.config.materialized ~ "' as materialized" %}
+    {% do out.append(line) %}
+  {% endfor %}
+
+  {% if out %}
+    {% set sql_statement = "create or replace table graph_table as\n" ~ (out | join('\nunion all\n')) %}
+    {{ return(sql_statement) }}
+  {% endif %}
+{% endif %}
+{% endmacro %}