Use a separate logger for unsafe thrift responses (databricks#153)

---------

Signed-off-by: Jesse Whitehouse <[email protected]>

Author: Jesse

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## 2.6.x (Unreleased)
 
+- Redact logged thrift responses by default
 - Add support for OAuth on Databricks Azure
 
 ## 2.6.2 (2023-06-14)

src/databricks/sql/thrift_backend.py

@@ -34,6 +34,16 @@
 
 logger = logging.getLogger(__name__)
 
+unsafe_logger = logging.getLogger("databricks.sql.unsafe")
+unsafe_logger.setLevel(logging.DEBUG)
+
+# To capture these logs in client code, add a non-NullHandler.
+# See our e2e test suite for an example with logging.FileHandler
+unsafe_logger.addHandler(logging.NullHandler())
+
+# Disable propagation so that handlers for `databricks.sql` don't pick up these messages
+unsafe_logger.propagate = False
+
 THRIFT_ERROR_MESSAGE_HEADER = "x-thriftserver-error-message"
 DATABRICKS_ERROR_OR_REDIRECT_HEADER = "x-databricks-error-or-redirect-message"
 DATABRICKS_REASON_HEADER = "x-databricks-reason-phrase"
@@ -318,13 +328,25 @@ def attempt_request(attempt):
 
             error, error_message, retry_delay = None, None, None
             try:
-                logger.debug("Sending request: {}".format(request))
+                # The MagicMocks in our unit tests have a `name` property instead of `__name__`.
+                logger.debug(
+                    "Sending request: {}(<REDACTED>)".format(
+                        getattr(
+                            method, "__name__", getattr(method, "name", "UnknownMethod")
+                        )
+                    )
+                )
+                unsafe_logger.debug("Sending request: {}".format(request))
                 response = method(request)
 
                 # Calling `close()` here releases the active HTTP connection back to the pool
                 self._transport.close()
 
-                logger.debug("Received response: {}".format(response))
+                # We need to call type(response) here because thrift doesn't implement __name__ attributes for thrift responses
+                logger.debug(
+                    "Received response: {}(<REDACTED>)".format(type(response).__name__)
+                )
+                unsafe_logger.debug("Received response: {}".format(response))
                 return response
 
             except urllib3.exceptions.HTTPError as err:

tests/e2e/driver_tests.py

@@ -28,6 +28,10 @@
 
 log = logging.getLogger(__name__)
 
+unsafe_logger = logging.getLogger("databricks.sql.unsafe")
+unsafe_logger.setLevel(logging.DEBUG)
+unsafe_logger.addHandler(logging.FileHandler("./tests-unsafe.log"))
+
 # manually decorate DecimalTestsMixin to need arrow support
 for name in loader.getTestCaseNames(DecimalTestsMixin, 'test_'):
     fn = getattr(DecimalTestsMixin, name)