Skip to content

Commit

Permalink
Add new invocation context class. (#52)
Browse files Browse the repository at this point in the history
* Add new invocation context class.

* Add changelog entry.

* Add unit tests for invocation context.

* Cleanup.

* Fix string constant.
  • Loading branch information
peterallenwebb authored Feb 1, 2024
1 parent 06668d1 commit caaab94
Show file tree
Hide file tree
Showing 3 changed files with 60 additions and 0 deletions.
6 changes: 6 additions & 0 deletions .changes/unreleased/Features-20240201-104437.yaml
Original file line number Diff line number Diff line change
@@ -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"
35 changes: 35 additions & 0 deletions dbt_common/context.py
Original file line number Diff line number Diff line change
@@ -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
19 changes: 19 additions & 0 deletions tests/unit/test_invocation_context.py
Original file line number Diff line number Diff line change
@@ -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"])

0 comments on commit caaab94

Please sign in to comment.