From 1203a62dfe17ee8f089a7eebf9d03e919dd7322e Mon Sep 17 00:00:00 2001
From: Jothi Prakash <jothi.prakash@databricks.com>
Date: Mon, 17 Jun 2024 14:30:02 +0530
Subject: [PATCH] Incorporated the 401 status code error into the
 NonRecoverableNetworkError error handling

---
 .gitignore                       | 2 ++
 src/databricks/sql/auth/retry.py | 9 ++-------
 src/databricks/sql/exc.py        | 4 ----
 3 files changed, 4 insertions(+), 11 deletions(-)

diff --git a/.gitignore b/.gitignore
index d89a4116a..682431225 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,6 +11,8 @@
 # Icon must end with two \r
 Icon
 
+# IntelliJ Files
+.idea
 
 # Thumbnails
 ._*
diff --git a/src/databricks/sql/auth/retry.py b/src/databricks/sql/auth/retry.py
index 24415490b..1faf1222b 100644
--- a/src/databricks/sql/auth/retry.py
+++ b/src/databricks/sql/auth/retry.py
@@ -19,7 +19,6 @@
     MaxRetryDurationError,
     NonRecoverableNetworkError,
     OperationalError,
-    AuthenticationFailureError,
     SessionAlreadyClosedError,
     UnsafeToRetryError,
 )
@@ -340,12 +339,8 @@ def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
         if status_code == 200:
             return False, "200 codes are not retried"
 
-        # Don't retry as there is an authentication error
-        if status_code == 401:
-            raise AuthenticationFailureError(
-                "Received 401 - UNAUTHORIZED. Authentication is required and has failed or has not yet been provided."
-            )
-        if status_code == 403:
+        # Invalid Credentials error. Don't retry
+        if status_code == 403 or status_code == 401:
             raise NonRecoverableNetworkError(
                 "Received 403 - FORBIDDEN. Confirm your authentication credentials."
             )
diff --git a/src/databricks/sql/exc.py b/src/databricks/sql/exc.py
index eff98f38b..3b27283a4 100644
--- a/src/databricks/sql/exc.py
+++ b/src/databricks/sql/exc.py
@@ -105,10 +105,6 @@ class NonRecoverableNetworkError(RequestError):
     """Thrown if an HTTP code 501 is received"""
 
 
-class AuthenticationFailureError(RequestError):
-    """Thrown if an HTTP code 401 is received"""
-
-
 class UnsafeToRetryError(RequestError):
     """Thrown if ExecuteStatement request receives a code other than 200, 429, or 503"""