From caaab94b4e3f633363d88a0780372d82e9e83c59 Mon Sep 17 00:00:00 2001
From: Peter Webb <peter.webb@dbtlabs.com>
Date: Thu, 1 Feb 2024 12:28:43 -0500
Subject: [PATCH] Add new invocation context class. (#52)

* Add new invocation context class.

* Add changelog entry.

* Add unit tests for invocation context.

* Cleanup.

* Fix string constant.
---
 .../unreleased/Features-20240201-104437.yaml  |  6 ++++
 dbt_common/context.py                         | 35 +++++++++++++++++++
 tests/unit/test_invocation_context.py         | 19 ++++++++++
 3 files changed, 60 insertions(+)
 create mode 100644 .changes/unreleased/Features-20240201-104437.yaml
 create mode 100644 dbt_common/context.py
 create mode 100644 tests/unit/test_invocation_context.py

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"])