From 9a1f20ae1d898a7faf9d6dcaaed1800c494e939a Mon Sep 17 00:00:00 2001
From: Gerda Shank <gerda.shank@gmail.com>
Date: Thu, 22 Feb 2024 09:45:23 -0500
Subject: [PATCH] Fix unset_contextvars to correctly unset the context vars
 (#83)

* Fix unset_contextvars to actually unset the right vars

* Changie

* Add test
---
 .../unreleased/Fixes-20240221-200204.yaml     |  6 ++++++
 dbt_common/events/contextvars.py              | 18 ++++++++--------
 tests/unit/test_contextvars.py                | 21 +++++++++++++++++++
 3 files changed, 36 insertions(+), 9 deletions(-)
 create mode 100644 .changes/unreleased/Fixes-20240221-200204.yaml
 create mode 100644 tests/unit/test_contextvars.py

diff --git a/.changes/unreleased/Fixes-20240221-200204.yaml b/.changes/unreleased/Fixes-20240221-200204.yaml
new file mode 100644
index 00000000..8df4eb14
--- /dev/null
+++ b/.changes/unreleased/Fixes-20240221-200204.yaml
@@ -0,0 +1,6 @@
+kind: Fixes
+body: Fix the unset_contextvars function in events/contextvars.py
+time: 2024-02-21T20:02:04.341681-05:00
+custom:
+  Author: gshank
+  Issue: "82"
diff --git a/dbt_common/events/contextvars.py b/dbt_common/events/contextvars.py
index 5bdb78fe..1508cdcc 100644
--- a/dbt_common/events/contextvars.py
+++ b/dbt_common/events/contextvars.py
@@ -58,12 +58,12 @@ def set_task_contextvars(**kwargs: Any) -> Mapping[str, contextvars.Token]:
 def set_contextvars(prefix: str, **kwargs: Any) -> Mapping[str, contextvars.Token]:
     cvar_tokens = {}
     for k, v in kwargs.items():
-        log_key = f"{prefix}{k}"
+        prefix_key = f"{prefix}{k}"
         try:
-            var = _context_vars[log_key]
+            var = _context_vars[prefix_key]
         except KeyError:
-            var = contextvars.ContextVar(log_key, default=Ellipsis)
-            _context_vars[log_key] = var
+            var = contextvars.ContextVar(prefix_key, default=Ellipsis)
+            _context_vars[prefix_key] = var
 
         cvar_tokens[k] = var.set(v)
 
@@ -73,17 +73,17 @@ def set_contextvars(prefix: str, **kwargs: Any) -> Mapping[str, contextvars.Toke
 # reset by Tokens
 def reset_contextvars(prefix: str, **kwargs: contextvars.Token) -> None:
     for k, v in kwargs.items():
-        log_key = f"{prefix}{k}"
-        var = _context_vars[log_key]
+        prefix_key = f"{prefix}{k}"
+        var = _context_vars[prefix_key]
         var.reset(v)
 
 
 # remove from contextvars
 def unset_contextvars(prefix: str, *keys: str) -> None:
     for k in keys:
-        if k in _context_vars:
-            log_key = f"{prefix}{k}"
-            _context_vars[log_key].set(Ellipsis)
+        prefix_key = f"{prefix}{k}"
+        if prefix_key in _context_vars:
+            _context_vars[prefix_key].set(Ellipsis)
 
 
 # Context manager or decorator to set and unset the context vars
diff --git a/tests/unit/test_contextvars.py b/tests/unit/test_contextvars.py
new file mode 100644
index 00000000..4eb58e6c
--- /dev/null
+++ b/tests/unit/test_contextvars.py
@@ -0,0 +1,21 @@
+from dbt_common.events.contextvars import log_contextvars, get_node_info, set_log_contextvars
+
+
+def test_contextvars():
+    node_info = {
+        "unique_id": "model.test.my_model",
+        "started_at": None,
+        "finished_at": None,
+    }
+    with log_contextvars(node_info=node_info):
+        assert node_info == get_node_info()
+        new_node_info = {
+            "unique_id": "model.test.my_model",
+            "started_at": "01-01-2024",
+            "finished_at": None,
+        }
+        set_log_contextvars(node_info=new_node_info)
+        assert get_node_info() == new_node_info
+
+    # Ensure that after the context manager ends, the node_info is gone
+    assert get_node_info() == {}