Clean up authentication

trino/auth.py

@@ -12,29 +12,19 @@
 
 import abc
 import os
+from typing import Optional, Sequence, Type
 
-from typing import Optional
-from requests.auth import AuthBase
+import requests
+import requests.auth
 
 
 class Authentication(metaclass=abc.ABCMeta):
     @abc.abstractmethod
-    def set_http_session(self, http_session):
+    def set_http_session(self, http_session: requests.Session) -> None:
         pass
 
-    @abc.abstractmethod
-    def set_client_session(self, client_session):
-        pass
-
-    @abc.abstractmethod
-    def setup(self):
-        pass
-
-    def get_exceptions(self):
-        return tuple()
-
-    def handle_err(self, error):
-        pass
+    def get_exceptions(self) -> Sequence[Type[Exception]]:
+        return ()
 
 
 class KerberosAuthentication(Authentication):
@@ -60,10 +50,7 @@ def __init__(
         self._delegate = delegate
         self._ca_bundle = ca_bundle
 
-    def set_client_session(self, client_session):
-        pass
-
-    def set_http_session(self, http_session):
+    def set_http_session(self, http_session: requests.Session) -> None:
         try:
             import requests_kerberos
         except ImportError:
@@ -83,57 +70,31 @@ def set_http_session(self, http_session):
         )
         if self._ca_bundle:
             http_session.verify = self._ca_bundle
-        return http_session
 
-    def setup(self, trino_client):
-        self.set_client_session(trino_client.client_session)
-        self.set_http_session(trino_client.http_session)
-
-    def get_exceptions(self):
+    def get_exceptions(self) -> Sequence[Type[Exception]]:
         try:
             from requests_kerberos.exceptions import KerberosExchangeError
 
             return (KerberosExchangeError,)
         except ImportError:
             raise RuntimeError("unable to import requests_kerberos")
 
-    def handle_error(self, handle_error):
-        pass
-
 
 class BasicAuthentication(Authentication):
-    def __init__(self, username, password):
+    def __init__(self, username: str, password: str) -> None:
         self._username = username
         self._password = password
 
-    def set_client_session(self, client_session):
-        pass
-
-    def set_http_session(self, http_session):
-        try:
-            import requests.auth
-        except ImportError:
-            raise RuntimeError("unable to import requests.auth")
-
+    def set_http_session(self, http_session: requests.Session) -> None:
         http_session.auth = requests.auth.HTTPBasicAuth(self._username, self._password)
-        return http_session
-
-    def setup(self, trino_client):
-        self.set_client_session(trino_client.client_session)
-        self.set_http_session(trino_client.http_session)
-
-    def get_exceptions(self):
-        return ()
-
-    def handle_error(self, handle_error):
-        pass
 
 
-class _BearerAuth(AuthBase):
+class _BearerAuth(requests.auth.AuthBase):
     """
-    Custom implementation of Authentication class for bearer token
+    Custom implementation of AuthBase class for bearer token
     """
-    def __init__(self, token):
+
+    def __init__(self, token: str) -> None:
         self.token = token
 
     def __call__(self, r):
@@ -142,23 +103,8 @@ def __call__(self, r):
 
 
 class JWTAuthentication(Authentication):
-
-    def __init__(self, token):
+    def __init__(self, token: str) -> None:
         self.token = token
 
-    def set_client_session(self, client_session):
-        pass
-
-    def set_http_session(self, http_session):
+    def set_http_session(self, http_session: requests.Session) -> None:
         http_session.auth = _BearerAuth(self.token)
-        return http_session
-
-    def setup(self, trino_client):
-        self.set_client_session(trino_client.client_session)
-        self.set_http_session(trino_client.http_session)
-
-    def get_exceptions(self):
-        return ()
-
-    def handle_error(self, handle_error):
-        pass

trino/client.py

@@ -41,6 +41,7 @@
 
 import trino.logging
 from trino import constants, exceptions
+from trino.auth import Authentication
 from trino.transaction import NO_TRANSACTION
 
 __all__ = ["TrinoQuery", "TrinoRequest"]
@@ -201,7 +202,7 @@ def __init__(
         http_headers: Optional[Dict[str, str]] = None,
         transaction_id: Optional[str] = NO_TRANSACTION,
         http_scheme: str = constants.HTTP,
-        auth: Optional[Any] = constants.DEFAULT_AUTH,
+        auth: Optional[Authentication] = constants.DEFAULT_AUTH,
         redirect_handler: Any = None,
         max_attempts: int = MAX_ATTEMPTS,
         request_timeout: Union[float, Tuple[float, float]] = constants.DEFAULT_REQUEST_TIMEOUT,

trino/constants.py

@@ -10,14 +10,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Optional
+from typing import Optional
 
+from trino.auth import Authentication
 
 DEFAULT_PORT = 8080
 DEFAULT_SOURCE = "trino-python-client"
 DEFAULT_CATALOG: Optional[str] = None
 DEFAULT_SCHEMA: Optional[str] = None
-DEFAULT_AUTH: Optional[Any] = None
+DEFAULT_AUTH: Optional[Authentication] = None
 DEFAULT_MAX_ATTEMPTS = 3
 DEFAULT_REQUEST_TIMEOUT: float = 30.0