fix: fill every elementary graph entry even when the mapping already exists (dbt v2)

integration_tests/dbt_project/macros/test_init_elementary_graph.sql

@@ -0,0 +1,8 @@
+{% macro test_init_elementary_graph(existing) %}
+    {#- `existing` is what an earlier cache read can leave on the graph; none means nothing yet -#}
+    {% if existing is none %} {% do graph.pop("elementary", none) %}
+    {% else %} {% do graph.update({"elementary": existing}) %}
+    {% endif %}
+    {% do elementary.init_elementary_graph() %}
+    {{ return(graph["elementary"]) }}
+{% endmacro %}

integration_tests/tests/test_init_elementary_graph.py

@@ -0,0 +1,48 @@
+import json
+from typing import Optional
+
+import pytest
+from dbt_project import DbtProject
+
+DEFAULTS = {
+    "elementary_test_results": {},
+    "elementary_test_failed_row_counts": {},
+    "tests_schema_name": None,
+    "tables": {
+        "metrics": {"relations": [], "rows": []},
+        "schema_snapshots": [],
+    },
+    "temp_test_table_relations_map": {},
+    "duration_context_stack": {},
+    "microbatch_compiled_code_by_unique_id": {},
+}
+
+
+@pytest.mark.parametrize(
+    "existing",
+    [
+        pytest.param(None, id="no_entry_yet"),
+        pytest.param({}, id="empty_entry_left_by_an_earlier_read"),
+        pytest.param({"tests_schema_name": "kept"}, id="entry_with_a_value"),
+        pytest.param(
+            {"elementary_test_results": {"test.pkg.not_null_id": [{"status": "pass"}]}},
+            id="cache_entry_with_results",
+        ),
+    ],
+)
+def test_init_elementary_graph_fills_every_missing_entry(
+    dbt_project: DbtProject, existing: Optional[dict]
+):
+    """init_elementary_graph fills each missing entry and keeps the ones present.
+
+    `existing` is what an earlier read can leave under graph["elementary"] (None: no
+    entry yet, as on dbt v1): dbt v2
+    shares `graph` between parsing and the run, so a get_cache rendered at parse
+    time leaves an empty mapping before the on-run-start hook runs.
+    """
+    result = dbt_project.dbt_runner.run_operation(
+        "elementary_tests.test_init_elementary_graph",
+        macro_args={"existing": existing},
+    )
+    elementary_graph = json.loads(result[0])
+    assert elementary_graph == {**DEFAULTS, **(existing or {})}

macros/edr/tests/on_run_start/init_elementary_graph.sql

@@ -1,17 +1,21 @@
 {% macro init_elementary_graph() %}
-    {% do graph.setdefault(
-        "elementary",
-        {
-            "elementary_test_results": {},
-            "elementary_test_failed_row_counts": {},
-            "tests_schema_name": none,
-            "tables": {
-                "metrics": {"relations": [], "rows": []},
-                "schema_snapshots": [],
-            },
-            "temp_test_table_relations_map": {},
-            "duration_context_stack": {},
-            "microbatch_compiled_code_by_unique_id": {},
+    {% set defaults = {
+        "elementary_test_results": {},
+        "elementary_test_failed_row_counts": {},
+        "tests_schema_name": none,
+        "tables": {
+            "metrics": {"relations": [], "rows": []},
+            "schema_snapshots": [],
         },
-    ) %}
+        "temp_test_table_relations_map": {},
+        "duration_context_stack": {},
+        "microbatch_compiled_code_by_unique_id": {},
+    } %}
+    {#- Each entry is filled on its own: a cache read before this hook (dbt v2 shares
+        `graph` between parsing and the run) can leave an empty "elementary" mapping,
+        which a single setdefault on the whole mapping would keep as it is. -#}
+    {% set elementary_graph = graph.setdefault("elementary", {}) %}
+    {% for entry, default in defaults.items() %}
+        {% do elementary_graph.setdefault(entry, default) %}
+    {% endfor %}
 {% endmacro %}