diff --git a/README.md b/README.md index 72d1f58..a501371 100644 --- a/README.md +++ b/README.md @@ -118,11 +118,55 @@ SpiceClient client = SpiceClient.builder() ``` Retries are performed for connection and system internal errors. It is the SDK user's responsibility to properly -handle other errors, for example RESOURCE_EXHAUSTED (HTTP 429). +handle other errors, for example RESOURCE_EXHAUSTED (HTTP 429). Retries use exponential backoff with jitter +(~250ms, 500ms, 1s, ... capped at 10s). If the server reports an expired authentication token +(UNAUTHENTICATED), the client automatically re-handshakes and retries. + +### Performance Tuning + +```java +SpiceClient client = SpiceClient.builder() + // Number of gRPC connections; queries are distributed round-robin. + // Increase for highly concurrent workloads with large result streams. + .withChannelCount(4) + // Deadline for query planning and statement preparation RPCs + // (result streaming is not limited by this timeout). + .withQueryTimeout(Duration.ofSeconds(30)) + // Max idle prepared statements reused by queryWithParams (default 64, 0 disables). + .withPreparedStatementCacheSize(128) + .build(); +``` + +### Connection Pooling and HikariCP + +`SpiceClient` is thread-safe and multiplexes concurrent queries over shared HTTP/2 connections — use **one client instance per application** and share it across threads. It does not need an external connection pool; for high concurrency, size the built-in pool with `withChannelCount(n)`. + +Applications that want JDBC semantics (ORMs, existing [HikariCP](https://github.com/brettwooldridge/HikariCP) infrastructure) can connect to Spice through the [Arrow Flight SQL JDBC driver](https://arrow.apache.org/docs/java/flight_sql_jdbc_driver.html) and pool those connections with HikariCP — this combination is exercised in this repo's test suite: + +```java +HikariConfig config = new HikariConfig(); +config.setJdbcUrl("jdbc:arrow-flight-sql://localhost:50051/?useEncryption=false"); +// For Spice Cloud: jdbc:arrow-flight-sql://flight.spiceai.io:443/?useEncryption=true +// with setUsername(appId) and setPassword(apiKey) — the same credentials the SDK uses. +config.setMaximumPoolSize(4); +// The Flight SQL JDBC driver does not implement Connection.isValid(): +config.setConnectionTestQuery("SELECT 1"); + +try (HikariDataSource pool = new HikariDataSource(config); + Connection conn = pool.getConnection(); + PreparedStatement ps = conn.prepareStatement("SELECT * FROM taxi_trips WHERE total_amount > ?")) { + ps.setDouble(1, 10.0); + try (ResultSet rs = ps.executeQuery()) { + while (rs.next()) { /* ... */ } + } +} +``` + +Requires `org.apache.arrow:flight-sql-jdbc-driver` (or `flight-sql-jdbc-core`) and `com.zaxxer:HikariCP` on your classpath. Keep the pool small — each JDBC connection opens its own Flight channel. Prefer the native `SpiceClient` where possible: it streams Arrow data without JDBC row conversion and is significantly faster for analytical results. ### Parameterized Queries (Recommended) -The SDK supports parameterized queries using ADBC (Arrow Database Connectivity), which is the recommended approach for queries with user input to prevent SQL injection: +The SDK supports parameterized queries using Arrow Flight SQL prepared statements, which is the recommended approach for queries with user input to prevent SQL injection. Prepared statements are cached and reused across repeated executions of the same SQL, and run on the same tuned connection as regular queries (DNS re-resolution, keep-alive, mTLS): ```java import org.apache.arrow.vector.VectorSchemaRoot; diff --git a/docs/parameterized_queries.md b/docs/parameterized_queries.md index a3f95a1..84937bb 100644 --- a/docs/parameterized_queries.md +++ b/docs/parameterized_queries.md @@ -294,7 +294,12 @@ try { 2. **Explicit Types**: No overhead, types are directly specified 3. **Arrow Conversion**: Zero-copy when possible, efficient serialization 4. **Large Data**: Use Large variants (LargeString, LargeBinary) for > 2GB data -5. **Connection Reuse**: ADBC connection is initialized lazily and reused +5. **Statement Reuse**: Prepared statements are cached per SQL string and reused, + removing the create/close round trips from repeated queries. Tune with + `withPreparedStatementCacheSize(n)` (default 64; 0 disables caching). +6. **Connection Reuse**: Parameterized queries share the same gRPC channels as + regular queries — one connection pool, one set of keep-alive/DNS/mTLS settings. + Use `withChannelCount(n)` for highly concurrent workloads. ## Security Benefits diff --git a/docs/release_notes/v0.7.0.md b/docs/release_notes/v0.7.0.md new file mode 100644 index 0000000..150c974 --- /dev/null +++ b/docs/release_notes/v0.7.0.md @@ -0,0 +1,77 @@ +# Spice Java SDK v0.7.0 + +## Highlights + +This release is a **performance overhaul**. Parameterized queries move off ADBC onto native Arrow Flight SQL prepared statements running on the same tuned connection as regular queries, prepared statements are cached and reused, retries use real exponential backoff, and the Netty/gRPC dependency stack is aligned onto supported versions. + +## What's New + +### ⚡ Native Flight SQL Prepared Statements (ADBC removed) + +`queryWithParams` is now implemented directly on Arrow Flight SQL prepared statements instead of the ADBC driver. This is a pure win with no API change — the method still returns an `ArrowReader`: + +- **One connection instead of two.** Parameterized queries previously ran on a separate ADBC-managed gRPC channel that had none of the SDK's transport tuning. They now share the primary channels, inheriting DNS re-resolution (`dns:///`), HTTP/2 keep-alive, and — new in this release for parameterized queries — **mTLS/custom CA support**. +- **Prepared statement caching.** Statements are cached per SQL string and reused across executions, removing the `CreatePreparedStatement` and `ClosePreparedStatement` round trips from every repeated query (5 → 3 round trips per call). Configure with `withPreparedStatementCacheSize(n)` (default 64, `0` disables). +- **Transparent recovery.** If the server no longer recognizes a cached statement handle (e.g. after a server restart), the SDK re-prepares and retries automatically. +- **Multi-endpoint results.** The `ArrowReader` consumes every endpoint of a partitioned result. +- Two dependencies removed (`adbc-driver-flight-sql`, `adbc-core`). + +### 🚀 Performance Tuning Options + +```java +SpiceClient client = SpiceClient.builder() + .withChannelCount(4) // round-robin gRPC connection pool + .withQueryTimeout(Duration.ofSeconds(30)) // deadline for query planning RPCs + .withPreparedStatementCacheSize(128) // prepared statement reuse + .build(); +``` + +- **`withChannelCount(n)`** — opens `n` gRPC connections and distributes queries round-robin, lifting the single-connection ceiling (server `MAX_CONCURRENT_STREAMS`, single-TCP throughput) for highly concurrent workloads. Default 1. +- **`withQueryTimeout(d)`** — bounds query planning, statement preparation, and parameter binding RPCs. Result streaming is intentionally not limited; dead connections during streaming are detected by keep-alive. Default: no timeout. +- **`withPreparedStatementCacheSize(n)`** — see above. + +### 🔁 Real Retry Backoff + +Retries previously waited 1–2ms between attempts — useless against real outages. They now use exponential backoff with jitter (~250ms, 500ms, 1s, … capped at 10s), so the default 3 retries cover ~2s outages such as load-balancer failovers. + +### 🔐 Automatic Re-Authentication + +If the server reports `UNAUTHENTICATED` (e.g. an expired handshake bearer token on a long-lived client), the SDK re-handshakes once and retries the query automatically. Previously this required a manual `reset()`. + +### 🧹 Smaller Fixes + +- Fixed a race where `reset()` concurrent with `queryWithParams` could throw `NullPointerException`. +- Parameter bind roots are sized for their single row instead of Arrow's ~3970-slot default (~48KB saved per string parameter per query). +- `refreshDataset` now sets a 60s request timeout (previously it could hang forever) and the HTTP client is created lazily. +- The user agent string is computed once instead of per request; fixed the OS version reported on Windows. +- `query()` logs a warning if the server returns a multi-endpoint result (only the first endpoint is consumed by the `FlightStream` API; use `queryWithParams` to consume all). + +## Dependency Changes + +| Dependency | v0.6.0 | v0.7.0 | +| ----------------------------------- | ------------- | ------------ | +| Apache Arrow ADBC Driver Flight SQL | 0.23.0 | **removed** | +| Apache Arrow ADBC Core | 0.23.0 | **removed** | +| netty-all | 4.2.12.Final | **removed** | +| Netty (via netty-bom) | 4.2.12.Final (forced) | 4.1.130.Final | +| gRPC (via grpc-bom) | 1.79.0 | 1.79.0 (pinned) | +| netty-transport-native-epoll | — | 4.1.130.Final (Linux runtime) | +| BouncyCastle (bcprov/bcpkix) | transitive | 1.80 (explicit) | + +The `netty-all 4.2.12` pin forced every Netty module to the 4.2 line, which `grpc-netty 1.79` (built against Netty 4.1.130) does not support, and dragged unused codecs (HTTP/3, MQTT, SMTP, Redis, …) onto consumers' classpaths. The stack is now aligned on Netty 4.1.130.Final via `netty-bom`, with native epoll transports added for Linux. BouncyCastle (used for mTLS PEM parsing) is now an explicit dependency instead of riding on ADBC's transitives. + +### 🐝 HikariCP / JDBC Interop + +`SpiceClient` is thread-safe and needs no external pool (use `withChannelCount`), but applications with JDBC infrastructure can pool connections to Spice with HikariCP via the Arrow Flight SQL JDBC driver. This combination — including authentication and prepared statements — is now verified by `HikariCpCompatTest`, and the README documents the recommended configuration. + +## Testing + +- New in-process mock Flight SQL server (`TestFlightSqlServer`) with per-RPC counters and failure injection — the full query, prepared-statement, retry, timeout, re-auth, and multi-endpoint behavior is now covered without an external runtime. +- New test suites: `LocalFlightServerTest`, `StatementCacheTest`, `ResilienceTest`, `FlightInfoReaderTest`, `ParamRootTest`, `RefreshHttpTest`, `HikariCpCompatTest`. +- New `PerfBenchmarkTest` printing latency percentiles and asserting the round-trip contract (cached path: zero prepares per query; uncached: one prepare + one close per query). + +## Compatibility + +- Public API unchanged: `query`, `queryWithParams`, `refreshDataset`, `reset`, `close`, and all existing builder methods behave as before. +- `queryWithParams` error causes are now `FlightRuntimeException` instead of `AdbcException` (both surfaced inside `ExecutionException`). +- To restore the previous prepare-per-query behavior: `withPreparedStatementCacheSize(0)`. diff --git a/pom.xml b/pom.xml index 60108fb..be77441 100644 --- a/pom.xml +++ b/pom.xml @@ -4,7 +4,7 @@ ai.spice spiceai jar - 0.6.0 + 0.7.0 Java Spice SDK https://github.com/spiceai/spice-java/ Spice provides a unified SQL query interface and portable runtime to locally materialize, accelerate, and query datasets from any database, data warehouse, or data lake. @@ -31,26 +31,42 @@ 11 11 + + 4.1.130.Final + 1.79.0 + + + + + io.netty + netty-bom + ${netty.version} + pom + import + + + + io.grpc + grpc-bom + ${grpc.version} + pom + import + + + org.apache.arrow flight-sql 19.0.0 - - org.apache.arrow.adbc - adbc-driver-flight-sql - 0.23.0 - - - org.apache.arrow.adbc - adbc-core - 0.23.0 - com.google.code.gson gson @@ -61,10 +77,33 @@ guava-retrying 2.0.0 + + + org.bouncycastle + bcprov-jdk18on + 1.80 + + + org.bouncycastle + bcpkix-jdk18on + 1.80 + + io.netty - netty-all - 4.2.12.Final + netty-transport-native-epoll + ${netty.version} + linux-x86_64 + runtime + + + io.netty + netty-transport-native-epoll + ${netty.version} + linux-aarch_64 + runtime org.slf4j @@ -83,6 +122,21 @@ 4.13.2 test + + + com.zaxxer + HikariCP + 5.1.0 + test + + + org.apache.arrow + flight-sql-jdbc-core + 19.0.0 + test + diff --git a/src/main/java/ai/spice/Config.java b/src/main/java/ai/spice/Config.java index 9063c0a..c73fc88 100644 --- a/src/main/java/ai/spice/Config.java +++ b/src/main/java/ai/spice/Config.java @@ -93,13 +93,21 @@ public static URI getCloudHttpAddressUri() throws URISyntaxException { return new URI(CLOUD_HTTP_ADDRESS); } + // Computed once: system properties are effectively immutable for the + // process lifetime, and this used to be recomputed on every request. + private static final String USER_AGENT = computeUserAgent(); + /** * Returns the Spice SDK user agent for this system, including the package * version, system OS, version and arch. - * + * * @return the Spice SDK user agent string for this system. */ public static String getUserAgent() { + return USER_AGENT; + } + + private static String computeUserAgent() { // change the os arch to match the pattern set in other SDKs String osArch = System.getProperty("os.arch"); if (osArch.equals("amd64")) { @@ -124,7 +132,7 @@ public static String getUserAgent() { } return "spice-java/" + Version.SPICE_JAVA_VERSION + " (" + osName + "/" - + System.getProperty("os.version") + " " + + osVersion + " " + osArch + ")"; } } diff --git a/src/main/java/ai/spice/FlightInfoReader.java b/src/main/java/ai/spice/FlightInfoReader.java new file mode 100644 index 0000000..f6ecf6e --- /dev/null +++ b/src/main/java/ai/spice/FlightInfoReader.java @@ -0,0 +1,163 @@ +/* +Copyright 2026 The Spice.ai OSS Authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +*/ + +package ai.spice; + +import java.io.IOException; +import java.util.Collections; +import java.util.List; + +import org.apache.arrow.flight.CallOption; +import org.apache.arrow.flight.FlightEndpoint; +import org.apache.arrow.flight.FlightInfo; +import org.apache.arrow.flight.FlightStream; +import org.apache.arrow.flight.sql.FlightSqlClient; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.VectorUnloader; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; +import org.apache.arrow.vector.types.pojo.Schema; + +/** + * An {@link ArrowReader} that streams the results described by a + * {@link FlightInfo}, consuming every endpoint in order on the connection + * that produced it. + * + *

Record batches are transferred from the underlying {@link FlightStream} + * into this reader's root via unload/load, which moves buffer ownership + * without copying the data.

