From 2dd2e993c7ac87409375e360ca079631f844c6a9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tor=20H=C3=B8dneb=C3=B8?= Date: Wed, 3 Jul 2024 10:46:14 +0200 Subject: [PATCH] Do not retry failing requests with status code 401 - Raises NonRecoverableNetworkError when request results in 401 status code MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Tor Hødnebø Signed-off-by: Tor Hødnebø --- CHANGELOG.md | 4 ++++ src/databricks/sql/auth/retry.py | 8 +++++++- tests/e2e/common/retry_test_mixins.py | 15 ++++++++++++++- 3 files changed, 25 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 487f0c32..03e8b836 100644 --- 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 index eee6f839..2581745d 100644 --- 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 index 9a49870e..9856a4bd 100644 --- 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) \ No newline at end of file + 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)