Cluster API: Address suggestions from code review by CodeRabbit

Author: amotl

cratedb_toolkit/cluster/croud.py

@@ -385,9 +385,8 @@ def get_jwt_token(self) -> t.Dict[str, str]:
         """
         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)
+            if errors:
+                raise IOError(f"Getting JWT token failed: {errors}")
+            raise IOError("Getting JWT token failed: Unknown error")
         return data

cratedb_toolkit/cluster/model.py

@@ -114,14 +114,9 @@ def close(self):
         """
 
         try:
-            self.adapter.connection.close()
+            self.adapter.close()
         except Exception as e:
-            logger.warning(f"Error closing adapter connection: {e}")
-
-        try:
-            self.adapter.engine.dispose()
-        except Exception as e:
-            logger.warning(f"Error disposing SQLAlchemy engine: {e}")
+            logger.warning(f"Error closing database adapter: {e}")
 
         try:
             self.dbapi.close()

cratedb_toolkit/shell/cli.py

@@ -31,7 +31,7 @@ def help_cli():
     - https://cratedb-toolkit.readthedocs.io/util/shell.html
     - https://cratedb.com/docs/crate/crash/
 
-    TODO: Learn/forward more options of `crash`.
+    TODO: Implement/forward more options of `crash`.
     """
 
 
@@ -84,19 +84,19 @@ def cli(
     if cluster.address is None:
         raise click.UsageError("Inquiring cluster address failed.")
 
-    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
+            jwt_token_obj = getattr(cluster.info, "jwt", None)
+            if jwt_token_obj:
+                jwt_token = jwt_token_obj.token
 
     run_crash(
-        hosts=http_url,
+        hosts=cluster.address.httpuri,
         username=username,
         password=password,
         jwt_token=jwt_token,

cratedb_toolkit/util/client.py

@@ -1,23 +1,24 @@
 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)):
+    original_request = crate.client.http.Client._request
+    crate.client.http.Client._request = _mk_crate_client_request(original_request, jwt_token)
+    try:
         yield
+    finally:
+        crate.client.http.Client._request = original_request
 
 
-def _mk_crate_client_server_request(jwt_token: str = None):
+def _mk_crate_client_request(orig, jwt_token: str = None):
     """
-    Create a monkey patched Server.request method to add the Authorization header for JWT token-based authentication.
+    Create a monkey patched Client._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):
+    def _crate_client_request(self, *args, **kwargs):
         """
         Monkey patch the Server.request method to add the Authorization header for JWT token-based authentication.
 
@@ -26,6 +27,6 @@ def _crate_client_server_request(self, *args, **kwargs):
         if jwt_token:
             kwargs.setdefault("headers", {})
             kwargs["headers"].update({"Authorization": "Bearer " + jwt_token})
-        return _crate_client_server_request_dist(self, *args, **kwargs)
+        return orig(self, *args, **kwargs)
 
-    return _crate_client_server_request
+    return _crate_client_request

cratedb_toolkit/util/database.py

@@ -123,13 +123,18 @@ def run_sql_real(self, sql: str, parameters: t.Mapping[str, str] = None, records
         for statement in sqlparse.split(sql):
             if self.internal:
                 statement += self.internal_tag
+            # FIXME: Persistent self.connection risks leaks & thread-unsafety.
+            #        https://github.com/crate/cratedb-toolkit/pull/81#discussion_r2071499204
             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]
+            if result.returns_rows:
+                if records:
+                    rows = result.mappings().fetchall()
+                    data = [dict(row.items()) for row in rows]
+                else:
+                    data = result.fetchall()
             else:
-                data = result.fetchall()
+                data = None
             results.append(data)
 
         # Backward-compatibility.
@@ -387,6 +392,25 @@ def describe_table_columns(self, table_name: str):
         table_address = TableAddress.from_string(table_name)
         return inspector.get_columns(table_name=t.cast(str, table_address.table), schema=table_address.schema)
 
+    def close(self):
+        """
+        Close all database connections created to this cluster.
+        Should be called when the cluster handle is no longer needed.
+        """
+
+        try:
+            self.connection.close()
+        except Exception as e:
+            logger.warning(f"Error closing adapter connection: {e}")
+
+        try:
+            self.engine.dispose()
+        except Exception as e:
+            logger.warning(f"Error disposing SQLAlchemy engine: {e}")
+
+    def __del__(self):
+        self.close()
+
 
 def sa_is_empty(thing):
     """

doc/util/shell.md

@@ -76,12 +76,8 @@ ctk shell --command "SELECT 42;"
 ```
 
 
-Include authentication credentials into the
-SQLAlchemy or HTTP connection URLs.
-
-
 :::{seealso}
-{ref}`cluster-api-tutorial` includes a full end-to-end tutorial which also includes
+{ref}`cluster-api-tutorial` demonstrates a full end-to-end tutorial, which also includes
 `ctk shell`.
 :::
 

examples/python/cloud_cluster.py

@@ -95,6 +95,8 @@ def workload_procedural():
     pprint(results, stream=sys.stderr)  # noqa: T201
 
     # Stop the cluster again.
+    # Note: We intentionally don't stop the cluster here to prevent accidental
+    # shutdown of resources that might be in use for demonstration purposes.
     # cluster.stop()
 
 

tests/cluster/test_cli.py

@@ -3,7 +3,7 @@
 from cratedb_toolkit.cluster.cli import cli
 
 
-def test_managed_cluster_info_empty(cloud_environment):
+def test_managed_cluster_info_default(cloud_environment):
     """
     Verify `ctk cluster info` on a managed cluster works when a valid environment is provided.
     """

tests/conftest.py

@@ -89,7 +89,10 @@ def prune_environment():
 @pytest.fixture(scope="session", autouse=True)
 def reset_environment():
     """
-    Reset all environment variables in use, so that they do not pollute the test suite.
+    Reset environment variables defined in ManagedClusterSettings.
+
+    This complements the `prune_environment` fixture by specifically targeting
+    variables from ManagedClusterSettings rather than all CRATEDB_* variables.
 
     TODO: Possibly synchronize with `prune_environment()`, as suggested.
     """

tests/shell/test_cli.py

@@ -54,7 +54,7 @@ def test_shell_managed_jwt(mocker, cloud_cluster_name):
 
 def test_shell_managed_username_password(mocker, cloud_cluster_name):
     """
-    Verify the successful incantation of `ctk shell` against CrateDB Cloud, using username/password authentication..
+    Verify the successful incantation of `ctk shell` against CrateDB Cloud, using username/password authentication.
     """
 
     settings = {