Incorporated the 401 status code error into the NonRecoverableNetworkError error handling

Author: jprakash-db

.gitignore

@@ -11,6 +11,8 @@
 # Icon must end with two \r
 Icon
 
+# IntelliJ Files
+.idea
 
 # Thumbnails
 ._*

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."
             )

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"""