Skip to content

Commit

Permalink
Fixed the Issue of TCP not closing
Browse files Browse the repository at this point in the history
  • Loading branch information
jprakash-db committed Dec 15, 2024
1 parent 680b3b6 commit a477dc8
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 1 deletion.
1 change: 1 addition & 0 deletions src/databricks/sql/auth/thrift_http_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
7 changes: 6 additions & 1 deletion src/databricks/sql/thrift_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__")
Expand All @@ -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...
Expand All @@ -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
Expand Down Expand Up @@ -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()

Expand Down

0 comments on commit a477dc8

Please sign in to comment.