diff --git a/.changes/unreleased/Features-20220909-204643.yaml b/.changes/unreleased/Features-20220909-204643.yaml new file mode 100644 index 00000000000..b1f2ada6d20 --- /dev/null +++ b/.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" diff --git a/core/dbt/context/providers.py b/core/dbt/context/providers.py index 0db9b0c07d0..fef7bd66faf 100644 --- a/core/dbt/context/providers.py +++ b/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) diff --git a/test/unit/test_context.py b/test/unit/test_context.py index 2aeba229bf4..534a999ef8e 100644 --- a/test/unit/test_context.py +++ b/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} diff --git a/tests/functional/context_methods/test_builtin_functions.py b/tests/functional/context_methods/test_builtin_functions.py index e5edda50d89..55524178005 100644 --- a/tests/functional/context_methods/test_builtin_functions.py +++ b/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,6 +46,26 @@ """ +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): @@ -46,6 +73,7 @@ def macros(self): "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