Fix unset_contextvars to correctly unset the context vars (#83)

* Fix unset_contextvars to actually unset the right vars

* Changie

* Add test

.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"

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

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() == {}