Give non-streamable result sets a route of their own

chdb-integration-tests/src/test/java/org/chdb/it/StreamingLifecycleIT.java

@@ -199,6 +199,50 @@ void queryTimeout() throws SQLException {
         }
     }
 
+    /**
+     * The same timeout, but expiring inside the call that opens the result set rather than
+     * during the fetches.
+     *
+     * <p>{@link #queryTimeout()} above covers a SELECT that emits as it scans, where the open
+     * costs milliseconds and the clock runs out in {@code next()}. A full aggregate has no
+     * first batch until the whole scan is done, so the entire query happens inside {@code
+     * streamOpen} -- and nothing can be cancelled there, because the handle {@code cancel()}
+     * would need is what the call is still producing. Measured on v26.7.0 before this was
+     * handled: {@code setQueryTimeout(1)} returned a working result set after 12.65 seconds.
+     *
+     * <p>Asserted as an invariant rather than against the clock: a result set is the right
+     * answer only if the statement beat its deadline.
+     */
+    @Test
+    @DisplayName("a timeout that expires while the result set is opening is still reported")
+    @Timeout(value = 5, unit = TimeUnit.MINUTES)
+    void queryTimeoutDuringTheOpen() throws SQLException {
+        try (Connection connection = openMemory();
+                Statement statement = connection.createStatement()) {
+            statement.setQueryTimeout(1);
+            long start = System.nanoTime();
+            // Sized so the scan runs for a few seconds against a 1 s deadline: the timeout
+            // cannot stop it, so the test pays the whole scan either way and there is no
+            // reason to make it longer than the margin needs.
+            try (ResultSet rs =
+                    statement.executeQuery(
+                            "SELECT max(sipHash64(number)) FROM numbers(500000000)")) {
+                long millis = (System.nanoTime() - start) / 1_000_000;
+                assertTrue(
+                        millis < 1_000,
+                        () -> "executeQuery returned a result set " + millis
+                                + " ms after a 1 s query timeout");
+                assertTrue(rs.next());
+            } catch (SQLTimeoutException expected) {
+                assertEquals("57014", expected.getSQLState());
+            }
+            // The slot has to have come back, or nothing else could run on this connection.
+            try (ResultSet rs = statement.executeQuery("SELECT 1")) {
+                assertTrue(rs.next());
+            }
+        }
+    }
+
     @Test
     @DisplayName("cancel from another thread stops an in-flight query")
     @Timeout(value = 2, unit = TimeUnit.MINUTES)

chdb-jdbc/src/main/java/org/chdb/internal/ChdbNative.java