+ */ +final class FlightInfoReader extends ArrowReader { + + private final FlightSqlClient client; + private final CallOption[] streamOptions; + private final List endpoints; + private final Schema schema; + private int nextEndpointIndex; + private FlightStream currentStream; + private long bytesRead; + + /** + * Creates a reader over all endpoints of the given FlightInfo. Opens the + * first stream eagerly so transient failures surface here (and can be + * retried by the caller) rather than mid-consumption. + */ + FlightInfoReader(BufferAllocator allocator, FlightSqlClient client, CallOption[] streamOptions, + FlightInfo flightInfo) throws IOException { + super(allocator); + this.client = client; + this.streamOptions = streamOptions; + this.endpoints = flightInfo.getEndpoints(); + try { + if (endpoints.isEmpty()) { + this.schema = flightInfo.getSchemaOptional() + .orElseGet(() -> new Schema(Collections.emptyList())); + } else { + this.currentStream = client.getStream(endpoints.get(0).getTicket(), streamOptions); + this.nextEndpointIndex = 1; + this.schema = currentStream.getSchema(); + } + ensureInitialized(); + } catch (IOException | RuntimeException e) { + closeStreamQuietly(e); + throw e; + } + } + + @Override + public boolean loadNextBatch() throws IOException { + while (currentStream != null) { + if (currentStream.next()) { + VectorSchemaRoot streamRoot = currentStream.getRoot(); + VectorUnloader unloader = new VectorUnloader(streamRoot); + try (ArrowRecordBatch batch = unloader.getRecordBatch()) { + bytesRead += batch.computeBodyLength(); + loadRecordBatch(batch); + } + return true; + } + advanceToNextEndpoint(); + } + return false; + } + + private void advanceToNextEndpoint() throws IOException { + try { + currentStream.close(); + } catch (Exception e) { + throw asIOException(e); + } + currentStream = null; + if (nextEndpointIndex < endpoints.size()) { + FlightEndpoint endpoint = endpoints.get(nextEndpointIndex++); + currentStream = client.getStream(endpoint.getTicket(), streamOptions); + if (!schema.equals(currentStream.getSchema())) { + Schema streamSchema = currentStream.getSchema(); + closeStreamQuietly(null); + throw new IOException( + "Endpoint returned inconsistent schema. Expected: " + schema + ", got: " + streamSchema); + } + } + } + + @Override + public long bytesRead() { + return bytesRead; + } + + @Override + protected void closeReadSource() throws IOException { + if (currentStream != null) { + try { + currentStream.close(); + } catch (Exception e) { + throw asIOException(e); + } finally { + currentStream = null; + } + } + } + + @Override + protected Schema readSchema() throws IOException { + return schema; + } + + private void closeStreamQuietly(Exception pending) { + if (currentStream != null) { + try { + currentStream.close(); + } catch (Exception closeEx) { + if (pending != null) { + pending.addSuppressed(closeEx); + } + } finally { + currentStream = null; + } + } + } + + private static IOException asIOException(Exception e) { + return (e instanceof IOException) ? (IOException) e : new IOException(e); + } +} diff --git a/src/main/java/ai/spice/SpiceClient.java b/src/main/java/ai/spice/SpiceClient.java index 03501a4..e9032af 100644 --- a/src/main/java/ai/spice/SpiceClient.java +++ b/src/main/java/ai/spice/SpiceClient.java @@ -22,6 +22,7 @@ of this software and associated documentation files (the "Software"), to deal package ai.spice; +import java.io.IOException; import java.math.BigDecimal; import java.net.ConnectException; import java.net.URI; @@ -34,32 +35,31 @@ of this software and associated documentation files (the "Software"), to deal import java.time.LocalDateTime; import java.time.LocalTime; import java.time.ZoneOffset; +import java.util.ArrayDeque; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicInteger; -import org.apache.arrow.adbc.core.AdbcConnection; -import org.apache.arrow.adbc.core.AdbcDatabase; -import org.apache.arrow.adbc.core.AdbcDriver; -import org.apache.arrow.adbc.core.AdbcException; -import org.apache.arrow.adbc.core.AdbcStatement; -import org.apache.arrow.adbc.core.AdbcStatusCode; -import org.apache.arrow.adbc.driver.flightsql.FlightSqlDriver; +import org.apache.arrow.flight.CallOption; +import org.apache.arrow.flight.CallOptions; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightClientMiddleware; +import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightGrpcUtils; +import org.apache.arrow.flight.FlightInfo; +import org.apache.arrow.flight.FlightRuntimeException; +import org.apache.arrow.flight.FlightStatusCode; import org.apache.arrow.flight.FlightStream; -import org.apache.arrow.flight.Location; import org.apache.arrow.flight.Ticket; import org.apache.arrow.flight.auth2.BasicAuthCredentialWriter; import org.apache.arrow.flight.auth2.ClientBearerHeaderHandler; import org.apache.arrow.flight.auth2.ClientIncomingAuthHeaderMiddleware; import org.apache.arrow.flight.grpc.CredentialCallOption; -import org.apache.arrow.flight.FlightInfo; -import org.apache.arrow.flight.FlightRuntimeException; import io.grpc.ManagedChannel; import io.grpc.netty.GrpcSslContexts; @@ -126,24 +126,37 @@ of this software and associated documentation files (the "Software"), to deal /** * Client to execute SQL queries against Spice.ai Cloud and Spice.ai OSS. - * Supports both regular queries and parameterized queries using ADBC. + * Supports both regular queries and parameterized queries using Arrow Flight + * SQL prepared statements. */ public class SpiceClient implements AutoCloseable { private static final Logger logger = LoggerFactory.getLogger(SpiceClient.class); - + private static final long BYTES_PER_MB = 1024L * 1024L; - + // Cached Gson instance for JSON serialization (thread-safe) private static final Gson GSON = new Gson(); - + // Cap for large dataset results and metadata (~2 GiB, max safe signed-int value) private static final int MAX_INBOUND_MESSAGE_SIZE = Integer.MAX_VALUE; private static final int MAX_INBOUND_METADATA_SIZE = Integer.MAX_VALUE; - // HttpClient for refresh operations (thread-safe, connection pooling) - private final HttpClient httpClient; - + /** Default number of gRPC channels (HTTP/2 connections) per client. */ + static final int DEFAULT_CHANNEL_COUNT = 1; + /** Default maximum number of idle prepared statements kept for reuse. */ + static final int DEFAULT_STATEMENT_CACHE_SIZE = 64; + + // Retry backoff: exponential (multiplier * 2^attempt) capped at a maximum, + // plus a random jitter so a fleet of clients does not retry in lockstep. + // First waits are ~250ms, 500ms, 1s, ... capped at 10s. + private static final long RETRY_BACKOFF_MULTIPLIER_MS = 125; + private static final long RETRY_BACKOFF_MAX_MS = 10_000; + private static final long RETRY_JITTER_MAX_MS = 250; + + private static final Duration HTTP_CONNECT_TIMEOUT = Duration.ofSeconds(15); + private static final Duration HTTP_REQUEST_TIMEOUT = Duration.ofSeconds(60); + // Pre-computed parameter field names to avoid string concatenation in hot path private static final String[] PARAM_NAMES = new String[64]; static { @@ -152,28 +165,45 @@ public class SpiceClient implements AutoCloseable { } } - private String appId; - private String apiKey; - private String userAgent; - private URI flightAddress; - private URI httpAddress; - private int maxRetries; - private String tlsClientCertFile; - private String tlsClientKeyFile; - private String tlsRootCertFile; - private FlightSqlClient flightClient; - private CredentialCallOption authCallOptions = null; - private BufferAllocator allocator; + private final String appId; + private final String apiKey; + private final String userAgent; + private final URI flightAddress; + private final URI httpAddress; + private final int maxRetries; + private final int channelCount; + private final Duration queryTimeout; + private final String tlsClientCertFile; + private final String tlsClientKeyFile; + private final String tlsRootCertFile; + private final BufferAllocator allocator; + private final PreparedStatementCache statementCache; + private final AtomicInteger channelSelector = new AtomicInteger(); + + /** + * The active connection generation. Volatile so the query hot path can + * snapshot it without locking; rebuilt under the client monitor by + * {@link #reset()}, lazy rebuild, and UNAUTHENTICATED recovery. + */ + private volatile FlightChannel[] channels; + + /** + * Channels replaced by {@link #reset()} or an auth rebuild whose transport + * has been gracefully shut down but whose Flight client (and buffer + * allocator) cannot be closed yet because RPCs or result streams may still + * be in flight on them. Swept once the transport reports terminated. + * Guarded by the client monitor. + */ + private final List retiredChannels = new ArrayList<>(); private volatile boolean closed = false; - + + // HttpClient for refresh operations, created lazily on first use so clients + // that never call refreshDataset() don't pay for its selector thread. + private volatile HttpClient httpClient; + // Cached retryers (immutable, thread-safe) - private Retryer adbcRetryer; private Retryer flightRetryer; - - // ADBC resources for parameterized queries - private volatile boolean adbcInitialized = false; - private AdbcDatabase adbcDatabase; - private AdbcConnection adbcConnection; + private Retryer readerRetryer; /** * Returns a new instance of SpiceClientBuilder @@ -187,7 +217,7 @@ public static SpiceClientBuilder builder() throws URISyntaxException { /** * Constructs a new SpiceClient instance with the specified parameters - * + * * @param appId the application ID used to identify the client * application * @param apiKey the API key used for authentication with Spice.ai @@ -196,7 +226,7 @@ public static SpiceClientBuilder builder() throws URISyntaxException { * Spice.ai * services * @param httpAddress the URI of the Spice.ai runtime HTTP address - * + * * @param maxRetries the maximum number of connection retries for the * client * @param userAgent the user agent string @@ -210,11 +240,20 @@ public SpiceClient(String appId, String apiKey, URI flightAddress, URI httpAddre public SpiceClient(String appId, String apiKey, URI flightAddress, URI httpAddress, int maxRetries, String userAgent, long memoryLimitMB, String tlsClientCertFile, String tlsClientKeyFile) { - this(appId, apiKey, flightAddress, httpAddress, maxRetries, userAgent, memoryLimitMB, tlsClientCertFile, tlsClientKeyFile, null); + this(appId, apiKey, flightAddress, httpAddress, maxRetries, userAgent, memoryLimitMB, tlsClientCertFile, + tlsClientKeyFile, null); } public SpiceClient(String appId, String apiKey, URI flightAddress, URI httpAddress, int maxRetries, - String userAgent, long memoryLimitMB, String tlsClientCertFile, String tlsClientKeyFile, String tlsRootCertFile) { + String userAgent, long memoryLimitMB, String tlsClientCertFile, String tlsClientKeyFile, + String tlsRootCertFile) { + this(appId, apiKey, flightAddress, httpAddress, maxRetries, userAgent, memoryLimitMB, tlsClientCertFile, + tlsClientKeyFile, tlsRootCertFile, DEFAULT_CHANNEL_COUNT, null, DEFAULT_STATEMENT_CACHE_SIZE); + } + + SpiceClient(String appId, String apiKey, URI flightAddress, URI httpAddress, int maxRetries, + String userAgent, long memoryLimitMB, String tlsClientCertFile, String tlsClientKeyFile, + String tlsRootCertFile, int channelCount, Duration queryTimeout, int statementCacheSize) { this.appId = appId; this.apiKey = apiKey; this.maxRetries = maxRetries; @@ -223,6 +262,9 @@ public SpiceClient(String appId, String apiKey, URI flightAddress, URI httpAddre this.tlsClientCertFile = tlsClientCertFile; this.tlsClientKeyFile = tlsClientKeyFile; this.tlsRootCertFile = tlsRootCertFile; + this.channelCount = channelCount; + this.queryTimeout = queryTimeout; + this.statementCache = new PreparedStatementCache(statementCacheSize); // Arrow Flight requires URI to be grpc protocol, convert http/https for // convinience @@ -241,22 +283,9 @@ public SpiceClient(String appId, String apiKey, URI flightAddress, URI httpAddre : memoryLimitMB * BYTES_PER_MB; this.allocator = new RootAllocator(memoryLimitBytes); - // Build the HTTP client with optional mTLS support - HttpClient.Builder httpBuilder = HttpClient.newBuilder() - .connectTimeout(Duration.ofSeconds(15)); - if (this.tlsRootCertFile != null || (this.tlsClientCertFile != null && this.tlsClientKeyFile != null)) { - try { - javax.net.ssl.SSLContext sslContext = buildSslContext(); - httpBuilder.sslContext(sslContext); - } catch (Exception e) { - throw new RuntimeException("Failed to configure TLS for HTTP client", e); - } - } - this.httpClient = httpBuilder.build(); - try { - // Build the Flight client (channel + auth handshake) - buildFlightClient(); + // Build the Flight channels (gRPC channel + auth handshake each) + buildFlightChannels(); // Initialize cached retryers (immutable, built once) initRetryers(); @@ -269,20 +298,164 @@ public SpiceClient(String appId, String apiKey, URI flightAddress, URI httpAddre throw e; } - logger.debug("SpiceClient initialized - flightAddress={}, appId={}", this.flightAddress, this.appId); + logger.debug("SpiceClient initialized - flightAddress={}, appId={}, channels={}", + this.flightAddress, this.appId, this.channelCount); + } + + /** + * A single gRPC connection to the Flight endpoint together with its + * per-connection auth token and pre-computed call options. + */ + private static final class FlightChannel { + /** The underlying transport, retained for graceful retirement. */ + final ManagedChannel grpcChannel; + final FlightSqlClient client; + /** Options for control-plane RPCs (GetFlightInfo, prepare, DoPut): auth + optional timeout. */ + final CallOption[] callOptions; + /** Options for DoGet streams: auth only — a deadline would kill long-running result streams. */ + final CallOption[] streamOptions; + + FlightChannel(ManagedChannel grpcChannel, FlightSqlClient client, CallOption[] callOptions, + CallOption[] streamOptions) { + this.grpcChannel = grpcChannel; + this.client = client; + this.callOptions = callOptions; + this.streamOptions = streamOptions; + } + } + + /** + * A server-side prepared statement bound to the channel that created it and + * tagged with the connection generation it belongs to. + */ + private static final class CachedStatement { + final FlightSqlClient.PreparedStatement statement; + final FlightChannel channel; + final Object generation; + + CachedStatement(FlightSqlClient.PreparedStatement statement, FlightChannel channel, Object generation) { + this.statement = statement; + this.channel = channel; + this.generation = generation; + } + } + + /** + * Bounded cache of idle prepared statements keyed by SQL text. Reusing a + * prepared statement saves the CreatePreparedStatement and + * ClosePreparedStatement round trips on every repeated query. + * + *

All methods are cheap map operations under the cache monitor — no RPC + * is ever performed while holding the lock. Statements from a previous + * connection generation are rejected on {@link #give} and drained by + * {@link #swapGeneration} when the transport is rebuilt.

+ */ + private static final class PreparedStatementCache { + private final int maxIdleStatements; + private final HashMap> idleBySql = new HashMap<>(); + private int idleCount; + private Object generation; + + PreparedStatementCache(int maxIdleStatements) { + this.maxIdleStatements = maxIdleStatements; + } + + synchronized CachedStatement borrow(String sql) { + ArrayDeque deque = idleBySql.get(sql); + if (deque == null) { + return null; + } + CachedStatement statement = deque.poll(); + if (statement != null) { + idleCount--; + if (deque.isEmpty()) { + idleBySql.remove(sql); + } + } + return statement; + } + + /** + * Offers a statement back for reuse. Returns false when the statement + * must be closed by the caller instead (cache full, disabled, or the + * statement belongs to a previous connection generation). + */ + synchronized boolean give(String sql, CachedStatement statement) { + if (statement.generation != generation || idleCount >= maxIdleStatements) { + return false; + } + idleBySql.computeIfAbsent(sql, k -> new ArrayDeque<>(2)).push(statement); + idleCount++; + return true; + } + + /** + * Starts a new connection generation, returning all drained statements + * for the caller to close outside the lock. + */ + synchronized List swapGeneration(Object newGeneration) { + this.generation = newGeneration; + if (idleCount == 0) { + return Collections.emptyList(); + } + List drained = new ArrayList<>(idleCount); + for (ArrayDeque deque : idleBySql.values()) { + drained.addAll(deque); + } + idleBySql.clear(); + idleCount = 0; + return drained; + } + } + + /** + * A view over a parameter root whose close() is a no-op. + * {@link FlightSqlClient.PreparedStatement#clearParameters()} closes the + * root it was given; passing this wrapper keeps ownership of the real root + * (and its buffers) with the SDK so it can be closed exactly once. + */ + private static final class NonOwningRoot extends VectorSchemaRoot { + NonOwningRoot(VectorSchemaRoot delegate) { + super(delegate.getSchema(), delegate.getFieldVectors(), delegate.getRowCount()); + } + + @Override + public void close() { + // The delegate owns the vectors; SpiceClient closes it. + } } - + /** - * Builds or rebuilds the Flight client, including the gRPC channel and auth handshake. - * This method is called during construction and after {@link #reset()}. + * Builds (or rebuilds) all Flight channels, including the gRPC channels and + * auth handshakes. Called during construction, after {@link #reset()}, and + * on UNAUTHENTICATED recovery. * - *

The gRPC channel is configured with:

+ *

Each gRPC channel is configured with:

*
    *
  • {@code dns:///} target scheme for periodic DNS re-resolution behind load balancers
  • *
  • HTTP/2 keep-alive (30s interval, 10s timeout) to detect dead connections quickly
  • *
*/ - private synchronized void buildFlightClient() { + private synchronized void buildFlightChannels() { + FlightChannel[] newChannels = new FlightChannel[channelCount]; + try { + for (int i = 0; i < channelCount; i++) { + newChannels[i] = buildFlightChannel(); + } + } catch (RuntimeException | Error e) { + for (FlightChannel channel : newChannels) { + if (channel != null) { + closeChannelQuietly(channel, e); + } + } + throw e; + } + this.channels = newChannels; + // Any cached prepared statements belong to the previous connections. + closeStatementsQuietly(statementCache.swapGeneration(newChannels)); + } + + private FlightChannel buildFlightChannel() { // Build a gRPC channel using forTarget() with the "dns:///" scheme so that // gRPC's DnsNameResolver periodically re-resolves the hostname. This is critical // for long-lived clients connecting to load-balanced endpoints (e.g. AWS ALBs) @@ -332,39 +505,50 @@ private synchronized void buildFlightClient() { ManagedChannel channel = channelBuilder.build(); try { + CredentialCallOption auth = null; + FlightClient client; if (Strings.isNullOrEmpty(apiKey)) { - FlightClient client = FlightGrpcUtils.createFlightClient(allocator, channel); - this.flightClient = new FlightSqlClient(client); - logger.debug("Flight client built (unauthenticated) - target={}", target); - return; - } - - // prepare additional headers to insert into Flight requests - Map headers = new HashMap<>(); - String uaString; - if (Strings.isNullOrEmpty(userAgent)) { - uaString = Config.getUserAgent(); + client = FlightGrpcUtils.createFlightClient(allocator, channel); + logger.debug("Flight channel built (unauthenticated) - target={}", target); } else { - // Prepend the user-supplied user agent string with the Spice.ai user agent - uaString = userAgent + " " + Config.getUserAgent(); - } - headers.put("User-Agent", uaString); + // prepare additional headers to insert into Flight requests + Map headers = new HashMap<>(); + String uaString; + if (Strings.isNullOrEmpty(userAgent)) { + uaString = Config.getUserAgent(); + } else { + // Prepend the user-supplied user agent string with the Spice.ai user agent + uaString = userAgent + " " + Config.getUserAgent(); + } + headers.put("User-Agent", uaString); - final ClientIncomingAuthHeaderMiddleware.Factory authFactory = new ClientIncomingAuthHeaderMiddleware.Factory( - new ClientBearerHeaderHandler()); + final ClientIncomingAuthHeaderMiddleware.Factory authFactory = + new ClientIncomingAuthHeaderMiddleware.Factory(new ClientBearerHeaderHandler()); - // Combine auth and custom header middleware into a single factory - final HeaderAuthMiddlewareFactory combinedFactory = new HeaderAuthMiddlewareFactory(authFactory, headers); + // Combine auth and custom header middleware into a single factory + final HeaderAuthMiddlewareFactory combinedFactory = new HeaderAuthMiddlewareFactory(authFactory, + headers); - List middleware = new ArrayList<>(); - middleware.add(combinedFactory); + List middleware = new ArrayList<>(); + middleware.add(combinedFactory); - final FlightClient client = FlightGrpcUtils.createFlightClient(allocator, channel, middleware); - client.handshake(new CredentialCallOption(new BasicAuthCredentialWriter(this.appId, this.apiKey))); - this.authCallOptions = authFactory.getCredentialCallOption(); - this.flightClient = new FlightSqlClient(client); + client = FlightGrpcUtils.createFlightClient(allocator, channel, middleware); + client.handshake(new CredentialCallOption(new BasicAuthCredentialWriter(this.appId, this.apiKey))); + auth = authFactory.getCredentialCallOption(); + logger.debug("Flight channel built (authenticated) - target={}, appId={}", target, this.appId); + } - logger.debug("Flight client built (authenticated) - target={}, appId={}", target, this.appId); + List options = new ArrayList<>(2); + if (auth != null) { + options.add(auth); + } + CallOption[] streamOptions = options.toArray(new CallOption[0]); + if (queryTimeout != null) { + options.add(CallOptions.timeout(queryTimeout.toMillis(), java.util.concurrent.TimeUnit.MILLISECONDS)); + } + CallOption[] callOptions = options.toArray(new CallOption[0]); + + return new FlightChannel(channel, new FlightSqlClient(client), callOptions, streamOptions); } catch (Exception e) { // Ensure the channel is shut down if client creation or handshake fails // to avoid leaking threads and file descriptors on repeated rebuild attempts. @@ -378,21 +562,36 @@ private synchronized void buildFlightClient() { } /** - * Ensures the Flight client is connected, rebuilding it if necessary - * (e.g. after a {@link #reset()} call). + * Returns the current channels, lazily rebuilding them if necessary + * (e.g. after a failed rebuild). Lock-free on the hot path. */ - private synchronized void ensureFlightClient() { - if (this.closed) { - throw new IllegalStateException("SpiceClient is closed"); + private FlightChannel[] currentChannels() { + FlightChannel[] snapshot = this.channels; + if (snapshot != null) { + return snapshot; } - if (this.flightClient == null) { - buildFlightClient(); + synchronized (this) { + if (closed) { + throw new IllegalStateException("SpiceClient is closed"); + } + if (this.channels == null) { + buildFlightChannels(); + } + return this.channels; } } + private FlightChannel selectChannel(FlightChannel[] snapshot) { + if (snapshot.length == 1) { + return snapshot[0]; + } + return snapshot[Math.floorMod(channelSelector.getAndIncrement(), snapshot.length)]; + } + /** - * Resets the underlying gRPC transport by closing the current Flight client and ADBC connections, - * then immediately establishes a fresh connection with a new DNS lookup and TLS handshake. + * Resets the underlying gRPC transport by closing the current Flight channels and + * cached prepared statements, then immediately establishes fresh connections with + * a new DNS lookup and TLS handshake. * This ensures the next {@link #query(String)} or {@link #queryWithParams(String, Object...)} * call does not incur connection setup overhead. * @@ -422,28 +621,46 @@ public synchronized void reset() { } logger.info("Resetting SpiceClient transport"); - // Close ADBC resources (they maintain a separate Flight connection) - closeADBC(); - - // Close Flight client (this also shuts down the underlying gRPC channel) - if (this.flightClient != null) { - try { - this.flightClient.close(); - } catch (Exception e) { - logger.warn("Error closing Flight client during reset: {}", e.getMessage()); - } - this.flightClient = null; - } - this.authCallOptions = null; + // Close cached prepared statements while their channels still accept + // RPCs, then retire the channels (graceful shutdown; deferred close so + // concurrent in-flight queries can finish on them). + closeStatementsQuietly(statementCache.swapGeneration(null)); + retireChannelsLocked(); + sweepRetiredChannels(false); - // Eagerly re-establish the connection so the next query has no setup overhead - buildFlightClient(); + // Eagerly re-establish the connections so the next query has no setup overhead + buildFlightChannels(); logger.info("SpiceClient transport reset and reconnected."); } /** - * Initializes the cached retryer instances. + * Rebuilds the transport once when the server reports UNAUTHENTICATED — + * typically an expired handshake bearer token on a long-lived client. The + * generation check makes concurrent failures trigger a single rebuild, and + * the retryer then re-attempts the query on the fresh connection. + */ + private void maybeRebuildOnAuthError(FlightRuntimeException e, FlightChannel[] expected) { + if (e.status().code() != FlightStatusCode.UNAUTHENTICATED || Strings.isNullOrEmpty(apiKey)) { + return; + } + synchronized (this) { + if (closed || this.channels != expected) { + return; + } + logger.info("Received UNAUTHENTICATED from server; re-authenticating with a fresh handshake"); + closeStatementsQuietly(statementCache.swapGeneration(null)); + retireChannelsLocked(); + sweepRetiredChannels(false); + try { + buildFlightChannels(); + } catch (RuntimeException rebuildError) { + // Leave channels null: the next query attempt rebuilds lazily. + logger.warn("Re-authentication failed: {}", rebuildError.getMessage()); + } + } + } + /** * Builds an SSLContext configured with the custom CA and/or client certificate * for the JDK HTTP client. @@ -467,15 +684,17 @@ private javax.net.ssl.SSLContext buildSslContext() throws Exception { // Parse the PEM private key using BouncyCastle java.security.PrivateKey privateKey; - try (java.io.FileReader keyReader = new java.io.FileReader(this.tlsClientKeyFile, java.nio.charset.StandardCharsets.UTF_8); - org.bouncycastle.openssl.PEMParser pemParser = new org.bouncycastle.openssl.PEMParser(keyReader)) { + try (java.io.FileReader keyReader = new java.io.FileReader(this.tlsClientKeyFile, + java.nio.charset.StandardCharsets.UTF_8); + org.bouncycastle.openssl.PEMParser pemParser = new org.bouncycastle.openssl.PEMParser(keyReader)) { Object parsed = pemParser.readObject(); org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter converter = new org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter().setProvider("BC"); if (parsed instanceof org.bouncycastle.asn1.pkcs.PrivateKeyInfo) { privateKey = converter.getPrivateKey((org.bouncycastle.asn1.pkcs.PrivateKeyInfo) parsed); } else if (parsed instanceof org.bouncycastle.openssl.PEMKeyPair) { - privateKey = converter.getPrivateKey(((org.bouncycastle.openssl.PEMKeyPair) parsed).getPrivateKeyInfo()); + privateKey = converter + .getPrivateKey(((org.bouncycastle.openssl.PEMKeyPair) parsed).getPrivateKeyInfo()); } else { throw new IllegalArgumentException("Unsupported PEM key format in " + this.tlsClientKeyFile); } @@ -485,7 +704,7 @@ private javax.net.ssl.SSLContext buildSslContext() throws Exception { java.security.KeyStore keyStore = java.security.KeyStore.getInstance("PKCS12"); keyStore.load(null, null); keyStore.setKeyEntry("client", privateKey, new char[0], - new java.security.cert.Certificate[]{clientCert}); + new java.security.cert.Certificate[] { clientCert }); javax.net.ssl.KeyManagerFactory kmf = javax.net.ssl.KeyManagerFactory.getInstance( javax.net.ssl.KeyManagerFactory.getDefaultAlgorithm()); kmf.init(keyStore, new char[0]); @@ -494,7 +713,8 @@ private javax.net.ssl.SSLContext buildSslContext() throws Exception { if (this.tlsRootCertFile != null) { java.security.cert.CertificateFactory cf = java.security.cert.CertificateFactory.getInstance("X.509"); - java.security.KeyStore trustStore = java.security.KeyStore.getInstance(java.security.KeyStore.getDefaultType()); + java.security.KeyStore trustStore = java.security.KeyStore + .getInstance(java.security.KeyStore.getDefaultType()); trustStore.load(null, null); try (java.io.FileInputStream fis = new java.io.FileInputStream(this.tlsRootCertFile)) { int i = 0; @@ -514,39 +734,49 @@ private javax.net.ssl.SSLContext buildSslContext() throws Exception { } /** + * Returns the HTTP client, creating it on first use. + */ + private HttpClient httpClient() { + HttpClient client = this.httpClient; + if (client != null) { + return client; + } + synchronized (this) { + if (this.httpClient == null) { + HttpClient.Builder httpBuilder = HttpClient.newBuilder() + .connectTimeout(HTTP_CONNECT_TIMEOUT); + if (this.tlsRootCertFile != null + || (this.tlsClientCertFile != null && this.tlsClientKeyFile != null)) { + try { + httpBuilder.sslContext(buildSslContext()); + } catch (Exception e) { + throw new RuntimeException("Failed to configure TLS for HTTP client", e); + } + } + this.httpClient = httpBuilder.build(); + } + return this.httpClient; + } + } + + /** + * Initializes the cached retryer instances. * Called from constructor and must be called after maxRetries is set. */ private void initRetryers() { - this.adbcRetryer = RetryerBuilder.newBuilder() - .retryIfException(throwable -> { - if (throwable instanceof AdbcException) { - AdbcStatusCode status = ((AdbcException) throwable).getStatus(); - switch (status) { - case IO: // maps to gRPC UNAVAILABLE - case UNKNOWN: - case TIMEOUT: // maps to gRPC DEADLINE_EXCEEDED - case INTERNAL: - return true; - default: - return false; - } - } - return false; - }) - .withWaitStrategy(WaitStrategies.fibonacciWait()) - .withStopStrategy(StopStrategies.stopAfterAttempt(this.maxRetries + 1)) - .build(); - - this.flightRetryer = RetryerBuilder.newBuilder() - .retryIfException(throwable -> { - if (throwable instanceof FlightRuntimeException) { - FlightRuntimeException flightException = (FlightRuntimeException) throwable; - CallStatus status = flightException.status(); - return shouldRetry(status); - } - return false; - }) - .withWaitStrategy(WaitStrategies.fibonacciWait()) + this.flightRetryer = buildRetryer(); + this.readerRetryer = buildRetryer(); + } + + private Retryer buildRetryer() { + return RetryerBuilder.newBuilder() + .retryIfException(throwable -> throwable instanceof FlightRuntimeException + && shouldRetry(((FlightRuntimeException) throwable).status())) + .withWaitStrategy(WaitStrategies.join( + WaitStrategies.exponentialWait(RETRY_BACKOFF_MULTIPLIER_MS, RETRY_BACKOFF_MAX_MS, + java.util.concurrent.TimeUnit.MILLISECONDS), + WaitStrategies.randomWait(1, java.util.concurrent.TimeUnit.MILLISECONDS, + RETRY_JITTER_MAX_MS, java.util.concurrent.TimeUnit.MILLISECONDS))) .withStopStrategy(StopStrategies.stopAfterAttempt(this.maxRetries + 1)) .build(); } @@ -562,6 +792,9 @@ public FlightStream query(String sql) throws ExecutionException { if (Strings.isNullOrEmpty(sql)) { throw new IllegalArgumentException("No SQL query provided"); } + if (closed) { + throw new IllegalStateException("Cannot query with a closed SpiceClient"); + } logger.debug("Executing query: {}", sql); try { @@ -576,13 +809,19 @@ public FlightStream query(String sql) throws ExecutionException { } /** - * Executes a parameterized SQL query using ADBC. + * Executes a parameterized SQL query using Arrow Flight SQL prepared + * statements. * This is the recommended method for queries with user input to prevent SQL * injection. * Parameters should use positional placeholders ($1, $2, etc.) in the SQL * query. * *

+ * Prepared statements are cached and reused for repeated executions of the + * same SQL, saving the create/close round trips on every call. + *

+ * + *

* Parameters can be: *

*
    @@ -595,13 +834,13 @@ public FlightStream query(String sql) throws ExecutionException { *

    * Example usage: *

    - * + * *
          * // With automatic type inference
          * ArrowReader reader = client.queryWithParams(
          *     "SELECT * FROM table WHERE id = $1 AND name = $2",
          *     123, "test");
    -     * 
    +     *
          * // With explicit types
          * ArrowReader reader = client.queryWithParams(
          *     "SELECT * FROM table WHERE id = $1 AND amount = $2",
    @@ -625,180 +864,172 @@ public ArrowReader queryWithParams(String sql, Object... params) throws Executio
     
             logger.debug("Executing parameterized query with {} parameters: {}", params != null ? params.length : 0, sql);
             try {
    -            initADBCIfNeeded();
    -            ArrowReader result = queryWithParamsInternal(sql, params);
    +            ArrowReader result = readerRetryer.call(() -> executeParameterizedQuery(sql, params));
                 logger.debug("Parameterized query executed successfully");
                 return result;
    -        } catch (AdbcException e) {
    -            logger.error("Parameterized query failed: {}", e.getMessage());
    -            throw new ExecutionException("Failed to execute parameterized query: " + e.getMessage(), e);
             } catch (RetryException e) {
                 Throwable err = e.getLastFailedAttempt().getExceptionCause();
    -            logger.error("Parameterized query failed after {} attempts: {}", e.getNumberOfFailedAttempts(), err.getMessage());
    +            logger.error("Parameterized query failed after {} attempts: {}", e.getNumberOfFailedAttempts(),
    +                    err.getMessage());
                 throw new ExecutionException("Failed to execute parameterized query due to error: " + err.toString(), err);
             }
         }
     
         /**
    -     * Initializes the ADBC connection if not already initialized.
    -     * Uses double-checked locking to avoid monitor contention on the hot path.
    +     * Executes a single attempt of a parameterized query: bind + execute on a
    +     * cached prepared statement when available, falling back to a freshly
    +     * prepared statement when the cached one fails (e.g. the server restarted
    +     * and no longer knows the handle).
          */
    -    private void initADBCIfNeeded() throws AdbcException {
    -        if (adbcInitialized) {
    -            return;
    -        }
    -        synchronized (this) {
    -            if (adbcInitialized) {
    -                return;
    +    private ArrowReader executeParameterizedQuery(String sql, Object[] params) throws IOException {
    +        final FlightChannel[] snapshot = currentChannels();
    +        try (VectorSchemaRoot paramRoot = (params != null && params.length > 0) ? createParameterRoot(params)
    +                : null) {
    +            CachedStatement statement = statementCache.borrow(sql);
    +            FlightInfo info = null;
    +            if (statement != null) {
    +                try {
    +                    info = bindAndExecute(statement, paramRoot);
    +                } catch (FlightRuntimeException e) {
    +                    logger.debug("Cached prepared statement failed ({}); preparing a fresh statement",
    +                            e.status().code());
    +                    closeStatementQuietly(statement);
    +                    statement = null;
    +                }
    +            }
    +            if (info == null) {
    +                FlightChannel channel = selectChannel(snapshot);
    +                try {
    +                    statement = new CachedStatement(channel.client.prepare(sql, channel.callOptions), channel,
    +                            snapshot);
    +                    info = bindAndExecute(statement, paramRoot);
    +                } catch (FlightRuntimeException e) {
    +                    if (statement != null) {
    +                        closeStatementQuietly(statement);
    +                    }
    +                    maybeRebuildOnAuthError(e, snapshot);
    +                    throw e;
    +                }
                 }
     
    -        logger.debug("Initializing ADBC connection");
    -        
    -        // Format the URI for ADBC FlightSQL driver
    -        String uri = this.flightAddress.toString();
    +            // The statement is healthy: return it to the cache before opening the
    +            // result stream so other threads can reuse it immediately.
    +            //
    +            // Safe because only bind+execute mutates statement state — and the
    +            // cache hands a statement to at most one thread at a time for that
    +            // phase. The result stream is served by the self-contained ticket
    +            // minted by execute(), not by live statement state: the previous
    +            // implementation closed the server-side statement outright at this
    +            // point (reader still open) and reads were unaffected, so a later
    +            // re-bind on the same handle cannot alter an already-issued ticket.
    +            if (!statementCache.give(sql, statement)) {
    +                closeStatementQuietly(statement);
    +            }
     
    -        // Convert grpc+tls:// to grpc+tls:// format expected by ADBC
    -        // and grpc+tcp:// to grpc:// format
    -        if (uri.startsWith("grpc+tcp://")) {
    -            uri = "grpc://" + uri.substring("grpc+tcp://".length());
    +            try {
    +                return new FlightInfoReader(allocator, statement.channel.client, statement.channel.streamOptions,
    +                        info);
    +            } catch (FlightRuntimeException e) {
    +                maybeRebuildOnAuthError(e, snapshot);
    +                throw e;
    +            }
             }
    +    }
     
    -        // Build driver options
    -        Map options = new HashMap<>();
    -        AdbcDriver.PARAM_URI.set(options, uri);
    +    /**
    +     * Binds the parameters (if any) and executes the prepared statement.
    +     * The parameter root is passed as a non-owning view: the statement's
    +     * clearParameters() closes only the view, never the caller's root.
    +     */
    +    private FlightInfo bindAndExecute(CachedStatement statement, VectorSchemaRoot paramRoot) {
    +        if (paramRoot != null) {
    +            statement.statement.setParameters(new NonOwningRoot(paramRoot));
    +        }
    +        try {
    +            return statement.statement.execute(statement.channel.callOptions);
    +        } finally {
    +            try {
    +                statement.statement.clearParameters();
    +            } catch (RuntimeException ignored) {
    +                // Best-effort: closing the non-owning view cannot fail meaningfully.
    +            }
    +        }
    +    }
     
    -        // Add authentication if available
    -        if (!Strings.isNullOrEmpty(apiKey)) {
    -            AdbcDriver.PARAM_USERNAME.set(options, appId);
    -            AdbcDriver.PARAM_PASSWORD.set(options, apiKey);
    +    private void closeStatementQuietly(CachedStatement statement) {
    +        try {
    +            statement.statement.close();
    +        } catch (Exception e) {
    +            logger.debug("Error closing prepared statement: {}", e.getMessage());
             }
    +    }
     
    -        // Add user agent header
    -        String uaString;
    -        if (Strings.isNullOrEmpty(userAgent)) {
    -            uaString = Config.getUserAgent();
    -        } else {
    -            uaString = userAgent + " " + Config.getUserAgent();
    +    private void closeStatementsQuietly(List statements) {
    +        for (CachedStatement statement : statements) {
    +            closeStatementQuietly(statement);
             }
    -        options.put("adbc.flight.sql.rpc.call_header.user-agent", uaString);
    +    }
     
    -        // Create the driver and database using local temporaries.
    -        // Only assign to fields after both open+connect succeed to avoid
    -        // leaking partially created resources on failure.
    -        FlightSqlDriver driver = new FlightSqlDriver(allocator);
    -        AdbcDatabase db = driver.open(options);
    -        AdbcConnection conn;
    +    private void closeChannelQuietly(FlightChannel channel, Throwable pending) {
             try {
    -            conn = db.connect();
    -        } catch (AdbcException e) {
    -            try { db.close(); } catch (Exception suppressed) { e.addSuppressed(suppressed); }
    -            throw e;
    -        }
    -        adbcDatabase = db;
    -        adbcConnection = conn;
    -        adbcInitialized = true;
    -        
    -        logger.debug("ADBC connection established - uri={}", uri);
    +            channel.client.close();
    +        } catch (Exception e) {
    +            if (pending != null) {
    +                pending.addSuppressed(e);
    +            } else {
    +                logger.warn("Error closing Flight channel: {}", e.getMessage());
    +            }
             }
         }
     
         /**
    -     * Closes the ADBC resources.
    +     * Retires the current channels: the gRPC transports are shut down
    +     * gracefully — in-flight RPCs and open result streams complete, new RPCs
    +     * fail with a retryable UNAVAILABLE — but the Flight clients (and their
    +     * buffer allocators) stay open until {@link #sweepRetiredChannels} sees
    +     * the transport terminate. Closing them inline would tear the allocator
    +     * out from under concurrent queries mid-RPC.
    +     * Must be called while holding the client monitor.
          */
    -    private void closeADBC() {
    -        adbcInitialized = false;
    -        if (adbcConnection != null) {
    -            try {
    -                adbcConnection.close();
    -                logger.debug("ADBC connection closed");
    -            } catch (Exception e) {
    -                logger.warn("Error closing ADBC connection: {}", e.getMessage());
    -            }
    -            adbcConnection = null;
    +    private void retireChannelsLocked() {
    +        FlightChannel[] snapshot = this.channels;
    +        if (snapshot == null) {
    +            return;
             }
    -        if (adbcDatabase != null) {
    +        for (FlightChannel channel : snapshot) {
                 try {
    -                adbcDatabase.close();
    -                logger.debug("ADBC database closed");
    -            } catch (Exception e) {
    -                logger.warn("Error closing ADBC database: {}", e.getMessage());
    +                channel.grpcChannel.shutdown();
    +            } catch (RuntimeException e) {
    +                logger.warn("Error shutting down Flight transport: {}", e.getMessage());
                 }
    -            adbcDatabase = null;
    +            retiredChannels.add(channel);
             }
    +        this.channels = null;
         }
     
         /**
    -     * Internal implementation of parameterized query execution.
    -     */
    -    private ArrowReader queryWithParamsInternal(String sql, Object... params)
    -            throws AdbcException, RetryException, ExecutionException {
    -        return adbcRetryer.call(() -> executeParameterizedQuery(sql, params));
    -    }
    -
    -    /**
    -     * Executes a single parameterized query using ADBC prepare/bind/execute
    -     * pattern.
    +     * Closes retired channels whose transport has fully terminated (no RPCs or
    +     * streams remain). With {@code force}, closes them regardless — used by
    +     * {@link #close()}, where interrupting anything still in flight is the
    +     * documented behavior.
    +     * Must be called while holding the client monitor.
          */
    -    private ArrowReader executeParameterizedQuery(String sql, Object... params) throws AdbcException {
    -        AdbcStatement stmt = adbcConnection.createStatement();
    -        VectorSchemaRoot paramRoot = null;
    -
    -        try {
    -            // Set the query
    -            stmt.setSqlQuery(sql);
    -
    -            // Prepare the statement
    -            stmt.prepare();
    -
    -            // Bind parameters if provided
    -            if (params != null && params.length > 0) {
    -                paramRoot = createParameterRoot(params);
    -                stmt.bind(paramRoot);
    -            }
    -
    -            // Execute the query - at this point parameters have been serialized
    -            AdbcStatement.QueryResult result = stmt.executeQuery();
    -            ArrowReader reader = result.getReader();
    -            
    -            // Now we can safely close the parameter root since it has been sent to server
    -            if (paramRoot != null) {
    -                paramRoot.close();
    -                paramRoot = null;
    -            }
    -
    -            // Close the statement eagerly — the reader holds its own Flight stream
    -            // and no longer needs the statement. This frees server-side resources
    -            // immediately rather than waiting for slow consumers to close the reader.
    -            try {
    -                stmt.close();
    -            } catch (Exception closeEx) {
    -                logger.warn("Error closing ADBC statement: {}", closeEx.getMessage());
    -            }
    -            
    -            return reader;
    -        } catch (AdbcException e) {
    -            // Clean up on error
    -            if (paramRoot != null) {
    -                try {
    -                    paramRoot.close();
    -                } catch (Exception closeEx) {
    -                    // Ignore close exception
    -                }
    -            }
    -            try {
    -                stmt.close();
    -            } catch (Exception closeEx) {
    -                // Ignore close exception
    +    private void sweepRetiredChannels(boolean force) {
    +        for (java.util.Iterator it = retiredChannels.iterator(); it.hasNext();) {
    +            FlightChannel channel = it.next();
    +            if (force || channel.grpcChannel.isTerminated()) {
    +                closeChannelQuietly(channel, null);
    +                it.remove();
                 }
    -            throw e;
             }
         }
     
         /**
          * Creates a VectorSchemaRoot containing the parameter values.
          * The caller is responsible for closing the returned root.
    +     * Package-private for tests.
          */
    -    private VectorSchemaRoot createParameterRoot(Object... params) throws AdbcException {
    +    VectorSchemaRoot createParameterRoot(Object... params) {
             final int numParams = params.length;
     
             // Single pass: build schema fields directly (no intermediate arrays)
    @@ -817,8 +1048,13 @@ private VectorSchemaRoot createParameterRoot(Object... params) throws AdbcExcept
             }
             Schema schema = new Schema(fields);
     
    -        // Create a VectorSchemaRoot and populate it
    +        // Create a VectorSchemaRoot sized for a single row of parameters —
    +        // without setInitialCapacity(1), allocateNew() reserves Arrow's default
    +        // ~3970-slot buffers per vector for what is always a 1-row binding.
             VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator);
    +        for (FieldVector vector : root.getFieldVectors()) {
    +            vector.setInitialCapacity(1);
    +        }
             root.allocateNew();
     
             // Populate vectors — read value from original params to avoid intermediate arrays
    @@ -831,9 +1067,9 @@ private VectorSchemaRoot createParameterRoot(Object... params) throws AdbcExcept
             }
     
             root.setRowCount(1);
    -        
    -        logger.debug("Created parameter root: rowCount={}, schema={}", 
    -            root.getRowCount(), root.getSchema());
    +
    +        logger.debug("Created parameter root: rowCount={}, schema={}",
    +                root.getRowCount(), root.getSchema());
     
             return root;
         }
    @@ -933,8 +1169,7 @@ private ArrowType inferArrowType(Object value) {
          * Uses setSafe methods to properly handle validity buffer and auto-expansion.
          */
         @SuppressWarnings("deprecation")
    -    private void appendValueToVector(FieldVector vector, int index, Object value, ArrowType type)
    -            throws AdbcException {
    +    private void appendValueToVector(FieldVector vector, int index, Object value, ArrowType type) {
             if (value == null) {
                 vector.setNull(index);
                 return;
    @@ -1073,14 +1308,13 @@ else if (vector instanceof DecimalVector) {
                     BigDecimal bd = (BigDecimal) value;
                     decVector.setSafe(index, bd);
                 } else {
    -                throw new AdbcException("Unsupported vector type: " + vector.getClass().getName(),
    -                        null, AdbcStatusCode.UNKNOWN, null, 0);
    +                throw new IllegalArgumentException("Unsupported vector type: " + vector.getClass().getName());
                 }
             } catch (ClassCastException e) {
    -            throw new AdbcException(
    +            throw new IllegalArgumentException(
                         "Cannot convert value of type " + value.getClass().getName() +
                                 " to " + vector.getClass().getSimpleName(),
    -                    e, AdbcStatusCode.INVALID_ARGUMENT, null, 0);
    +                    e);
             }
         }
     
    @@ -1134,7 +1368,7 @@ private boolean probe(String path, String expectedBody, boolean authenticate) {
                 HttpRequest request = buildProbeRequest(this.httpAddress, path,
                         authenticate ? this.apiKey : null);
     
    -            HttpResponse response = httpClient.send(request,
    +            HttpResponse response = httpClient().send(request,
                         HttpResponse.BodyHandlers.ofString());
     
                 if (response.statusCode() != 200) {
    @@ -1204,7 +1438,7 @@ public List runtimeStatus() throws ExecutionException {
             try {
                 HttpRequest request = buildProbeRequest(this.httpAddress, "/v1/status", this.apiKey);
     
    -            HttpResponse response = httpClient.send(request,
    +            HttpResponse response = httpClient().send(request,
                         HttpResponse.BodyHandlers.ofString());
     
                 if (response.statusCode() != 200) {
    @@ -1300,7 +1534,7 @@ public void refreshDataset(String dataset) throws ExecutionException {
         /**
          * Refreshes an accelerated dataset using the configured dataset acceleration
          * settings
    -     * 
    +     *
          * @param dataset        the name of the dataset to refresh
          * @param refreshOptions the refresh options to use when refreshing the dataset
          * @throws ExecutionException if there is an error refreshing the dataset
    @@ -1314,6 +1548,7 @@ public void refreshDataset(String dataset, RefreshOptions refreshOptions) throws
             try {
                 HttpRequest.Builder builder = HttpRequest.newBuilder()
                         .uri(new URI(String.format("%s/v1/datasets/%s/acceleration/refresh", this.httpAddress, dataset)))
    +                    .timeout(HTTP_REQUEST_TIMEOUT)
                         .header("Content-Type", "application/json")
                         .header("X-Spice-User-Agent", Config.getUserAgent());
     
    @@ -1325,10 +1560,11 @@ public void refreshDataset(String dataset, RefreshOptions refreshOptions) throws
                 }
     
                 HttpRequest request = builder.build();
    -            HttpResponse response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
    +            HttpResponse response = httpClient().send(request, HttpResponse.BodyHandlers.ofString());
     
                 if (response.statusCode() != 201) {
    -                logger.error("Dataset refresh failed - dataset={}, statusCode={}, response={}", dataset, response.statusCode(), response.body());
    +                logger.error("Dataset refresh failed - dataset={}, statusCode={}, response={}", dataset,
    +                        response.statusCode(), response.body());
                     throw new ExecutionException(
                             String.format("Failed to trigger dataset refresh. Status Code: %d, Response: %s",
                                     response.statusCode(),
    @@ -1350,18 +1586,26 @@ public void refreshDataset(String dataset, RefreshOptions refreshOptions) throws
         }
     
         private FlightStream queryInternal(String sql) {
    -        // Snapshot the client and auth under the lock, then release it
    -        // so concurrent queries can execute RPCs in parallel.
    -        final FlightSqlClient client;
    -        final CredentialCallOption auth;
    -        synchronized (this) {
    -            ensureFlightClient();
    -            client = this.flightClient;
    -            auth = this.authCallOptions;
    +        FlightChannel[] snapshot = currentChannels();
    +        FlightChannel channel = selectChannel(snapshot);
    +        try {
    +            FlightInfo flightInfo = channel.client.execute(sql, channel.callOptions);
    +            List endpoints = flightInfo.getEndpoints();
    +            if (endpoints.isEmpty()) {
    +                throw CallStatus.INTERNAL
    +                        .withDescription("Server returned a FlightInfo with no endpoints for the query")
    +                        .toRuntimeException();
    +            }
    +            if (endpoints.size() > 1) {
    +                logger.warn("Server returned {} endpoints; query() consumes only the first. "
    +                        + "Use queryWithParams() (ArrowReader) to consume all endpoints.", endpoints.size());
    +            }
    +            Ticket ticket = endpoints.get(0).getTicket();
    +            return channel.client.getStream(ticket, channel.streamOptions);
    +        } catch (FlightRuntimeException e) {
    +            maybeRebuildOnAuthError(e, snapshot);
    +            throw e;
             }
    -        FlightInfo flightInfo = client.execute(sql, auth);
    -        Ticket ticket = flightInfo.getEndpoints().get(0).getTicket();
    -        return client.getStream(ticket, auth);
         }
     
         private FlightStream queryInternalWithRetry(String sql) throws ExecutionException, RetryException {
    @@ -1375,6 +1619,10 @@ private boolean shouldRetry(CallStatus status) {
                 case TIMED_OUT:
                 case INTERNAL:
                     return true;
    +            case UNAUTHENTICATED:
    +                // Retried after maybeRebuildOnAuthError() re-handshakes; only
    +                // meaningful for authenticated clients.
    +                return !Strings.isNullOrEmpty(apiKey);
                 default:
                     return false;
             }
    @@ -1389,24 +1637,33 @@ public synchronized void close() throws Exception {
             logger.debug("Closing SpiceClient");
             List exceptions = new ArrayList<>();
     
    -        // Close ADBC resources first
    +        // Close cached prepared statements first (best-effort RPCs that need the
    +        // channels to still be open).
             try {
    -            closeADBC();
    +            closeStatementsQuietly(statementCache.swapGeneration(null));
             } catch (Exception e) {
    -            logger.warn("Error during ADBC cleanup: {}", e.getMessage());
    +            logger.warn("Error during prepared statement cleanup: {}", e.getMessage());
                 exceptions.add(e);
             }
     
    -        // Close Flight client
    -        if (this.flightClient != null) {
    +        // Close Flight channels: current ones and any retired channels whose
    +        // in-flight work never finished.
    +        FlightChannel[] snapshot = this.channels;
    +        this.channels = null;
    +        List toClose = new ArrayList<>(retiredChannels);
    +        retiredChannels.clear();
    +        if (snapshot != null) {
    +            toClose.addAll(java.util.Arrays.asList(snapshot));
    +        }
    +        for (FlightChannel channel : toClose) {
                 try {
    -                this.flightClient.close();
    -                logger.debug("Flight client closed");
    +                channel.client.close();
                 } catch (Exception e) {
    -                logger.warn("Error closing Flight client: {}", e.getMessage());
    +                logger.warn("Error closing Flight channel: {}", e.getMessage());
                     exceptions.add(e);
                 }
             }
    +        logger.debug("Flight channels closed");
     
             // Close allocator
             try {
    @@ -1427,7 +1684,7 @@ public synchronized void close() throws Exception {
                 logger.error("SpiceClient closed with {} error(s)", exceptions.size());
                 throw first;
             }
    -        
    +
             logger.debug("SpiceClient closed successfully");
         }
     }
    diff --git a/src/main/java/ai/spice/SpiceClientBuilder.java b/src/main/java/ai/spice/SpiceClientBuilder.java
    index cbba1fc..b2f95cc 100644
    --- a/src/main/java/ai/spice/SpiceClientBuilder.java
    +++ b/src/main/java/ai/spice/SpiceClientBuilder.java
    @@ -24,6 +24,7 @@ of this software and associated documentation files (the "Software"), to deal
     
     import java.net.URI;
     import java.net.URISyntaxException;
    +import java.time.Duration;
     
     import com.google.common.base.Strings;
     
    @@ -32,6 +33,11 @@ of this software and associated documentation files (the "Software"), to deal
      */
     public class SpiceClientBuilder {
     
    +    /** Maximum number of gRPC channels a single client may open. */
    +    private static final int MAX_CHANNEL_COUNT = 16;
    +    /** Maximum number of idle prepared statements the cache may hold. */
    +    private static final int MAX_STATEMENT_CACHE_SIZE = 1024;
    +
         private String appId;
         private String apiKey;
         private String userAgent;
    @@ -42,6 +48,9 @@ public class SpiceClientBuilder {
         private String tlsClientCertFile;
         private String tlsClientKeyFile;
         private String tlsRootCertFile;
    +    private int channelCount = SpiceClient.DEFAULT_CHANNEL_COUNT;
    +    private Duration queryTimeout;
    +    private int statementCacheSize = SpiceClient.DEFAULT_STATEMENT_CACHE_SIZE;
     
         /**
          * Constructs a new SpiceClientBuilder instance
    @@ -206,6 +215,70 @@ public SpiceClientBuilder withTlsRootCertFile(String caFile) {
             return this;
         }
     
    +    /**
    +     * Sets the number of gRPC channels (HTTP/2 connections) the client opens to
    +     * the Flight endpoint. Queries are distributed round-robin across channels.
    +     *
    +     * 

    A single HTTP/2 connection multiplexes all concurrent queries and is + * limited by the server's MAX_CONCURRENT_STREAMS and the throughput of one + * TCP connection. Increase this for highly concurrent workloads with large + * result streams. The default of 1 is appropriate for most applications.

    + * + * @param channelCount Number of connections, between 1 and 16. + * @return The current instance of SpiceClientBuilder for method chaining. + */ + public SpiceClientBuilder withChannelCount(int channelCount) { + if (channelCount < 1 || channelCount > MAX_CHANNEL_COUNT) { + throw new IllegalArgumentException( + "channelCount must be between 1 and " + MAX_CHANNEL_COUNT + ", got: " + channelCount); + } + this.channelCount = channelCount; + return this; + } + + /** + * Sets a deadline for query control-plane RPCs: query planning + * (GetFlightInfo), statement preparation, and parameter binding. Without a + * timeout, a hung server can block the calling thread indefinitely. + * + *

    The timeout intentionally does not apply to result streaming (DoGet) — + * large results may legitimately stream for longer than any planning + * deadline. Dead connections during streaming are detected by HTTP/2 + * keep-alive instead.

    + * + * @param queryTimeout The timeout, must be positive. + * @return The current instance of SpiceClientBuilder for method chaining. + */ + public SpiceClientBuilder withQueryTimeout(Duration queryTimeout) { + if (queryTimeout == null || queryTimeout.isZero() || queryTimeout.isNegative()) { + throw new IllegalArgumentException("queryTimeout must be positive, got: " + queryTimeout); + } + this.queryTimeout = queryTimeout; + return this; + } + + /** + * Sets the maximum number of idle prepared statements cached for reuse by + * {@link SpiceClient#queryWithParams(String, Object...)}. + * + *

    Reusing a prepared statement removes the CreatePreparedStatement and + * ClosePreparedStatement round trips from every repeated parameterized + * query. Set to 0 to disable caching (each query then prepares and closes + * its own statement, matching the pre-0.7 behavior).

    + * + * @param statementCacheSize Maximum idle statements, between 0 and 1024. Default is 64. + * @return The current instance of SpiceClientBuilder for method chaining. + */ + public SpiceClientBuilder withPreparedStatementCacheSize(int statementCacheSize) { + if (statementCacheSize < 0 || statementCacheSize > MAX_STATEMENT_CACHE_SIZE) { + throw new IllegalArgumentException( + "statementCacheSize must be between 0 and " + MAX_STATEMENT_CACHE_SIZE + ", got: " + + statementCacheSize); + } + this.statementCacheSize = statementCacheSize; + return this; + } + /** * Creates SpiceClient with provided parameters. * @@ -220,6 +293,7 @@ public SpiceClient build() { "Both tlsClientCertFile and tlsClientKeyFile must be provided together for mTLS. " + (hasCert ? "tlsClientKeyFile is missing." : "tlsClientCertFile is missing.")); } - return new SpiceClient(appId, apiKey, flightAddress, httpAddress, maxRetries, userAgent, memoryLimitMB, tlsClientCertFile, tlsClientKeyFile, tlsRootCertFile); + return new SpiceClient(appId, apiKey, flightAddress, httpAddress, maxRetries, userAgent, memoryLimitMB, + tlsClientCertFile, tlsClientKeyFile, tlsRootCertFile, channelCount, queryTimeout, statementCacheSize); } } diff --git a/src/main/java/ai/spice/Version.java b/src/main/java/ai/spice/Version.java index 23d79a7..2964cda 100644 --- a/src/main/java/ai/spice/Version.java +++ b/src/main/java/ai/spice/Version.java @@ -30,6 +30,6 @@ public class Version { public static final String SPICE_JAVA_VERSION; static { - SPICE_JAVA_VERSION = "0.6.0"; + SPICE_JAVA_VERSION = "0.7.0"; } } diff --git a/src/test/java/ai/spice/FlightInfoReaderTest.java b/src/test/java/ai/spice/FlightInfoReaderTest.java new file mode 100644 index 0000000..1b66157 --- /dev/null +++ b/src/test/java/ai/spice/FlightInfoReaderTest.java @@ -0,0 +1,87 @@ +/* +Copyright 2026 The Spice.ai OSS Authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +*/ + +package ai.spice; + +import java.util.Collections; + +import org.apache.arrow.flight.CallOption; +import org.apache.arrow.flight.FlightClient; +import org.apache.arrow.flight.FlightDescriptor; +import org.apache.arrow.flight.FlightInfo; +import org.apache.arrow.flight.Location; +import org.apache.arrow.flight.sql.FlightSqlClient; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; + +import junit.framework.TestCase; + +/** + * Unit tests for the SDK's FlightInfoReader (edge cases not reachable through + * the public query API). + */ +public class FlightInfoReaderTest extends TestCase { + + /** + * A FlightInfo with no endpoints yields an empty reader with the schema + * taken from the FlightInfo itself. + */ + public void testEmptyEndpointsYieldsEmptyReader() throws Exception { + try (TestFlightSqlServer server = new TestFlightSqlServer(); + BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + FlightClient flightClient = FlightClient.builder(allocator, + Location.forGrpcInsecure("localhost", server.getPort())).build(); + try (FlightSqlClient client = new FlightSqlClient(flightClient)) { + FlightInfo info = new FlightInfo(TestFlightSqlServer.RESULT_SCHEMA, + FlightDescriptor.command(new byte[0]), Collections.emptyList(), -1, -1); + try (FlightInfoReader reader = new FlightInfoReader( + allocator, client, new CallOption[0], info)) { + assertEquals(TestFlightSqlServer.RESULT_SCHEMA, + reader.getVectorSchemaRoot().getSchema()); + assertFalse("no batches expected", reader.loadNextBatch()); + assertEquals(0, reader.bytesRead()); + } + } + } + } + + /** + * The reader tracks bytesRead across batches and endpoints. + */ + public void testBytesReadAccumulatesAcrossEndpoints() throws Exception { + try (TestFlightSqlServer server = new TestFlightSqlServer()) { + server.endpointCount = 2; + server.batchesPerEndpoint = 2; + try (SpiceClient client = SpiceClient.builder().withFlightAddress(server.flightUri()).build(); + org.apache.arrow.vector.ipc.ArrowReader reader = client.queryWithParams("SELECT 1", 1)) { + long previous = 0; + int batches = 0; + while (reader.loadNextBatch()) { + batches++; + assertTrue("bytesRead must be monotonically increasing", reader.bytesRead() > previous); + previous = reader.bytesRead(); + } + assertEquals(4, batches); + } + } + } +} diff --git a/src/test/java/ai/spice/FlightQueryTest.java b/src/test/java/ai/spice/FlightQueryTest.java index 036a20d..9d4f3f2 100644 --- a/src/test/java/ai/spice/FlightQueryTest.java +++ b/src/test/java/ai/spice/FlightQueryTest.java @@ -78,7 +78,9 @@ public void testQuerySpiceCloudPlatform() throws ExecutionException, Interrupted public void testQuerySpiceOSS() throws ExecutionException, InterruptedException { try { + // maxRetries=0: keep the no-local-server skip path fast under real backoff SpiceClient spiceClient = SpiceClient.builder() + .withMaxRetries(0) .build(); String sql = "SELECT tpep_pickup_datetime, total_amount, passenger_count from taxi_trips limit 10;"; diff --git a/src/test/java/ai/spice/HikariCpCompatTest.java b/src/test/java/ai/spice/HikariCpCompatTest.java new file mode 100644 index 0000000..64394af --- /dev/null +++ b/src/test/java/ai/spice/HikariCpCompatTest.java @@ -0,0 +1,190 @@ +/* +Copyright 2026 The Spice.ai OSS Authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +*/ + +package ai.spice; + +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.Statement; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; + +import com.zaxxer.hikari.HikariConfig; +import com.zaxxer.hikari.HikariDataSource; + +import junit.framework.TestCase; + +/** + * Verifies that Spice works with HikariCP-pooled JDBC connections. + * + *

    {@link SpiceClient} itself is thread-safe and multiplexes queries over + * shared gRPC channels, so it does not need — and cannot use — a JDBC pool. + * Applications that want HikariCP (ORMs, existing JDBC infrastructure) + * connect to Spice through the Arrow Flight SQL JDBC driver instead; these + * tests prove that combination works end-to-end against a Flight SQL server, + * including authentication and prepared statements.

    + */ +public class HikariCpCompatTest extends TestCase { + + private TestFlightSqlServer server; + + @Override + protected void setUp() throws Exception { + super.setUp(); + server = new TestFlightSqlServer(); + } + + @Override + protected void tearDown() throws Exception { + server.close(); + super.tearDown(); + } + + private HikariDataSource newPool(int maxPoolSize, String user, String password) throws Exception { + HikariConfig config = new HikariConfig(); + config.setJdbcUrl("jdbc:arrow-flight-sql://localhost:" + server.getPort() + "/?useEncryption=false"); + if (user != null) { + config.setUsername(user); + config.setPassword(password); + } + config.setMaximumPoolSize(maxPoolSize); + config.setMinimumIdle(1); + // The Flight SQL JDBC driver does not implement Connection.isValid(), + // so give Hikari an explicit validation query. + config.setConnectionTestQuery("SELECT 1"); + return new HikariDataSource(config); + } + + private static long countRows(ResultSet resultSet) throws Exception { + long rows = 0; + while (resultSet.next()) { + assertNotNull(resultSet.getString("name")); + rows++; + } + return rows; + } + + public void testPlainStatementThroughPool() throws Exception { + try (HikariDataSource pool = newPool(2, null, null)) { + try (Connection connection = pool.getConnection(); + Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery("SELECT * FROM test")) { + assertEquals(server.expectedTotalRows(), countRows(resultSet)); + } + } + } + + public void testPreparedStatementThroughPool() throws Exception { + try (HikariDataSource pool = newPool(2, null, null)) { + try (Connection connection = pool.getConnection(); + PreparedStatement statement = connection + .prepareStatement("SELECT * FROM test WHERE id > ?")) { + statement.setLong(1, 5L); + try (ResultSet resultSet = statement.executeQuery()) { + assertEquals(server.expectedTotalRows(), countRows(resultSet)); + } + } + assertTrue("JDBC prepared statement should bind parameters via DoPut", + server.doPutParameterCalls.get() > 0); + } + } + + /** + * Connections are pooled and reused: cycling through the pool repeatedly + * must not accumulate server-side connections or fail. + */ + public void testConnectionReuseAcrossBorrows() throws Exception { + try (HikariDataSource pool = newPool(1, null, null)) { + for (int i = 0; i < 10; i++) { + try (Connection connection = pool.getConnection(); + Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery("SELECT " + i)) { + assertEquals(server.expectedTotalRows(), countRows(resultSet)); + } + } + } + } + + public void testConcurrentQueriesThroughPool() throws Exception { + final int threads = 4; + final int queriesPerThread = 5; + final AtomicInteger failures = new AtomicInteger(); + final AtomicLong totalRows = new AtomicLong(); + + try (HikariDataSource pool = newPool(threads, null, null)) { + ExecutorService executor = Executors.newFixedThreadPool(threads); + final CountDownLatch start = new CountDownLatch(1); + for (int t = 0; t < threads; t++) { + executor.submit(() -> { + try { + start.await(); + for (int i = 0; i < queriesPerThread; i++) { + try (Connection connection = pool.getConnection(); + Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery("SELECT * FROM test")) { + totalRows.addAndGet(countRows(resultSet)); + } + } + } catch (Throwable e) { + failures.incrementAndGet(); + } + }); + } + start.countDown(); + executor.shutdown(); + assertTrue(executor.awaitTermination(60, TimeUnit.SECONDS)); + + assertEquals("no pooled query should fail", 0, failures.get()); + assertEquals(threads * queriesPerThread * server.expectedTotalRows(), totalRows.get()); + } + } + + /** + * The same basic-auth handshake credentials the native client uses + * (user = app id, password = full API key) work through Hikari's + * username/password configuration. + */ + public void testAuthenticatedPool() throws Exception { + try (TestFlightSqlServer authServer = new TestFlightSqlServer("testapp", "testapp|secret")) { + HikariConfig config = new HikariConfig(); + config.setJdbcUrl("jdbc:arrow-flight-sql://localhost:" + authServer.getPort() + + "/?useEncryption=false"); + config.setUsername("testapp"); + config.setPassword("testapp|secret"); + config.setMaximumPoolSize(1); + config.setConnectionTestQuery("SELECT 1"); + try (HikariDataSource pool = new HikariDataSource(config); + Connection connection = pool.getConnection(); + Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery("SELECT * FROM test")) { + assertEquals(authServer.expectedTotalRows(), countRows(resultSet)); + } + assertTrue("pool connections must authenticate via handshake", + authServer.basicAuthValidations.get() >= 1); + } + } +} diff --git a/src/test/java/ai/spice/LocalFlightServerTest.java b/src/test/java/ai/spice/LocalFlightServerTest.java new file mode 100644 index 0000000..6bfa6bd --- /dev/null +++ b/src/test/java/ai/spice/LocalFlightServerTest.java @@ -0,0 +1,223 @@ +/* +Copyright 2026 The Spice.ai OSS Authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +*/ + +package ai.spice; + +import java.math.BigDecimal; +import java.time.LocalDate; +import java.util.List; +import java.util.concurrent.ExecutionException; + +import org.apache.arrow.flight.FlightStream; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.apache.arrow.vector.util.Text; + +import junit.framework.TestCase; + +/** + * End-to-end tests for both query paths against an in-process Flight SQL + * server. These run without any external Spice runtime. + */ +public class LocalFlightServerTest extends TestCase { + + private TestFlightSqlServer server; + private SpiceClient client; + + @Override + protected void setUp() throws Exception { + super.setUp(); + server = new TestFlightSqlServer(); + client = SpiceClient.builder() + .withFlightAddress(server.flightUri()) + .build(); + } + + @Override + protected void tearDown() throws Exception { + client.close(); + server.close(); + super.tearDown(); + } + + static long countRows(FlightStream stream) throws Exception { + long rows = 0; + while (stream.next()) { + rows += stream.getRoot().getRowCount(); + } + return rows; + } + + static long countRows(ArrowReader reader) throws Exception { + long rows = 0; + while (reader.loadNextBatch()) { + rows += reader.getVectorSchemaRoot().getRowCount(); + } + return rows; + } + + public void testPlainQueryReturnsAllRows() throws Exception { + try (FlightStream stream = client.query("SELECT * FROM test")) { + assertNotNull(stream.getSchema().findField("id")); + assertNotNull(stream.getSchema().findField("name")); + assertEquals(server.expectedTotalRows(), countRows(stream)); + } + assertEquals(1, server.getFlightInfoCalls.get()); + assertEquals(1, server.doGetCalls.get()); + } + + public void testQueryWithParamsReturnsAllRows() throws Exception { + try (ArrowReader reader = client.queryWithParams("SELECT * FROM test WHERE id > $1", 5L)) { + assertNotNull(reader.getVectorSchemaRoot().getSchema().findField("id")); + assertEquals(server.expectedTotalRows(), countRows(reader)); + assertTrue("bytesRead should be positive", reader.bytesRead() > 0); + } + assertEquals(1, server.createPreparedStatementCalls.get()); + assertEquals(1, server.doPutParameterCalls.get()); + assertEquals(1, server.doGetCalls.get()); + } + + public void testQueryWithParamsWithoutParameters() throws Exception { + try (ArrowReader reader = client.queryWithParams("SELECT 1")) { + assertEquals(server.expectedTotalRows(), countRows(reader)); + } + // No parameters bound: no DoPut should have happened. + assertEquals(0, server.doPutParameterCalls.get()); + } + + /** + * queryWithParams (ArrowReader) consumes every endpoint of a partitioned + * result. + */ + public void testQueryWithParamsReadsAllEndpoints() throws Exception { + server.endpointCount = 3; + server.batchesPerEndpoint = 2; + server.rowsPerBatch = 7; + try (ArrowReader reader = client.queryWithParams("SELECT * FROM test", 1)) { + assertEquals(3 * 2 * 7, countRows(reader)); + } + assertEquals("one DoGet per endpoint", 3, server.doGetCalls.get()); + } + + /** + * Documents the known limitation of the FlightStream-returning query() + * API: only the first endpoint of a partitioned result is consumed. + */ + public void testPlainQueryConsumesOnlyFirstEndpoint() throws Exception { + server.endpointCount = 3; + try (FlightStream stream = client.query("SELECT * FROM test")) { + assertEquals((long) server.batchesPerEndpoint * server.rowsPerBatch, countRows(stream)); + } + assertEquals(1, server.doGetCalls.get()); + } + + public void testParameterValuesArriveAtServer() throws Exception { + try (ArrowReader reader = client.queryWithParams( + "SELECT * FROM test WHERE a=$1 AND b=$2 AND c=$3 AND d=$4 AND e=$5 AND f=$6 AND g=$7 AND h=$8", + 42, 42L, "hello", 3.5, true, new byte[] { 1, 2, 3 }, + LocalDate.of(2026, 7, 30), new BigDecimal("12.34"))) { + countRows(reader); + } + List bound = server.lastBoundParameters; + assertNotNull(bound); + assertEquals(8, bound.size()); + assertEquals(42, bound.get(0)); + assertEquals(42L, bound.get(1)); + assertEquals(new Text("hello"), bound.get(2)); + assertEquals(3.5, bound.get(3)); + assertEquals(Boolean.TRUE, bound.get(4)); + assertTrue(bound.get(5) instanceof byte[]); + assertEquals(3, ((byte[]) bound.get(5)).length); + // DateDayVector surfaces the raw day count since the Unix epoch. + assertEquals((int) LocalDate.of(2026, 7, 30).toEpochDay(), bound.get(6)); + assertEquals(new BigDecimal("12.34"), bound.get(7)); + } + + public void testExplicitParamTypesArriveAtServer() throws Exception { + try (ArrowReader reader = client.queryWithParams( + "SELECT * FROM test WHERE a=$1 AND b=$2 AND c=$3", + Param.int32(7), Param.string("typed"), Param.float64(2.25))) { + countRows(reader); + } + List bound = server.lastBoundParameters; + assertNotNull(bound); + assertEquals(3, bound.size()); + assertEquals(7, bound.get(0)); + assertEquals(new Text("typed"), bound.get(1)); + assertEquals(2.25, bound.get(2)); + } + + public void testRepeatedQueriesReturnConsistentResults() throws Exception { + for (int i = 0; i < 5; i++) { + try (ArrowReader reader = client.queryWithParams("SELECT * FROM test WHERE id = $1", (long) i)) { + assertEquals(server.expectedTotalRows(), countRows(reader)); + } + try (FlightStream stream = client.query("SELECT " + i)) { + assertEquals(server.expectedTotalRows(), countRows(stream)); + } + } + } + + public void testEmptySqlThrowsIllegalArgument() throws Exception { + try { + client.query(""); + fail("Expected IllegalArgumentException"); + } catch (IllegalArgumentException expected) { + // expected + } + try { + client.queryWithParams("", 1); + fail("Expected IllegalArgumentException"); + } catch (IllegalArgumentException expected) { + // expected + } + } + + public void testQueryAfterCloseThrowsIllegalState() throws Exception { + SpiceClient shortLived = SpiceClient.builder().withFlightAddress(server.flightUri()).build(); + shortLived.close(); + try { + shortLived.query("SELECT 1"); + fail("Expected IllegalStateException"); + } catch (IllegalStateException expected) { + assertTrue(expected.getMessage().contains("closed")); + } + try { + shortLived.queryWithParams("SELECT $1", 1); + fail("Expected IllegalStateException"); + } catch (IllegalStateException expected) { + assertTrue(expected.getMessage().contains("closed")); + } + } + + public void testUnsupportedParameterTypeFailsWithoutRpc() throws Exception { + long infoCallsBefore = server.getFlightInfoCalls.get(); + try { + client.queryWithParams("SELECT $1", new Object()); + fail("Expected ExecutionException"); + } catch (ExecutionException e) { + assertTrue("cause should be IllegalArgumentException, got: " + e.getCause(), + e.getCause() instanceof IllegalArgumentException); + } + assertEquals("type inference failures must not reach the server", + infoCallsBefore, server.getFlightInfoCalls.get()); + } +} diff --git a/src/test/java/ai/spice/ParamRootTest.java b/src/test/java/ai/spice/ParamRootTest.java new file mode 100644 index 0000000..cd6d231 --- /dev/null +++ b/src/test/java/ai/spice/ParamRootTest.java @@ -0,0 +1,80 @@ +/* +Copyright 2026 The Spice.ai OSS Authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +*/ + +package ai.spice; + +import java.math.BigDecimal; +import java.time.LocalDate; + +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VectorSchemaRoot; + +import junit.framework.TestCase; + +/** + * Tests for parameter root construction: allocation sizing and type mapping. + */ +public class ParamRootTest extends TestCase { + + /** + * The parameter root always holds exactly one row; its vectors must be + * sized for that, not for Arrow's default ~3970-value capacity (which + * previously allocated ~48KB per string parameter per query). + */ + public void testParameterRootIsRightSizedForOneRow() throws Exception { + try (SpiceClient client = SpiceClient.builder().build()) { + try (VectorSchemaRoot root = client.createParameterRoot( + 42, 42L, "hello", 3.5, true, LocalDate.of(2026, 7, 30), new BigDecimal("9.99"))) { + assertEquals(1, root.getRowCount()); + for (FieldVector vector : root.getFieldVectors()) { + assertTrue( + "vector " + vector.getField() + " should be sized for ~1 value, capacity was " + + vector.getValueCapacity(), + vector.getValueCapacity() < 512); + } + } + } + } + + public void testParameterRootValues() throws Exception { + try (SpiceClient client = SpiceClient.builder().build()) { + try (VectorSchemaRoot root = client.createParameterRoot(7, "abc")) { + assertEquals(1, root.getRowCount()); + assertEquals("$1", root.getSchema().getFields().get(0).getName()); + assertEquals("$2", root.getSchema().getFields().get(1).getName()); + assertEquals(7, root.getVector(0).getObject(0)); + assertEquals("abc", root.getVector(1).getObject(0).toString()); + } + } + } + + public void testUnsupportedTypeThrowsIllegalArgument() throws Exception { + try (SpiceClient client = SpiceClient.builder().build()) { + try { + client.createParameterRoot(new Object()); + fail("Expected IllegalArgumentException"); + } catch (IllegalArgumentException expected) { + assertTrue(expected.getMessage().contains("Unsupported parameter type")); + } + } + } +} diff --git a/src/test/java/ai/spice/ParameterizedQueryTest.java b/src/test/java/ai/spice/ParameterizedQueryTest.java index 08e0fa2..b06d1c1 100644 --- a/src/test/java/ai/spice/ParameterizedQueryTest.java +++ b/src/test/java/ai/spice/ParameterizedQueryTest.java @@ -82,7 +82,7 @@ public void testParameterizedQuerySpiceCloud() throws Exception { * Test parameterized query with local Spice OSS runtime. */ public void testParameterizedQuerySpiceOSS() throws Exception { - try (SpiceClient spiceClient = SpiceClient.builder().build()) { + try (SpiceClient spiceClient = SpiceClient.builder().withMaxRetries(0).build()) { // Test with float parameter on tpch.orders String sql = "SELECT o_orderkey, o_totalprice FROM tpch.orders WHERE o_totalprice > $1 ORDER BY o_totalprice LIMIT 5"; @@ -113,7 +113,7 @@ public void testParameterizedQuerySpiceOSS() throws Exception { * Test parameterized query with multiple parameters. */ public void testMultipleParameters() throws Exception { - try (SpiceClient spiceClient = SpiceClient.builder().build()) { + try (SpiceClient spiceClient = SpiceClient.builder().withMaxRetries(0).build()) { String sql = "SELECT o_orderkey, o_totalprice FROM tpch.orders WHERE o_totalprice > $1 AND o_custkey > $2 LIMIT 5"; try (ArrowReader reader = spiceClient.queryWithParams(sql, 5000.0, 100)) { @@ -140,7 +140,7 @@ public void testMultipleParameters() throws Exception { * Test parameterized query with string parameter. */ public void testStringParameter() throws Exception { - try (SpiceClient spiceClient = SpiceClient.builder().build()) { + try (SpiceClient spiceClient = SpiceClient.builder().withMaxRetries(0).build()) { // Use c_mktsegment which is a string column in tpch.customer String sql = "SELECT c_custkey, c_mktsegment FROM tpch.customer WHERE c_mktsegment = $1 LIMIT 5"; @@ -168,7 +168,7 @@ public void testStringParameter() throws Exception { * Test parameterized query with explicit Param types. */ public void testExplicitParamTypes() throws Exception { - try (SpiceClient spiceClient = SpiceClient.builder().build()) { + try (SpiceClient spiceClient = SpiceClient.builder().withMaxRetries(0).build()) { // Use explicit int64 type on tpch.customer String sql = "SELECT c_custkey, c_name, c_nationkey FROM tpch.customer WHERE c_nationkey = $1 LIMIT 5"; @@ -195,7 +195,7 @@ public void testExplicitParamTypes() throws Exception { * Test parameterized query with mixed parameter types. */ public void testMixedParameterTypes() throws Exception { - try (SpiceClient spiceClient = SpiceClient.builder().build()) { + try (SpiceClient spiceClient = SpiceClient.builder().withMaxRetries(0).build()) { String sql = "SELECT o_orderkey, o_totalprice FROM tpch.orders WHERE o_totalprice > $1 AND o_orderstatus = $2 LIMIT 5"; try (ArrowReader reader = spiceClient.queryWithParams(sql, @@ -290,7 +290,7 @@ public void testParamFactoryMethods() { * Test that null SQL throws IllegalArgumentException. */ public void testNullSqlThrows() throws Exception { - try (SpiceClient spiceClient = SpiceClient.builder().build()) { + try (SpiceClient spiceClient = SpiceClient.builder().withMaxRetries(0).build()) { try { spiceClient.queryWithParams(null, 1); fail("Expected IllegalArgumentException"); @@ -310,7 +310,7 @@ public void testNullSqlThrows() throws Exception { * Test that empty SQL throws IllegalArgumentException. */ public void testEmptySqlThrows() throws Exception { - try (SpiceClient spiceClient = SpiceClient.builder().build()) { + try (SpiceClient spiceClient = SpiceClient.builder().withMaxRetries(0).build()) { try { spiceClient.queryWithParams("", 1); fail("Expected IllegalArgumentException"); diff --git a/src/test/java/ai/spice/PerfBenchmarkTest.java b/src/test/java/ai/spice/PerfBenchmarkTest.java new file mode 100644 index 0000000..24d6ab3 --- /dev/null +++ b/src/test/java/ai/spice/PerfBenchmarkTest.java @@ -0,0 +1,173 @@ +/* +Copyright 2026 The Spice.ai OSS Authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +*/ + +package ai.spice; + +import java.util.Arrays; +import java.util.concurrent.Callable; + +import org.apache.arrow.flight.FlightStream; +import org.apache.arrow.vector.ipc.ArrowReader; + +import junit.framework.TestCase; + +/** + * In-process micro-benchmarks for the two query paths. + * + *

    Latency numbers are printed for humans (loopback RTT makes wall-clock + * savings look small; against a real network every saved round trip is a full + * RTT). The assertions are on server-side RPC counts, which are deterministic + * and are the metric that dominates real-world latency.

    + */ +public class PerfBenchmarkTest extends TestCase { + + private static final int WARMUP_ITERATIONS = 50; + private static final int MEASURED_ITERATIONS = 300; + private static final String SQL = "SELECT * FROM bench WHERE id > $1"; + + private TestFlightSqlServer server; + + @Override + protected void setUp() throws Exception { + super.setUp(); + server = new TestFlightSqlServer(); + server.rowsPerBatch = 100; + } + + @Override + protected void tearDown() throws Exception { + server.close(); + super.tearDown(); + } + + private static long[] measure(int iterations, Callable op) throws Exception { + long[] samples = new long[iterations]; + for (int i = 0; i < iterations; i++) { + long start = System.nanoTime(); + op.call(); + samples[i] = System.nanoTime() - start; + } + Arrays.sort(samples); + return samples; + } + + private static String stats(String label, long[] sortedNanos) { + long p50 = sortedNanos[sortedNanos.length / 2]; + long p95 = sortedNanos[(int) (sortedNanos.length * 0.95)]; + long p99 = sortedNanos[(int) (sortedNanos.length * 0.99)]; + return String.format("%-28s p50=%6dus p95=%6dus p99=%6dus", label, + p50 / 1_000, p95 / 1_000, p99 / 1_000); + } + + public void testBenchmarkPlainQuery() throws Exception { + try (SpiceClient client = SpiceClient.builder().withFlightAddress(server.flightUri()).build()) { + Callable op = () -> { + try (FlightStream stream = client.query("SELECT * FROM bench")) { + return LocalFlightServerTest.countRows(stream); + } + }; + measure(WARMUP_ITERATIONS, op); + long getFlightInfoBefore = server.getFlightInfoCalls.get(); + long[] samples = measure(MEASURED_ITERATIONS, op); + System.out.println("[bench] " + stats("query()", samples)); + + // The plain query path is exactly 2 RPCs: GetFlightInfo + DoGet. + assertEquals(MEASURED_ITERATIONS, server.getFlightInfoCalls.get() - getFlightInfoBefore); + } + } + + public void testBenchmarkParameterizedQueryCachedVsUncached() throws Exception { + try (SpiceClient cachedClient = SpiceClient.builder() + .withFlightAddress(server.flightUri()) + .build(); + SpiceClient uncachedClient = SpiceClient.builder() + .withFlightAddress(server.flightUri()) + .withPreparedStatementCacheSize(0) + .build()) { + Callable cachedOp = () -> { + try (ArrowReader reader = cachedClient.queryWithParams(SQL, 5L)) { + return LocalFlightServerTest.countRows(reader); + } + }; + Callable uncachedOp = () -> { + try (ArrowReader reader = uncachedClient.queryWithParams(SQL, 5L)) { + return LocalFlightServerTest.countRows(reader); + } + }; + + // Warm both paths before measuring either, so JIT compilation of the + // shared code doesn't bias whichever block runs second. + measure(WARMUP_ITERATIONS, cachedOp); + measure(WARMUP_ITERATIONS, uncachedOp); + measure(WARMUP_ITERATIONS, cachedOp); + measure(WARMUP_ITERATIONS, uncachedOp); + + long preparesBeforeCached = server.createPreparedStatementCalls.get(); + long[] cachedSamples = measure(MEASURED_ITERATIONS, cachedOp); + long cachedPrepares = server.createPreparedStatementCalls.get() - preparesBeforeCached; + + long preparesBeforeUncached = server.createPreparedStatementCalls.get(); + long closesBeforeUncached = server.closePreparedStatementCalls.get(); + long[] uncachedSamples = measure(MEASURED_ITERATIONS, uncachedOp); + long uncachedPrepares = server.createPreparedStatementCalls.get() - preparesBeforeUncached; + long uncachedCloses = server.closePreparedStatementCalls.get() - closesBeforeUncached; + + System.out.println("[bench] " + stats("queryWithParams (cached)", cachedSamples)); + System.out.println("[bench] " + stats("queryWithParams (uncached)", uncachedSamples)); + System.out.printf( + "[bench] RPCs per %d queries: cached=%d prepares, uncached=%d prepares + %d closes " + + "(2 round trips saved per query on a real network)%n", + MEASURED_ITERATIONS, cachedPrepares, uncachedPrepares, uncachedCloses); + + // The RPC counts are the deterministic performance contract. + assertEquals("cached path must reuse the prepared statement", 0, cachedPrepares); + assertEquals("uncached path prepares on every query", + MEASURED_ITERATIONS, uncachedPrepares); + assertEquals("uncached path closes on every query", + MEASURED_ITERATIONS, uncachedCloses); + } + } + + /** + * Guards the parameter-root allocation fix: binding parameters for a + * query must stay in the low-kilobyte range, not Arrow's ~48KB-per-string + * default-capacity allocation. + */ + public void testParameterBindAllocationStaysSmall() throws Exception { + try (SpiceClient client = SpiceClient.builder().build()) { + long totalCapacityBytes = 0; + for (int i = 0; i < 100; i++) { + try (org.apache.arrow.vector.VectorSchemaRoot root = + client.createParameterRoot(42L, "hello-" + i, 3.14)) { + for (org.apache.arrow.vector.FieldVector vector : root.getFieldVectors()) { + totalCapacityBytes += vector.getBufferSize(); + } + } + } + System.out.printf("[bench] param-root buffer bytes for 100 binds of (long,string,double): %d%n", + totalCapacityBytes); + // Old behavior: >48KB per string vector alone → many MB over 100 binds. + assertTrue("parameter roots should stay small, used " + totalCapacityBytes + " bytes", + totalCapacityBytes < 200_000); + } + } +} diff --git a/src/test/java/ai/spice/RefreshHttpTest.java b/src/test/java/ai/spice/RefreshHttpTest.java new file mode 100644 index 0000000..71dc0eb --- /dev/null +++ b/src/test/java/ai/spice/RefreshHttpTest.java @@ -0,0 +1,117 @@ +/* +Copyright 2026 The Spice.ai OSS Authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +*/ + +package ai.spice; + +import java.io.IOException; +import java.io.InputStream; +import java.net.InetSocketAddress; +import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicReference; + +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpServer; + +import junit.framework.TestCase; + +/** + * Tests for refreshDataset against a local HTTP server. + */ +public class RefreshHttpTest extends TestCase { + + private HttpServer httpServer; + private final AtomicReference lastPath = new AtomicReference<>(); + private final AtomicReference lastBody = new AtomicReference<>(); + private final AtomicReference lastUserAgentHeader = new AtomicReference<>(); + private volatile int responseCode = 201; + + @Override + protected void setUp() throws Exception { + super.setUp(); + httpServer = HttpServer.create(new InetSocketAddress("localhost", 0), 0); + httpServer.createContext("/", this::handle); + httpServer.start(); + } + + @Override + protected void tearDown() throws Exception { + httpServer.stop(0); + super.tearDown(); + } + + private void handle(HttpExchange exchange) throws IOException { + lastPath.set(exchange.getRequestURI().getPath()); + lastUserAgentHeader.set(exchange.getRequestHeaders().getFirst("X-Spice-User-Agent")); + try (InputStream body = exchange.getRequestBody()) { + lastBody.set(new String(body.readAllBytes(), StandardCharsets.UTF_8)); + } + exchange.sendResponseHeaders(responseCode, -1); + exchange.close(); + } + + private SpiceClient newClient() throws Exception { + return SpiceClient.builder() + .withHttpAddress(new URI("http://localhost:" + httpServer.getAddress().getPort())) + .build(); + } + + public void testRefreshDataset() throws Exception { + try (SpiceClient client = newClient()) { + client.refreshDataset("taxi_trips"); + assertEquals("/v1/datasets/taxi_trips/acceleration/refresh", lastPath.get()); + assertEquals("{}", lastBody.get()); + assertNotNull(lastUserAgentHeader.get()); + assertTrue(lastUserAgentHeader.get().startsWith("spice-java/")); + } + } + + public void testRefreshDatasetWithOptions() throws Exception { + try (SpiceClient client = newClient()) { + client.refreshDataset("taxi_trips", new RefreshOptions() + .withRefreshMode("full") + .withRefreshSql("SELECT * FROM taxi_trips")); + assertTrue(lastBody.get().contains("\"refresh_mode\":\"full\"")); + assertTrue(lastBody.get().contains("\"refresh_sql\":\"SELECT * FROM taxi_trips\"")); + } + } + + public void testRefreshDatasetNon201Throws() throws Exception { + responseCode = 500; + try (SpiceClient client = newClient()) { + client.refreshDataset("taxi_trips"); + fail("Expected ExecutionException"); + } catch (ExecutionException e) { + assertTrue(e.getMessage().contains("500")); + } + } + + public void testRefreshDatasetEmptyNameThrows() throws Exception { + try (SpiceClient client = newClient()) { + client.refreshDataset(""); + fail("Expected IllegalArgumentException"); + } catch (IllegalArgumentException expected) { + // expected + } + } +} diff --git a/src/test/java/ai/spice/ResetTest.java b/src/test/java/ai/spice/ResetTest.java index e57728b..75fe4f0 100644 --- a/src/test/java/ai/spice/ResetTest.java +++ b/src/test/java/ai/spice/ResetTest.java @@ -53,7 +53,7 @@ protected void setUp() throws Exception { if (!serverAvailabilityChecked) { synchronized (ResetTest.class) { if (!serverAvailabilityChecked) { - try (SpiceClient probe = SpiceClient.builder().build()) { + try (SpiceClient probe = SpiceClient.builder().withMaxRetries(0).build()) { // Probe with taxi_trips (not SELECT 1) to ensure // the dataset is loaded and ready, not just that // the server is up. @@ -119,9 +119,10 @@ public void testCloseAfterReset() throws Exception { * After reset(), the client should have eagerly rebuilt its Flight connection. * A query should work without any NullPointerException. * If no local server is available, a connection error is expected. + * maxRetries=0 keeps the no-server failure path fast (real backoff otherwise). */ public void testQueryAfterResetRebuildsClient() throws Exception { - SpiceClient client = SpiceClient.builder().build(); + SpiceClient client = SpiceClient.builder().withMaxRetries(0).build(); client.reset(); try { @@ -145,11 +146,11 @@ public void testQueryAfterResetRebuildsClient() throws Exception { } /** - * After reset(), queryWithParams should work because the Flight client was - * eagerly rebuilt (ADBC is still lazily initialized on first parameterized query). + * After reset(), queryWithParams should work because the Flight channels were + * eagerly rebuilt (prepared statements are re-created on first use). */ public void testQueryWithParamsAfterResetRebuildsClient() throws Exception { - SpiceClient client = SpiceClient.builder().build(); + SpiceClient client = SpiceClient.builder().withMaxRetries(0).build(); client.reset(); try { @@ -184,7 +185,7 @@ public void testMultipleResetsAreIdempotent() throws Exception { * Each reset discards the transport; each query rebuilds it. */ public void testResetQueryResetQueryCycle() throws Exception { - SpiceClient client = SpiceClient.builder().build(); + SpiceClient client = SpiceClient.builder().withMaxRetries(0).build(); for (int i = 0; i < 3; i++) { client.reset(); @@ -269,7 +270,7 @@ public void testConcurrentResetDoesNotThrow() throws Exception { * (query may fail with connection errors, but not NPE or IllegalStateException.) */ public void testConcurrentResetAndQuery() throws Exception { - final SpiceClient client = SpiceClient.builder().build(); + final SpiceClient client = SpiceClient.builder().withMaxRetries(0).build(); final int iterations = 5; final CountDownLatch startLatch = new CountDownLatch(1); final List unexpectedErrors = new CopyOnWriteArrayList<>(); @@ -508,7 +509,7 @@ public void testResetThenQueryWithParamsIntegration() throws Exception { // Reset client.reset(); - // Second query (re-initializes both Flight and ADBC) + // Second query (rebuilds channels and re-prepares the statement) try (ArrowReader reader2 = client.queryWithParams( "SELECT total_amount FROM taxi_trips WHERE total_amount > $1 LIMIT 2", 0.0)) { diff --git a/src/test/java/ai/spice/ResilienceTest.java b/src/test/java/ai/spice/ResilienceTest.java new file mode 100644 index 0000000..8e87c3a --- /dev/null +++ b/src/test/java/ai/spice/ResilienceTest.java @@ -0,0 +1,264 @@ +/* +Copyright 2026 The Spice.ai OSS Authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +*/ + +package ai.spice; + +import java.time.Duration; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.List; + +import org.apache.arrow.flight.CallStatus; +import org.apache.arrow.flight.FlightRuntimeException; +import org.apache.arrow.flight.FlightStatusCode; +import org.apache.arrow.flight.FlightStream; +import org.apache.arrow.vector.ipc.ArrowReader; + +import junit.framework.TestCase; + +/** + * Tests for retry with backoff, query timeouts, automatic re-authentication, + * the channel pool, and reset()/query race safety — all against the + * in-process Flight SQL server. + */ +public class ResilienceTest extends TestCase { + + /** + * A transient UNAVAILABLE is retried and succeeds, and the retry waits a + * real backoff (the previous 1-2ms fibonacci wait made retries useless + * against real outages). + */ + public void testTransientUnavailableIsRetriedWithBackoff() throws Exception { + try (TestFlightSqlServer server = new TestFlightSqlServer(); + SpiceClient client = SpiceClient.builder().withFlightAddress(server.flightUri()).build()) { + server.failNextGetFlightInfo(1, CallStatus.UNAVAILABLE); + + long startNanos = System.nanoTime(); + try (FlightStream stream = client.query("SELECT 1")) { + assertEquals(server.expectedTotalRows(), LocalFlightServerTest.countRows(stream)); + } + long elapsedMs = (System.nanoTime() - startNanos) / 1_000_000; + + assertEquals("one failed and one successful attempt", 2, server.getFlightInfoCalls.get()); + assertTrue("retry should wait a real backoff (>=200ms), waited " + elapsedMs + "ms", + elapsedMs >= 200); + } + } + + public void testRetriesExhaustedSurfaceLastError() throws Exception { + try (TestFlightSqlServer server = new TestFlightSqlServer(); + SpiceClient client = SpiceClient.builder() + .withFlightAddress(server.flightUri()) + .withMaxRetries(1) + .build()) { + server.failNextGetFlightInfo(10, CallStatus.UNAVAILABLE); + try { + client.query("SELECT 1"); + fail("Expected ExecutionException"); + } catch (ExecutionException e) { + assertTrue(e.getCause() instanceof FlightRuntimeException); + assertEquals(FlightStatusCode.UNAVAILABLE, + ((FlightRuntimeException) e.getCause()).status().code()); + } + assertEquals("initial attempt plus one retry", 2, server.getFlightInfoCalls.get()); + } + } + + /** + * Non-retryable errors (e.g. INVALID_ARGUMENT for bad SQL) fail fast + * without burning retry attempts or backoff time. + */ + public void testNonRetryableErrorFailsFast() throws Exception { + try (TestFlightSqlServer server = new TestFlightSqlServer(); + SpiceClient client = SpiceClient.builder().withFlightAddress(server.flightUri()).build()) { + server.failNextGetFlightInfo(1, CallStatus.INVALID_ARGUMENT); + + long startNanos = System.nanoTime(); + try { + client.query("SELECT invalid"); + fail("Expected ExecutionException"); + } catch (ExecutionException e) { + assertEquals(FlightStatusCode.INVALID_ARGUMENT, + ((FlightRuntimeException) e.getCause()).status().code()); + } + long elapsedMs = (System.nanoTime() - startNanos) / 1_000_000; + + assertEquals("no retry for non-retryable status", 1, server.getFlightInfoCalls.get()); + assertTrue("should fail fast, took " + elapsedMs + "ms", elapsedMs < 200); + } + } + + /** + * withQueryTimeout bounds how long query planning may hang. Without it, + * this query would block for the full server delay. + */ + public void testQueryTimeoutBoundsPlanning() throws Exception { + try (TestFlightSqlServer server = new TestFlightSqlServer(); + SpiceClient client = SpiceClient.builder() + .withFlightAddress(server.flightUri()) + .withQueryTimeout(Duration.ofMillis(200)) + .withMaxRetries(0) + .build()) { + server.getFlightInfoDelayMs = 1_500; + + long startNanos = System.nanoTime(); + try { + client.query("SELECT 1"); + fail("Expected ExecutionException"); + } catch (ExecutionException e) { + assertEquals(FlightStatusCode.TIMED_OUT, + ((FlightRuntimeException) e.getCause()).status().code()); + } + long elapsedMs = (System.nanoTime() - startNanos) / 1_000_000; + assertTrue("timeout should fire well before the 1.5s server delay, took " + elapsedMs + "ms", + elapsedMs < 1_200); + // server.close() waits for the still-sleeping handler via awaitTermination. + } + } + + /** + * A long-lived client whose handshake bearer token expires re-handshakes + * automatically and the query succeeds — no manual reset() required. + */ + public void testExpiredBearerTokenRecoversAutomatically() throws Exception { + try (TestFlightSqlServer server = new TestFlightSqlServer("testapp", "testapp|secret"); + SpiceClient client = SpiceClient.builder() + .withFlightAddress(server.flightUri()) + .withApiKey("testapp|secret") + .build()) { + assertEquals("constructor performs the initial handshake", 1, server.basicAuthValidations.get()); + + try (FlightStream stream = client.query("SELECT 1")) { + assertEquals(server.expectedTotalRows(), LocalFlightServerTest.countRows(stream)); + } + + server.rejectNextBearerToken(); + + try (FlightStream stream = client.query("SELECT 2")) { + assertEquals(server.expectedTotalRows(), LocalFlightServerTest.countRows(stream)); + } + assertEquals("expired token must trigger exactly one re-handshake", 2, + server.basicAuthValidations.get()); + } + } + + /** + * Parameterized queries work against an authenticated server (the + * prepared-statement path shares the same authenticated channels; with + * ADBC it used a second unauthenticated-configured connection). + */ + public void testQueryWithParamsOnAuthenticatedClient() throws Exception { + try (TestFlightSqlServer server = new TestFlightSqlServer("testapp", "testapp|secret"); + SpiceClient client = SpiceClient.builder() + .withFlightAddress(server.flightUri()) + .withApiKey("testapp|secret") + .build()) { + try (ArrowReader reader = client.queryWithParams("SELECT * FROM t WHERE id=$1", 7L)) { + assertEquals(server.expectedTotalRows(), LocalFlightServerTest.countRows(reader)); + } + assertEquals("no extra handshake beyond the constructor's", 1, + server.basicAuthValidations.get()); + } + } + + public void testChannelPoolServesQueries() throws Exception { + try (TestFlightSqlServer server = new TestFlightSqlServer(); + SpiceClient client = SpiceClient.builder() + .withFlightAddress(server.flightUri()) + .withChannelCount(4) + .build()) { + for (int i = 0; i < 8; i++) { + try (FlightStream stream = client.query("SELECT " + i)) { + assertEquals(server.expectedTotalRows(), LocalFlightServerTest.countRows(stream)); + } + } + try (ArrowReader reader = client.queryWithParams("SELECT * FROM t WHERE id=$1", 1L)) { + assertEquals(server.expectedTotalRows(), LocalFlightServerTest.countRows(reader)); + } + assertEquals(9, server.getFlightInfoCalls.get()); + + // reset() rebuilds all channels and queries still work. + client.reset(); + try (FlightStream stream = client.query("SELECT after_reset")) { + assertEquals(server.expectedTotalRows(), LocalFlightServerTest.countRows(stream)); + } + } + } + + /** + * Regression test for the reset()-vs-queryWithParams race: concurrent + * resets while parameterized queries are in flight must never surface + * NullPointerException (previously the connection field could be nulled + * mid-query), and against a healthy server every query must succeed via + * the retry/fallback path. + */ + public void testConcurrentResetAndQueryWithParams() throws Exception { + try (TestFlightSqlServer server = new TestFlightSqlServer(); + SpiceClient client = SpiceClient.builder().withFlightAddress(server.flightUri()).build()) { + final int queries = 20; + final List failures = new CopyOnWriteArrayList<>(); + final CountDownLatch start = new CountDownLatch(1); + ExecutorService executor = Executors.newFixedThreadPool(2); + + executor.submit(() -> { + try { + start.await(); + for (int i = 0; i < 8; i++) { + client.reset(); + Thread.sleep(15); + } + } catch (Throwable e) { + failures.add(e); + } + }); + executor.submit(() -> { + try { + start.await(); + for (int i = 0; i < queries; i++) { + try (ArrowReader reader = client.queryWithParams( + "SELECT * FROM t WHERE id=$1", (long) i)) { + LocalFlightServerTest.countRows(reader); + } + Thread.sleep(5); + } + } catch (Throwable e) { + failures.add(e); + } + }); + + start.countDown(); + executor.shutdown(); + assertTrue(executor.awaitTermination(120, TimeUnit.SECONDS)); + + for (Throwable failure : failures) { + assertFalse("NPE means the reset race is back: " + failure, + failure instanceof NullPointerException); + } + assertTrue("all queries and resets should succeed against a healthy server, failures: " + + failures, failures.isEmpty()); + } + } +} diff --git a/src/test/java/ai/spice/SpiceClientBuilderTest.java b/src/test/java/ai/spice/SpiceClientBuilderTest.java index 0c0bdcb..5bbc296 100644 --- a/src/test/java/ai/spice/SpiceClientBuilderTest.java +++ b/src/test/java/ai/spice/SpiceClientBuilderTest.java @@ -312,6 +312,81 @@ public void testMultipleConfigOverrides() throws Exception { client.close(); } + // ==================== Channel Count / Timeout / Statement Cache ==================== + + public void testWithChannelCountValid() throws Exception { + try (SpiceClient client = SpiceClient.builder().withChannelCount(2).build()) { + assertNotNull(client); + } + } + + public void testWithChannelCountInvalid() throws Exception { + try { + SpiceClient.builder().withChannelCount(0); + fail("Expected IllegalArgumentException for channelCount=0"); + } catch (IllegalArgumentException expected) { + // expected + } + try { + SpiceClient.builder().withChannelCount(17); + fail("Expected IllegalArgumentException for channelCount=17"); + } catch (IllegalArgumentException expected) { + // expected + } + } + + public void testWithQueryTimeoutValid() throws Exception { + try (SpiceClient client = SpiceClient.builder() + .withQueryTimeout(java.time.Duration.ofSeconds(30)).build()) { + assertNotNull(client); + } + } + + public void testWithQueryTimeoutInvalid() throws Exception { + try { + SpiceClient.builder().withQueryTimeout(null); + fail("Expected IllegalArgumentException for null timeout"); + } catch (IllegalArgumentException expected) { + // expected + } + try { + SpiceClient.builder().withQueryTimeout(java.time.Duration.ZERO); + fail("Expected IllegalArgumentException for zero timeout"); + } catch (IllegalArgumentException expected) { + // expected + } + try { + SpiceClient.builder().withQueryTimeout(java.time.Duration.ofSeconds(-1)); + fail("Expected IllegalArgumentException for negative timeout"); + } catch (IllegalArgumentException expected) { + // expected + } + } + + public void testWithPreparedStatementCacheSizeValid() throws Exception { + try (SpiceClient client = SpiceClient.builder().withPreparedStatementCacheSize(0).build()) { + assertNotNull(client); + } + try (SpiceClient client = SpiceClient.builder().withPreparedStatementCacheSize(1024).build()) { + assertNotNull(client); + } + } + + public void testWithPreparedStatementCacheSizeInvalid() throws Exception { + try { + SpiceClient.builder().withPreparedStatementCacheSize(-1); + fail("Expected IllegalArgumentException for negative cache size"); + } catch (IllegalArgumentException expected) { + // expected + } + try { + SpiceClient.builder().withPreparedStatementCacheSize(1025); + fail("Expected IllegalArgumentException for oversized cache"); + } catch (IllegalArgumentException expected) { + // expected + } + } + // ==================== Close/Resource Management Tests ==================== public void testMultipleClose() throws Exception { diff --git a/src/test/java/ai/spice/StatementCacheTest.java b/src/test/java/ai/spice/StatementCacheTest.java new file mode 100644 index 0000000..14a521d --- /dev/null +++ b/src/test/java/ai/spice/StatementCacheTest.java @@ -0,0 +1,196 @@ +/* +Copyright 2026 The Spice.ai OSS Authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +*/ + +package ai.spice; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; + +import org.apache.arrow.vector.ipc.ArrowReader; + +import junit.framework.TestCase; + +/** + * Tests for the prepared statement cache: reuse, invalidation, and recovery. + * The server-side RPC counters make the round-trip savings directly + * observable and deterministic. + */ +public class StatementCacheTest extends TestCase { + + private static final String SQL = "SELECT * FROM test WHERE id = $1"; + + private TestFlightSqlServer server; + + @Override + protected void setUp() throws Exception { + super.setUp(); + server = new TestFlightSqlServer(); + } + + @Override + protected void tearDown() throws Exception { + server.close(); + super.tearDown(); + } + + private SpiceClient newClient() throws Exception { + return SpiceClient.builder().withFlightAddress(server.flightUri()).build(); + } + + private static void runQuery(SpiceClient client, String sql, Object... params) throws Exception { + try (ArrowReader reader = client.queryWithParams(sql, params)) { + long rows = LocalFlightServerTest.countRows(reader); + assertTrue("query should return rows", rows > 0); + } + } + + /** + * The core saving: repeated executions of the same SQL prepare exactly + * once instead of once per query (2 fewer round trips per call). + */ + public void testSameSqlPreparesOnce() throws Exception { + try (SpiceClient client = newClient()) { + for (int i = 0; i < 5; i++) { + runQuery(client, SQL, (long) i); + } + assertEquals("statement should be prepared exactly once", 1, + server.createPreparedStatementCalls.get()); + assertEquals("statement should not be closed between queries", 0, + server.closePreparedStatementCalls.get()); + assertEquals("every query still binds its own parameters", 5, + server.doPutParameterCalls.get()); + assertEquals(5, server.getFlightInfoCalls.get()); + assertEquals(5, server.doGetCalls.get()); + } + // close() drains the cache and closes the statement on the server. + assertEquals(1, server.closePreparedStatementCalls.get()); + } + + public void testDistinctSqlPreparesDistinctStatements() throws Exception { + try (SpiceClient client = newClient()) { + runQuery(client, "SELECT * FROM a WHERE id = $1", 1L); + runQuery(client, "SELECT * FROM b WHERE id = $1", 1L); + runQuery(client, "SELECT * FROM c WHERE id = $1", 1L); + assertEquals(3, server.createPreparedStatementCalls.get()); + // Re-run the first: cache hit, no new prepare. + runQuery(client, "SELECT * FROM a WHERE id = $1", 2L); + assertEquals(3, server.createPreparedStatementCalls.get()); + } + assertEquals(3, server.closePreparedStatementCalls.get()); + } + + /** + * Cache disabled restores the prepare/close-per-query behavior. + */ + public void testCacheDisabled() throws Exception { + try (SpiceClient client = SpiceClient.builder() + .withFlightAddress(server.flightUri()) + .withPreparedStatementCacheSize(0) + .build()) { + for (int i = 0; i < 3; i++) { + runQuery(client, SQL, (long) i); + } + } + assertEquals(3, server.createPreparedStatementCalls.get()); + assertEquals(3, server.closePreparedStatementCalls.get()); + } + + /** + * When the server no longer recognizes a cached handle (e.g. it + * restarted), the query transparently re-prepares and succeeds. + */ + public void testStaleHandleRePreparesTransparently() throws Exception { + try (SpiceClient client = newClient()) { + runQuery(client, SQL, 1L); + assertEquals(1, server.createPreparedStatementCalls.get()); + + server.invalidatePreparedStatements(); + + runQuery(client, SQL, 2L); + assertEquals("stale handle must trigger exactly one re-prepare", 2, + server.createPreparedStatementCalls.get()); + } + } + + /** + * reset() invalidates cached statements along with the transport. + */ + public void testResetInvalidatesCache() throws Exception { + try (SpiceClient client = newClient()) { + runQuery(client, SQL, 1L); + assertEquals(1, server.createPreparedStatementCalls.get()); + + client.reset(); + assertEquals("reset must close the cached statement", 1, + server.closePreparedStatementCalls.get()); + + runQuery(client, SQL, 2L); + assertEquals(2, server.createPreparedStatementCalls.get()); + } + } + + /** + * Concurrent queries on the same SQL are safe: a statement is used by one + * thread at a time, extra statements are prepared on cache misses, and + * everything is closed by client.close() — no server-side leaks. + */ + public void testConcurrentSameSqlQueries() throws Exception { + final int threads = 8; + final int queriesPerThread = 10; + final AtomicInteger failures = new AtomicInteger(); + final AtomicLong totalRows = new AtomicLong(); + + try (SpiceClient client = newClient()) { + ExecutorService executor = Executors.newFixedThreadPool(threads); + final CountDownLatch start = new CountDownLatch(1); + for (int t = 0; t < threads; t++) { + executor.submit(() -> { + try { + start.await(); + for (int i = 0; i < queriesPerThread; i++) { + try (ArrowReader reader = client.queryWithParams(SQL, (long) i)) { + totalRows.addAndGet(LocalFlightServerTest.countRows(reader)); + } + } + } catch (Throwable e) { + failures.incrementAndGet(); + } + }); + } + start.countDown(); + executor.shutdown(); + assertTrue(executor.awaitTermination(60, TimeUnit.SECONDS)); + + assertEquals("no query should fail", 0, failures.get()); + assertEquals(threads * queriesPerThread * server.expectedTotalRows(), totalRows.get()); + assertTrue("at most one prepared statement per thread, got " + + server.createPreparedStatementCalls.get(), + server.createPreparedStatementCalls.get() <= threads); + } + assertEquals("every prepared statement must be closed on client.close()", + server.createPreparedStatementCalls.get(), server.closePreparedStatementCalls.get()); + } +} diff --git a/src/test/java/ai/spice/TestFlightSqlServer.java b/src/test/java/ai/spice/TestFlightSqlServer.java new file mode 100644 index 0000000..0d47562 --- /dev/null +++ b/src/test/java/ai/spice/TestFlightSqlServer.java @@ -0,0 +1,362 @@ +/* +Copyright 2026 The Spice.ai OSS Authors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +*/ + +package ai.spice; + +import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; + +import org.apache.arrow.flight.CallHeaders; +import org.apache.arrow.flight.CallStatus; +import org.apache.arrow.flight.FlightDescriptor; +import org.apache.arrow.flight.FlightEndpoint; +import org.apache.arrow.flight.FlightInfo; +import org.apache.arrow.flight.FlightServer; +import org.apache.arrow.flight.FlightStream; +import org.apache.arrow.flight.Location; +import org.apache.arrow.flight.PutResult; +import org.apache.arrow.flight.Result; +import org.apache.arrow.flight.Ticket; +import org.apache.arrow.flight.auth2.BasicCallHeaderAuthenticator; +import org.apache.arrow.flight.auth2.CallHeaderAuthenticator; +import org.apache.arrow.flight.auth2.GeneratedBearerTokenAuthenticator; +import org.apache.arrow.flight.sql.NoOpFlightSqlProducer; +import org.apache.arrow.flight.sql.impl.FlightSql; +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; + +import com.google.protobuf.Any; +import com.google.protobuf.ByteString; + +/** + * In-process Flight SQL server for tests: serves deterministic data, counts + * every RPC, and supports failure injection (transient errors, prepared + * statement invalidation, artificial delays, bearer-token expiry). + */ +final class TestFlightSqlServer implements AutoCloseable { + + static final Schema RESULT_SCHEMA = new Schema(Arrays.asList( + Field.nullable("id", new ArrowType.Int(64, true)), + Field.nullable("name", ArrowType.Utf8.INSTANCE))); + + // ==================== RPC counters ==================== + final AtomicInteger createPreparedStatementCalls = new AtomicInteger(); + final AtomicInteger closePreparedStatementCalls = new AtomicInteger(); + final AtomicInteger getFlightInfoCalls = new AtomicInteger(); + final AtomicInteger doPutParameterCalls = new AtomicInteger(); + final AtomicInteger doGetCalls = new AtomicInteger(); + final AtomicInteger basicAuthValidations = new AtomicInteger(); + + // ==================== failure injection ==================== + private final AtomicInteger failNextGetFlightInfo = new AtomicInteger(); + private volatile CallStatus injectedFailureStatus = CallStatus.UNAVAILABLE; + private final AtomicBoolean rejectNextBearer = new AtomicBoolean(); + volatile long getFlightInfoDelayMs = 0; + + // ==================== result shape ==================== + volatile int endpointCount = 1; + volatile int batchesPerEndpoint = 1; + volatile int rowsPerBatch = 10; + + /** Parameter values (row 0 of each vector) captured from the last DoPut. */ + volatile List lastBoundParameters; + + private final Set validHandles = ConcurrentHashMap.newKeySet(); + private final AtomicLong handleCounter = new AtomicLong(); + + private final BufferAllocator allocator; + private final FlightServer server; + + /** Starts an unauthenticated server on an ephemeral port. */ + TestFlightSqlServer() throws Exception { + this(null, null); + } + + /** + * Starts a server on an ephemeral port. When expectedUser/expectedPassword + * are non-null, the server requires a basic-auth handshake and issues + * bearer tokens for subsequent calls (mirroring Spice's auth flow). + */ + TestFlightSqlServer(String expectedUser, String expectedPassword) throws Exception { + this.allocator = new RootAllocator(Long.MAX_VALUE); + FlightServer.Builder builder = FlightServer.builder( + allocator, Location.forGrpcInsecure("localhost", 0), new Producer()); + if (expectedUser != null) { + CallHeaderAuthenticator inner = new GeneratedBearerTokenAuthenticator( + new BasicCallHeaderAuthenticator((username, password) -> { + basicAuthValidations.incrementAndGet(); + if (!expectedUser.equals(username) || !expectedPassword.equals(password)) { + throw CallStatus.UNAUTHENTICATED.withDescription("invalid credentials") + .toRuntimeException(); + } + return () -> username; + })); + builder.headerAuthenticator(headers -> { + String authorization = headers.get("authorization"); + if (authorization != null && authorization.startsWith("Bearer ") + && rejectNextBearer.compareAndSet(true, false)) { + throw CallStatus.UNAUTHENTICATED.withDescription("token expired").toRuntimeException(); + } + return inner.authenticate(headers); + }); + } + this.server = builder.build(); + this.server.start(); + } + + int getPort() { + return server.getPort(); + } + + URI flightUri() throws Exception { + return new URI("grpc://localhost:" + getPort()); + } + + /** The next N GetFlightInfo calls (plain or prepared) fail with the given status. */ + void failNextGetFlightInfo(int count, CallStatus status) { + this.injectedFailureStatus = status; + this.failNextGetFlightInfo.set(count); + } + + /** Simulates a server restart: all existing prepared statement handles become unknown. */ + void invalidatePreparedStatements() { + validHandles.clear(); + } + + /** The next call presenting a bearer token is rejected with UNAUTHENTICATED. */ + void rejectNextBearerToken() { + rejectNextBearer.set(true); + } + + long expectedTotalRows() { + return (long) endpointCount * batchesPerEndpoint * rowsPerBatch; + } + + @Override + public void close() throws Exception { + server.shutdown(); + server.awaitTermination(); + allocator.close(); + } + + private void maybeInjectGetFlightInfoFailure() { + if (getFlightInfoDelayMs > 0) { + try { + Thread.sleep(getFlightInfoDelayMs); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + int remaining = failNextGetFlightInfo.get(); + while (remaining > 0) { + if (failNextGetFlightInfo.compareAndSet(remaining, remaining - 1)) { + throw injectedFailureStatus.withDescription("injected failure").toRuntimeException(); + } + remaining = failNextGetFlightInfo.get(); + } + } + + private List buildEndpoints(java.util.function.Function ticketForEndpoint) { + List endpoints = new ArrayList<>(endpointCount); + for (int i = 0; i < endpointCount; i++) { + endpoints.add(new FlightEndpoint(ticketForEndpoint.apply(i))); + } + return endpoints; + } + + private void serveData(int endpointIndex, org.apache.arrow.flight.FlightProducer.ServerStreamListener listener) { + doGetCalls.incrementAndGet(); + try (BufferAllocator streamAllocator = allocator.newChildAllocator("doget", 0, Long.MAX_VALUE); + VectorSchemaRoot root = VectorSchemaRoot.create(RESULT_SCHEMA, streamAllocator)) { + listener.start(root); + BigIntVector idVector = (BigIntVector) root.getVector("id"); + VarCharVector nameVector = (VarCharVector) root.getVector("name"); + for (int batch = 0; batch < batchesPerEndpoint; batch++) { + root.allocateNew(); + for (int row = 0; row < rowsPerBatch; row++) { + long id = endpointIndex * 1_000_000L + batch * 1_000L + row; + idVector.setSafe(row, id); + nameVector.setSafe(row, ("row-" + id).getBytes(StandardCharsets.UTF_8)); + } + root.setRowCount(rowsPerBatch); + listener.putNext(); + } + listener.completed(); + } + } + + private final class Producer extends NoOpFlightSqlProducer { + + @Override + public void createPreparedStatement(FlightSql.ActionCreatePreparedStatementRequest request, + CallContext context, StreamListener listener) { + createPreparedStatementCalls.incrementAndGet(); + String handle = "ps-" + handleCounter.incrementAndGet(); + validHandles.add(handle); + FlightSql.ActionCreatePreparedStatementResult result = FlightSql.ActionCreatePreparedStatementResult + .newBuilder() + .setPreparedStatementHandle(ByteString.copyFromUtf8(handle)) + // The dataset schema tells clients this statement is a query. + // JDBC (Avatica) clients treat an empty dataset schema as an + // UPDATE and would route execution to DoPut(...Update). + .setDatasetSchema(ByteString.copyFrom(RESULT_SCHEMA.serializeAsMessage())) + .setParameterSchema(ByteString.copyFrom(parameterSchemaFor(request.getQuery()) + .serializeAsMessage())) + .build(); + listener.onNext(new Result(Any.pack(result).toByteArray())); + listener.onCompleted(); + } + + @Override + public void closePreparedStatement(FlightSql.ActionClosePreparedStatementRequest request, + CallContext context, StreamListener listener) { + closePreparedStatementCalls.incrementAndGet(); + validHandles.remove(request.getPreparedStatementHandle().toStringUtf8()); + listener.onCompleted(); + } + + @Override + public Runnable acceptPutPreparedStatementQuery(FlightSql.CommandPreparedStatementQuery command, + CallContext context, FlightStream flightStream, StreamListener ackStream) { + return () -> { + doPutParameterCalls.incrementAndGet(); + String handle = command.getPreparedStatementHandle().toStringUtf8(); + if (!validHandles.contains(handle)) { + ackStream.onError(CallStatus.NOT_FOUND + .withDescription("unknown prepared statement handle").toRuntimeException()); + return; + } + List captured = new ArrayList<>(); + while (flightStream.next()) { + VectorSchemaRoot root = flightStream.getRoot(); + if (root.getRowCount() > 0) { + captured.clear(); + for (FieldVector vector : root.getFieldVectors()) { + captured.add(vector.getObject(0)); + } + } + } + lastBoundParameters = captured; + FlightSql.DoPutPreparedStatementResult metadata = FlightSql.DoPutPreparedStatementResult.newBuilder() + .setPreparedStatementHandle(command.getPreparedStatementHandle()) + .build(); + byte[] bytes = metadata.toByteArray(); + try (ArrowBuf buffer = allocator.buffer(bytes.length)) { + buffer.writeBytes(bytes); + ackStream.onNext(PutResult.metadata(buffer)); + } + ackStream.onCompleted(); + }; + } + + @Override + public FlightInfo getFlightInfoPreparedStatement(FlightSql.CommandPreparedStatementQuery command, + CallContext context, FlightDescriptor descriptor) { + getFlightInfoCalls.incrementAndGet(); + maybeInjectGetFlightInfoFailure(); + String handle = command.getPreparedStatementHandle().toStringUtf8(); + if (!validHandles.contains(handle)) { + throw CallStatus.NOT_FOUND.withDescription("unknown prepared statement handle").toRuntimeException(); + } + List endpoints = buildEndpoints(i -> new Ticket( + Any.pack(FlightSql.CommandPreparedStatementQuery.newBuilder() + .setPreparedStatementHandle(ByteString.copyFromUtf8(handle + "#" + i)) + .build()).toByteArray())); + return new FlightInfo(RESULT_SCHEMA, descriptor, endpoints, -1, -1); + } + + @Override + public void getStreamPreparedStatement(FlightSql.CommandPreparedStatementQuery command, + CallContext context, ServerStreamListener listener) { + serveData(parseEndpointIndex(command.getPreparedStatementHandle().toStringUtf8()), listener); + } + + @Override + public FlightInfo getFlightInfoStatement(FlightSql.CommandStatementQuery command, + CallContext context, FlightDescriptor descriptor) { + getFlightInfoCalls.incrementAndGet(); + maybeInjectGetFlightInfoFailure(); + List endpoints = buildEndpoints(i -> new Ticket( + Any.pack(FlightSql.TicketStatementQuery.newBuilder() + .setStatementHandle(ByteString.copyFromUtf8("q#" + i)) + .build()).toByteArray())); + return new FlightInfo(RESULT_SCHEMA, descriptor, endpoints, -1, -1); + } + + @Override + public void getStreamStatement(FlightSql.TicketStatementQuery ticket, + CallContext context, ServerStreamListener listener) { + serveData(parseEndpointIndex(ticket.getStatementHandle().toStringUtf8()), listener); + } + } + + /** + * Advertises one nullable BIGINT parameter per JDBC-style '?' placeholder. + * JDBC clients build their bind roots from this schema; the native SDK uses + * $N placeholders (no '?') and constructs its own parameter roots, so it + * receives an empty schema here — matching real servers' optionality. + */ + private static Schema parameterSchemaFor(String query) { + List fields = new ArrayList<>(); + for (int i = 0; i < query.length(); i++) { + if (query.charAt(i) == '?') { + fields.add(Field.nullable("$" + (fields.size() + 1), new ArrowType.Int(64, true))); + } + } + return new Schema(fields); + } + + /** + * Extracts the endpoint index from a "handle#index" ticket, converting any + * malformed input into a protocol-appropriate INVALID_ARGUMENT instead of + * an uncaught NumberFormatException. + */ + private static int parseEndpointIndex(String handle) { + int separator = handle.lastIndexOf('#'); + if (separator < 0 || separator == handle.length() - 1) { + throw CallStatus.INVALID_ARGUMENT + .withDescription("malformed ticket handle: " + handle).toRuntimeException(); + } + try { + return Integer.parseInt(handle.substring(separator + 1)); + } catch (NumberFormatException e) { + throw CallStatus.INVALID_ARGUMENT + .withDescription("malformed ticket handle: " + handle).toRuntimeException(); + } + } +} diff --git a/src/test/java/ai/spice/TpchIntegrationTest.java b/src/test/java/ai/spice/TpchIntegrationTest.java index 33b38e9..9c2fb6d 100644 --- a/src/test/java/ai/spice/TpchIntegrationTest.java +++ b/src/test/java/ai/spice/TpchIntegrationTest.java @@ -48,23 +48,31 @@ public class TpchIntegrationTest extends TestCase { private SpiceClient client; private boolean tpchAvailable = true; + // Availability is probed once for the whole class (with retries disabled) + // so a missing local runtime doesn't cost retry backoff per test method. + private static volatile Boolean tpchAvailableCached; + @Override protected void setUp() throws Exception { super.setUp(); - try { - client = SpiceClient.builder().build(); - // Check if TPC-H tables are available by querying tpch.customer - try (FlightStream stream = client.query("SELECT c_custkey FROM tpch.customer LIMIT 1")) { - stream.next(); - } - } catch (Exception e) { - // TPC-H tables not available (either no server or no TPC-H data) - tpchAvailable = false; - if (client != null) { - try { client.close(); } catch (Exception ex) {} - client = null; + if (tpchAvailableCached == null) { + synchronized (TpchIntegrationTest.class) { + if (tpchAvailableCached == null) { + try (SpiceClient probe = SpiceClient.builder().withMaxRetries(0).build(); + FlightStream stream = probe.query("SELECT c_custkey FROM tpch.customer LIMIT 1")) { + stream.next(); + tpchAvailableCached = Boolean.TRUE; + } catch (Exception e) { + // TPC-H tables not available (either no server or no TPC-H data) + tpchAvailableCached = Boolean.FALSE; + } + } } } + tpchAvailable = tpchAvailableCached; + if (tpchAvailable) { + client = SpiceClient.builder().build(); + } } @Override