Cluster API: Naming things and refactoring

Author: amotl

cratedb_toolkit/cli.py

@@ -10,7 +10,6 @@
 from .docs.cli import cli as docs_cli
 from .info.cli import cli as info_cli
 from .io.cli import cli as io_cli
-from .job.cli import cli_list_jobs
 from .query.cli import cli as query_cli
 from .settings.cli import cli as settings_cli
 from .shell.cli import cli as shell_cli
@@ -37,4 +36,3 @@ def cli(ctx: click.Context, verbose: bool, debug: bool):
 cli.add_command(shell_cli, name="shell")
 cli.add_command(settings_cli, name="settings")
 cli.add_command(tail_cli, name="tail")
-cli.add_command(cli_list_jobs)

cratedb_toolkit/cluster/cli.py

@@ -1,7 +1,10 @@
 import contextlib
+import json
 import logging
+import sys
 
 import click
+import yaml
 from click import ClickException
 from click_aliases import ClickAliasedGroup
 
@@ -45,6 +48,15 @@ def help_stop():
     """
 
 
+def help_list_jobs():
+    """
+    List jobs on cluster.
+
+    ctk cluster list-jobs
+    croud clusters import-jobs
+    """
+
+
 @click.group(cls=ClickAliasedGroup)  # type: ignore[arg-type]
 @click.option("--verbose", is_flag=True, required=False, help="Turn on logging")
 @click.option("--debug", is_flag=True, required=False, help="Turn on logging with debug level")
@@ -106,6 +118,39 @@ def stop(ctx: click.Context, cluster_id: str, cluster_name: str):
         jd(cluster.info.asdict())
 
 
+@make_command(cli, name="list-jobs", help=help_list_jobs)
+@option_cluster_id
+@option_cluster_name
+@click.option(
+    "--format",
+    "output_format",
+    type=click.Choice(["json", "yaml"], case_sensitive=False),
+    default="json",
+    required=False,
+    help="The output format for the result of the operation",
+)
+@click.pass_context
+def list_jobs(ctx: click.Context, cluster_id: str, cluster_name: str, output_format: str):
+    """
+    List jobs on cluster.
+    """
+
+    with handle_command_errors("list cluster jobs"):
+        # Acquire the database cluster handle and acquire job information.
+        cluster = ManagedCluster(cluster_id=cluster_id, cluster_name=cluster_name).probe()
+        if not cluster.operation:
+            raise CroudException("Cluster does not support job operations")
+        data = cluster.operation.list_jobs()
+
+        # Display job information.
+        if output_format == "json":
+            print(json.dumps(data, indent=2), file=sys.stdout)  # noqa: T201
+        elif output_format == "yaml":
+            print(yaml.dump(data), file=sys.stdout)  # noqa: T201
+        else:
+            raise ValueError(f"Unknown output format: {output_format}")
+
+
 @contextlib.contextmanager
 def handle_command_errors(operation_name):
     """Handle common command errors and exit with appropriate error messages."""

cratedb_toolkit/cluster/core.py

@@ -10,7 +10,7 @@
 import click
 from boltons.urlutils import URL
 
-from cratedb_toolkit.cluster.croud import CloudManager
+from cratedb_toolkit.cluster.croud import CloudClusterServices, CloudRootServices
 from cratedb_toolkit.cluster.guide import DataImportGuide
 from cratedb_toolkit.cluster.model import ClientBundle, ClusterBase, ClusterInformation
 from cratedb_toolkit.config import CONFIG
@@ -132,7 +132,8 @@ def __init__(
                 "Failed to address cluster: Either cluster identifier or name needs to be specified"
             )
 
-        self.cm = CloudManager()
+        self.root = CloudRootServices()
+        self.operation: t.Optional[CloudClusterServices] = None
         self._jwt_ctx: t.ContextManager = nullcontext()
         self._client_bundle: t.Optional[ClientBundle] = None
 
@@ -148,15 +149,18 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
         try:
             self.close_connections()
-            if self.stop_on_exit:
-                self.stop()
-                logger.info(f"Successfully stopped cluster: id={self.cluster_id}, name={self.cluster_name}")
         except Exception as ex:
-            logger.error(f"Failed to stop cluster: {ex}")
-            # Don't swallow the original exception
-            return False
+            logger.error(f"Failed to close connections: {ex}")
+        finally:
+            if self.stop_on_exit:
+                try:
+                    self.stop()
+                    logger.info(f"Successfully stopped cluster: id={self.cluster_id}, name={self.cluster_name}")
+                except Exception as ex:
+                    logger.error(f"Failed to stop cluster: {ex}")
 
-        return False  # Don't suppress any exceptions
+        # Don't suppress the original exception.
+        return False
 
     @classmethod
     @flexfun(domain="settings")
@@ -202,13 +206,19 @@ def probe(self) -> "ManagedCluster":
         Probe a CrateDB Cloud cluster, API-wise.
 
         TODO: Investigate callers, and reduce number of invocations.
+        TODO: self._jwt_ctx is created once in probe() and reused for every query.
+              If the per-cluster JWT expires (default 1 h) subsequent queries will fail.
+              Consider refreshing self._jwt_ctx (or simply calling jwt_token_patch()
+              inside query()) when probe() is older than e.g. 30 min.
         """
         try:
             self.info = ClusterInformation.from_id_or_name(cluster_id=self.cluster_id, cluster_name=self.cluster_name)
             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)
