Add support for JWT Authentication

README.md

@@ -58,6 +58,25 @@ cur.execute('SELECT * FROM system.runtime.nodes')
 rows = cur.fetchall()
 ```
 
+# JWT Authentication
+The `JWTAuthentication` can be used to connect to a JWT token authentication enabled Trino cluster:
+```python
+import trino
+JWT_TOKEN = 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJuYW1lIjoiSm9obiBEb2UifQ.DjwRE2jZhren2Wt37t5hlVru6Myq4AhpGLiiefF69u8'
+conn = trino.dbapi.connect(
+ host='coordinator-url',
+ port=8443,
+ user='the-user',
+ catalog='the-catalog',
+ schema='the-schema',
+ http_scheme='https',
+ auth=trino.auth.JWTAuthentication(JWT_TOKEN),
+)
+cur = conn.cursor()
+cur.execute('SELECT * FROM system.runtime.nodes')
+rows = cur.fetchall()
+```
+
 # Transactions
 The client runs by default in *autocommit* mode. To enable transactions, set
 *isolation_level* to a value different than `IsolationLevel.AUTOCOMMIT`:

trino/auth.py

@@ -18,6 +18,7 @@
 
 from six import with_metaclass
 from typing import Any, Optional, Text # NOQA
+from requests.auth import AuthBase
 
 
 class Authentication(with_metaclass(abc.ABCMeta)): # type: ignore
@@ -131,3 +132,38 @@ def get_exceptions(self):
 
  def handle_error(self, handle_error):
  pass
+
+
+class _BearerAuth(AuthBase):
+ """
+ Custom implementation of Authentication class for bearer token
+ """
+ def __init__(self, token):
+ self.token = token
+
+ def __call__(self, r):
+ r.headers["Authorization"] = "Bearer " + self.token
+ return r
+
+
+class JWTAuthentication(Authentication):
+
+ def __init__(self, token):
+ self.token = token
+
+ def set_client_session(self, client_session):
+ pass
+
+ def set_http_session(self, http_session):
+ 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