From 0a559eecae07965a7b2a650ef2f97f97750db8d7 Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Wed, 19 Feb 2025 14:27:22 +0530 Subject: [PATCH 1/7] Added some logging support --- src/databricks/sql/auth/retry.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/databricks/sql/auth/retry.py b/src/databricks/sql/auth/retry.py index eedcc773f..2d04c56d3 100755 --- a/src/databricks/sql/auth/retry.py +++ b/src/databricks/sql/auth/retry.py @@ -345,6 +345,8 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]: if a retry would violate the configured policy. """ + logger.info(f"Received status code {status_code} for {method} request") + # Request succeeded. Don't retry. if status_code == 200: return False, "200 codes are not retried" From 9944f78023e3d6f17e4fd5a096e346f2b501c533 Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Thu, 20 Feb 2025 12:47:25 +0530 Subject: [PATCH 2/7] Logs are based on what the server returns --- src/databricks/sql/auth/retry.py | 26 ++++++++++++++------------ 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/src/databricks/sql/auth/retry.py b/src/databricks/sql/auth/retry.py index 2d04c56d3..1bacbbdbe 100755 --- a/src/databricks/sql/auth/retry.py +++ b/src/databricks/sql/auth/retry.py @@ -352,18 +352,17 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]: return False, "200 codes are not retried" if status_code == 401: - raise NonRecoverableNetworkError( - "Received 401 - UNAUTHORIZED. Confirm your authentication credentials." + return ( + False, + "Received 401 - UNAUTHORIZED. Confirm your authentication credentials.", ) if status_code == 403: - raise NonRecoverableNetworkError( - "Received 403 - FORBIDDEN. Confirm your authentication credentials." - ) + return False, "403 codes are not retried" # Request failed and server said NotImplemented. This isn't recoverable. Don't retry. if status_code == 501: - raise NonRecoverableNetworkError("Received code 501 from server.") + return False, "Received code 501 from server." # Request failed and this method is not retryable. We only retry POST requests. if not self._is_method_retryable(method): @@ -382,8 +381,9 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]: and self.command_type == CommandType.CLOSE_SESSION and len(self.history) > 0 ): - raise SessionAlreadyClosedError( - "CloseSession received 404 code from Databricks. Session is already closed." + return ( + False, + "CloseSession received 404 code from Databricks. Session is already closed.", ) # Request failed with 404 because CloseOperation returns 404 if you repeat the request. @@ -392,8 +392,9 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]: and self.command_type == CommandType.CLOSE_OPERATION and len(self.history) > 0 ): - raise CursorAlreadyClosedError( - "CloseOperation received 404 code from Databricks. Cursor is already closed." + return ( + False, + "CloseOperation received 404 code from Databricks. Cursor is already closed.", ) # Request failed, was an ExecuteStatement and the command may have reached the server @@ -402,8 +403,9 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]: and status_code not in self.status_forcelist and status_code not in self.force_dangerous_codes ): - raise UnsafeToRetryError( - "ExecuteStatement command can only be retried for codes 429 and 503" + return ( + False, + "ExecuteStatement command can only be retried for codes 429 and 503", ) # Request failed with a dangerous code, was an ExecuteStatement, but user forced retries for this From 9ffaee4d37ae6ea902a96403bf1d370c8d249055 Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Fri, 21 Feb 2025 23:00:43 +0530 Subject: [PATCH 3/7] Added back some of the better errors --- src/databricks/sql/auth/retry.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/databricks/sql/auth/retry.py b/src/databricks/sql/auth/retry.py index 1bacbbdbe..57cfeed58 100755 --- a/src/databricks/sql/auth/retry.py +++ b/src/databricks/sql/auth/retry.py @@ -381,9 +381,8 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]: and self.command_type == CommandType.CLOSE_SESSION and len(self.history) > 0 ): - return ( - False, - "CloseSession received 404 code from Databricks. Session is already closed.", + raise SessionAlreadyClosedError( + "CloseSession received 404 code from Databricks. Session is already closed." ) # Request failed with 404 because CloseOperation returns 404 if you repeat the request. @@ -392,9 +391,8 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]: and self.command_type == CommandType.CLOSE_OPERATION and len(self.history) > 0 ): - return ( - False, - "CloseOperation received 404 code from Databricks. Cursor is already closed.", + raise CursorAlreadyClosedError( + "CloseOperation received 404 code from Databricks. Cursor is already closed." ) # Request failed, was an ExecuteStatement and the command may have reached the server From 547d55b4153ff82a5e7bd00ec9d2908f49347937 Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Sat, 22 Feb 2025 00:25:02 +0530 Subject: [PATCH 4/7] Print status if not detected as error --- 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 9beab0371..969ec44c3 100644 --- a/src/databricks/sql/thrift_backend.py +++ b/src/databricks/sql/thrift_backend.py @@ -382,7 +382,7 @@ def attempt_request(attempt): # 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__) + "Received response: {}() with status {}".format(type(response).__name__, getattr(response, "status", None)) ) unsafe_logger.debug("Received response: {}".format(response)) return response From f305198dc50ed2897b1fb2a7552a78aad2869ef6 Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Sat, 22 Feb 2025 11:45:10 +0530 Subject: [PATCH 5/7] Added logging in the flush step --- src/databricks/sql/auth/thrift_http_client.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/auth/thrift_http_client.py b/src/databricks/sql/auth/thrift_http_client.py index 6273ab284..e1f894b06 100644 --- a/src/databricks/sql/auth/thrift_http_client.py +++ b/src/databricks/sql/auth/thrift_http_client.py @@ -192,12 +192,14 @@ def flush(self): timeout=self.__timeout, retries=self.retry_policy, ) - + # Get reply to flush the request self.code = self.__resp.status self.message = self.__resp.reason self.headers = self.__resp.headers + logger.info("HTTP Response with status code {}, message: {}".format(self.code, self.message)) + @staticmethod def basic_proxy_auth_headers(proxy): if proxy is None or not proxy.username: From c4aad86433ff19497b6b5ee94914e7248c9e5608 Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Sat, 22 Feb 2025 11:46:13 +0530 Subject: [PATCH 6/7] Reformatted --- src/databricks/sql/auth/thrift_http_client.py | 8 ++++++-- src/databricks/sql/thrift_backend.py | 4 +++- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/src/databricks/sql/auth/thrift_http_client.py b/src/databricks/sql/auth/thrift_http_client.py index e1f894b06..f0daae162 100644 --- a/src/databricks/sql/auth/thrift_http_client.py +++ b/src/databricks/sql/auth/thrift_http_client.py @@ -192,13 +192,17 @@ def flush(self): timeout=self.__timeout, retries=self.retry_policy, ) - + # Get reply to flush the request self.code = self.__resp.status self.message = self.__resp.reason self.headers = self.__resp.headers - logger.info("HTTP Response with status code {}, message: {}".format(self.code, self.message)) + logger.info( + "HTTP Response with status code {}, message: {}".format( + self.code, self.message + ) + ) @staticmethod def basic_proxy_auth_headers(proxy): diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py index 969ec44c3..0505b0bac 100644 --- a/src/databricks/sql/thrift_backend.py +++ b/src/databricks/sql/thrift_backend.py @@ -382,7 +382,9 @@ def attempt_request(attempt): # We need to call type(response) here because thrift doesn't implement __name__ attributes for thrift responses logger.debug( - "Received response: {}() with status {}".format(type(response).__name__, getattr(response, "status", None)) + "Received response: {}() with status {}".format( + type(response).__name__, getattr(response, "status", None) + ) ) unsafe_logger.debug("Received response: {}".format(response)) return response From 17838eed3b425b5185ad5aabb90d05573cb7f3ca Mon Sep 17 00:00:00 2001 From: Jothi Prakash Date: Sat, 22 Feb 2025 18:03:27 +0530 Subject: [PATCH 7/7] Removed some unnecessary --- src/databricks/sql/thrift_backend.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py index 0505b0bac..9beab0371 100644 --- a/src/databricks/sql/thrift_backend.py +++ b/src/databricks/sql/thrift_backend.py @@ -382,9 +382,7 @@ def attempt_request(attempt): # We need to call type(response) here because thrift doesn't implement __name__ attributes for thrift responses logger.debug( - "Received response: {}() with status {}".format( - type(response).__name__, getattr(response, "status", None) - ) + "Received response: {}()".format(type(response).__name__) ) unsafe_logger.debug("Received response: {}".format(response)) return response