diff --git a/.changes/unreleased/Features-20240201-104437.yaml b/.changes/unreleased/Features-20240201-104437.yaml new file mode 100644 index 00000000..ff251faf --- /dev/null +++ b/.changes/unreleased/Features-20240201-104437.yaml @@ -0,0 +1,6 @@ +kind: Features +body: Added a new InvocationContext class for application context management. +time: 2024-02-01T10:44:37.161298-05:00 +custom: + Author: peterallenwebb + Issue: "57" diff --git a/dbt_common/context.py b/dbt_common/context.py new file mode 100644 index 00000000..07434432 --- /dev/null +++ b/dbt_common/context.py @@ -0,0 +1,35 @@ +from contextvars import ContextVar +from typing import List, Mapping, Optional + +from dbt_common.constants import SECRET_ENV_PREFIX + + +class InvocationContext: + def __init__(self, env: Mapping[str, str]): + self._env = env + self._env_secrets: Optional[List[str]] = None + # This class will also eventually manage the invocation_id, flags, event manager, etc. + + @property + def env(self) -> Mapping[str, str]: + return self._env + + @property + def env_secrets(self) -> List[str]: + if self._env_secrets is None: + self._env_secrets = [ + v for k, v in self.env.items() if k.startswith(SECRET_ENV_PREFIX) and v.strip() + ] + return self._env_secrets + + +_INVOCATION_CONTEXT_VAR: ContextVar[InvocationContext] = ContextVar("DBT_INVOCATION_CONTEXT_VAR") + + +def set_invocation_context(env: Mapping[str, str]) -> None: + _INVOCATION_CONTEXT_VAR.set(InvocationContext(env)) + + +def get_invocation_context() -> InvocationContext: + ctx = _INVOCATION_CONTEXT_VAR.get() + return ctx diff --git a/tests/unit/test_invocation_context.py b/tests/unit/test_invocation_context.py new file mode 100644 index 00000000..fa63cbc4 --- /dev/null +++ b/tests/unit/test_invocation_context.py @@ -0,0 +1,19 @@ +from dbt_common.constants import SECRET_ENV_PREFIX +from dbt_common.context import InvocationContext + + +def test_invocation_context_env(): + test_env = {"VAR_1": "value1", "VAR_2": "value2"} + ic = InvocationContext(env=test_env) + assert ic.env == test_env + + +def test_invocation_context_secrets(): + test_env = { + f"{SECRET_ENV_PREFIX}_VAR_1": "secret1", + f"{SECRET_ENV_PREFIX}VAR_2": "secret2", + "NON_SECRET": "non-secret", + f"foo{SECRET_ENV_PREFIX}": "non-secret", + } + ic = InvocationContext(env=test_env) + assert set(ic.env_secrets) == set(["secret1", "secret2"])