Updated the non retying to only inline case

Author: jprakash-db

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.

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
                             ]

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)