Auth: Add JWT authentication to client API and ctk shell

Author: amotl

CHANGES.md

@@ -9,6 +9,7 @@
   deploy/start/resume and data import procedures using fluent API and CLI.
 - CLI naming things: Rename `--cratedb-sqlalchemy-url` to `--sqlalchemy-url`
   and `--cratedb-http-url` to `--http-url`.
+- Auth: Added JWT authentication to client API and `ctk shell`.
 
 ## 2025/04/23 v0.0.32
 - MCP: Add subsystem providing a few server and client utilities through

cratedb_toolkit/cluster/core.py

@@ -3,6 +3,7 @@
 import logging
 import time
 import typing as t
+from contextlib import nullcontext
 from copy import deepcopy
 from pathlib import Path
 
@@ -20,6 +21,7 @@
     OperationFailed,
 )
 from cratedb_toolkit.model import DatabaseAddress, InputOutputResource, TableAddress
+from cratedb_toolkit.util.client import jwt_token_patch
 from cratedb_toolkit.util.data import asbool
 from cratedb_toolkit.util.database import DatabaseAdapter
 from cratedb_toolkit.util.runtime import flexfun
@@ -131,6 +133,8 @@ def __init__(
             )
 
         self.cm = CloudManager()
+        self._jwt_ctx: t.ContextManager = nullcontext()
+        self._client_bundle: t.Optional[ClientBundle] = None
 
     def __enter__(self):
         """Enter the context manager, ensuring the cluster is running."""
@@ -204,6 +208,7 @@ def probe(self) -> "ManagedCluster":
             self.cluster_id = self.info.cloud["id"]
             self.cluster_name = self.info.cloud["name"]
             self.address = DatabaseAddress.from_httpuri(self.info.cloud["url"])
+            self._jwt_ctx = jwt_token_patch(self.info.jwt.token)
 
         except (CroudException, DatabaseAddressMissingError) as ex:
             self.exists = False
@@ -415,8 +420,9 @@ def query(self, sql: str):
         # Ensure we have cluster connection details.
         if not self.info or not self.info.cloud.get("url"):
             self.probe()
-        client_bundle = self.get_client_bundle()
-        return client_bundle.adapter.run_sql(sql, records=True)
+        with self._jwt_ctx:
+            client_bundle = self.get_client_bundle()
+            return client_bundle.adapter.run_sql(sql, records=True)
 
 
 @dataclasses.dataclass
@@ -429,6 +435,7 @@ class StandaloneCluster(ClusterBase):
     info: t.Optional[ClusterInformation] = None
     exists: bool = False
     _load_table_result: t.Optional[bool] = None
+    _client_bundle: t.Optional[ClientBundle] = None
 
     def __post_init__(self):
         super().__init__()

cratedb_toolkit/cluster/croud.py

@@ -1,17 +1,21 @@
 import json
+import logging
 import os
 import typing as t
 from pathlib import Path
 
 from cratedb_toolkit.exception import CroudException
 from cratedb_toolkit.model import InputOutputResource, TableAddress
-from cratedb_toolkit.util.croud import CroudCall, CroudWrapper
+from cratedb_toolkit.util.croud import CroudCall, CroudClient, CroudWrapper
 
 # Default to a stable version if not specified in the environment.
 # TODO: Use `latest` CrateDB by default, or even `nightly`?
 DEFAULT_CRATEDB_VERSION = "5.10.4"
 
 
+logger = logging.getLogger(__name__)
+
+
 class CloudManager:
     """
     A wrapper around the CrateDB Cloud API through the `croud` package, providing common methods.
@@ -374,3 +378,16 @@ def create_import_job(self, resource: InputOutputResource, target: TableAddress)
 
         wr = CroudWrapper(call=call)
         return wr.invoke()
+
+    def get_jwt_token(self) -> t.Dict[str, str]:
+        """
+        Retrieve per-cluster JWT token.
+        """
+        client = CroudClient.create()
+        data, errors = client.get(f"/api/v2/clusters/{self.cluster_id}/jwt/")
+        errmsg = "Getting JWT token failed: Unknown error"
+        if errors:
+            errmsg = f"Getting JWT token failed: {errors}"
+        if data is None:
+            raise IOError(errmsg)
+        return data

cratedb_toolkit/cluster/model.py

@@ -16,6 +16,18 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclasses.dataclass
+class JwtResponse:
+    expiry: str
+    refresh: str
+    token: str
+
+    def get_token(self):
+        # TODO: Persist token across sessions.
+        # TODO: Refresh automatically when expired.
+        return self.token
+
+
 @dataclasses.dataclass
 class ClusterInformation:
     """
@@ -76,6 +88,14 @@ def from_name(cls, cluster_name: str) -> "ClusterInformation":
     def asdict(self) -> t.Dict[str, t.Any]:
         return deepcopy(dataclasses.asdict(self))
 
+    @property
+    def jwt(self) -> JwtResponse:
+        """
+        Return per-cluster JWT token response.
+        """
+        cc = CloudCluster(cluster_id=self.cloud_id)
+        return JwtResponse(**cc.get_jwt_token())
+
 
 @dataclasses.dataclass
 class ClientBundle:

