Add cratedb_toolkit.shell.run_sql utility primitive

CHANGES.md

@@ -3,6 +3,7 @@
 
 ## Unreleased
 
+- Add `cratedb_toolkit.shell.run_sql` utility primitive
 
 
 ## 2023/11/06 v0.0.2

cratedb_toolkit/shell/__init__.py

@@ -0,0 +1 @@
+from .crash import run_sql # noqa: F401

cratedb_toolkit/shell/crash.py

@@ -0,0 +1,54 @@
+import contextlib
+import io
+import json
+import sys
+import typing as t
+from pathlib import Path
+from unittest import mock
+
+
+def run_sql(statement: t.Union[str, Path, io.IOBase], hosts: str = None, schema: str = None):
+ """
+ Run SQL from string or file, using `crash`.
+
+ TODO: Validate it works well in different scenarios, both on CrateDB and CrateDB Cloud.
+ TODO: Returning stderr/logging from crash does not work yet.
+ """
+ import crate.crash.command
+
+ sys.argv = ["crash"]
+ if hosts:
+ sys.argv += ["--hosts", hosts]
+ if schema:
+ sys.argv += ["--schema", schema]
+ if isinstance(statement, str):
+ sys.argv += ["--command", statement]
+ elif isinstance(statement, Path):
+ sys.stdin = io.StringIO(statement.read_text())
+ elif isinstance(statement, io.IOBase):
+ sys.stdin = statement # type: ignore[assignment]
+ else:
+ raise ValueError("Either statement or filepath must be given")
+
+ sys.argv += ["--format", "json"]
+
+ # Temporarily patch some shortcomings of `crash`, when used programmatically.
+ # TODO: See what can be done over in `crash` in a later iteration.
+ with mock.patch("crate.crash.repl.SQLCompleter._populate_keywords"), mock.patch(
+ "crate.crash.command.CrateShell.close"
+ ):
+ buffer_out = io.StringIO()
+ buffer_err = io.StringIO()
+ with contextlib.redirect_stdout(buffer_out), contextlib.redirect_stderr(buffer_err):
+ # Invoke `crash`, to execute the SQL statement.
+ try:
+ crate.crash.command.main()
+ except SystemExit as ex:
+ if ex.code != 0:
+ raise
+
+ buffer_out.seek(0)
+ buffer_err.seek(0)
+ out = buffer_out.read()
+ err = buffer_err.read()
+ return json.loads(out), err

tests/test_shell.py

@@ -0,0 +1,45 @@
+import io
+
+import pytest
+
+from cratedb_toolkit.shell import run_sql
+
+
+def test_run_sql_from_string():
+ sql = "SELECT 1;"
+ data, messages = run_sql(sql)
+ assert data == [{"1": 1}]
+
+ # TODO: Returning stderr/logging from crash does not work yet.
+ assert messages == ""
+
+
+def test_run_sql_from_file(tmp_path):
+ sql_file = tmp_path / "temp.sql"
+ sql_file.write_text("SELECT 1;")
+ data, messages = run_sql(sql_file)
+ assert data == [{"1": 1}]
+
+ # TODO: Returning stderr/logging from crash does not work yet.
+ assert messages == ""
+
+
+def test_run_sql_from_buffer():
+ sql_buffer = io.StringIO("SELECT 1;")
+ data, messages = run_sql(sql_buffer)
+ assert data == [{"1": 1}]
+
+ # TODO: Returning stderr/logging from crash does not work yet.
+ assert messages == ""
+
+
+def test_run_sql_invalid_host(capsys):
+ sql = "SELECT 1;"
+ with pytest.raises(SystemExit) as ex:
+ run_sql(sql, hosts="localhost:12345")
+
+ # TODO: Returning stderr/logging from crash does not work yet.
+ out, err = capsys.readouterr()
+ assert out == ""
+ assert err == ""
+ assert ex.match("1")