diff --git a/.github/workflows/MainDistributionPipeline.yml b/.github/workflows/MainDistributionPipeline.yml index eb26aac..e040845 100644 --- a/.github/workflows/MainDistributionPipeline.yml +++ b/.github/workflows/MainDistributionPipeline.yml @@ -11,11 +11,19 @@ concurrency: group: ${{ github.workflow }}-${{ github.ref }}-${{ github.head_ref || '' }}-${{ github.base_ref || '' }}-${{ github.ref != 'refs/heads/main' && github.sha || '' }} cancel-in-progress: true +env: + # Single source of truth for the DuckDB/extension-ci-tools version. + # The `uses:` refs below still pin the tag literally (expressions are not + # allowed there), so bump those alongside this. + DUCKDB_VERSION: v1.5.6 + jobs: duckdb-stable-build: name: Build extension binaries uses: duckdb/extension-ci-tools/.github/workflows/_extension_distribution.yml@v1.5.6 with: + # env context is not allowed in reusable-workflow inputs — keep in sync + # with DUCKDB_VERSION above duckdb_version: v1.5.6 ci_tools_version: v1.5.6 extension_name: ggsql @@ -26,7 +34,32 @@ jobs: name: Code Quality Check uses: duckdb/extension-ci-tools/.github/workflows/_extension_code_quality.yml@v1.5.6 with: + # env context is not allowed in reusable-workflow inputs — keep in sync + # with DUCKDB_VERSION above duckdb_version: v1.5.6 ci_tools_version: v1.5.6 extension_name: ggsql format_checks: 'format;tidy' + + interrupt-tests: + name: Embedded connection interruption tests + needs: duckdb-stable-build + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + - uses: actions/download-artifact@v4 + with: + name: ggsql-${{ env.DUCKDB_VERSION }}-extension-linux_amd64 + path: build/interrupt-tests + - name: Install test dependencies + # Strip the leading "v" for the pip version specifier + run: python -m pip install "duckdb==${DUCKDB_VERSION#v}" numpy + - name: Test interruption through the Python client + env: + GGSQL_EXTENSION_PATH: ${{ github.workspace }}/build/interrupt-tests/ggsql.duckdb_extension + GGSQL_NO_OPEN_BROWSER: '1' + run: python -m unittest discover -s test/python -v diff --git a/CLAUDE.md b/CLAUDE.md index 00ebd7a..19bf7eb 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -30,6 +30,7 @@ duckdb/, extension-ci-tools/ Submodules; versions bumped per release (see docs/ - CMake drives `cargo build --release` via a custom command; Rust sources + `assets/*` are listed as dependencies so edits trigger rebuilds. - macOS links `CoreFoundation`, `Security`, `SystemConfiguration` (needed by Rust std / tiny_http TLS bits). - DuckDB target version is set in three places — keep them in sync: `.github/workflows/MainDistributionPipeline.yml` (`duckdb_version`, `ci_tools_version`, workflow `@` tag) and the `duckdb` + `extension-ci-tools` submodule refs. `docs/UPDATING.md` has the full checklist. +- The workflow's embedded interruption tests must use the matching Python DuckDB version and distribution artifact name when bumping DuckDB. ## User-visible surface @@ -61,6 +62,15 @@ Every ggsql call opens a **sibling `Connection` on the same `DatabaseInstance`** The inner `Connection` is created lazily and **persists for the whole `ggsql_execute` call**, so temp tables created by one `exec_sql` call (e.g. ggsql's CTE materialisation, which now goes through `execute_sql(create_or_replace_temp_table_sql(...))`) remain visible to subsequent `exec_sql` calls within the same invocation. +`BridgeCtx` also owns an interrupt-forwarding thread, started with the inner connection. +It polls the outer context every 10 ms and calls the inner connection's `Interrupt()`; +these operations only access DuckDB's atomic interrupt flags, never the context locks. +Keep forwarding after cancellation because each new inner query resets its flag. +The forwarder must stop and join before the connection is destroyed. SQL callbacks +and Arrow stream callbacks also check cancellation, and `RunGgsqlQuery` restores the +DuckDB `InterruptException` type after freeing the Rust result buffer. Regression +tests in `test/python/test_interrupt.py` exercise cancellation from another thread. + ## HTTP server - `tiny_http` bound to `127.0.0.1:0` (random port), stored in a `OnceCell` — one server for the process lifetime. diff --git a/README.md b/README.md index ef2ec5b..a7d6e4a 100644 --- a/README.md +++ b/README.md @@ -131,6 +131,19 @@ SELECT * FROM flights VISUALISE dep_delay, arr_delay DRAW point; -- ✅ The reason is structural: calling `Query` back into the outer `ClientContext` from inside an executing table function deadlocks on the context's mutex, so we open a sibling `Connection` — which by DuckDB's design has its own temp catalog. +### Query interruption + +Interrupting the calling connection (for example, Python's `connection.interrupt()`) +also interrupts SQL running on ggsql's sibling connection. This applies to both +`SELECT ggsql('...')` and direct `VISUALISE` statements. Cancellation is reported as +a DuckDB interruption error, and the connection can be used for subsequent queries. + +The bridge checks for cancellation between SQL calls and Arrow batches, and forwards +the caller's interrupt flag every 10 ms while the inner connection is alive. Actual +termination still depends on the running DuckDB operation reaching a cancellation +check; the bridge cannot preempt code inside a blocking external function or ggsql's +Rust rendering code. + ## Running the tests SQL logic tests under `test/sql/` are the primary test surface: @@ -140,3 +153,16 @@ GGSQL_NO_OPEN_BROWSER=1 make test ``` `GGSQL_NO_OPEN_BROWSER=1` prevents a browser tab from opening for every test query. + +The embedded interruption tests use the Python client to interrupt an active query +from another thread. Install a Python DuckDB version matching the extension build: + +```sh +python -m pip install duckdb==1.5.4 numpy +GGSQL_EXTENSION_PATH="$PWD/build/release/extension/ggsql/ggsql.duckdb_extension" \ + python -m unittest discover -s test/python -v +``` + +These tests cover both entry points, single- and multi-threaded execution, repeated +interruptions, connection reuse, and isolation from unrelated queries. Each scenario runs in a subprocess with a +timeout so a cancellation regression fails instead of hanging the test runner. diff --git a/src/ggsql_bridge.cpp b/src/ggsql_bridge.cpp index 2bf1580..e957292 100644 --- a/src/ggsql_bridge.cpp +++ b/src/ggsql_bridge.cpp @@ -3,23 +3,89 @@ #include "duckdb/common/arrow/result_arrow_wrapper.hpp" #include "duckdb/main/client_context.hpp" +#include +#include +#include #include #include +#include #include +#include namespace duckdb { +// Query() blocks during binding and execution, so checking only between bridge +// calls cannot cancel a long-running inner query. Interrupt() and IsInterrupted() +// use DuckDB's atomic flag and do not acquire either connection's context lock. +class InterruptForwarder { +public: + InterruptForwarder(ClientContext &outer, Connection &inner) + : outer(outer), inner(inner), worker([this] { Run(); }) { + } + + ~InterruptForwarder() { + { + std::lock_guard lock(mutex); + stopped = true; + } + condition.notify_one(); + worker.join(); + } + +private: + void Run() { + std::unique_lock lock(mutex); + while (!stopped) { + if (outer.IsInterrupted()) { + inner.Interrupt(); + } + // Keep forwarding: starting another inner query clears its interrupt flag. + condition.wait_for(lock, std::chrono::milliseconds(10), [this] { return stopped; }); + } + } + + ClientContext &outer; + Connection &inner; + std::mutex mutex; + std::condition_variable condition; + bool stopped = false; + std::thread worker; +}; + +BridgeCtx::BridgeCtx(ClientContext &outer) : outer(outer) { +} + +BridgeCtx::~BridgeCtx() = default; + +bool BridgeCtx::IsInterrupted() const { + return outer.IsInterrupted(); +} + +Connection &BridgeCtx::GetInnerConnection() { + if (IsInterrupted()) { + throw InterruptException(); + } + if (!inner) { + inner = make_uniq(*outer.db); + interrupt_forwarder = make_uniq(outer, *inner); + } + return *inner; +} + namespace { constexpr idx_t ARROW_STREAM_BATCH_SIZE = 100000; // --------------------------------------------------------------------------- -// Owner block for a stream produced by the inner Connection's Query(). -// Kept alive (via the Arrow stream's release callback) until Rust finishes -// consuming the stream, so client_properties.client_context stays valid. +// Own the result until Rust releases the stream. The bridge and its connection +// outlive all streams consumed during ggsql_execute. // --------------------------------------------------------------------------- struct InnerStream { - unique_ptr connection; + explicit InnerStream(BridgeCtx &bridge) : bridge(bridge) { + } + + BridgeCtx &bridge; + bool interrupted = false; unique_ptr inner; }; @@ -57,16 +123,29 @@ extern "C" void CppFreeBuffer(ggsql_byte_buffer_t *buf) { extern "C" int InnerGetSchema(ArrowArrayStream *stream, ArrowSchema *out) { auto *self = static_cast(stream->private_data); + if (self->bridge.IsInterrupted()) { + self->interrupted = true; + out->release = nullptr; + return EINTR; + } return self->inner->stream.get_schema(&self->inner->stream, out); } extern "C" int InnerGetNext(ArrowArrayStream *stream, ArrowArray *out) { auto *self = static_cast(stream->private_data); + if (self->bridge.IsInterrupted()) { + self->interrupted = true; + out->release = nullptr; + return EINTR; + } return self->inner->stream.get_next(&self->inner->stream, out); } extern "C" const char *InnerGetLastError(ArrowArrayStream *stream) { auto *self = static_cast(stream->private_data); + if (self->interrupted) { + return "Interrupted!"; + } return self->inner->stream.get_last_error(&self->inner->stream); } @@ -80,12 +159,6 @@ extern "C" void InnerRelease(ArrowArrayStream *stream) { delete self; } -void EnsureInnerConnection(BridgeCtx &bctx) { - if (!bctx.inner) { - bctx.inner = make_uniq(*bctx.outer->db); - } -} - // --- Bridge callbacks ------------------------------------------------------ extern "C" int32_t ExecSqlCallback(void *ctx, const char *sql, size_t sql_len, struct ArrowArrayStream *out_stream, @@ -107,11 +180,13 @@ extern "C" int32_t ExecSqlCallback(void *ctx, const char *sql, size_t sql_len, s auto *bctx = static_cast(ctx); std::string query(sql, sql_len); - auto holder = make_uniq(); + auto holder = make_uniq(*bctx); unique_ptr result; try { - EnsureInnerConnection(*bctx); - result = bctx->inner->Query(query); + result = bctx->GetInnerConnection().Query(query); + if (bctx->IsInterrupted()) { + throw InterruptException(); + } } catch (const std::exception &ex) { WriteError(out_err, ex.what()); return 1; @@ -129,10 +204,8 @@ extern "C" int32_t ExecSqlCallback(void *ctx, const char *sql, size_t sql_len, s return 1; } - // The inner Connection must outlive the stream (QueryResult holds a raw pointer to - // its ClientContext). Own both inside `holder` and hand off via the stream's - // release callback. We use the inner `bctx->inner` directly — it already lives for - // the full ggsql_execute call. + // QueryResult holds a raw ClientContext pointer. BridgeCtx keeps the inner + // connection alive until Rust has consumed and released every stream. try { holder->inner = make_uniq(std::move(result), ARROW_STREAM_BATCH_SIZE); } catch (const std::exception &ex) { diff --git a/src/ggsql_exec.cpp b/src/ggsql_exec.cpp index f70a9a2..fcb665b 100644 --- a/src/ggsql_exec.cpp +++ b/src/ggsql_exec.cpp @@ -60,8 +60,10 @@ OutputConfig ResolveOutputConfig(ClientContext &context) { // payload on success (text for url/spec/html/svg, raw bytes for pdf/hep, empty // for silent). string RunGgsqlQuery(ClientContext &context, const string &query, const OutputConfig &cfg) { - BridgeCtx bctx; - bctx.outer = &context; + if (context.IsInterrupted()) { + throw InterruptException(); + } + BridgeCtx bctx(context); auto bridge = BuildReaderBridge(bctx); ggsql_byte_buffer_t out; @@ -77,6 +79,12 @@ string RunGgsqlQuery(ClientContext &context, const string &query, const OutputCo } ggsql_free_buffer(&out); + // The Rust/Arrow bridge transports errors as strings. Restore DuckDB's + // interruption type after releasing the FFI buffer, rather than reporting it + // as invalid input (or returning a result after cancellation). + if (context.IsInterrupted()) { + throw InterruptException(); + } if (rc != 0) { throw InvalidInputException(payload.empty() ? "ggsql: unknown error" : payload); } diff --git a/src/ggsql_extension.cpp b/src/ggsql_extension.cpp index 31cf9d2..7473998 100644 --- a/src/ggsql_extension.cpp +++ b/src/ggsql_extension.cpp @@ -39,8 +39,7 @@ static void LoadInternal(ExtensionLoader &loader) { ggsql_save_desc.description = "Renders a ggsql query straight to a file; the writer is inferred from the file extension " "(.svg, .pdf, .hep, .html, .json) and the output path is returned."; - ggsql_save_desc.examples = { - "ggsql_save('SELECT range AS x FROM range(10) VISUALISE x DRAW line', 'plot.svg')"}; + ggsql_save_desc.examples = {"ggsql_save('SELECT range AS x FROM range(10) VISUALISE x DRAW line', 'plot.svg')"}; ggsql_save_desc.categories = {"plotting"}; ggsql_save_info.descriptions.push_back(std::move(ggsql_save_desc)); loader.RegisterFunction(std::move(ggsql_save_info)); @@ -55,8 +54,7 @@ static void LoadInternal(ExtensionLoader &loader) { ggsql_run_desc.description = "Executes a ggsql query and returns the plot as a one-row table; used for binary output modes " "('pdf', 'hep') where the result column is typed as BLOB."; - ggsql_run_desc.examples = { - "SELECT plot FROM ggsql_run('SELECT range AS x FROM range(10) VISUALISE x DRAW line')"}; + ggsql_run_desc.examples = {"SELECT plot FROM ggsql_run('SELECT range AS x FROM range(10) VISUALISE x DRAW line')"}; ggsql_run_desc.categories = {"plotting"}; ggsql_run_info.descriptions.push_back(std::move(ggsql_run_desc)); loader.RegisterFunction(std::move(ggsql_run_info)); diff --git a/src/include/ggsql_bridge.hpp b/src/include/ggsql_bridge.hpp index 19f7885..9bb2def 100644 --- a/src/include/ggsql_bridge.hpp +++ b/src/include/ggsql_bridge.hpp @@ -10,14 +10,25 @@ extern "C" { namespace duckdb { +class InterruptForwarder; + // Per-invocation bridge state. Allocated on the stack in the caller (GgsqlRunExec or // GgsqlScalarFun), passed as `ctx` through the Rust FFI. Holds the outer ClientContext // for reference, and a lazily-created inner Connection that persists across every // exec_sql callback within a single ggsql_execute call — so temp tables created by one // callback (e.g. ggsql's CTE materialisation) are visible to subsequent ones. struct BridgeCtx { - ClientContext *outer = nullptr; + explicit BridgeCtx(ClientContext &outer); + ~BridgeCtx(); + + bool IsInterrupted() const; + Connection &GetInnerConnection(); + +private: + ClientContext &outer; unique_ptr inner; + // Destroy the forwarder (and join its thread) before destroying the connection. + unique_ptr interrupt_forwarder; }; // Build a ggsql_reader_bridge_t whose function pointers dispatch back through `bctx`. diff --git a/test/python/test_interrupt.py b/test/python/test_interrupt.py new file mode 100644 index 0000000..dc311b9 --- /dev/null +++ b/test/python/test_interrupt.py @@ -0,0 +1,133 @@ +"""Embedded cancellation tests; pass GGSQL_EXTENSION_PATH to a built extension. + +Run with: python -m unittest discover -s test/python -v +Each scenario runs in a subprocess so a regression cannot hang the test runner. +""" + +import os +from pathlib import Path +import subprocess +import sys +import threading +import traceback +import unittest + + +def run_scenario(form, threads): + import duckdb + + extension = Path(os.environ['GGSQL_EXTENSION_PATH']).resolve() + con = duckdb.connect(config={'allow_unsigned_extensions': 'true', 'threads': threads}) + con.execute("LOAD '" + str(extension).replace("'", "''") + "'") + con.execute("SET ggsql_output = 'spec'") + started = threading.Event() + + def mark_started(): + started.set() + return 0 + + # Registered functions live in the shared catalog, so the sibling connection + # can call this. The event proves we interrupt during inner SQL execution. + con.create_function('ggsql_mark_started', mark_started, [], 'BIGINT', side_effects=True) + query = ( + 'SELECT sum(i + (SELECT ggsql_mark_started())) AS x, 1 AS y ' + 'FROM range(1000000000000) t(i) VISUALISE x, y DRAW point' + ) + if form == 'scalar': + query = "SELECT ggsql('" + query + "')" + + # Keep unrelated work active on the same database while cancelling ggsql. + other = con.cursor() + other_started = threading.Event() + other_errors = [] + + def mark_other_started(): + other_started.set() + return 0 + + con.create_function('ggsql_other_started', mark_other_started, [], 'BIGINT', side_effects=True) + + def execute_other(): + try: + other.execute( + 'WITH started AS MATERIALIZED (SELECT ggsql_other_started() AS marker) ' + 'SELECT sum(i + marker) FROM range(1000000000000) t(i), started' + ).fetchall() + except Exception as exc: + other_errors.append(exc) + + other_worker = threading.Thread(target=execute_other, daemon=True) + other_worker.start() + assert other_started.wait(10), f'unrelated query did not start: {other_errors}' + + for _ in range(3): + started.clear() + errors = [] + + def execute(): + try: + con.execute(query).fetchall() + except Exception as exc: + errors.append(exc) + + worker = threading.Thread(target=execute, daemon=True) + worker.start() + assert started.wait(10), f'inner query did not start: {errors}' + con.interrupt() + worker.join(5) + assert not worker.is_alive(), 'outer interrupt did not stop the inner query within 5 seconds' + assert len(errors) == 1 and isinstance(errors[0], duckdb.InterruptException), errors + assert other_worker.is_alive() and not other_errors, 'interrupt reached an unrelated connection' + + # Cancellation must not poison the next invocation or its inner connection. + assert con.execute('SELECT 42').fetchone() == (42,) + spec = con.execute( + "SELECT ggsql('WITH data AS (SELECT 1 AS x, 2 AS y) " "SELECT * FROM data VISUALISE x, y DRAW point')" + ).fetchone()[0] + assert 'vega-lite' in spec + + other.interrupt() + other_worker.join(5) + assert not other_worker.is_alive(), 'unrelated query failed to stop on its own interrupt' + assert len(other_errors) == 1 and isinstance(other_errors[0], duckdb.InterruptException), other_errors + other.close() + + # An ordinary inner SQL failure also needs to stop and join the forwarder. + try: + con.execute("SELECT ggsql('SELECT missing_column AS x VISUALISE x DRAW point')") + except duckdb.InvalidInputException: + pass + else: + raise AssertionError('expected the inner SQL binding error') + assert con.execute('SELECT 42').fetchone() == (42,) + con.close() + + +class InterruptTests(unittest.TestCase): + def test_interrupt(self): + for form in ('scalar', 'parser'): + for threads in (1, 4): + with self.subTest(form=form, threads=threads): + result = subprocess.run( + [sys.executable, __file__, form, str(threads)], + capture_output=True, + text=True, + timeout=60, + env={**os.environ, 'GGSQL_NO_OPEN_BROWSER': '1'}, + ) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + + +if __name__ == '__main__': + if len(sys.argv) == 3 and sys.argv[1] in ('scalar', 'parser'): + try: + run_scenario(sys.argv[1], int(sys.argv[2])) + except BaseException: + traceback.print_exc() + # A failing cancellation test may still have a blocked native worker. + # Exit without waiting for Python/DuckDB shutdown to join that worker. + sys.stdout.flush() + sys.stderr.flush() + os._exit(1) + else: + unittest.main()