From a10f4cac6c3b89e54ec01b7251b6997d58f44edb Mon Sep 17 00:00:00 2001 From: Josiah Lee Date: Tue, 20 Feb 2024 12:13:14 -0800 Subject: [PATCH] fix(weave): Make traces work again and still remove logging (#1222) * Make traces work again and still remove logging * lint * add helper * lint --- weave/engine_trace.py | 13 +++++++------ weave/environment.py | 4 ++++ weave/weave_server.py | 9 +++++++-- 3 files changed, 18 insertions(+), 8 deletions(-) diff --git a/weave/engine_trace.py b/weave/engine_trace.py index d06988257211..d1048d6eb1d0 100644 --- a/weave/engine_trace.py +++ b/weave/engine_trace.py @@ -18,6 +18,7 @@ import json import dataclasses +from . import environment from . import logs from . import stream_data_interfaces @@ -155,7 +156,7 @@ def attributes(self): def set_tag(self, key, unredacted_val, redacted_val=None): if "tags" not in self.attributes: self.attributes["tags"] = {} - if not os.getenv("DISABLE_WEAVE_PII"): + if not environment.disable_weave_pii(): self.attributes["tags"][key] = unredacted_val elif redacted_val is not None: self.attributes["tags"][key] = redacted_val @@ -163,7 +164,7 @@ def set_tag(self, key, unredacted_val, redacted_val=None): def set_meta(self, key, unredacted_val, redacted_val=None): if "metadata" not in self.attributes: self.attributes["metadata"] = {} - if not os.getenv("DISABLE_WEAVE_PII"): + if not environment.disable_weave_pii(): self.attributes["metadata"][key] = unredacted_val elif redacted_val is not None: self.attributes["metadata"][key] = redacted_val @@ -171,7 +172,7 @@ def set_meta(self, key, unredacted_val, redacted_val=None): def set_metric(self, key, unredacted_val, redacted_val=None): if "metrics" not in self.attributes: self.attributes["metrics"] = {} - if not os.getenv("DISABLE_WEAVE_PII"): + if not environment.disable_weave_pii(): self.attributes["metrics"][key] = unredacted_val elif redacted_val is not None: self.attributes["metrics"][key] = redacted_val @@ -285,18 +286,18 @@ def patch_ddtrace_set_tag(): # Only logged redacted values if flag is on def set_tag(self, key, unredacted_val=None, redacted_val=None): - if redacted_val is not None and os.getenv("DISABLE_WEAVE_PII"): + if redacted_val is not None and environment.disable_weave_pii(): old_set_tag(self, key, redacted_val) elif ( unredacted_val is not None - and not os.getenv("DISABLE_WEAVE_PII") + and not environment.disable_weave_pii() or "_dd." in key ): old_set_tag(self, key, unredacted_val) # Log metric if flag is off or if flag is on and redacted def set_metric(self, key, val, is_pii_redacted=False): - if not os.getenv("DISABLE_WEAVE_PII") or is_pii_redacted or "_dd." in key: + if not environment.disable_weave_pii() or is_pii_redacted or "_dd." in key: old_set_metric(self, key, val) ddtrace_span.Span.set_metric = set_metric diff --git a/weave/environment.py b/weave/environment.py index f7a18c46ddc5..26481c213027 100644 --- a/weave/environment.py +++ b/weave/environment.py @@ -259,3 +259,7 @@ def dd_env() -> str: def env_is_ci() -> bool: return util.parse_boolean_env_var("WEAVE_CI") + + +def disable_weave_pii() -> bool: + return os.getenv("DISABLE_WEAVE_PII", "false").strip().lower() == "true" diff --git a/weave/weave_server.py b/weave/weave_server.py index 9001e9efe481..92b5fdbbcbb2 100644 --- a/weave/weave_server.py +++ b/weave/weave_server.py @@ -71,8 +71,13 @@ def execute(self, document, *args, **kwargs): # crashes import ddtrace - if not os.environ.get("DISABLE_WEAVE_PII"): - ddtrace.patch_all(logging=True) + # patch all libraries to datadog except logging + ddtrace.patch_all(logging=False) + + if not environment.disable_weave_pii(): + # patch logging to datadog, if not disabled + ddtrace.patch(logging=True) + custom_dd_patch()