+            if self.cluster_id:
+                self.operation = CloudClusterServices(cluster_id=self.cluster_id)
 
         except (CroudException, DatabaseAddressMissingError) as ex:
             self.exists = False
@@ -269,14 +279,13 @@ def deploy(self) -> "ManagedCluster":
 
         Command: ctk cluster start
         """
-        # FIXME: Accept id or name.
         if self.cluster_name is None:
             raise DatabaseAddressMissingError("Need cluster name to deploy")
 
         # Find the existing project by name (equals cluster name).
         project_id = None
         try:
-            projects = self.cm.list_projects()
+            projects = self.root.list_projects()
             for project in projects:
                 if project["name"] == self.cluster_name:
                     project_id = project["id"]
@@ -287,11 +296,11 @@ def deploy(self) -> "ManagedCluster":
 
         # Create a new project if none exists.
         if not project_id:
-            project = self.cm.create_project(name=self.cluster_name, organization_id=self.settings.organization_id)
+            project = self.root.create_project(name=self.cluster_name, organization_id=self.settings.organization_id)
             project_id = project["id"]
             logger.info(f"Created project: {project_id}")
 
-        cluster_info = self.cm.deploy_cluster(
+        cluster_info = self.root.deploy_cluster(
             name=self.cluster_name, project_id=project_id, subscription_id=self.settings.subscription_id
         )
 
@@ -309,7 +318,8 @@ def resume(self) -> "ManagedCluster":
         if self.cluster_id is None:
             raise DatabaseAddressMissingError("Need cluster identifier to resume cluster")
         logger.info(f"Resuming CrateDB Cloud Cluster: id={self.cluster_id}, name={self.cluster_name}")
-        self.cm.resume_cluster(identifier=self.cluster_id)
+        if self.operation:
+            self.operation.resume()
         self.probe()
         return self
 
@@ -322,7 +332,8 @@ def stop(self) -> "ManagedCluster":
         if self.cluster_id is None:
             raise DatabaseAddressMissingError("Need cluster identifier to stop cluster")
         logger.info(f"Stopping CrateDB Cloud Cluster: id={self.cluster_id}, name={self.cluster_name}")
-        self.cm.suspend_cluster(identifier=self.cluster_id)
+        if self.operation:
+            self.operation.suspend()
         self.probe()
         return self
 
@@ -348,7 +359,6 @@ def load_table(
         self.probe()
         target = target or TableAddress()
 
-        # FIXME: Accept id or name.
         if self.cluster_id is None:
             raise DatabaseAddressMissingError("Need cluster identifier to load table")