cratedb_toolkit/shell/cli.py

@@ -1,3 +1,5 @@
+import logging
+
 import click
 
 from cratedb_toolkit import DatabaseCluster
@@ -13,6 +15,8 @@
 from cratedb_toolkit.util.cli import boot_click, docstring_format_verbatim
 from cratedb_toolkit.util.crash import get_crash_output_formats, run_crash
 
+logger = logging.getLogger(__name__)
+
 output_formats = get_crash_output_formats()
 
 
@@ -39,7 +43,7 @@ def help_cli():
 @option_username
 @option_password
 @option_schema
-@click.option("--command", type=str, required=False, help="SQL command")
+@click.option("--command", "-c", type=str, required=False, help="SQL command")
 @click.option(
     "--format",
     "format_",
@@ -82,10 +86,20 @@ def cli(
 
     http_url = cluster.address.httpuri
 
+    is_cloud = cluster_id is not None or cluster_name is not None
+    jwt_token = None
+    if is_cloud:
+        if username is not None:
+            logger.info("Using username/password credentials for authentication")
+        else:
+            logger.info("Using JWT token for authentication")
+            jwt_token = cluster.info.jwt.token
+
     run_crash(
         hosts=http_url,
         username=username,
         password=password,
+        jwt_token=jwt_token,
         schema=schema,
         command=command,
         output_format=format_,

cratedb_toolkit/util/client.py

@@ -0,0 +1,31 @@
+import contextlib
+from unittest import mock
+
+import crate.client.http
+
+
+@contextlib.contextmanager
+def jwt_token_patch(jwt_token: str = None):
+    with mock.patch.object(crate.client.http.Client, "_request", _mk_crate_client_server_request(jwt_token)):
+        yield
+
+
+def _mk_crate_client_server_request(jwt_token: str = None):
+    """
+    Create a monkey patched Server.request method to add the Authorization header for JWT token-based authentication.
+    """
+
+    _crate_client_server_request_dist = crate.client.http.Client._request
+
+    def _crate_client_server_request(self, *args, **kwargs):
+        """
+        Monkey patch the Server.request method to add the Authorization header for JWT token-based authentication.
+
+        TODO: Submit to upstream libraries and programs `crate.client`, `crate.crash`, and `sqlalchemy-cratedb`.
+        """
+        if jwt_token:
+            kwargs.setdefault("headers", {})
+            kwargs["headers"].update({"Authorization": "Bearer " + jwt_token})
+        return _crate_client_server_request_dist(self, *args, **kwargs)
+
+    return _crate_client_server_request

cratedb_toolkit/util/crash.py

@@ -6,9 +6,17 @@
 
 from crate.crash.command import main
 
+from cratedb_toolkit.util.client import jwt_token_patch
+
 
 def run_crash(
-    hosts: str, command: str, output_format: str = None, schema: str = None, username: str = None, password: str = None
+    hosts: str,
+    command: str,
+    output_format: str = None,
+    schema: str = None,
+    username: str = None,
+    password: str = None,
+    jwt_token: str = None,
 ):
     """
     Run the interactive CrateDB database shell using `crash`.
@@ -23,12 +31,11 @@ def run_crash(
         cmd += ["--command", command]
     if output_format:
         cmd += ["--format", output_format]
-    with mock.patch.object(sys, "argv", cmd):
-        password_context: contextlib.AbstractContextManager = contextlib.nullcontext()
-        if password:
-            password_context = mock.patch.dict(os.environ, {"CRATEPW": password})
-        with password_context:
-            main()
+    password_context: contextlib.AbstractContextManager = contextlib.nullcontext()
+    if password:
+        password_context = mock.patch.dict(os.environ, {"CRATEPW": password})
+    with mock.patch.object(sys, "argv", cmd), jwt_token_patch(jwt_token=jwt_token), password_context:
+        main()
 
 
 def get_crash_output_formats() -> t.List[str]:

cratedb_toolkit/util/croud.py

@@ -14,6 +14,7 @@
 
 import croud.api
 import yaml
+from boltons.typeutils import classproperty
 from croud.config.configuration import Configuration
 
 import cratedb_toolkit
@@ -151,24 +152,6 @@ def invoke_capturing(self, fun: t.Callable, *args: t.List[t.Any], **kwargs: t.Di
         buffer.seek(0)
         return buffer.read()
 
-    @property
-    def headless_config(self) -> Configuration:
-        cfg = Configuration("croud.yaml")
-        if cfg._file_path.exists() and "CRATEDB_CLOUD_API_KEY" not in os.environ:
-            return cfg
-
-        tmp_file = NamedTemporaryFile()
-        tmp_path = Path(tmp_file.name)
-        config = Configuration("headless.yaml", tmp_path)
-
-        # Get credentials from the environment.
-        config.profile["key"] = os.environ.get("CRATEDB_CLOUD_API_KEY")
-        config.profile["secret"] = os.environ.get("CRATEDB_CLOUD_API_SECRET")
-        config.profile["organization-id"] = os.environ.get("CRATEDB_CLOUD_ORGANIZATION_ID")
-        # config.profile["endpoint"] = os.environ.get("CRATEDB_CLOUD_ENDPOINT")  # noqa: ERA001
-
-        return config
-
     def run_croud_fun(self, fun: t.Callable, with_exceptions: bool = True):
         """
         Wrapper function to call into `croud`, for catching and converging error messages.
@@ -200,14 +183,13 @@ def print_fun(levelname: str, *args, **kwargs):
         # https://stackoverflow.com/a/46481946
         levels = ["debug", "info", "warning", "error", "success"]
         with contextlib.ExitStack() as stack:
-            # Patch all `print_*` functions.
+            # Patch all `print_*` functions, as they would obstruct the output.
             for level in levels:
                 p = patch(f"croud.printer.print_{level}", functools.partial(print_fun, level))
                 stack.enter_context(p)
 
             # Patch configuration.
-            p = patch("croud.config._CONFIG", self.headless_config)
-            stack.enter_context(p)
+            stack.enter_context(headless_config())
 
             # TODO: When aiming to disable wait-for-completion.
             """
@@ -219,6 +201,15 @@ def print_fun(levelname: str, *args, **kwargs):
             return fun()
 
 
+@contextlib.contextmanager
+def headless_config():
+    """
+    Patch the `croud.config` module to use a headless configuration.
+    """
+    with patch("croud.config._CONFIG", CroudClient.get_headless_config):
+        yield
+
+
 class CroudClient(croud.api.Client):
     """
     A slightly modified `croud.api.Client` class, to inject a custom User-Agent header.
@@ -229,6 +220,42 @@ def __init__(self, *args, **kwargs):
         ua = f"{cratedb_toolkit.__appname__}/{cratedb_toolkit.__version__} Python/{python_version()}"
         self.session.headers["User-Agent"] = ua
 
+    @staticmethod
+    def create() -> "croud.api.Client":
+        """
+        Canonical factory method for creating a `croud.api.Client` instance.
+        """
+        with headless_config():
+            from croud.config import CONFIG
+
+            return croud.api.Client(
+                CONFIG.endpoint,
+                token=CONFIG.token,
+                on_token=CONFIG.set_current_auth_token,
+                key=CONFIG.key,
+                secret=CONFIG.secret,
+                region=CONFIG.region,
+                sudo=False,
+            )
+
+    @classproperty
+    def get_headless_config(cls) -> Configuration:
+        cfg = Configuration("croud.yaml")
+        if cfg._file_path.exists() and "CRATEDB_CLOUD_API_KEY" not in os.environ:
+            return cfg
+
+        tmp_file = NamedTemporaryFile()
+        tmp_path = Path(tmp_file.name)
+        config = Configuration("headless.yaml", tmp_path)
+
+        # Get credentials from the environment.
+        config.profile["key"] = os.environ.get("CRATEDB_CLOUD_API_KEY")
+        config.profile["secret"] = os.environ.get("CRATEDB_CLOUD_API_SECRET")
+        config.profile["organization-id"] = os.environ.get("CRATEDB_CLOUD_ORGANIZATION_ID")
+        # config.profile["endpoint"] = os.environ.get("CRATEDB_CLOUD_ENDPOINT")  # noqa: ERA001
+
+        return config
+
 
 croud.api.Client = CroudClient
 

cratedb_toolkit/util/database.py

@@ -120,18 +120,17 @@ def run_sql_real(self, sql: str, parameters: t.Mapping[str, str] = None, records
         Invoke SQL statement, and return results.
         """
         results = []
-        with self.engine.connect() as connection:
-            for statement in sqlparse.split(sql):
-                if self.internal:
-                    statement += self.internal_tag
-                result = connection.execute(sa.text(statement), parameters)
-                data: t.Any
-                if records:
-                    rows = result.mappings().fetchall()
-                    data = [dict(row.items()) for row in rows]
-                else:
-                    data = result.fetchall()
-                results.append(data)
+        for statement in sqlparse.split(sql):
+            if self.internal:
+                statement += self.internal_tag
+            result = self.connection.execute(sa.text(statement), parameters)
+            data: t.Any
+            if records:
+                rows = result.mappings().fetchall()
+                data = [dict(row.items()) for row in rows]
+            else:
+                data = result.fetchall()
+            results.append(data)
 
         # Backward-compatibility.
         if len(results) == 1:

doc/cluster/tutorial.md

@@ -21,10 +21,6 @@ CRATEDB_CLOUD_API_SECRET='<YOUR_API_SECRET_HERE>'
 # CrateDB Cloud cluster identifier (id or name).
 # CRATEDB_CLUSTER_ID='<YOUR_CLUSTER_ID_HERE>'
 CRATEDB_CLUSTER_NAME='<YOUR_CLUSTER_NAME_HERE>'
-
-# Database credentials.
-CRATEDB_USERNAME='<YOUR_DATABASE_USERNAME_HERE>'
-CRATEDB_PASSWORD='<YOUR_DATABASE_PASSWORD_HERE>'
 EOF
 ```