From 8e738cd0ad0e7dce644f66bb749a7b46770badee Mon Sep 17 00:00:00 2001
From: Maciej Obuchowski <obuchowski.maciej@gmail.com>
Date: Tue, 15 Aug 2023 11:10:03 +0200
Subject: [PATCH] openlineage: do not try to redact Proxy objects from
 deprecated config (#33393)

Signed-off-by: Maciej Obuchowski <obuchowski.maciej@gmail.com>
---
 airflow/providers/openlineage/utils/utils.py  | 60 +++++++++++--------
 .../openlineage/plugins/test_utils.py         |  6 ++
 2 files changed, 42 insertions(+), 24 deletions(-)

diff --git a/airflow/providers/openlineage/utils/utils.py b/airflow/providers/openlineage/utils/utils.py
index 20b9afef49294..3ff7610c4f1e7 100644
--- a/airflow/providers/openlineage/utils/utils.py
+++ b/airflow/providers/openlineage/utils/utils.py
@@ -38,6 +38,7 @@
     AirflowMappedTaskRunFacet,
     AirflowRunFacet,
 )
+from airflow.utils.context import AirflowContextDeprecationWarning
 from airflow.utils.log.secrets_masker import Redactable, Redacted, SecretsMasker, should_hide_value_for_key
 
 if TYPE_CHECKING:
@@ -345,29 +346,40 @@ def _redact(self, item: Redactable, name: str | None, depth: int, max_depth: int
         if depth > max_depth:
             return item
         try:
-            if name and should_hide_value_for_key(name):
-                return self._redact_all(item, depth, max_depth)
-            if attrs.has(type(item)):
-                # TODO: fixme when mypy gets compatible with new attrs
-                for dict_key, subval in attrs.asdict(item, recurse=False).items():  # type: ignore[arg-type]
-                    if _is_name_redactable(dict_key, item):
-                        setattr(
-                            item,
-                            dict_key,
-                            self._redact(subval, name=dict_key, depth=(depth + 1), max_depth=max_depth),
-                        )
-                return item
-            elif is_json_serializable(item) and hasattr(item, "__dict__"):
-                for dict_key, subval in item.__dict__.items():
-                    if _is_name_redactable(dict_key, item):
-                        setattr(
-                            item,
-                            dict_key,
-                            self._redact(subval, name=dict_key, depth=(depth + 1), max_depth=max_depth),
-                        )
-                return item
-            else:
-                return super()._redact(item, name, depth, max_depth)
+            # It's impossible to check the type of variable in a dict without accessing it, and
+            # this already causes warning - so suppress it
+            with suppress(AirflowContextDeprecationWarning):
+                if type(item).__name__ == "Proxy":
+                    # Those are deprecated values in _DEPRECATION_REPLACEMENTS
+                    # in airflow.utils.context.Context
+                    return "<<non-redactable: Proxy>>"
+                if name and should_hide_value_for_key(name):
+                    return self._redact_all(item, depth, max_depth)
+                if attrs.has(type(item)):
+                    # TODO: fixme when mypy gets compatible with new attrs
+                    for dict_key, subval in attrs.asdict(
+                        item, recurse=False  # type: ignore[arg-type]
+                    ).items():
+                        if _is_name_redactable(dict_key, item):
+                            setattr(
+                                item,
+                                dict_key,
+                                self._redact(subval, name=dict_key, depth=(depth + 1), max_depth=max_depth),
+                            )
+                    return item
+                elif is_json_serializable(item) and hasattr(item, "__dict__"):
+                    for dict_key, subval in item.__dict__.items():
+                        if type(subval).__name__ == "Proxy":
+                            return "<<non-redactable: Proxy>>"
+                        if _is_name_redactable(dict_key, item):
+                            setattr(
+                                item,
+                                dict_key,
+                                self._redact(subval, name=dict_key, depth=(depth + 1), max_depth=max_depth),
+                            )
+                    return item
+                else:
+                    return super()._redact(item, name, depth, max_depth)
         except Exception as e:
             log.warning(
                 "Unable to redact %s. Error was: %s: %s",
@@ -375,7 +387,7 @@ def _redact(self, item: Redactable, name: str | None, depth: int, max_depth: int
                 type(e).__name__,
                 str(e),
             )
-            return item
+        return item
 
 
 def is_json_serializable(item):
diff --git a/tests/providers/openlineage/plugins/test_utils.py b/tests/providers/openlineage/plugins/test_utils.py
index afbc40ee4c692..59abb43e094b7 100644
--- a/tests/providers/openlineage/plugins/test_utils.py
+++ b/tests/providers/openlineage/plugins/test_utils.py
@@ -171,6 +171,9 @@ class NotMixin:
         def __init__(self):
             self.password = "passwd"
 
+    class Proxy:
+        pass
+
     def default(self, o):
         if isinstance(o, NotMixin):
             return o.__dict__
@@ -180,6 +183,9 @@ def default(self, o):
     monkeypatch.setattr(JSONEncoder, "default", default)
     assert redactor.redact(NotMixin()).password == "***"
 
+    assert redactor.redact(Proxy()) == "<<non-redactable: Proxy>>"
+    assert redactor.redact({"a": "a", "b": Proxy()}) == {"a": "a", "b": "<<non-redactable: Proxy>>"}
+
     class Mixined(RedactMixin):
         _skip_redact = ["password"]