Add JWT token authentication

Author: amotl

CHANGES.rst

@@ -4,6 +4,7 @@ Changes for crate
 
 Unreleased
 ==========
+- Added JWT token authentication
 
 2025/01/30 2.0.0
 ================

docs/connect.rst

@@ -246,6 +246,12 @@ and password.
     authenticate as the CrateDB superuser, which is ``crate``. The superuser
     does not have a password, so you can omit the ``password`` argument.
 
+Alternatively, authenticate using a JWT token:
+
+    >>> connection = client.connect(..., jwt_token="<JWT_TOKEN>")
+
+Here, replace ``<JWT_TOKEN>`` with the appropriate JWT token.
+
 .. _schema-selection:
 
 Schema selection

src/crate/client/connection.py

@@ -42,6 +42,7 @@ def __init__(
         ssl_relax_minimum_version=False,
         username=None,
         password=None,
+        jwt_token=None,
         schema=None,
         pool_size=None,
         socket_keepalive=True,
@@ -81,6 +82,8 @@ def __init__(
             the username in the database.
         :param password:
             the password of the user in the database.
+        :param jwt_token:
+            the JWT token to authenticate with the server.
         :param pool_size:
             (optional)
             Number of connections to save that can be reused.
@@ -148,6 +151,7 @@ def __init__(
                 ssl_relax_minimum_version=ssl_relax_minimum_version,
                 username=username,
                 password=password,
+                jwt_token=jwt_token,
                 schema=schema,
                 pool_size=pool_size,
                 socket_keepalive=socket_keepalive,

src/crate/client/http.py

@@ -158,6 +158,7 @@ def request(
         headers=None,
         username=None,
         password=None,
+        jwt_token=None,
         schema=None,
         backoff_factor=0,
         **kwargs,
@@ -173,6 +174,10 @@ def request(
             if length is not None:
                 headers["Content-Length"] = length
 
+        # Authentication token
+        if jwt_token is not None and "Authorization" not in headers:
+            headers["Authorization"] = "Bearer %s" % jwt_token
+
         # Authentication credentials
         if username is not None:
             if "Authorization" not in headers and username is not None:
@@ -418,6 +423,7 @@ def __init__(
         ssl_relax_minimum_version=False,
         username=None,
         password=None,
+        jwt_token=None,
         schema=None,
         pool_size=None,
         socket_keepalive=True,
@@ -473,6 +479,7 @@ def __init__(
         self._local = threading.local()
         self.username = username
         self.password = password
+        self.jwt_token = jwt_token
         self.schema = schema
 
         self.path = self.SQL_PATH
@@ -589,6 +596,7 @@ def _request(self, method, path, server=None, **kwargs):
                     path,
                     username=self.username,
                     password=self.password,
+                    jwt_token=self.jwt_token,
                     backoff_factor=self.backoff_factor,
                     schema=self.schema,
                     **kwargs,

tests/client/test_http.py

@@ -571,14 +571,20 @@ def do_POST(self):
         self.server.SHARED["schema"] = self.headers.get("Default-Schema")
 
         if self.headers.get("Authorization") is not None:
-            auth_header = self.headers["Authorization"].replace("Basic ", "")
-            credentials = b64decode(auth_header).decode("utf-8").split(":", 1)
-            self.server.SHARED["username"] = credentials[0]
-            if len(credentials) > 1 and credentials[1]:
-                self.server.SHARED["password"] = credentials[1]
-            else:
-                self.server.SHARED["password"] = None
+            auth_header = self.headers["Authorization"]
+            if "Basic" in auth_header:
+                auth_header = auth_header.replace("Basic ", "")
+                credentials = b64decode(auth_header).decode("utf-8").split(":", 1)
+                self.server.SHARED["username"] = credentials[0]
+                if len(credentials) > 1 and credentials[1]:
+                    self.server.SHARED["password"] = credentials[1]
+                else:
+                    self.server.SHARED["password"] = None
+            elif "Bearer" in auth_header:
+                jwt_token = auth_header.replace("Bearer ", "")
+                self.server.SHARED["jwt_token"] = jwt_token
         else:
+            self.server.SHARED["jwt_token"] = None
             self.server.SHARED["username"] = None
 
         if self.headers.get("X-User") is not None:
@@ -604,6 +610,7 @@ class TestingHTTPServer(HTTPServer):
     SHARED = manager.dict()
     SHARED["count"] = 0
     SHARED["usernameFromXUser"] = None
+    SHARED["jwt_token"] = None
     SHARED["username"] = None
     SHARED["password"] = None
     SHARED["schema"] = None
@@ -689,13 +696,15 @@ class TestUsernameSentAsHeader(TestingHttpServerTestCase):
     def setUp(self):
         super().setUp()
         self.clientWithoutUsername = self.clientWithKwargs()
+        self.clientWithJwtToken = self.clientWithKwargs(jwt_token="testJwtToken")
         self.clientWithUsername = self.clientWithKwargs(username="testDBUser")
         self.clientWithUsernameAndPassword = self.clientWithKwargs(
             username="testDBUser", password="test:password"
         )
 
     def tearDown(self):
         self.clientWithoutUsername.close()
+        self.clientWithJwtToken.close()
         self.clientWithUsername.close()
         self.clientWithUsernameAndPassword.close()
         super().tearDown()
@@ -720,6 +729,13 @@ def test_username(self):
         self.assertEqual(TestingHTTPServer.SHARED["username"], "testDBUser")
         self.assertEqual(TestingHTTPServer.SHARED["password"], "test:password")
 
+    def test_jwt_token(self):
+        self.clientWithoutUsername.sql("select * from fake")
+        self.assertEqual(TestingHTTPServer.SHARED["jwt_token"], None)
+
+        self.clientWithJwtToken.sql("select * from fake")
+        self.assertEqual(TestingHTTPServer.SHARED["jwt_token"], "testJwtToken")
+
 
 class TestCrateJsonEncoder(TestCase):
     def test_naive_datetime(self):