[GitHub] [airflow] alex-astronomer commented on a change in pull request #21281: Mask secrets in stdout for `airflow tasks test ...` CLI command (#17476)

[GitBox <gi...@apache.org>]

alex-astronomer commented on a change in pull request #21281:
URL: https://github.com/apache/airflow/pull/21281#discussion_r798629986



##########
File path: airflow/utils/log/secrets_masker.py
##########
@@ -252,3 +253,44 @@ def add_mask(self, secret: Union[str, dict, Iterable], name: Optional[str] = Non
         elif isinstance(secret, collections.abc.Iterable):
             for v in secret:
                 self.add_mask(v, name)
+
+
+class StdoutRedactContext:
+    """
+    This class redirects all stdout within the context through the logger specified in the init.
+    Print statements get turned into log lines through the logger member.
+
+    """
+
+    def __init__(self):
+        self.stdout_lines = []
+        self._redirector = contextlib.redirect_stdout(self)
+
+    def write(self, stdout_line):
+        """Write stdout to the internal buffer when stdout is written to."""
+        self.stdout_lines.append(redact(stdout_line))
+
+    def flush(self):
+        """
+        No flush functionality needed here, but this method is needed to make this class a "file-like" which
+        can be called by the context lib in self._redirector.
+
+        """
+
+    def __enter__(self):
+        """
+        At the beginning of the context, enter the "redirect mode" which is initiated by entering the context
+        of `contextlib.redirect_stdout`.
+
+        """
+        self._redirector.__enter__()
+
+    def __exit__(self, *args):
+        """
+        On exit from the context, close the `contextlib.redirect_stdout` context to stop capturing lines
+        from stdout, then route all lines one by one through self.redirect_to.
+
+        """
+        self._redirector.__exit__(*args)
+        for log in self.stdout_lines:
+            sys.stdout.write(log)

Review comment:
       I think the problem here is that every time it writes to ‘stdout’ you’ll need to exit the internal redirect context manager because otherwise it’ll get stuck in an infinite recursive loop. I thought about doing that but since ‘__exit__’ expects a few arguments that only exist during the external context manager wrapper’s exit function. Do you have any thoughts about being able to disable and reenable the internal redirect context manager within the ‘write’ function?




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org