@@ -28,7 +28,7 @@ public final class ChdbNative {
      * Must equal {@code chdb_jni::kJniAbiVersion} in the shim and {@code
      * chdb.jni.abi.version} in the poms. Verified at load time before any other call.
      */
-    public static final int JNI_ABI_VERSION = 1;
+    public static final int JNI_ABI_VERSION = 2;
 
     /** Handle kinds, for {@link #openHandleCount(int)}. */
     public static final int KIND_CONNECTION = 1;
@@ -172,6 +172,43 @@ public static native long streamOpen(
             boolean unsupportedAsBinary,
             boolean stringAsString);
 
+    /**
+     * {@code chdb_query_arrow_n}: runs the statement to completion and exports the whole
+     * result set through the same Arrow C Data Interface the streaming path consumes.
+     *
+     * <p>The route for a statement that has a result set the engine refuses to stream --
+     * {@code SHOW}, {@code DESCRIBE}, {@code EXPLAIN}, {@code EXISTS}, {@code CHECK}. The
+     * engine's streaming door admits only a SELECT pipeline, so without this every non-SELECT
+     * read failed with {@code Streaming query is not supported} (issue #12).
+     *
+     * <p>Returns a handle of the same kind as {@link #streamOpen}, driven by the same {@code
+     * stream*} methods: the batches come from the exported Arrow stream instead of from the
+     * engine's pipeline, and nothing above the shim can tell the difference.
+     *
+     * <p>Differences a caller has to know about:
+     *
+     * <ul>
+     *   <li>Memory is not bounded by one batch. The result is fully materialized in the engine
+     *       before this returns, which is what makes it safe only for the statements above:
+     *       their results are bounded by the schema, not by the data.
+     *   <li>No server-side parameters. The engine exports {@code chdb_query_arrow_n} but no
+     *       {@code _with_params_n} variant of it, so a parameterized non-streamable statement
+     *       has no route at all; {@code ChdbStatement} reports that rather than silently
+     *       dropping the bindings.
+     *   <li>{@link #streamCancel(long, long)} cannot stop the execution, only the iteration:
+     *       the statement has already run by the time there is a handle to cancel.
+     * </ul>
+     *
+     * @return a stream handle
+     * @throws ChdbNativeException carrying the engine's error text if the statement failed
+     */
+    public static native long streamOpenMaterialized(
+            long connection,
+            byte[] sql,
+            boolean lowCardinalityAsDictionary,
+            boolean unsupportedAsBinary,
+            boolean stringAsString);
+
     /** Column names from the Arrow schema, in ordinal order. */
     public static native String[] streamColumnNames(long stream);
 

chdb-jdbc/src/main/java/org/chdb/jdbc/ChdbStatement.java

@@ -3,6 +3,7 @@
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.SQLTimeoutException;
 import java.sql.SQLWarning;
 import java.sql.Statement;
 import java.util.Collections;
@@ -15,12 +16,19 @@
 /**
  * Executes one statement at a time on a {@link ChdbConnection}.
  *
- * <h2>Streaming or materialized</h2>
- * A statement with a result set goes through the Arrow streaming API and yields a {@link
- * ChdbResultSet} that holds one batch at a time. Everything else -- DDL, DML, session control
- * -- goes through {@code chdb_query_n}, because the streaming API only accepts statements
- * that have a result schema. {@link StatementShape} decides which, from the engine's
- * classifier where available.
+ * <h2>Three routes, not two</h2>
+ * A statement with a result set the engine will stream goes through {@code
+ * chdb_stream_query_arrow_n} and yields a {@link ChdbResultSet} that holds one batch at a
+ * time. A statement with a result set the engine refuses to stream -- {@code SHOW}, {@code
+ * DESCRIBE}, {@code EXPLAIN}, {@code EXISTS}, {@code CHECK}, none of which parse as a SELECT
+ * pipeline -- goes through {@code chdb_query_arrow_n}, which delivers the same Arrow C Data
+ * Interface materialized, and produces a {@link ChdbResultSet} indistinguishable from the
+ * first. Everything else -- DDL, DML, session control -- goes through {@code chdb_query_n}.
+ * {@link StatementShape} decides which, from the engine's classifier where available.
+ *
+ * <p>The route is decided once, before execution, and no statement is ever executed twice:
+ * {@link StatementShape} documents why a failed stream open is reported rather than retried
+ * on the materialized route.
  *
  * <h2>Cancellation</h2>
  * {@link #cancel()} is the one method that may be called from another thread while this one
@@ -31,6 +39,11 @@
  * <p>{@link #setQueryTimeout(int)} is built on the same mechanism: a timer thread calls the
  * same cancel path, so a timeout actually stops the engine rather than only abandoning the
  * Java-side wait (work plan section 5.8).
+ *
+ * <p>There is one window in which neither can interrupt anything, because there is no handle
+ * to cancel yet: the call that opens the result set. {@link
+ * #checkDeadlineSurvivedTheOpen(String)} is what keeps that window from turning a timeout into
+ * a late success.
  */
 public class ChdbStatement implements Statement {
 
@@ -132,7 +145,8 @@ boolean executeInternal(
         connection.statementSlot().acquire(sql);
         boolean handedOff = false;
         try {
-            boolean producesResultSet = producesResultSet(sql);
+            StatementShape.Route route = route(sql);
+            boolean producesResultSet = route != StatementShape.Route.NO_RESULT_SET;
             if (expectResultSet != null && expectResultSet && !producesResultSet) {
                 return false;
             }
@@ -143,9 +157,12 @@ boolean executeInternal(
                 // where the time goes.
                 activeTimeout = QueryTimeout.start(this, queryTimeoutSeconds);
                 try {
-                    boolean opened = openStream(sql, parameterNames, parameterValues);
-                    // The result set now owns the slot and releases it when it closes.
-                    handedOff = opened;
+                    boolean opened = openStream(route, sql, parameterNames, parameterValues);
+                    if (opened) {
+                        // The result set now owns the slot and releases it when it closes.
+                        handedOff = true;
+                        checkDeadlineSurvivedTheOpen(sql);
+                    }
                     return opened;
                 } catch (RuntimeException | SQLException e) {
                     stopTimeout();
@@ -164,7 +181,71 @@ boolean executeInternal(
         }
     }
 
-    private boolean producesResultSet(String sql) throws SQLException {
+    /**
+     * Fails the execution if the query timeout expired, or {@link #cancel()} was called, while
+     * the open was in flight.
+     *
+     * <h2>Why this is needed at all</h2>
+     * {@link #cancel()} works by handing the in-flight stream handle to the engine -- and
+     * there is no handle until the call that creates it returns. Every exported cancel in the
+     * C ABI takes a result or stream handle ({@code chdb_stream_cancel_query}, {@code
+     * chdb_streaming_cancel_query}, {@code chdb_stream_cancel_insert}); there is no
+     * connection-level cancel. So the open is uninterruptible on both routes, and a timer that
+     * fires inside it finds {@code inFlightStream} still null and returns having done nothing.
+     *
+     * <p>Left there, the caller got the worst of both worlds: {@code executeQuery} came back
+     * <em>successfully</em>, with a usable result set, long after the deadline it set. Measured
+     * on v26.7.0 before this check: {@code setQueryTimeout(1)} on {@code SELECT
+     * max(sipHash64(number)) FROM numbers(2000000000)} returned a working result set after
+     * 12.65 seconds.
+     *
+     * <h2>How long that window is</h2>
+     * On the streaming route it is the pipeline init plus the first batch, which for a SELECT
+     * that emits as it scans is milliseconds -- but for a full aggregate, a {@code GROUP BY} or
+     * an {@code ORDER BY} without a {@code LIMIT} there is no first batch until the whole scan
+     * is done, so the window is the entire query. On the materialized route it is always the
+     * entire statement, because {@code chdb_query_arrow_n} runs it to completion before
+     * returning.
+     *
+     * <h2>What the caller gets instead</h2>
+     * A deadline, not an interrupt. The statement has already run to completion by the time
+     * this is reached and the engine cannot be told to stop, so the work is spent either way --
+     * but reporting {@link SQLTimeoutException} is what a caller can act on, and it keeps
+     * {@code setQueryTimeout} from being a setting that silently does nothing. The result set
+     * is closed first, so the stream and the connection's statement slot are released rather
+     * than pinned by a result nobody can reach.
+     */
+    private void checkDeadlineSurvivedTheOpen(String sql) throws SQLException {
+        boolean expired = activeTimeout.expired();
+        if (!expired && !cancelled.get()) {
+            return;
+        }
+        // Releases the stream and the statement slot; the caller set handedOff before calling
+        // here, so this close is the one release of that slot.
+        closeCurrentResultSet();
+        if (expired) {
+            throw new SQLTimeoutException(
+                    "The query exceeded the statement's query timeout of "
+                            + queryTimeoutSeconds
+                            + "s while "
+                            + describeStatement(sql)
+                            + " was still being executed. The engine could not be interrupted:"
+                            + " the chDB C ABI has no cancel that applies before a query returns"
+                            + " a handle, so the statement ran to completion and its result was"
+                            + " discarded.",
+                    "57014");
+        }
+        throw new SQLException(
+                "The statement was cancelled while "
+                        + describeStatement(sql)
+                        + " was still being executed. The engine could not be interrupted before"
+                        + " the query returned a handle, so it ran to completion and its result"
+                        + " was discarded.",
+                "57014",
+                394);
+    }
+
+    private StatementShape.Route route(String sql) throws SQLException {
         int[] analysis;
         try {
             // Null when the engine predates chdb_classify_query_n, which the v26.7.2-rc.2
@@ -176,23 +257,59 @@ private boolean producesResultSet(String sql) throws SQLException {
             // message than anything this classification step could produce.
             analysis = null;
         }
-        return StatementShape.producesResultSet(analysis, sql);
+        return StatementShape.route(analysis, sql);
     }
 
-    private boolean openStream(String sql, List<String> parameterNames, List<String> parameterValues)
+    private boolean openStream(
+            StatementShape.Route route,
+            String sql,
+            List<String> parameterNames,
+            List<String> parameterValues)
             throws SQLException {
         ChdbUrl url = connection.chdbUrl();
+        boolean lowCardinalityAsDictionary =
+                url.booleanProperty(ChdbUrl.PROP_LOW_CARDINALITY_AS_DICTIONARY, false);
+        boolean unsupportedAsBinary = url.booleanProperty(ChdbUrl.PROP_UNSUPPORTED_AS_BINARY, false);
+        boolean stringAsString = url.booleanProperty(ChdbUrl.PROP_STRING_AS_STRING, true);
+
+        boolean materialize = route == StatementShape.Route.MATERIALIZED_RESULT_SET;
+        if (materialize && !parameterNames.isEmpty()) {
+            // The engine exports chdb_query_arrow_n but no _with_params_n variant of it, so
+            // there is no entry point that both accepts this statement and binds parameters.
+            // Reported rather than worked around: interpolating the values into the SQL is
+            // the injection this driver's server-side binding exists to avoid, and running
+            // the statement with the bindings dropped would answer the wrong question.
+            throw new SQLException(
+                    "Server-side parameters are not supported for "
+                            + describeStatement(sql)
+                            + ", because the engine has no parameter-binding form of the Arrow"
+                            + " entry point that accepts it (chdb_query_arrow_with_params_n is"
+                            + " not exported by engine "
+                            + ChdbNative.engineVersion()
+                            + "). Use a Statement with the value written into the SQL, or a"
+                            + " query over system.tables / system.columns, which is a SELECT and"
+                            + " does take parameters.",
+                    "0A000");
+        }
+
         long stream;
         try {
             stream =
-                    ChdbNative.streamOpen(
-                            connection.handle(),
-                            Utf8.encode(sql),
-                            Utf8.encodeAll(parameterNames),
-                            Utf8.encodeAll(parameterValues),
-                            url.booleanProperty(ChdbUrl.PROP_LOW_CARDINALITY_AS_DICTIONARY, false),
-                            url.booleanProperty(ChdbUrl.PROP_UNSUPPORTED_AS_BINARY, false),
-                            url.booleanProperty(ChdbUrl.PROP_STRING_AS_STRING, true));
+                    materialize
+                            ? ChdbNative.streamOpenMaterialized(
+                                    connection.handle(),
+                                    Utf8.encode(sql),
+                                    lowCardinalityAsDictionary,
+                                    unsupportedAsBinary,
+                                    stringAsString)
+                            : ChdbNative.streamOpen(
+                                    connection.handle(),
+                                    Utf8.encode(sql),
+                                    Utf8.encodeAll(parameterNames),
+                                    Utf8.encodeAll(parameterValues),
+                                    lowCardinalityAsDictionary,
+                                    unsupportedAsBinary,
+                                    stringAsString);
         } catch (ChdbNativeException e) {
             throw ChdbExceptions.wrap("Query failed", e);
         }

chdb-jdbc/src/main/java/org/chdb/jdbc/QueryTimeout.java

@@ -23,6 +23,16 @@
  * Statements with no result set run through {@code chdb_query_n}, which is synchronous and
  * has no cancellation handle in the C ABI. There is nothing to interrupt, so the timeout does
  * not apply to DDL or DML, and {@code ChdbStatement} does not arm one for them.
+ *
+ * <p>Nor can it interrupt the call that <em>opens</em> a result set, for the same reason: every
+ * cancel the C ABI exports takes a result or stream handle, and that handle is what the open is
+ * still producing. So the cancel this fires during an open reaches nothing. How long that
+ * window lasts depends on the statement, not on the driver -- milliseconds for a SELECT that
+ * emits as it scans, the whole query for a full aggregate or an {@code ORDER BY} without a
+ * {@code LIMIT}, and the whole statement for anything on the materialized Arrow route. {@code
+ * ChdbStatement} therefore also treats an expired timeout as a deadline once the open returns:
+ * the work is already spent, but the caller is told, rather than handed a result set that
+ * arrived after the time it allowed.
  */
 final class QueryTimeout {