Do not retry failing requests with status code 401

Hodnebo · Hodnebo · commit 37a1a15e4e88 · 2024-07-03T10:46:14.000+02:00
- Raises NonRecoverableNetworkError when request results in 401 status code

Signed-off-by: Tor Hødnebø &lt;thodnebo@gmail.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # Release History
 
+# x.x.x (TBD)
+
+- Don't retry requests that fail with code 401
+
 # 3.2.0 (2024-06-06)
 
 - Update proxy authentication (databricks/databricks-sql-python#354 by @amir-haroun)
diff --git a/src/databricks/sql/auth/retry.py b/src/databricks/sql/auth/retry.py
@@ -325,7 +325,8 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
                default, this means ExecuteStatement is only retried for codes 429 and 503.
                This limit prevents automatically retrying non-idempotent commands that could
                be destructive.
-            5. The request received a 403 response, because this can never succeed.
+            5. The request received a 401 response, because this can never succeed.
+            6. The request received a 403 response, because this can never succeed.
 
 
         Q: What about OSErrors and Redirects?
@@ -339,6 +340,11 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
         if status_code == 200:
             return False, "200 codes are not retried"
 
+        if status_code == 401:
+            raise NonRecoverableNetworkError(
+                "Received 401 - UNAUTHORIZED. Confirm your authentication credentials."
+            )
+
         if status_code == 403:
             raise NonRecoverableNetworkError(
                 "Received 403 - FORBIDDEN. Confirm your authentication credentials."
diff --git a/tests/e2e/common/retry_test_mixins.py b/tests/e2e/common/retry_test_mixins.py
@@ -420,4 +420,17 @@ def test_403_not_retried(self):
             with pytest.raises(RequestError) as cm:
                 with self.connection(extra_params=self._retry_policy) as conn:
                     pass
-                assert isinstance(cm.value.args[1], NonRecoverableNetworkError)
+                assert isinstance(cm.value.args[1], NonRecoverableNetworkError)
+
+    def test_401_not_retried(self):
+        """GIVEN the server returns a code 401
+        WHEN the connector receives this response
+        THEN nothing is retried and an exception is raised
+        """
+
+        # Code 401 is an Unauthorized error
+        with mocked_server_response(status=401):
+            with pytest.raises(RequestError) as cm:
+                with self.connection(extra_params=self._retry_policy):
+                    pass
+                assert isinstance(cm.value.args[1], NonRecoverableNetworkError)
