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();
+ /**
+ * Ensures the Flight client is connected, rebuilding it if necessary
+ * (e.g. after a {@link #reset()} call).
+ */
+ private synchronized void ensureFlightClient() {
+ if (this.closed) {
+ throw new IllegalStateException("SpiceClient is closed");
+ }
+ if (this.flightClient == null) {
+ buildFlightClient();
+ }
+ }
+
+ /**
+ * 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.
+ * This ensures the next {@link #query(String)} or {@link #queryWithParams(String, Object...)}
+ * call does not incur connection setup overhead.
+ *
+ * Use this method to recover from unrecoverable transport failures such as:
+ *
+ * - SSLHandshakeException with mismatched certificates (e.g. load-balancer routing to wrong backend)
+ * - Persistent UNAVAILABLE errors after exhausting retries
+ * - Stale connections pinned to decommissioned backend IPs
+ *
+ *
+ * Example usage for long-lived clients:
+ * {@code
+ * try {
+ * return client.query(sql);
+ * } catch (ExecutionException e) {
+ * if (isTransportFailure(e.getCause())) {
+ * client.reset();
+ * return client.query(sql); // retry with fresh connection
+ * }
+ * throw e;
+ * }
+ * }
+ */
+ public synchronized void reset() {
+ if (closed) {
+ throw new IllegalStateException("Cannot reset a closed SpiceClient");
}
- headers.put("User-Agent", uaString);
+ logger.info("Resetting SpiceClient transport");
- final ClientIncomingAuthHeaderMiddleware.Factory authFactory = new ClientIncomingAuthHeaderMiddleware.Factory(
- new ClientBearerHeaderHandler());
+ // Close ADBC resources (they maintain a separate Flight connection)
+ closeADBC();
- // Combine auth and custom header middleware into a single factory
- final HeaderAuthMiddlewareFactory combinedFactory = new HeaderAuthMiddlewareFactory(authFactory, headers);
+ // 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;
- List middleware = new ArrayList<>();
- middleware.add(combinedFactory);
+ // Eagerly re-establish the connection so the next query has no setup overhead
+ buildFlightClient();
- 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);
-
- // Initialize cached retryers (immutable, built once)
- initRetryers();
-
- logger.debug("SpiceClient initialized (authenticated) - flightAddress={}, appId={}, target={}", this.flightAddress, this.appId, target);
+ logger.info("SpiceClient transport reset and reconnected.");
}
-
+
/**
* Initializes the cached retryer instances.
* Called from constructor and must be called after maxRetries is set.
@@ -297,11 +410,16 @@ private void initRetryers() {
this.adbcRetryer = RetryerBuilder.newBuilder()
.retryIfException(throwable -> {
if (throwable instanceof AdbcException) {
- String message = throwable.getMessage();
- return message != null && (message.contains("UNAVAILABLE") ||
- message.contains("UNKNOWN") ||
- message.contains("DEADLINE_EXCEEDED") ||
- message.contains("INTERNAL"));
+ 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;
})
@@ -391,6 +509,9 @@ public ArrowReader queryWithParams(String sql, Object... params) throws Executio
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 parameterized query with {} parameters: {}", params != null ? params.length : 0, sql);
try {
@@ -410,12 +531,16 @@ public ArrowReader queryWithParams(String sql, Object... params) throws Executio
/**
* Initializes the ADBC connection if not already initialized.
- * This is called lazily on the first parameterized query.
+ * Uses double-checked locking to avoid monitor contention on the hot path.
*/
- private synchronized void initADBCIfNeeded() throws AdbcException {
- if (adbcDatabase != null && adbcConnection != null) {
+ private void initADBCIfNeeded() throws AdbcException {
+ if (adbcInitialized) {
return;
}
+ synchronized (this) {
+ if (adbcInitialized) {
+ return;
+ }
logger.debug("Initializing ADBC connection");
@@ -447,18 +572,31 @@ private synchronized void initADBCIfNeeded() throws AdbcException {
}
options.put("adbc.flight.sql.rpc.call_header.user-agent", uaString);
- // Create the driver and database
+ // 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 = driver.open(options);
- adbcConnection = adbcDatabase.connect();
+ AdbcDatabase db = driver.open(options);
+ AdbcConnection conn;
+ 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);
+ }
}
/**
* Closes the ADBC resources.
*/
private void closeADBC() {
+ adbcInitialized = false;
if (adbcConnection != null) {
try {
adbcConnection.close();
@@ -515,6 +653,16 @@ private ArrowReader executeParameterizedQuery(String sql, Object... params) thro
// 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;
@@ -534,8 +682,6 @@ private ArrowReader executeParameterizedQuery(String sql, Object... params) thro
}
throw e;
}
- // Note: We don't close the statement here because the reader needs it
- // The statement will be closed when the reader is closed
}
/**
@@ -543,34 +689,21 @@ private ArrowReader executeParameterizedQuery(String sql, Object... params) thro
* The caller is responsible for closing the returned root.
*/
private VectorSchemaRoot createParameterRoot(Object... params) throws AdbcException {
- // Extract values and determine types
final int numParams = params.length;
- Object[] values = new Object[numParams];
- ArrowType[] types = new ArrowType[numParams];
+ // Single pass: build schema fields directly (no intermediate arrays)
+ List fields = new ArrayList<>(numParams);
for (int i = 0; i < numParams; i++) {
Object param = params[i];
-
+ ArrowType type;
if (param instanceof Param) {
Param p = (Param) param;
- values[i] = p.getValue();
- if (p.hasExplicitType()) {
- types[i] = p.getType();
- } else {
- types[i] = inferArrowType(p.getValue());
- }
+ type = p.hasExplicitType() ? p.getType() : inferArrowType(p.getValue());
} else {
- values[i] = param;
- types[i] = inferArrowType(param);
+ type = inferArrowType(param);
}
- }
-
- // Build the schema for parameters with pre-sized list
- List fields = new ArrayList<>(numParams);
- for (int i = 0; i < numParams; i++) {
- // Use cached field names for common cases (up to 64 params)
String fieldName = (i < PARAM_NAMES.length) ? PARAM_NAMES[i] : "$" + (i + 1);
- fields.add(new Field(fieldName, FieldType.nullable(types[i]), null));
+ fields.add(new Field(fieldName, FieldType.nullable(type), null));
}
Schema schema = new Schema(fields);
@@ -578,10 +711,12 @@ private VectorSchemaRoot createParameterRoot(Object... params) throws AdbcExcept
VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator);
root.allocateNew();
- // Append values to the vectors and set value count for each
+ // Populate vectors — read value from original params to avoid intermediate arrays
for (int i = 0; i < numParams; i++) {
+ Object param = params[i];
+ Object value = (param instanceof Param) ? ((Param) param).getValue() : param;
FieldVector vector = root.getVector(i);
- appendValueToVector(vector, 0, values[i], types[i]);
+ appendValueToVector(vector, 0, value, vector.getField().getType());
vector.setValueCount(1);
}
@@ -907,9 +1042,18 @@ public void refreshDataset(String dataset, RefreshOptions refreshOptions) throws
}
private FlightStream queryInternal(String sql) {
- FlightInfo flightInfo = this.flightClient.execute(sql, authCallOptions);
+ // 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;
+ }
+ FlightInfo flightInfo = client.execute(sql, auth);
Ticket ticket = flightInfo.getEndpoints().get(0).getTicket();
- return this.flightClient.getStream(ticket, authCallOptions);
+ return client.getStream(ticket, auth);
}
private FlightStream queryInternalWithRetry(String sql) throws ExecutionException, RetryException {
@@ -929,7 +1073,11 @@ private boolean shouldRetry(CallStatus status) {
}
@Override
- public void close() throws Exception {
+ public synchronized void close() throws Exception {
+ if (closed) {
+ return;
+ }
+ closed = true;
logger.debug("Closing SpiceClient");
List exceptions = new ArrayList<>();
@@ -942,12 +1090,14 @@ public void close() throws Exception {
}
// Close Flight client
- try {
- this.flightClient.close();
- logger.debug("Flight client closed");
- } catch (Exception e) {
- logger.warn("Error closing Flight client: {}", e.getMessage());
- exceptions.add(e);
+ if (this.flightClient != null) {
+ try {
+ this.flightClient.close();
+ logger.debug("Flight client closed");
+ } catch (Exception e) {
+ logger.warn("Error closing Flight client: {}", e.getMessage());
+ exceptions.add(e);
+ }
}
// Close allocator
diff --git a/src/test/java/ai/spice/ResetTest.java b/src/test/java/ai/spice/ResetTest.java
new file mode 100644
index 0000000..c92904a
--- /dev/null
+++ b/src/test/java/ai/spice/ResetTest.java
@@ -0,0 +1,493 @@
+/*
+Copyright 2024-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.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+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 org.apache.arrow.flight.FlightStream;
+import org.apache.arrow.vector.ipc.ArrowReader;
+
+import junit.framework.TestCase;
+
+/**
+ * Tests for SpiceClient.reset() and transport resilience features
+ * (keep-alive, dns:/// resolution, lazy rebuild).
+ */
+public class ResetTest extends TestCase {
+
+ private static volatile boolean serverAvailable = false;
+ private static volatile boolean serverAvailabilityChecked = false;
+
+ @Override
+ protected void setUp() throws Exception {
+ super.setUp();
+ if (!serverAvailabilityChecked) {
+ synchronized (ResetTest.class) {
+ if (!serverAvailabilityChecked) {
+ try (SpiceClient probe = SpiceClient.builder().build()) {
+ // Probe with taxi_trips (not SELECT 1) to ensure
+ // the dataset is loaded and ready, not just that
+ // the server is up.
+ try (FlightStream stream = probe.query(
+ "SELECT total_amount FROM taxi_trips LIMIT 1")) {
+ stream.next();
+ }
+ serverAvailable = true;
+ } catch (Exception e) {
+ serverAvailable = false;
+ } finally {
+ serverAvailabilityChecked = true;
+ }
+ }
+ }
+ }
+ }
+
+ // ==================== reset() Happy Path ====================
+
+ /**
+ * reset() on a freshly built unauthenticated client should not throw.
+ */
+ public void testResetOnFreshClient() throws Exception {
+ SpiceClient client = SpiceClient.builder().build();
+ client.reset(); // should not throw
+ client.close();
+ }
+
+ /**
+ * reset() on a freshly built TLS client should not throw.
+ */
+ public void testResetOnTlsClient() throws Exception {
+ SpiceClient client = SpiceClient.builder()
+ .withFlightAddress(new URI("grpc+tls://localhost:443"))
+ .build();
+ client.reset(); // should not throw
+ client.close();
+ }
+
+ /**
+ * reset() on a freshly built HTTPS client (auto-converted to grpc+tls) should not throw.
+ */
+ public void testResetOnHttpsClient() throws Exception {
+ SpiceClient client = SpiceClient.builder()
+ .withFlightAddress(new URI("https://localhost:443"))
+ .build();
+ client.reset(); // should not throw
+ client.close();
+ }
+
+ /**
+ * After reset(), close() should complete without errors.
+ * reset() eagerly rebuilds the Flight client, so close() tears down the new connection.
+ */
+ public void testCloseAfterReset() throws Exception {
+ SpiceClient client = SpiceClient.builder().build();
+ client.reset();
+ client.close(); // should not throw
+ }
+
+ /**
+ * 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.
+ */
+ public void testQueryAfterResetRebuildsClient() throws Exception {
+ SpiceClient client = SpiceClient.builder().build();
+ client.reset();
+
+ try {
+ try (FlightStream stream = client.query("SELECT 1")) {
+ // If a local Spice runtime is running, this succeeds
+ stream.next();
+ }
+ } catch (Exception e) {
+ // Connection errors are expected when no server is running.
+ // A NullPointerException here would indicate the rebuild failed.
+ assertFalse("Should not get NullPointerException after reset (rebuild failed)",
+ e instanceof NullPointerException);
+ String msg = e.getMessage() != null ? e.getMessage().toLowerCase() : "";
+ assertTrue("Expected a connection/transport error, got: " + e.getMessage(),
+ msg.contains("unavailable") || msg.contains("connection refused")
+ || msg.contains("not found") || msg.contains("io exception")
+ || msg.contains("failed to execute"));
+ } finally {
+ client.close();
+ }
+ }
+
+ /**
+ * After reset(), queryWithParams should work because the Flight client was
+ * eagerly rebuilt (ADBC is still lazily initialized on first parameterized query).
+ */
+ public void testQueryWithParamsAfterResetRebuildsClient() throws Exception {
+ SpiceClient client = SpiceClient.builder().build();
+ client.reset();
+
+ try {
+ try (ArrowReader reader = client.queryWithParams("SELECT $1", 42)) {
+ while (reader.loadNextBatch()) {
+ // consume
+ }
+ }
+ } catch (Exception e) {
+ assertFalse("Should not get NullPointerException after reset (rebuild failed)",
+ e instanceof NullPointerException);
+ } finally {
+ client.close();
+ }
+ }
+
+ // ==================== reset() Edge Cases ====================
+
+ /**
+ * Calling reset() multiple times in a row should be safe (idempotent).
+ */
+ public void testMultipleResetsAreIdempotent() throws Exception {
+ SpiceClient client = SpiceClient.builder().build();
+ client.reset();
+ client.reset();
+ client.reset();
+ client.close(); // should not throw
+ }
+
+ /**
+ * reset() → query → reset() → query cycle should work.
+ * Each reset discards the transport; each query rebuilds it.
+ */
+ public void testResetQueryResetQueryCycle() throws Exception {
+ SpiceClient client = SpiceClient.builder().build();
+
+ for (int i = 0; i < 3; i++) {
+ client.reset();
+ try {
+ FlightStream stream = client.query("SELECT 1");
+ stream.close();
+ } catch (Exception e) {
+ // Connection errors are fine — we're testing the reset/rebuild cycle,
+ // not server availability. NPE would indicate a broken rebuild.
+ assertFalse("Cycle " + i + ": NPE after reset means rebuild is broken",
+ e instanceof NullPointerException);
+ }
+ }
+
+ client.close();
+ }
+
+ /**
+ * reset() after close() should throw IllegalStateException.
+ */
+ public void testResetAfterClose() throws Exception {
+ SpiceClient client = SpiceClient.builder().build();
+ client.close();
+ try {
+ client.reset();
+ fail("Expected IllegalStateException when resetting a closed client");
+ } catch (IllegalStateException e) {
+ assertTrue(e.getMessage().contains("closed"));
+ }
+ }
+
+ /**
+ * try-with-resources with a reset in between should work cleanly.
+ */
+ public void testTryWithResourcesAndReset() throws Exception {
+ try (SpiceClient client = SpiceClient.builder().build()) {
+ assertNotNull(client);
+ client.reset();
+ // auto-close on scope exit
+ }
+ }
+
+ // ==================== Concurrent reset() ====================
+
+ /**
+ * Concurrent calls to reset() from multiple threads should not throw
+ * or corrupt the client's internal state (reset() is expected to be thread-safe).
+ */
+ public void testConcurrentResetDoesNotThrow() throws Exception {
+ final SpiceClient client = SpiceClient.builder().build();
+ final int threadCount = 8;
+ final CountDownLatch startLatch = new CountDownLatch(1);
+ final AtomicInteger errors = new AtomicInteger(0);
+ ExecutorService executor = Executors.newFixedThreadPool(threadCount);
+
+ try {
+ for (int i = 0; i < threadCount; i++) {
+ executor.submit(() -> {
+ try {
+ startLatch.await();
+ client.reset();
+ } catch (Exception e) {
+ errors.incrementAndGet();
+ }
+ });
+ }
+
+ startLatch.countDown(); // release all threads
+ executor.shutdown();
+ assertTrue("Concurrent reset should complete within 30s",
+ executor.awaitTermination(30, TimeUnit.SECONDS));
+
+ assertEquals("No threads should have encountered errors", 0, errors.get());
+ } finally {
+ executor.shutdownNow();
+ client.close();
+ }
+ }
+
+ /**
+ * Concurrent reset() and query() should not throw unexpected errors.
+ * (query may fail with connection errors, but not NPE or IllegalStateException.)
+ */
+ public void testConcurrentResetAndQuery() throws Exception {
+ final SpiceClient client = SpiceClient.builder().build();
+ final int iterations = 5;
+ final CountDownLatch startLatch = new CountDownLatch(1);
+ final List unexpectedErrors = new CopyOnWriteArrayList<>();
+ ExecutorService executor = Executors.newFixedThreadPool(2);
+
+ try {
+ // Thread 1: repeated resets
+ executor.submit(() -> {
+ try {
+ startLatch.await();
+ for (int i = 0; i < iterations; i++) {
+ client.reset();
+ Thread.sleep(10);
+ }
+ } catch (InterruptedException ignored) {
+ }
+ });
+
+ // Thread 2: repeated queries
+ executor.submit(() -> {
+ try {
+ startLatch.await();
+ for (int i = 0; i < iterations; i++) {
+ try {
+ FlightStream stream = client.query("SELECT 1");
+ stream.close();
+ } catch (NullPointerException e) {
+ unexpectedErrors.add(e);
+ } catch (Exception e) {
+ // Connection errors are expected
+ }
+ Thread.sleep(10);
+ }
+ } catch (InterruptedException ignored) {
+ }
+ });
+
+ startLatch.countDown();
+ executor.shutdown();
+ assertTrue("Concurrent reset+query should complete within 30s",
+ executor.awaitTermination(30, TimeUnit.SECONDS));
+
+ assertTrue("Should not get NullPointerException during concurrent reset+query: " + unexpectedErrors,
+ unexpectedErrors.isEmpty());
+ } finally {
+ executor.shutdownNow();
+ client.close();
+ }
+ }
+
+ // ==================== Construction / DNS / Keep-alive ====================
+
+ /**
+ * Client built with default (plaintext) address constructs successfully.
+ * This implicitly validates dns:/// target construction for grpc+tcp.
+ */
+ public void testDefaultPlaintextConstruction() throws Exception {
+ SpiceClient client = SpiceClient.builder().build();
+ assertNotNull(client);
+ client.close();
+ }
+
+ /**
+ * Client built with grpc+tls address constructs successfully.
+ * Validates dns:/// target + TLS + keep-alive configuration.
+ */
+ public void testGrpcTlsConstruction() throws Exception {
+ SpiceClient client = SpiceClient.builder()
+ .withFlightAddress(new URI("grpc+tls://localhost:443"))
+ .build();
+ assertNotNull(client);
+ client.close();
+ }
+
+ /**
+ * Client built with HTTP address (auto-converted to grpc+tcp).
+ * Validates scheme conversion + dns:/// target.
+ */
+ public void testHttpSchemeConversion() throws Exception {
+ SpiceClient client = SpiceClient.builder()
+ .withFlightAddress(new URI("http://localhost:50051"))
+ .build();
+ assertNotNull(client);
+ client.close();
+ }
+
+ /**
+ * Client built with HTTPS address (auto-converted to grpc+tls).
+ * Validates scheme conversion + dns:/// target + TLS.
+ */
+ public void testHttpsSchemeConversion() throws Exception {
+ SpiceClient client = SpiceClient.builder()
+ .withFlightAddress(new URI("https://localhost:443"))
+ .build();
+ assertNotNull(client);
+ client.close();
+ }
+
+ /**
+ * Client built with no explicit port should use default (443 for TLS, 80 for plaintext).
+ */
+ public void testDefaultPortForTls() throws Exception {
+ SpiceClient client = SpiceClient.builder()
+ .withFlightAddress(new URI("grpc+tls://somehost"))
+ .build();
+ assertNotNull(client);
+ client.close();
+ }
+
+ public void testDefaultPortForPlaintext() throws Exception {
+ SpiceClient client = SpiceClient.builder()
+ .withFlightAddress(new URI("grpc+tcp://somehost"))
+ .build();
+ assertNotNull(client);
+ client.close();
+ }
+
+ /**
+ * Client built with Spice Cloud configuration (production-like scenario).
+ * Validates the full TLS + dns:/// + keep-alive path.
+ */
+ public void testSpiceCloudConstruction() throws Exception {
+ SpiceClient client = SpiceClient.builder()
+ .withSpiceCloud()
+ .build();
+ assertNotNull(client);
+ client.close();
+ }
+
+ /**
+ * Client with custom memory limit and retries + reset cycle still works.
+ */
+ public void testResetWithCustomConfig() throws Exception {
+ SpiceClient client = SpiceClient.builder()
+ .withMaxRetries(5)
+ .withArrowMemoryLimitMB(256)
+ .withUserAgent("TestApp/1.0")
+ .build();
+ client.reset();
+
+ try {
+ FlightStream stream = client.query("SELECT 1");
+ stream.close();
+ } catch (Exception e) {
+ assertFalse("NPE after reset with custom config",
+ e instanceof NullPointerException);
+ }
+
+ client.close();
+ }
+
+ // ==================== Integration: reset then query (server required) ====================
+
+ /**
+ * If a local Spice runtime is running, verify that
+ * reset() followed by query() actually returns data.
+ * Uses taxi_trips which is available in the CI quickstart dataset.
+ */
+ public void testResetThenQueryIntegration() throws Exception {
+ if (!serverAvailable) return;
+
+ try (SpiceClient client = SpiceClient.builder().build()) {
+ // First query (establishes connection)
+ try (FlightStream stream1 = client.query(
+ "SELECT total_amount FROM taxi_trips LIMIT 1")) {
+ int rows1 = 0;
+ while (stream1.next()) {
+ rows1 += stream1.getRoot().getRowCount();
+ }
+ assertEquals("First query should return 1 row", 1, rows1);
+ }
+
+ // Reset (discards transport)
+ client.reset();
+
+ // Second query (lazy rebuild)
+ try (FlightStream stream2 = client.query(
+ "SELECT total_amount FROM taxi_trips LIMIT 2")) {
+ int rows2 = 0;
+ while (stream2.next()) {
+ rows2 += stream2.getRoot().getRowCount();
+ }
+ assertEquals("Second query after reset should return 2 rows", 2, rows2);
+ }
+ }
+ }
+
+ /**
+ * If a local Spice runtime is running, verify that
+ * reset() followed by queryWithParams() actually returns data.
+ * Uses taxi_trips which is available in the CI quickstart dataset.
+ */
+ public void testResetThenQueryWithParamsIntegration() throws Exception {
+ if (!serverAvailable) return;
+
+ try (SpiceClient client = SpiceClient.builder().build()) {
+ // First query
+ try (ArrowReader reader1 = client.queryWithParams(
+ "SELECT total_amount FROM taxi_trips WHERE total_amount > $1 LIMIT 1",
+ 0.0)) {
+ int rows1 = 0;
+ while (reader1.loadNextBatch()) {
+ rows1 += reader1.getVectorSchemaRoot().getRowCount();
+ }
+ assertTrue("First parameterized query should return rows", rows1 > 0);
+ }
+
+ // Reset
+ client.reset();
+
+ // Second query (re-initializes both Flight and ADBC)
+ try (ArrowReader reader2 = client.queryWithParams(
+ "SELECT total_amount FROM taxi_trips WHERE total_amount > $1 LIMIT 2",
+ 0.0)) {
+ int rows2 = 0;
+ while (reader2.loadNextBatch()) {
+ rows2 += reader2.getVectorSchemaRoot().getRowCount();
+ }
+ assertTrue("Second parameterized query after reset should return rows", rows2 > 0);
+ }
+ }
+ }
+}