Fixed the Issue of TCP not closing

Author: jprakash-db

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)

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()