From 3f86f7d4d7bf1ab7be056aafd957573658f15112 Mon Sep 17 00:00:00 2001
From: Andreas Motl <andreas.motl@crate.io>
Date: Mon, 16 Dec 2024 23:58:11 +0100
Subject: [PATCH] Improve `ctk shell` to also talk to CrateDB standalone
 databases

---
 CHANGES.md                   |  1 +
 cratedb_toolkit/shell/cli.py | 22 ++++++++++++++++----
 tests/test_shell.py          | 39 ++++++++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+), 4 deletions(-)
 create mode 100644 tests/test_shell.py

diff --git a/CHANGES.md b/CHANGES.md
index 84d59a62..d8eeafd8 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -4,6 +4,7 @@
 - Dependencies: Minimize dependencies of core installation,
   defer `polars` to `cratedb-toolkit[io]`.
 - Fixed `cratedb-wtf record` about too large values of `ulimit_hard`
+- Improved `ctk shell` to also talk to CrateDB standalone databases
 
 ## 2024/10/13 v0.0.29
 - MongoDB: Added Zyp transformations to the CDC subsystem,
diff --git a/cratedb_toolkit/shell/cli.py b/cratedb_toolkit/shell/cli.py
index 0df97bbb..3324432e 100644
--- a/cratedb_toolkit/shell/cli.py
+++ b/cratedb_toolkit/shell/cli.py
@@ -1,15 +1,21 @@
 import click
 
 from cratedb_toolkit.cluster.util import get_cluster_info
+from cratedb_toolkit.model import DatabaseAddress
 from cratedb_toolkit.util.cli import boot_click
 from cratedb_toolkit.util.crash import get_crash_output_formats, run_crash
 
 output_formats = get_crash_output_formats()
 
+cratedb_sqlalchemy_option = click.option(
+    "--cratedb-sqlalchemy-url", envvar="CRATEDB_SQLALCHEMY_URL", type=str, required=False, help="CrateDB SQLAlchemy URL"
+)
+
 
 @click.command()
+@cratedb_sqlalchemy_option
 @click.option(
-    "--cluster-id", envvar="CRATEDB_CLOUD_CLUSTER_ID", type=str, required=True, help="CrateDB Cloud cluster identifier"
+    "--cluster-id", envvar="CRATEDB_CLOUD_CLUSTER_ID", type=str, required=False, help="CrateDB Cloud cluster identifier"
 )
 @click.option("--username", envvar="CRATEDB_USERNAME", type=str, required=False, help="Username for CrateDB cluster")
 @click.option("--password", envvar="CRATEDB_PASSWORD", type=str, required=False, help="Password for CrateDB cluster")
@@ -34,6 +40,7 @@
 @click.pass_context
 def cli(
     ctx: click.Context,
+    cratedb_sqlalchemy_url: str,
     cluster_id: str,
     username: str,
     password: str,
@@ -46,13 +53,20 @@ def cli(
     """
     Start an interactive database shell, or invoke SQL commands.
 
-    TODO: Only talks to CrateDB Cloud for now. Also implement for standalone CrateDB servers.
     TODO: Learn/forward more options of `crash`.
     """
     boot_click(ctx, verbose, debug)
 
-    cluster_info = get_cluster_info(cluster_id=cluster_id)
-    cratedb_http_url = cluster_info.cloud["url"]
+    if cratedb_sqlalchemy_url:
+        address = DatabaseAddress.from_string(cratedb_sqlalchemy_url)
+        cratedb_http_url = address.httpuri
+    elif cluster_id:
+        cluster_info = get_cluster_info(cluster_id=cluster_id)
+        cratedb_http_url = cluster_info.cloud["url"]
+    else:
+        raise ValueError(
+            "Unknown database address, please specify either cluster id or database URI in SQLAlchemy format"
+        )
 
     run_crash(
         hosts=cratedb_http_url,
diff --git a/tests/test_shell.py b/tests/test_shell.py
new file mode 100644
index 00000000..de649971
--- /dev/null
+++ b/tests/test_shell.py
@@ -0,0 +1,39 @@
+import json
+
+import pytest
+from click.testing import CliRunner
+
+from cratedb_toolkit.shell.cli import cli
+from tests.conftest import TESTDRIVE_DATA_SCHEMA
+
+
+def test_shell_success(cratedb):
+    """
+    Verify successful incantation of `ctk shell`.
+    """
+    runner = CliRunner()
+
+    database_url = cratedb.get_connection_url() + "?schema=" + TESTDRIVE_DATA_SCHEMA
+
+    result = runner.invoke(
+        cli,
+        args=f"--cratedb-sqlalchemy-url='{database_url}' --command 'SELECT 42 AS answer;' --format=json",
+        catch_exceptions=False,
+    )
+    assert result.exit_code == 0
+    assert json.loads(result.output) == [{"answer": 42}]
+
+
+def test_shell_failure_no_address():
+    """
+    Verify `ctk shell` fails when invoked without database address.
+    """
+    runner = CliRunner()
+
+    with pytest.raises(ValueError) as ex:
+        runner.invoke(
+            cli,
+            args="--command 'SELECT 42 AS answer;' --format=json",
+            catch_exceptions=False,
+        )
+    assert ex.match("Unknown database address")