From dfd535f1338563773b7344d550ef8114ea1b6249 Mon Sep 17 00:00:00 2001
From: Andreas Motl <andreas.motl@crate.io>
Date: Tue, 7 Nov 2023 17:41:57 +0100
Subject: [PATCH] Add `cratedb_toolkit.shell.run_sql` utility primitive

---
 CHANGES.md                        |  1 +
 cratedb_toolkit/shell/__init__.py |  1 +
 cratedb_toolkit/shell/crash.py    | 54 +++++++++++++++++++++++++++++++
 tests/test_shell.py               | 45 ++++++++++++++++++++++++++
 4 files changed, 101 insertions(+)
 create mode 100644 cratedb_toolkit/shell/__init__.py
 create mode 100644 cratedb_toolkit/shell/crash.py
 create mode 100644 tests/test_shell.py

diff --git a/CHANGES.md b/CHANGES.md
index e032df26..2338b248 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -3,6 +3,7 @@
 
 ## Unreleased
 
+- Add `cratedb_toolkit.shell.run_sql` utility primitive
 
 
 ## 2023/11/06 v0.0.2
diff --git a/cratedb_toolkit/shell/__init__.py b/cratedb_toolkit/shell/__init__.py
new file mode 100644
index 00000000..252b9e20
--- /dev/null
+++ b/cratedb_toolkit/shell/__init__.py
@@ -0,0 +1 @@
+from .crash import run_sql  # noqa: F401
diff --git a/cratedb_toolkit/shell/crash.py b/cratedb_toolkit/shell/crash.py
new file mode 100644
index 00000000..aa193237
--- /dev/null
+++ b/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
diff --git a/tests/test_shell.py b/tests/test_shell.py
new file mode 100644
index 00000000..2cacc92d
--- /dev/null
+++ b/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")