generated from dbt-labs/dbt-oss-template
-
Notifications
You must be signed in to change notification settings - Fork 15
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add new invocation context class. (#52)
* Add new invocation context class. * Add changelog entry. * Add unit tests for invocation context. * Cleanup. * Fix string constant.
- Loading branch information
1 parent
06668d1
commit caaab94
Showing
3 changed files
with
60 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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"]) |