diff --git a/src/databricks/sql/auth/thrift_http_client.py b/src/databricks/sql/auth/thrift_http_client.py index 6273ab28..42a008b6 100644 --- a/src/databricks/sql/auth/thrift_http_client.py +++ b/src/databricks/sql/auth/thrift_http_client.py @@ -144,6 +144,7 @@ def close(self): self.__resp and self.__resp.drain_conn() self.__resp and self.__resp.release_conn() self.__resp = None + self.__pool = None def read(self, sz): return self.__resp.read(sz) diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py index 8c212c55..0fad24df 100644 --- a/src/databricks/sql/thrift_backend.py +++ b/src/databricks/sql/thrift_backend.py @@ -365,6 +365,7 @@ def attempt_request(attempt): # - non-None retry_delay -> sleep delay before retry # - error, error_message always set when available + exception_occurred = False error, error_message, retry_delay = None, None, None try: this_method_name = getattr(method, "__name__") @@ -388,6 +389,7 @@ def attempt_request(attempt): return response except urllib3.exceptions.HTTPError as err: + exception_occurred = True # retry on timeout. Happens a lot in Azure and it is safe as data has not been sent to server yet # TODO: don't use exception handling for GOS polling... @@ -406,6 +408,7 @@ def attempt_request(attempt): else: raise err except OSError as err: + exception_occurred = True error = err error_message = str(err) # fmt: off @@ -436,12 +439,14 @@ def attempt_request(attempt): else: logger.warning(log_string) except Exception as err: + exception_occurred = True error = err retry_delay = extract_retry_delay(attempt) error_message = ThriftBackend._extract_error_message_from_headers( getattr(self._transport, "headers", {}) ) - finally: + + if exception_occurred: # Calling `close()` here releases the active HTTP connection back to the pool self._transport.close()