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")