From 2785d0aef287fc646a8a270cff3272b6074db588 Mon Sep 17 00:00:00 2001
From: Jothi Prakash <jothi.prakash@databricks.com>
Date: Wed, 11 Dec 2024 11:48:17 +0530
Subject: [PATCH 01/12] Raised error when incorrect Row offset it returned

---
 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 8c212c554..e433e878a 100644
--- a/src/databricks/sql/thrift_backend.py
+++ b/src/databricks/sql/thrift_backend.py
@@ -1046,7 +1046,7 @@ def fetch_results(
 
         resp = self.make_request(self._client.FetchResults, req)
         if resp.results.startRowOffset > expected_row_start_offset:
-            logger.warning(
+            raise ValueError(
                 "Expected results to start from {} but they instead start at {}".format(
                     expected_row_start_offset, resp.results.startRowOffset
                 )

From 1f3bc1937bfe9f5fbb9f6e96b992d8efc9377c4d Mon Sep 17 00:00:00 2001
From: Jothi Prakash <jothi.prakash@databricks.com>
Date: Fri, 13 Dec 2024 00:03:08 +0530
Subject: [PATCH 02/12] Changed error type

---
 src/databricks/sql/thrift_backend.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py
index e433e878a..281d0468c 100644
--- a/src/databricks/sql/thrift_backend.py
+++ b/src/databricks/sql/thrift_backend.py
@@ -1046,8 +1046,8 @@ def fetch_results(
 
         resp = self.make_request(self._client.FetchResults, req)
         if resp.results.startRowOffset > expected_row_start_offset:
-            raise ValueError(
-                "Expected results to start from {} but they instead start at {}".format(
+            raise DataError(
+                "fetch_results failed due to inconsistency in the state between the client and the server. Expected results to start from {} but they instead start at {}, some result batch must have been skipped".format(
                     expected_row_start_offset, resp.results.startRowOffset
                 )
             )

From 8d20b79d665f47bc6067f62512d2fc709c440405 Mon Sep 17 00:00:00 2001
From: Jothi Prakash <jothi.prakash@databricks.com>
Date: Fri, 13 Dec 2024 00:04:41 +0530
Subject: [PATCH 03/12] grammar fix

---
 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 281d0468c..141d9e1f6 100644
--- a/src/databricks/sql/thrift_backend.py
+++ b/src/databricks/sql/thrift_backend.py
@@ -1047,7 +1047,7 @@ def fetch_results(
         resp = self.make_request(self._client.FetchResults, req)
         if resp.results.startRowOffset > expected_row_start_offset:
             raise DataError(
-                "fetch_results failed due to inconsistency in the state between the client and the server. Expected results to start from {} but they instead start at {}, some result batch must have been skipped".format(
+                "fetch_results failed due to inconsistency in the state between the client and the server. Expected results to start from {} but they instead start at {}, some result batches must have been skipped".format(
                     expected_row_start_offset, resp.results.startRowOffset
                 )
             )

From 887a9e0510b523f286f1cf65e7f2664c4a26085d Mon Sep 17 00:00:00 2001
From: Jothi Prakash <jothi.prakash@databricks.com>
Date: Sun, 15 Dec 2024 13:20:36 +0530
Subject: [PATCH 04/12] Added unit tests and modified the code

---
 src/databricks/sql/auth/retry.py     | 7 +++++++
 src/databricks/sql/thrift_backend.py | 9 +++++++++
 tests/unit/test_retry.py             | 6 ++++++
 3 files changed, 22 insertions(+)

diff --git a/src/databricks/sql/auth/retry.py b/src/databricks/sql/auth/retry.py
index 0243d0aa2..85c4386fd 100755
--- a/src/databricks/sql/auth/retry.py
+++ b/src/databricks/sql/auth/retry.py
@@ -32,6 +32,7 @@ class CommandType(Enum):
     CLOSE_SESSION = "CloseSession"
     CLOSE_OPERATION = "CloseOperation"
     GET_OPERATION_STATUS = "GetOperationStatus"
+    FETCH_RESULTS_ORIENTATION_FETCH_NEXT = "FetchResultsOrientation_FETCH_NEXT"
     OTHER = "Other"
 
     @classmethod
@@ -362,6 +363,12 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
         if status_code == 501:
             raise NonRecoverableNetworkError("Received code 501 from server.")
 
+        if self.command_type == CommandType.FETCH_RESULTS_ORIENTATION_FETCH_NEXT:
+            return (
+                False,
+                "FetchResults with FETCH_NEXT orientation are not idempotent and is not retried",
+            )
+
         # Request failed and this method is not retryable. We only retry POST requests.
         if not self._is_method_retryable(method):
             return False, "Only POST requests are retried"
diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py
index 141d9e1f6..711fa388b 100644
--- a/src/databricks/sql/thrift_backend.py
+++ b/src/databricks/sql/thrift_backend.py
@@ -374,6 +374,15 @@ def attempt_request(attempt):
 
                 # These three lines are no-ops if the v3 retry policy is not in use
                 if self.enable_v3_retries:
+                    # Not to retry when FetchResults has orientation as FETCH_NEXT as it is not idempotent
+                    if this_method_name == "FetchResults":
+                        this_method_name += (
+                            "Orientation_"
+                            + ttypes.TFetchOrientation._VALUES_TO_NAMES[
+                                request.orientation
+                            ]
+                        )
+
                     this_command_type = CommandType.get(this_method_name)
                     self._transport.set_retry_command_type(this_command_type)
                     self._transport.startRetryTimer()
diff --git a/tests/unit/test_retry.py b/tests/unit/test_retry.py
index 1e18e1f43..3261d178f 100644
--- a/tests/unit/test_retry.py
+++ b/tests/unit/test_retry.py
@@ -84,3 +84,9 @@ def test_sleep__retry_after_present(self, t_mock, retry_policy, error_history):
         retry_policy.history = [error_history, error_history, error_history]
         retry_policy.sleep(HTTPResponse(status=503, headers={"Retry-After": "3"}))
         t_mock.assert_called_with(3)
+
+    def test_not_retryable__fetch_results_orientation_fetch_next(self, retry_policy):
+        HTTP_STATUS_CODES = [200, 429, 503, 504]
+        retry_policy.command_type = CommandType.FETCH_RESULTS_ORIENTATION_FETCH_NEXT
+        for status_code in HTTP_STATUS_CODES:
+            assert not retry_policy.is_retry("METHOD_NAME", status_code=status_code)

From bf2688143eb1563e064359f99444bf90be24932b Mon Sep 17 00:00:00 2001
From: Jothi Prakash <jothi.prakash@databricks.com>
Date: Mon, 16 Dec 2024 22:59:45 +0530
Subject: [PATCH 05/12] Updated error message

---
 src/databricks/sql/auth/retry.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/databricks/sql/auth/retry.py b/src/databricks/sql/auth/retry.py
index 85c4386fd..08bb6846e 100755
--- a/src/databricks/sql/auth/retry.py
+++ b/src/databricks/sql/auth/retry.py
@@ -366,7 +366,7 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
         if self.command_type == CommandType.FETCH_RESULTS_ORIENTATION_FETCH_NEXT:
             return (
                 False,
-                "FetchResults with FETCH_NEXT orientation are not idempotent and is not retried",
+                "FetchResults with FETCH_NEXT orientation are not idempotent in inline mode and is not retried",
             )
 
         # Request failed and this method is not retryable. We only retry POST requests.

From b2202c7d87a6f1a46594eaaa284efb2a726ef8d1 Mon Sep 17 00:00:00 2001
From: Jothi Prakash <jothi.prakash@databricks.com>
Date: Tue, 17 Dec 2024 22:37:34 +0530
Subject: [PATCH 06/12] Updated the non retying to only inline case

---
 src/databricks/sql/auth/retry.py     |  6 +++---
 src/databricks/sql/thrift_backend.py | 11 ++++++++---
 tests/unit/test_retry.py             |  2 +-
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/src/databricks/sql/auth/retry.py b/src/databricks/sql/auth/retry.py
index 08bb6846e..109b6949e 100755
--- a/src/databricks/sql/auth/retry.py
+++ b/src/databricks/sql/auth/retry.py
@@ -32,7 +32,7 @@ class CommandType(Enum):
     CLOSE_SESSION = "CloseSession"
     CLOSE_OPERATION = "CloseOperation"
     GET_OPERATION_STATUS = "GetOperationStatus"
-    FETCH_RESULTS_ORIENTATION_FETCH_NEXT = "FetchResultsOrientation_FETCH_NEXT"
+    FETCH_RESULTS_INLINE_FETCH_NEXT = "FetchResultsInline_FETCH_NEXT"
     OTHER = "Other"
 
     @classmethod
@@ -363,10 +363,10 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
         if status_code == 501:
             raise NonRecoverableNetworkError("Received code 501 from server.")
 
-        if self.command_type == CommandType.FETCH_RESULTS_ORIENTATION_FETCH_NEXT:
+        if self.command_type == CommandType.FETCH_RESULTS_INLINE_FETCH_NEXT:
             return (
                 False,
-                "FetchResults with FETCH_NEXT orientation are not idempotent in inline mode and is not retried",
+                "FetchResults in INLINE mode with FETCH_NEXT orientation are not idempotent and is not retried",
             )
 
         # Request failed and this method is not retryable. We only retry POST requests.
diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py
index 711fa388b..a1f9815de 100644
--- a/src/databricks/sql/thrift_backend.py
+++ b/src/databricks/sql/thrift_backend.py
@@ -156,6 +156,8 @@ def __init__(
         )
 
         # Cloud fetch
+        self._use_cloud_fetch = kwargs.get("use_cloud_fetch", True)
+
         self.max_download_threads = kwargs.get("max_download_threads", 10)
 
         self._ssl_options = ssl_options
@@ -374,10 +376,13 @@ def attempt_request(attempt):
 
                 # These three lines are no-ops if the v3 retry policy is not in use
                 if self.enable_v3_retries:
-                    # Not to retry when FetchResults has orientation as FETCH_NEXT as it is not idempotent
-                    if this_method_name == "FetchResults":
+                    # Not to retry when FetchResults in INLINE mode when it has orientation as FETCH_NEXT as it is not idempotent
+                    if (
+                        this_method_name == "FetchResults"
+                        and self._use_cloud_fetch == False
+                    ):
                         this_method_name += (
-                            "Orientation_"
+                            "Inline_"
                             + ttypes.TFetchOrientation._VALUES_TO_NAMES[
                                 request.orientation
                             ]
diff --git a/tests/unit/test_retry.py b/tests/unit/test_retry.py
index 3261d178f..9e52df72a 100644
--- a/tests/unit/test_retry.py
+++ b/tests/unit/test_retry.py
@@ -87,6 +87,6 @@ def test_sleep__retry_after_present(self, t_mock, retry_policy, error_history):
 
     def test_not_retryable__fetch_results_orientation_fetch_next(self, retry_policy):
         HTTP_STATUS_CODES = [200, 429, 503, 504]
-        retry_policy.command_type = CommandType.FETCH_RESULTS_ORIENTATION_FETCH_NEXT
+        retry_policy.command_type = CommandType.FETCH_RESULTS_INLINE_FETCH_NEXT
         for status_code in HTTP_STATUS_CODES:
             assert not retry_policy.is_retry("METHOD_NAME", status_code=status_code)

From 8e759dbb6b6f63ce9516833a1b024da1f6d63804 Mon Sep 17 00:00:00 2001
From: Jothi Prakash <jothi.prakash@databricks.com>
Date: Thu, 19 Dec 2024 10:52:04 +0530
Subject: [PATCH 07/12] Updated fix

---
 src/databricks/sql/thrift_backend.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py
index a1f9815de..70b29d32a 100644
--- a/src/databricks/sql/thrift_backend.py
+++ b/src/databricks/sql/thrift_backend.py
@@ -156,8 +156,6 @@ def __init__(
         )
 
         # Cloud fetch
-        self._use_cloud_fetch = kwargs.get("use_cloud_fetch", True)
-
         self.max_download_threads = kwargs.get("max_download_threads", 10)
 
         self._ssl_options = ssl_options
@@ -900,6 +898,8 @@ def execute_command(
     ):
         assert session_handle is not None
 
+        self._use_cloud_fetch = use_cloud_fetch
+
         spark_arrow_types = ttypes.TSparkArrowTypes(
             timestampAsArrow=self._use_arrow_native_timestamps,
             decimalAsArrow=self._use_arrow_native_decimals,

From 1ac441a5c5d8c64aec196767702093231db92779 Mon Sep 17 00:00:00 2001
From: Jothi Prakash <jothi.prakash@databricks.com>
Date: Thu, 19 Dec 2024 23:49:23 +0530
Subject: [PATCH 08/12] Changed the flow

---
 src/databricks/sql/auth/retry.py              | 16 ++++++++------
 src/databricks/sql/auth/thrift_http_client.py |  9 ++++++++
 src/databricks/sql/client.py                  |  4 ++++
 src/databricks/sql/thrift_backend.py          | 21 +++++--------------
 tests/unit/test_retry.py                      |  2 +-
 5 files changed, 29 insertions(+), 23 deletions(-)

diff --git a/src/databricks/sql/auth/retry.py b/src/databricks/sql/auth/retry.py
index 109b6949e..56cd34127 100755
--- a/src/databricks/sql/auth/retry.py
+++ b/src/databricks/sql/auth/retry.py
@@ -32,7 +32,6 @@ class CommandType(Enum):
     CLOSE_SESSION = "CloseSession"
     CLOSE_OPERATION = "CloseOperation"
     GET_OPERATION_STATUS = "GetOperationStatus"
-    FETCH_RESULTS_INLINE_FETCH_NEXT = "FetchResultsInline_FETCH_NEXT"
     OTHER = "Other"
 
     @classmethod
@@ -242,6 +241,14 @@ def command_type(self) -> Optional[CommandType]:
     def command_type(self, value: CommandType) -> None:
         self._command_type = value
 
+    @property
+    def is_retryable(self) -> bool:
+        return self._is_retryable
+
+    @is_retryable.setter
+    def is_retryable(self, value: bool) -> None:
+        self._is_retryable = value
+
     @property
     def delay_default(self) -> float:
         """Time in seconds the connector will wait between requests polling a GetOperationStatus Request
@@ -363,11 +370,8 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
         if status_code == 501:
             raise NonRecoverableNetworkError("Received code 501 from server.")
 
-        if self.command_type == CommandType.FETCH_RESULTS_INLINE_FETCH_NEXT:
-            return (
-                False,
-                "FetchResults in INLINE mode with FETCH_NEXT orientation are not idempotent and is not retried",
-            )
+        if self.is_retryable == False:
+            return False, "Request is not retryable"
 
         # Request failed and this method is not retryable. We only retry POST requests.
         if not self._is_method_retryable(method):
diff --git a/src/databricks/sql/auth/thrift_http_client.py b/src/databricks/sql/auth/thrift_http_client.py
index 6273ab284..6381ac61f 100644
--- a/src/databricks/sql/auth/thrift_http_client.py
+++ b/src/databricks/sql/auth/thrift_http_client.py
@@ -216,3 +216,12 @@ def set_retry_command_type(self, value: CommandType):
             logger.warning(
                 "DatabricksRetryPolicy is currently bypassed. The CommandType cannot be set."
             )
+
+    def set_is_retryable(self, retryable: bool):
+        """Pass the provided retryable flag to the retry policy"""
+        if isinstance(self.retry_policy, DatabricksRetryPolicy):
+            self.retry_policy.is_retryable = retryable
+        else:
+            logger.warning(
+                "DatabricksRetryPolicy is currently bypassed. The is_retryable flag cannot be set."
+            )
diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py
index 8ea81e12b..341a6ebec 100755
--- a/src/databricks/sql/client.py
+++ b/src/databricks/sql/client.py
@@ -808,6 +808,7 @@ def execute(
             self.thrift_backend,
             self.buffer_size_bytes,
             self.arraysize,
+            self.connection.use_cloud_fetch,
         )
 
         if execute_response.is_staging_operation:
@@ -1202,6 +1203,7 @@ def __init__(
         thrift_backend: ThriftBackend,
         result_buffer_size_bytes: int = DEFAULT_RESULT_BUFFER_SIZE_BYTES,
         arraysize: int = 10000,
+        use_cloud_fetch: bool = True,
     ):
         """
         A ResultSet manages the results of a single command.
@@ -1223,6 +1225,7 @@ def __init__(
         self.description = execute_response.description
         self._arrow_schema_bytes = execute_response.arrow_schema_bytes
         self._next_row_index = 0
+        self.use_cloud_fetch = use_cloud_fetch
 
         if execute_response.arrow_queue:
             # In this case the server has taken the fast path and returned an initial batch of
@@ -1250,6 +1253,7 @@ def _fill_results_buffer(self):
             lz4_compressed=self.lz4_compressed,
             arrow_schema_bytes=self._arrow_schema_bytes,
             description=self.description,
+            use_cloud_fetch=self.use_cloud_fetch,
         )
         self.results = results
         self.has_more_rows = has_more_rows
diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py
index 70b29d32a..ab4da3154 100644
--- a/src/databricks/sql/thrift_backend.py
+++ b/src/databricks/sql/thrift_backend.py
@@ -321,7 +321,7 @@ def _handle_request_error(self, error_info, attempt, elapsed):
 
     # FUTURE: Consider moving to https://github.com/litl/backoff or
     # https://github.com/jd/tenacity for retry logic.
-    def make_request(self, method, request):
+    def make_request(self, method, request, retryable=True):
         """Execute given request, attempting retries when
             1. Receiving HTTP 429/503 from server
             2. OSError is raised during a GetOperationStatus
@@ -374,20 +374,9 @@ def attempt_request(attempt):
 
                 # These three lines are no-ops if the v3 retry policy is not in use
                 if self.enable_v3_retries:
-                    # Not to retry when FetchResults in INLINE mode when it has orientation as FETCH_NEXT as it is not idempotent
-                    if (
-                        this_method_name == "FetchResults"
-                        and self._use_cloud_fetch == False
-                    ):
-                        this_method_name += (
-                            "Inline_"
-                            + ttypes.TFetchOrientation._VALUES_TO_NAMES[
-                                request.orientation
-                            ]
-                        )
-
                     this_command_type = CommandType.get(this_method_name)
                     self._transport.set_retry_command_type(this_command_type)
+                    self._transport.set_is_retryable(retryable)
                     self._transport.startRetryTimer()
 
                 response = method(request)
@@ -898,8 +887,6 @@ def execute_command(
     ):
         assert session_handle is not None
 
-        self._use_cloud_fetch = use_cloud_fetch
-
         spark_arrow_types = ttypes.TSparkArrowTypes(
             timestampAsArrow=self._use_arrow_native_timestamps,
             decimalAsArrow=self._use_arrow_native_decimals,
@@ -1042,6 +1029,7 @@ def fetch_results(
         lz4_compressed,
         arrow_schema_bytes,
         description,
+        use_cloud_fetch=True,
     ):
         assert op_handle is not None
 
@@ -1058,7 +1046,8 @@ def fetch_results(
             includeResultSetMetadata=True,
         )
 
-        resp = self.make_request(self._client.FetchResults, req)
+        # Fetch results in Inline mode with FETCH_NEXT orientation are not idempotent and hence not retried
+        resp = self.make_request(self._client.FetchResults, req, use_cloud_fetch)
         if resp.results.startRowOffset > expected_row_start_offset:
             raise DataError(
                 "fetch_results failed due to inconsistency in the state between the client and the server. Expected results to start from {} but they instead start at {}, some result batches must have been skipped".format(
diff --git a/tests/unit/test_retry.py b/tests/unit/test_retry.py
index 9e52df72a..be8e1ecc9 100644
--- a/tests/unit/test_retry.py
+++ b/tests/unit/test_retry.py
@@ -87,6 +87,6 @@ def test_sleep__retry_after_present(self, t_mock, retry_policy, error_history):
 
     def test_not_retryable__fetch_results_orientation_fetch_next(self, retry_policy):
         HTTP_STATUS_CODES = [200, 429, 503, 504]
-        retry_policy.command_type = CommandType.FETCH_RESULTS_INLINE_FETCH_NEXT
+        retry_policy.is_retryable = False
         for status_code in HTTP_STATUS_CODES:
             assert not retry_policy.is_retry("METHOD_NAME", status_code=status_code)

From fba9759e4a4073052d3bab0bcdf77b455c822189 Mon Sep 17 00:00:00 2001
From: Jothi Prakash <jothi.prakash@databricks.com>
Date: Thu, 19 Dec 2024 23:57:07 +0530
Subject: [PATCH 09/12] Minor update

---
 src/databricks/sql/auth/retry.py | 1 -
 src/databricks/sql/client.py     | 4 ++--
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/databricks/sql/auth/retry.py b/src/databricks/sql/auth/retry.py
index 56cd34127..dcb489621 100755
--- a/src/databricks/sql/auth/retry.py
+++ b/src/databricks/sql/auth/retry.py
@@ -351,7 +351,6 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
         Returns True if the request should be retried. Returns False or raises an exception
         if a retry would violate the configured policy.
         """
-
         # Request succeeded. Don't retry.
         if status_code == 200:
             return False, "200 codes are not retried"
diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py
index 341a6ebec..aefed1ef0 100755
--- a/src/databricks/sql/client.py
+++ b/src/databricks/sql/client.py
@@ -1225,7 +1225,7 @@ def __init__(
         self.description = execute_response.description
         self._arrow_schema_bytes = execute_response.arrow_schema_bytes
         self._next_row_index = 0
-        self.use_cloud_fetch = use_cloud_fetch
+        self._use_cloud_fetch = use_cloud_fetch
 
         if execute_response.arrow_queue:
             # In this case the server has taken the fast path and returned an initial batch of
@@ -1253,7 +1253,7 @@ def _fill_results_buffer(self):
             lz4_compressed=self.lz4_compressed,
             arrow_schema_bytes=self._arrow_schema_bytes,
             description=self.description,
-            use_cloud_fetch=self.use_cloud_fetch,
+            use_cloud_fetch=self._use_cloud_fetch,
         )
         self.results = results
         self.has_more_rows = has_more_rows

From aefd44b9d76e1cdeb1b2cd7eaf9e78a034b4324d Mon Sep 17 00:00:00 2001
From: Jothi Prakash <jothi.prakash@databricks.com>
Date: Fri, 20 Dec 2024 22:02:54 +0530
Subject: [PATCH 10/12] Updated the retryable condition

---
 src/databricks/sql/auth/retry.py              | 11 -----------
 src/databricks/sql/auth/thrift_http_client.py |  9 ---------
 src/databricks/sql/thrift_backend.py          |  3 +--
 tests/unit/test_retry.py                      |  6 ------
 4 files changed, 1 insertion(+), 28 deletions(-)

diff --git a/src/databricks/sql/auth/retry.py b/src/databricks/sql/auth/retry.py
index dcb489621..606ac04cf 100755
--- a/src/databricks/sql/auth/retry.py
+++ b/src/databricks/sql/auth/retry.py
@@ -241,14 +241,6 @@ def command_type(self) -> Optional[CommandType]:
     def command_type(self, value: CommandType) -> None:
         self._command_type = value
 
-    @property
-    def is_retryable(self) -> bool:
-        return self._is_retryable
-
-    @is_retryable.setter
-    def is_retryable(self, value: bool) -> None:
-        self._is_retryable = value
-
     @property
     def delay_default(self) -> float:
         """Time in seconds the connector will wait between requests polling a GetOperationStatus Request
@@ -369,9 +361,6 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
         if status_code == 501:
             raise NonRecoverableNetworkError("Received code 501 from server.")
 
-        if self.is_retryable == False:
-            return False, "Request is not retryable"
-
         # Request failed and this method is not retryable. We only retry POST requests.
         if not self._is_method_retryable(method):
             return False, "Only POST requests are retried"
diff --git a/src/databricks/sql/auth/thrift_http_client.py b/src/databricks/sql/auth/thrift_http_client.py
index 6381ac61f..6273ab284 100644
--- a/src/databricks/sql/auth/thrift_http_client.py
+++ b/src/databricks/sql/auth/thrift_http_client.py
@@ -216,12 +216,3 @@ def set_retry_command_type(self, value: CommandType):
             logger.warning(
                 "DatabricksRetryPolicy is currently bypassed. The CommandType cannot be set."
             )
-
-    def set_is_retryable(self, retryable: bool):
-        """Pass the provided retryable flag to the retry policy"""
-        if isinstance(self.retry_policy, DatabricksRetryPolicy):
-            self.retry_policy.is_retryable = retryable
-        else:
-            logger.warning(
-                "DatabricksRetryPolicy is currently bypassed. The is_retryable flag cannot be set."
-            )
diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py
index ab4da3154..5fbd9f749 100644
--- a/src/databricks/sql/thrift_backend.py
+++ b/src/databricks/sql/thrift_backend.py
@@ -376,7 +376,6 @@ def attempt_request(attempt):
                 if self.enable_v3_retries:
                     this_command_type = CommandType.get(this_method_name)
                     self._transport.set_retry_command_type(this_command_type)
-                    self._transport.set_is_retryable(retryable)
                     self._transport.startRetryTimer()
 
                 response = method(request)
@@ -461,7 +460,7 @@ def attempt_request(attempt):
         #       return on success
         #       if available: bounded delay and retry
         #       if not: raise error
-        max_attempts = self._retry_stop_after_attempts_count
+        max_attempts = self._retry_stop_after_attempts_count if retryable else 1
 
         # use index-1 counting for logging/human consistency
         for attempt in range(1, max_attempts + 1):
diff --git a/tests/unit/test_retry.py b/tests/unit/test_retry.py
index be8e1ecc9..1e18e1f43 100644
--- a/tests/unit/test_retry.py
+++ b/tests/unit/test_retry.py
@@ -84,9 +84,3 @@ def test_sleep__retry_after_present(self, t_mock, retry_policy, error_history):
         retry_policy.history = [error_history, error_history, error_history]
         retry_policy.sleep(HTTPResponse(status=503, headers={"Retry-After": "3"}))
         t_mock.assert_called_with(3)
-
-    def test_not_retryable__fetch_results_orientation_fetch_next(self, retry_policy):
-        HTTP_STATUS_CODES = [200, 429, 503, 504]
-        retry_policy.is_retryable = False
-        for status_code in HTTP_STATUS_CODES:
-            assert not retry_policy.is_retry("METHOD_NAME", status_code=status_code)

From ea686c78ca2b2b2e60132947a18d43135616bc11 Mon Sep 17 00:00:00 2001
From: Jothi Prakash <jothi.prakash@databricks.com>
Date: Fri, 20 Dec 2024 22:11:09 +0530
Subject: [PATCH 11/12] Minor test fix

---
 tests/unit/test_fetches.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index e9a58acdd..89cedcfae 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -66,6 +66,7 @@ def fetch_results(
             lz4_compressed,
             arrow_schema_bytes,
             description,
+            use_cloud_fetch=True,
         ):
             nonlocal batch_index
             results = FetchTests.make_arrow_queue(batch_list[batch_index])

From e6e573c462fe961f69664de9306792e82e6f3c9d Mon Sep 17 00:00:00 2001
From: Jothi Prakash <jothi.prakash@databricks.com>
Date: Fri, 20 Dec 2024 22:50:12 +0530
Subject: [PATCH 12/12] Added extra space

---
 src/databricks/sql/auth/retry.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/databricks/sql/auth/retry.py b/src/databricks/sql/auth/retry.py
index 606ac04cf..0243d0aa2 100755
--- a/src/databricks/sql/auth/retry.py
+++ b/src/databricks/sql/auth/retry.py
@@ -343,6 +343,7 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
         Returns True if the request should be retried. Returns False or raises an exception
         if a retry would violate the configured policy.
         """
+
         # Request succeeded. Don't retry.
         if status_code == 200:
             return False, "200 codes are not retried"