From a477dc8e1d3e7d8122b1aa9cf981a007777c7628 Mon Sep 17 00:00:00 2001
From: Jothi Prakash <jothi.prakash@databricks.com>
Date: Sun, 15 Dec 2024 15:57:10 +0530
Subject: [PATCH] Fixed the Issue of TCP not closing

---
 src/databricks/sql/auth/thrift_http_client.py | 1 +
 src/databricks/sql/thrift_backend.py          | 7 ++++++-
 2 files changed, 7 insertions(+), 1 deletion(-)

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