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] 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 0243d0aa..85c4386f 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 141d9e1f..711fa388 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 1e18e1f4..3261d178 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)