Add metadata env method to provider context

.changes/unreleased/Features-20220909-204643.yaml

@@ -0,0 +1,7 @@
+kind: Features
+body: Add metadata env method to ProviderContext class
+time: 2022-09-09T20:46:43.889302+01:00
+custom:
+  Author: jared-rimmer
+  Issue: "5522"
+  PR: "5794"

core/dbt/context/providers.py

@@ -41,6 +41,7 @@
     ParsedSourceDefinition,
 )
 from dbt.contracts.graph.metrics import MetricReference, ResolvedMetricReference
+from dbt.contracts.util import get_metadata_env
 from dbt.exceptions import (
     CompilationException,
     ParsingException,
@@ -710,6 +711,10 @@ def _get_namespace_builder(self):
             self.model,
         )
 
+    @contextproperty
+    def dbt_metadata_envs(self) -> Dict[str, str]:
+        return get_metadata_env()
+
     @contextproperty
     def invocation_args_dict(self):
         return args_to_dict(self.config.args)

test/unit/test_context.py

@@ -200,7 +200,7 @@ def assert_has_keys(required_keys: Set[str], maybe_keys: Set[str], ctx: Dict[str
         "modules",
         "flags",
         "print",
-        "diff_of_two_dicts"
+        "diff_of_two_dicts",
     }
 )
 
@@ -239,6 +239,7 @@ def assert_has_keys(required_keys: Set[str], maybe_keys: Set[str], ctx: Dict[str
     "adapter_macro",
     "selected_resources",
     "invocation_args_dict",
+    "dbt_metadata_envs"
 }
 REQUIRED_MODEL_KEYS = REQUIRED_MACRO_KEYS | {"this", "compiled_code",}
 MAYBE_KEYS = frozenset({"debug"})
@@ -499,3 +500,21 @@ def test_macro_namespace(config_postgres, manifest_fx):
         assert result["dbt"]["some_macro"].macro is pg_macro
         assert result["root"]["some_macro"].macro is package_macro
         assert result["some_macro"].macro is package_macro
+
+def test_dbt_metadata_envs(monkeypatch, config_postgres, manifest_fx, get_adapter, get_include_paths):
+    envs = {
+        "DBT_ENV_CUSTOM_ENV_RUN_ID": 1234,
+        "DBT_ENV_CUSTOM_ENV_JOB_ID": 5678,
+        "DBT_ENV_RUN_ID": 91011,
+        "RANDOM_ENV": 121314
+    }
+    monkeypatch.setattr(os, 'environ', envs)
+
+    ctx = providers.generate_runtime_macro_context(
+        macro=manifest_fx.macros["macro.root.macro_a"],
+        config=config_postgres,
+        manifest=manifest_fx,
+        package_name="root",
+    ) 
+
+    assert ctx["dbt_metadata_envs"] == {'JOB_ID': 5678, 'RUN_ID': 1234}

tests/functional/context_methods/test_builtin_functions.py

@@ -1,5 +1,6 @@
 import pytest
 import json
+import os
 
 from dbt.tests.util import run_dbt, run_dbt_and_capture, write_file
 from dbt.exceptions import CompilationException
@@ -30,6 +31,12 @@
 {% endmacro %}
 """
 
+macros__validate_dbt_metadata_envs_sql = """
+{% macro validate_dbt_metadata_envs() %}
+    {{ log("dbt_metadata_envs_result:"~ dbt_metadata_envs) }}
+{% endmacro %}
+"""
+
 models__set_exception_sql = """
 {% set set_strict_result = set_strict(1) %}
 """
@@ -39,13 +46,34 @@
 """
 
 
+def parse_json_logs(json_log_output):
+    parsed_logs = []
+    for line in json_log_output.split("\n"):
+        try:
+            log = json.loads(line)
+        except ValueError:
+            continue
+
+        parsed_logs.append(log)
+
+    return parsed_logs
+
+
+def find_result_in_parsed_logs(parsed_logs, result_name):
+    return next(
+        (item for item in parsed_logs if result_name in item["data"].get("msg", "msg")),
+        False,
+    )
+
+
 class TestContextBuiltins:
     @pytest.fixture(scope="class")
     def macros(self):
         return {
             "validate_set.sql": macros__validate_set_sql,
             "validate_zip.sql": macros__validate_zip_sql,
             "validate_invocation.sql": macros__validate_invocation_sql,
+            "validate_dbt_metadata_envs.sql": macros__validate_dbt_metadata_envs_sql,
         }
 
     def test_builtin_set_function(self, project):
@@ -75,24 +103,8 @@ def test_builtin_invocation_args_dict_function(self, project):
             ]
         )
 
-        parsed_logs = []
-        for line in log_output.split("\n"):
-            try:
-                log = json.loads(line)
-            except ValueError:
-                continue
-
-            parsed_logs.append(log)
-
-        # Value of msg is a string
-        result = next(
-            (
-                item
-                for item in parsed_logs
-                if "invocation_result" in item["data"].get("msg", "msg")
-            ),
-            False,
-        )
+        parsed_logs = parse_json_logs(log_output)
+        result = find_result_in_parsed_logs(parsed_logs, "invocation_result")
 
         assert result
 
@@ -103,6 +115,27 @@ def test_builtin_invocation_args_dict_function(self, project):
         expected = "'send_anonymous_usage_stats': False, 'event_buffer_size': 100000, 'quiet': False, 'no_print': False, 'macro': 'validate_invocation', 'args': '{my_variable: test_variable}', 'which': 'run-operation', 'rpc_method': 'run-operation', 'indirect_selection': 'eager'}"
         assert expected in str(result)
 
+    def test_builtin_dbt_metadata_envs_function(self, project, monkeypatch):
+        envs = {
+            "DBT_ENV_CUSTOM_ENV_RUN_ID": 1234,
+            "DBT_ENV_CUSTOM_ENV_JOB_ID": 5678,
+            "DBT_ENV_RUN_ID": 91011,
+            "RANDOM_ENV": 121314,
+        }
+        monkeypatch.setattr(os, "environ", envs)
+
+        _, log_output = run_dbt_and_capture(
+            ["--debug", "--log-format=json", "run-operation", "validate_dbt_metadata_envs"]
+        )
+
+        parsed_logs = parse_json_logs(log_output)
+        result = find_result_in_parsed_logs(parsed_logs, "dbt_metadata_envs_result")
+
+        assert result
+
+        expected = "dbt_metadata_envs_result:{'RUN_ID': 1234, 'JOB_ID': 5678}"
+        assert expected in str(result)
+
 
 class TestContextBuiltinExceptions:
     # Assert compilation errors are raised with _strict equivalents