From 06130600d9a40fa5df620f5fbe1e87edeb4792db Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Wed, 21 Jun 2023 15:50:25 -0500 Subject: [PATCH 1/6] Use a separate logger for unsafe thrift responses Signed-off-by: Jesse Whitehouse --- CHANGELOG.md | 1 + src/databricks/sql/thrift_backend.py | 13 +++++++++++-- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a947be50..20d8f8a7 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py index c17da877..a1d66da1 100644 --- a/src/databricks/sql/thrift_backend.py +++ b/src/databricks/sql/thrift_backend.py @@ -34,6 +34,10 @@ logger = logging.getLogger(__name__) +# Disable propagation so that handlers for `databricks.sql` don't pick up these messages +unsafe_logger = logging.getLogger("databricks.sql.unsafe") +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 +322,18 @@ def attempt_request(attempt): error, error_message, retry_delay = None, None, None try: - logger.debug("Sending request: {}".format(request)) + logger.debug("Sending request: {}()".format(method.__name__)) + 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: {}()".format(type(response).__name__) + ) + unsafe_logger.debug("Received response: {}".format(response)) return response except urllib3.exceptions.HTTPError as err: From 7ff14ce3ff759665e0bd048ecea0dbd12c687c45 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Wed, 21 Jun 2023 16:24:19 -0500 Subject: [PATCH 2/6] Set up e2e tests to record unsafe logs This gives an example of how to actual view the unsafe logs when needed. Signed-off-by: Jesse Whitehouse --- src/databricks/sql/thrift_backend.py | 8 +++++++- tests/e2e/driver_tests.py | 4 ++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py index a1d66da1..e115b4c5 100644 --- a/src/databricks/sql/thrift_backend.py +++ b/src/databricks/sql/thrift_backend.py @@ -34,8 +34,14 @@ logger = logging.getLogger(__name__) -# Disable propagation so that handlers for `databricks.sql` don't pick up these messages unsafe_logger = logging.getLogger("databricks.sql.unsafe") +unsafe_logger.setLevel(logging.DEBUG) + +# To capture these logs in client code, add a non-NullHandler for example. +# 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" diff --git a/tests/e2e/driver_tests.py b/tests/e2e/driver_tests.py index 831ed21f..f8350475 100644 --- a/tests/e2e/driver_tests.py +++ b/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) From a98a42ffc6e28e2a2e257ce1879341362d5bc05f Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Wed, 21 Jun 2023 16:32:21 -0500 Subject: [PATCH 3/6] Fix unit tests because the MagicMocks don't have a __name__ property but do have a `name` property. Signed-off-by: Jesse Whitehouse --- src/databricks/sql/thrift_backend.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py index e115b4c5..f6d8bdf1 100644 --- a/src/databricks/sql/thrift_backend.py +++ b/src/databricks/sql/thrift_backend.py @@ -328,7 +328,13 @@ def attempt_request(attempt): error, error_message, retry_delay = None, None, None try: - logger.debug("Sending request: {}()".format(method.__name__)) + logger.debug( + "Sending request: {}()".format( + getattr( + method, "__name__", getattr(method, "name", "UnknownMethod") + ) + ) + ) unsafe_logger.debug("Sending request: {}".format(request)) response = method(request) From 06deb2a87323a2bebcbcc505906340ac62484595 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Wed, 21 Jun 2023 16:40:18 -0500 Subject: [PATCH 4/6] Update comments Signed-off-by: Jesse Whitehouse --- src/databricks/sql/thrift_backend.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py index f6d8bdf1..5ec2f1bf 100644 --- a/src/databricks/sql/thrift_backend.py +++ b/src/databricks/sql/thrift_backend.py @@ -37,7 +37,7 @@ unsafe_logger = logging.getLogger("databricks.sql.unsafe") unsafe_logger.setLevel(logging.DEBUG) -# To capture these logs in client code, add a non-NullHandler for example. +# 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()) @@ -328,6 +328,7 @@ def attempt_request(attempt): error, error_message, retry_delay = None, None, None try: + # The MagicMocks in our unit tests have a `name` property instead of `__name__. logger.debug( "Sending request: {}()".format( getattr( From 053011256f37f9d2277c79278790f357bd5b2e00 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Wed, 21 Jun 2023 16:40:42 -0500 Subject: [PATCH 5/6] Typo Signed-off-by: Jesse Whitehouse --- src/databricks/sql/thrift_backend.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py index 5ec2f1bf..a95b9acc 100644 --- a/src/databricks/sql/thrift_backend.py +++ b/src/databricks/sql/thrift_backend.py @@ -328,7 +328,7 @@ def attempt_request(attempt): error, error_message, retry_delay = None, None, None try: - # The MagicMocks in our unit tests have a `name` property instead of `__name__. + # The MagicMocks in our unit tests have a `name` property instead of `__name__`. logger.debug( "Sending request: {}()".format( getattr( From 4adcd713a6c44be9f114592368eb21a437883c26 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Wed, 21 Jun 2023 16:48:26 -0500 Subject: [PATCH 6/6] Fix linting error Signed-off-by: Jesse Whitehouse --- src/databricks/sql/thrift_backend.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py index a95b9acc..b4afeaa3 100644 --- a/src/databricks/sql/thrift_backend.py +++ b/src/databricks/sql/thrift_backend.py @@ -328,7 +328,7 @@ def attempt_request(attempt): error, error_message, retry_delay = None, None, None try: - # The MagicMocks in our unit tests have a `name` property instead of `__name__`. + # The MagicMocks in our unit tests have a `name` property instead of `__name__`. logger.debug( "Sending request: {}()".format( getattr(