From 2ab78064f7f8f2caf197b75f3ebe2631ef9a6151 Mon Sep 17 00:00:00 2001 From: unikdahal Date: Sat, 27 Jun 2026 03:18:16 +0530 Subject: [PATCH 01/10] feat(java/driver/flight-sql): implement Flight SQL session management MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires Flight SQL session options and CloseSession into the Java ADBC driver, bringing it to parity with the Go/Python drivers. What changed: FlightSqlConnectionProperties — five new string constants for session option key prefixes (adbc.flight.sql.session.option., adbc.flight.sql.session.optionbool., adbc.flight.sql.session.optionstringlist., adbc.flight.sql.session.optionerase., and the read-only adbc.flight.sql.session.options JSON blob key). FlightSqlClientWithCallOptions — three new delegation methods: setSessionOptions, getSessionOptions, closeSession, all forwarding connectionOptions via the existing combine() pattern. FlightSqlConnection — overrides getOption/setOption on AdbcConnection for all session key prefixes; each typed prefix (string, bool, stringlist, long, double) dispatches to the appropriate SessionOptionValue via a visitor. setOption null value is rejected early with INVALID_ARGUMENT. The erase prefix calls makeEmptySessionOptionValue. close() now sends CloseSession best-effort, catching only FlightRuntimeException so InterruptedException is not swallowed. FlightSqlSessionUtil — new package-private class holding all JSON helpers (sessionOptionsToJson, toJsonArray, parseJsonArray, escapeJson) and the six typed SessionOptionValue visitors. Extracted from FlightSqlConnection to keep that class focused on the ADBC surface. escapeJson handles the full U+0000-U+001F control-character range. parseJsonArray handles \uXXXX, \b, \f and all JSON whitespace. JSON_VALUE_VISITOR emits "null" for NaN/Infinity doubles to produce valid RFC 8259 output. FlightSqlSessionTest — new unit tests covering set/get for string, bool, and string-list option types, erase, the full SESSION_OPTIONS JSON blob, CloseSession being called on close(), read-only blob rejection, and close() not throwing when the server returns UNIMPLEMENTED. Relates to #745. Closes #4443. Co-Authored-By: Claude Sonnet 4.6 --- .../FlightSqlClientWithCallOptions.java | 20 + .../driver/flightsql/FlightSqlConnection.java | 179 +++++++ .../FlightSqlConnectionProperties.java | 6 + .../flightsql/FlightSqlSessionUtil.java | 211 ++++++++ .../flightsql/FlightSqlSessionTest.java | 466 ++++++++++++++++++ 5 files changed, 882 insertions(+) create mode 100644 java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java create mode 100644 java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionTest.java diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlClientWithCallOptions.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlClientWithCallOptions.java index 881013fc69..d53eee9762 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlClientWithCallOptions.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlClientWithCallOptions.java @@ -26,12 +26,18 @@ import org.apache.arrow.flight.CallOption; import org.apache.arrow.flight.CancelFlightInfoRequest; import org.apache.arrow.flight.CancelFlightInfoResult; +import org.apache.arrow.flight.CloseSessionRequest; +import org.apache.arrow.flight.CloseSessionResult; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightStream; +import org.apache.arrow.flight.GetSessionOptionsRequest; +import org.apache.arrow.flight.GetSessionOptionsResult; import org.apache.arrow.flight.RenewFlightEndpointRequest; import org.apache.arrow.flight.SchemaResult; +import org.apache.arrow.flight.SetSessionOptionsRequest; +import org.apache.arrow.flight.SetSessionOptionsResult; import org.apache.arrow.flight.Ticket; import org.apache.arrow.flight.sql.CancelResult; import org.apache.arrow.flight.sql.FlightSqlClient; @@ -287,6 +293,20 @@ public FlightEndpoint renewFlightEndpoint( return client.renewFlightEndpoint(request, combine(options)); } + public SetSessionOptionsResult setSessionOptions( + SetSessionOptionsRequest request, CallOption... options) { + return client.setSessionOptions(request, combine(options)); + } + + public GetSessionOptionsResult getSessionOptions( + GetSessionOptionsRequest request, CallOption... options) { + return client.getSessionOptions(request, combine(options)); + } + + public CloseSessionResult closeSession(CloseSessionRequest request, CallOption... options) { + return client.closeSession(request, combine(options)); + } + @Override public void close() throws Exception { AutoCloseables.close(client); diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java index f099cb64c8..e6b0b39750 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java @@ -38,13 +38,23 @@ import org.apache.arrow.adbc.core.AdbcStatement; import org.apache.arrow.adbc.core.AdbcStatusCode; import org.apache.arrow.adbc.core.BulkIngestMode; +import org.apache.arrow.adbc.core.TypedKey; import org.apache.arrow.adbc.sql.SqlQuirks; import org.apache.arrow.flight.CallOption; +import org.apache.arrow.flight.CloseSessionRequest; import org.apache.arrow.flight.FlightCallHeaders; import org.apache.arrow.flight.FlightClient; import org.apache.arrow.flight.FlightEndpoint; +import org.apache.arrow.flight.FlightRuntimeException; +import org.apache.arrow.flight.FlightStatusCode; +import org.apache.arrow.flight.GetSessionOptionsRequest; +import org.apache.arrow.flight.GetSessionOptionsResult; import org.apache.arrow.flight.HeaderCallOption; import org.apache.arrow.flight.Location; +import org.apache.arrow.flight.SessionOptionValue; +import org.apache.arrow.flight.SessionOptionValueFactory; +import org.apache.arrow.flight.SetSessionOptionsRequest; +import org.apache.arrow.flight.SetSessionOptionsResult; import org.apache.arrow.flight.Ticket; import org.apache.arrow.flight.auth2.BasicAuthCredentialWriter; import org.apache.arrow.flight.client.ClientCookieMiddleware; @@ -203,8 +213,143 @@ public void setAutoCommit(boolean enableAutoCommit) throws AdbcException { } } + @Override + @SuppressWarnings("unchecked") + public T getOption(TypedKey key) throws AdbcException { + final String k = key.getKey(); + + if (k.equals(FlightSqlConnectionProperties.SESSION_OPTIONS)) { + return key.cast(FlightSqlSessionUtil.sessionOptionsToJson(fetchSessionOptions())); + + } else if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_BOOL_PREFIX)) { + final String name = + k.substring(FlightSqlConnectionProperties.SESSION_OPTION_BOOL_PREFIX.length()); + final SessionOptionValue val = fetchSessionOptions().get(name); + if (val == null) { + throw new AdbcException( + "[Flight SQL] Session option not found: " + name, + null, + AdbcStatusCode.NOT_FOUND, + null, + 0); + } + if (key.getType() == Boolean.class) { + return key.cast(val.acceptVisitor(FlightSqlSessionUtil.BOOL_VISITOR)); + } + return key.cast(val.acceptVisitor(FlightSqlSessionUtil.STRING_VISITOR)); + + } else if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_STRING_LIST_PREFIX)) { + final String name = + k.substring(FlightSqlConnectionProperties.SESSION_OPTION_STRING_LIST_PREFIX.length()); + final SessionOptionValue val = fetchSessionOptions().get(name); + if (val == null) { + throw new AdbcException( + "[Flight SQL] Session option not found: " + name, + null, + AdbcStatusCode.NOT_FOUND, + null, + 0); + } + if (key.getType() != String[].class) { + return AdbcConnection.super.getOption(key); + } + return key.cast(val.acceptVisitor(FlightSqlSessionUtil.STRING_ARRAY_VISITOR)); + + } else if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_PREFIX)) { + final String name = + k.substring(FlightSqlConnectionProperties.SESSION_OPTION_PREFIX.length()); + final SessionOptionValue val = fetchSessionOptions().get(name); + if (val == null) { + throw new AdbcException( + "[Flight SQL] Session option not found: " + name, + null, + AdbcStatusCode.NOT_FOUND, + null, + 0); + } + try { + if (key.getType() == Long.class) { + return key.cast(val.acceptVisitor(FlightSqlSessionUtil.LONG_VISITOR)); + } else if (key.getType() == Double.class) { + return key.cast(val.acceptVisitor(FlightSqlSessionUtil.DOUBLE_VISITOR)); + } + } catch (NumberFormatException e) { + throw AdbcException.invalidArgument( + "[Flight SQL] Session option '" + name + "' value cannot be parsed as " + + key.getType().getSimpleName() + ": " + e.getMessage()); + } + return key.cast(val.acceptVisitor(FlightSqlSessionUtil.STRING_VISITOR)); + } + + return AdbcConnection.super.getOption(key); + } + + @Override + public void setOption(TypedKey key, T value) throws AdbcException { + final String k = key.getKey(); + + if (value == null) { + throw AdbcException.invalidArgument( + "[Flight SQL] null value not allowed for key: " + k + + " — use adbc.flight.sql.session.optionerase. to erase an option"); + } + + if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_ERASE_PREFIX)) { + final String name = + k.substring(FlightSqlConnectionProperties.SESSION_OPTION_ERASE_PREFIX.length()); + doSetSessionOption(name, SessionOptionValueFactory.makeEmptySessionOptionValue()); + + } else if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_BOOL_PREFIX)) { + final String name = + k.substring(FlightSqlConnectionProperties.SESSION_OPTION_BOOL_PREFIX.length()); + final boolean b; + if (value instanceof Boolean) { + b = (Boolean) value; + } else { + b = Boolean.parseBoolean(value.toString()); + } + doSetSessionOption(name, SessionOptionValueFactory.makeSessionOptionValue(b)); + + } else if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_STRING_LIST_PREFIX)) { + final String name = + k.substring(FlightSqlConnectionProperties.SESSION_OPTION_STRING_LIST_PREFIX.length()); + final String[] arr; + if (value instanceof String[]) { + arr = (String[]) value; + } else { + arr = FlightSqlSessionUtil.parseJsonArray(value.toString()); + } + doSetSessionOption(name, SessionOptionValueFactory.makeSessionOptionValue(arr)); + + } else if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_PREFIX)) { + final String name = + k.substring(FlightSqlConnectionProperties.SESSION_OPTION_PREFIX.length()); + final SessionOptionValue sv; + if (value instanceof Long) { + sv = SessionOptionValueFactory.makeSessionOptionValue((Long) value); + } else if (value instanceof Double) { + sv = SessionOptionValueFactory.makeSessionOptionValue((Double) value); + } else { + sv = SessionOptionValueFactory.makeSessionOptionValue(value.toString()); + } + doSetSessionOption(name, sv); + + } else if (k.equals(FlightSqlConnectionProperties.SESSION_OPTIONS)) { + throw AdbcException.notImplemented( + "[Flight SQL] adbc.flight.sql.session.options is read-only"); + + } else { + AdbcConnection.super.setOption(key, value); + } + } + @Override public void close() throws Exception { + try { + client.closeSession(new CloseSessionRequest()); + } catch (FlightRuntimeException ignored) { + // best-effort: UNIMPLEMENTED is normal; other errors are silently swallowed at close time + } clientCache.invalidateAll(); AutoCloseables.close(client, allocator); } @@ -214,6 +359,40 @@ public String toString() { return "FlightSqlConnection{" + "client=" + client + '}'; } + private Map fetchSessionOptions() throws AdbcException { + try { + GetSessionOptionsResult result = + client.getSessionOptions(new GetSessionOptionsRequest()); + return result.getSessionOptions(); + } catch (FlightRuntimeException e) { + if (e.status().code() == FlightStatusCode.UNIMPLEMENTED) { + return Collections.emptyMap(); + } + throw FlightSqlDriverUtil.fromFlightException(e); + } + } + + private void doSetSessionOption(String name, SessionOptionValue value) throws AdbcException { + if (name.isEmpty()) { + throw AdbcException.invalidArgument( + "[Flight SQL] Session option name must not be empty"); + } + final SetSessionOptionsResult result; + try { + result = + client.setSessionOptions( + new SetSessionOptionsRequest(Collections.singletonMap(name, value))); + } catch (FlightRuntimeException e) { + throw FlightSqlDriverUtil.fromFlightException(e); + } + if (result.hasErrors()) { + final SetSessionOptionsResult.Error err = result.getErrors().get(name); + final String errType = (err != null) ? err.value.name() : "UNKNOWN"; + throw AdbcException.invalidArgument( + "[Flight SQL] Failed to set session option '" + name + "': " + errType); + } + } + /** * Initialize cached data to share between connections and create, test, and authenticate the * first connection. diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnectionProperties.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnectionProperties.java index 4ab1955a1b..39da07a03f 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnectionProperties.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnectionProperties.java @@ -35,4 +35,10 @@ public interface FlightSqlConnectionProperties { TypedKey WITH_COOKIE_MIDDLEWARE = new TypedKey<>("adbc.flight.sql.rpc.with_cookie_middleware", Boolean.class); String RPC_CALL_HEADER_PREFIX = "adbc.flight.sql.rpc.call_header."; + + String SESSION_OPTIONS = "adbc.flight.sql.session.options"; + String SESSION_OPTION_PREFIX = "adbc.flight.sql.session.option."; + String SESSION_OPTION_BOOL_PREFIX = "adbc.flight.sql.session.optionbool."; + String SESSION_OPTION_STRING_LIST_PREFIX = "adbc.flight.sql.session.optionstringlist."; + String SESSION_OPTION_ERASE_PREFIX = "adbc.flight.sql.session.optionerase."; } diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java new file mode 100644 index 0000000000..df3d97c18c --- /dev/null +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java @@ -0,0 +1,211 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adbc.driver.flightsql; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import org.apache.arrow.adbc.core.AdbcException; +import org.apache.arrow.flight.SessionOptionValue; +import org.apache.arrow.flight.SessionOptionValueVisitor; + +/** Helpers for Flight SQL session option serialization and type conversion. */ +final class FlightSqlSessionUtil { + + private FlightSqlSessionUtil() {} + + // -- JSON helpers -- + + static String sessionOptionsToJson(Map opts) { + final StringBuilder sb = new StringBuilder("{"); + boolean first = true; + for (Map.Entry e : opts.entrySet()) { + if (!first) sb.append(','); + sb.append('"').append(escapeJson(e.getKey())).append("\":").append(sessionValueToJson(e.getValue())); + first = false; + } + return sb.append('}').toString(); + } + + static String sessionValueToJson(SessionOptionValue val) { + return val.acceptVisitor(JSON_VALUE_VISITOR); + } + + static String toJsonArray(String[] values) { + final StringBuilder sb = new StringBuilder("["); + for (int i = 0; i < values.length; i++) { + if (i > 0) sb.append(','); + sb.append('"').append(escapeJson(values[i])).append('"'); + } + return sb.append(']').toString(); + } + + static String[] parseJsonArray(String json) throws AdbcException { + final String trimmed = json.trim(); + if (!trimmed.startsWith("[") || !trimmed.endsWith("]")) { + throw AdbcException.invalidArgument( + "[Flight SQL] Expected JSON array for string list option, got: " + json); + } + final String inner = trimmed.substring(1, trimmed.length() - 1).trim(); + if (inner.isEmpty()) { + return new String[0]; + } + final List result = new ArrayList<>(); + int i = 0; + while (i < inner.length()) { + while (i < inner.length() && Character.isWhitespace(inner.charAt(i))) i++; + if (i >= inner.length()) break; + if (inner.charAt(i) != '"') { + throw AdbcException.invalidArgument( + "[Flight SQL] Invalid JSON string array element at position " + i + " in: " + json); + } + i++; // skip opening quote + final StringBuilder elem = new StringBuilder(); + while (i < inner.length() && inner.charAt(i) != '"') { + if (inner.charAt(i) == '\\' && i + 1 < inner.length()) { + i++; + switch (inner.charAt(i)) { + case '"': elem.append('"'); break; + case '\\': elem.append('\\'); break; + case '/': elem.append('/'); break; + case 'n': elem.append('\n'); break; + case 'r': elem.append('\r'); break; + case 't': elem.append('\t'); break; + case 'b': elem.append('\b'); break; + case 'f': elem.append('\f'); break; + case 'u': + if (i + 4 < inner.length()) { + final String hex = inner.substring(i + 1, i + 5); + try { + elem.append((char) Integer.parseInt(hex, 16)); + i += 4; + } catch (NumberFormatException e) { + elem.append('u'); + } + } else { + elem.append('u'); + } + break; + default: elem.append(inner.charAt(i)); + } + } else { + elem.append(inner.charAt(i)); + } + i++; + } + if (i >= inner.length()) { + throw AdbcException.invalidArgument( + "[Flight SQL] Unterminated string in JSON array: " + json); + } + i++; // skip closing quote + result.add(elem.toString()); + while (i < inner.length() && Character.isWhitespace(inner.charAt(i))) i++; + if (i < inner.length() && inner.charAt(i) == ',') i++; + } + return result.toArray(new String[0]); + } + + static String escapeJson(String s) { + final StringBuilder sb = new StringBuilder(s.length()); + for (int i = 0; i < s.length(); i++) { + final char c = s.charAt(i); + switch (c) { + case '"': sb.append("\\\""); break; + case '\\': sb.append("\\\\"); break; + case '\n': sb.append("\\n"); break; + case '\r': sb.append("\\r"); break; + case '\t': sb.append("\\t"); break; + case '\b': sb.append("\\b"); break; + case '\f': sb.append("\\f"); break; + default: + if (c < 0x20) { + sb.append(String.format("\\u%04x", (int) c)); + } else { + sb.append(c); + } + } + } + return sb.toString(); + } + + // -- SessionOptionValue visitors -- + + static final SessionOptionValueVisitor BOOL_VISITOR = + new SessionOptionValueVisitor() { + @Override public Boolean visit(String v) { return Boolean.parseBoolean(v); } + @Override public Boolean visit(boolean v) { return v; } + @Override public Boolean visit(long v) { return v != 0; } + @Override public Boolean visit(double v) { return v != 0; } + @Override public Boolean visit(String[] v) { return false; } + @Override public Boolean visit(Void v) { return false; } + }; + + static final SessionOptionValueVisitor STRING_ARRAY_VISITOR = + new SessionOptionValueVisitor() { + @Override public String[] visit(String v) { return new String[]{v}; } + @Override public String[] visit(boolean v) { return new String[]{String.valueOf(v)}; } + @Override public String[] visit(long v) { return new String[]{String.valueOf(v)}; } + @Override public String[] visit(double v) { return new String[]{String.valueOf(v)}; } + @Override public String[] visit(String[] v) { return v.clone(); } + @Override public String[] visit(Void v) { return new String[0]; } + }; + + static final SessionOptionValueVisitor STRING_VISITOR = + new SessionOptionValueVisitor() { + @Override public String visit(String v) { return v; } + @Override public String visit(boolean v) { return String.valueOf(v); } + @Override public String visit(long v) { return String.valueOf(v); } + @Override public String visit(double v) { return String.valueOf(v); } + @Override public String visit(String[] v) { return toJsonArray(v); } + @Override public String visit(Void v) { return ""; } + }; + + static final SessionOptionValueVisitor LONG_VISITOR = + new SessionOptionValueVisitor() { + @Override public Long visit(String v) { return Long.parseLong(v); } + @Override public Long visit(boolean v) { return v ? 1L : 0L; } + @Override public Long visit(long v) { return v; } + @Override public Long visit(double v) { return (long) v; } + @Override public Long visit(String[] v) { return 0L; } + @Override public Long visit(Void v) { return 0L; } + }; + + static final SessionOptionValueVisitor DOUBLE_VISITOR = + new SessionOptionValueVisitor() { + @Override public Double visit(String v) { return Double.parseDouble(v); } + @Override public Double visit(boolean v) { return v ? 1.0 : 0.0; } + @Override public Double visit(long v) { return (double) v; } + @Override public Double visit(double v) { return v; } + @Override public Double visit(String[] v) { return 0.0; } + @Override public Double visit(Void v) { return 0.0; } + }; + + static final SessionOptionValueVisitor JSON_VALUE_VISITOR = + new SessionOptionValueVisitor() { + @Override public String visit(String v) { return '"' + escapeJson(v) + '"'; } + @Override public String visit(boolean v) { return String.valueOf(v); } + @Override public String visit(long v) { return String.valueOf(v); } + @Override public String visit(double v) { + if (Double.isNaN(v) || Double.isInfinite(v)) { + return "null"; + } + return String.valueOf(v); + } + @Override public String visit(String[] v) { return toJsonArray(v); } + @Override public String visit(Void v) { return "null"; } + }; +} diff --git a/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionTest.java b/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionTest.java new file mode 100644 index 0000000000..571964a1db --- /dev/null +++ b/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionTest.java @@ -0,0 +1,466 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.adbc.driver.flightsql; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; +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.AdbcStatusCode; +import org.apache.arrow.adbc.core.TypedKey; +import org.apache.arrow.flight.CallStatus; +import org.apache.arrow.flight.CloseSessionRequest; +import org.apache.arrow.flight.Criteria; +import org.apache.arrow.flight.CloseSessionResult; +import org.apache.arrow.flight.FlightDescriptor; +import org.apache.arrow.flight.FlightInfo; +import org.apache.arrow.flight.FlightServer; +import org.apache.arrow.flight.FlightStream; +import org.apache.arrow.flight.GetSessionOptionsRequest; +import org.apache.arrow.flight.GetSessionOptionsResult; +import org.apache.arrow.flight.Location; +import org.apache.arrow.flight.PutResult; +import org.apache.arrow.flight.Result; +import org.apache.arrow.flight.SchemaResult; +import org.apache.arrow.flight.SessionOptionValue; +import org.apache.arrow.flight.SessionOptionValueFactory; +import org.apache.arrow.flight.SetSessionOptionsRequest; +import org.apache.arrow.flight.SetSessionOptionsResult; +import org.apache.arrow.flight.sql.FlightSqlProducer; +import org.apache.arrow.flight.sql.impl.FlightSql; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.util.AutoCloseables; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Tests for Flight SQL session management (get/set/erase options, CloseSession). */ +class FlightSqlSessionTest { + static BufferAllocator allocator; + static SessionProducer producer; + static FlightServer server; + static AdbcDriver driver; + static AdbcDatabase database; + AdbcConnection connection; + + @BeforeAll + static void beforeAll() throws Exception { + allocator = new RootAllocator(); + producer = new SessionProducer(); + server = + FlightServer.builder() + .allocator(allocator) + .producer(producer) + .location(Location.forGrpcInsecure("localhost", 0)) + .build(); + server.start(); + driver = new FlightSqlDriver(allocator); + Map parameters = new HashMap<>(); + AdbcDriver.PARAM_URI.set( + parameters, Location.forGrpcInsecure("localhost", server.getPort()).getUri().toString()); + database = driver.open(parameters); + } + + @BeforeEach + void beforeEach() throws Exception { + producer.reset(); + connection = database.connect(); + } + + @AfterEach + void afterEach() throws Exception { + AutoCloseables.close(connection); + } + + @AfterAll + static void afterAll() throws Exception { + AutoCloseables.close(database, server, allocator); + } + + @Test + void testSetGetStringOption() throws Exception { + connection.setOption( + new TypedKey<>( + FlightSqlConnectionProperties.SESSION_OPTION_PREFIX + "catalog", String.class), + "my_catalog"); + + String value = + connection.getOption( + new TypedKey<>( + FlightSqlConnectionProperties.SESSION_OPTION_PREFIX + "catalog", String.class)); + assertThat(value).isEqualTo("my_catalog"); + } + + @Test + void testSetGetBoolOption() throws Exception { + connection.setOption( + new TypedKey<>( + FlightSqlConnectionProperties.SESSION_OPTION_BOOL_PREFIX + "flag", Boolean.class), + true); + + Boolean value = + connection.getOption( + new TypedKey<>( + FlightSqlConnectionProperties.SESSION_OPTION_BOOL_PREFIX + "flag", Boolean.class)); + assertThat(value).isTrue(); + } + + @Test + void testSetGetStringListOptionAsArray() throws Exception { + connection.setOption( + new TypedKey<>( + FlightSqlConnectionProperties.SESSION_OPTION_STRING_LIST_PREFIX + "tags", + String[].class), + new String[] {"a", "b", "c"}); + + String[] value = + connection.getOption( + new TypedKey<>( + FlightSqlConnectionProperties.SESSION_OPTION_STRING_LIST_PREFIX + "tags", + String[].class)); + assertThat(value).containsExactly("a", "b", "c"); + } + + @Test + void testSetGetStringListOptionAsJson() throws Exception { + connection.setOption( + new TypedKey<>( + FlightSqlConnectionProperties.SESSION_OPTION_STRING_LIST_PREFIX + "tags", String.class), + "[\"x\",\"y\"]"); + + String[] value = + connection.getOption( + new TypedKey<>( + FlightSqlConnectionProperties.SESSION_OPTION_STRING_LIST_PREFIX + "tags", + String[].class)); + assertThat(value).containsExactly("x", "y"); + } + + @Test + void testEraseOption() throws Exception { + connection.setOption( + new TypedKey<>( + FlightSqlConnectionProperties.SESSION_OPTION_PREFIX + "toErase", String.class), + "value"); + + connection.setOption( + new TypedKey<>( + FlightSqlConnectionProperties.SESSION_OPTION_ERASE_PREFIX + "toErase", String.class), + ""); + + AdbcException ex = + assertThrows( + AdbcException.class, + () -> + connection.getOption( + new TypedKey<>( + FlightSqlConnectionProperties.SESSION_OPTION_PREFIX + "toErase", + String.class))); + assertThat(ex.getStatus()).isEqualTo(AdbcStatusCode.NOT_FOUND); + } + + @Test + void testGetSessionOptionsBlob() throws Exception { + connection.setOption( + new TypedKey<>( + FlightSqlConnectionProperties.SESSION_OPTION_PREFIX + "k1", String.class), + "v1"); + + String blob = + connection.getOption( + new TypedKey<>(FlightSqlConnectionProperties.SESSION_OPTIONS, String.class)); + assertThat(blob).contains("\"k1\""); + assertThat(blob).contains("\"v1\""); + } + + @Test + void testCloseSessionCalledOnClose() throws Exception { + producer.closeSessionCalled.set(false); + connection.close(); + connection = null; // prevent double-close in afterEach + assertThat(producer.closeSessionCalled.get()).isTrue(); + } + + @Test + void testSetOptionReadOnlyBlobThrows() { + AdbcException ex = + assertThrows( + AdbcException.class, + () -> + connection.setOption( + new TypedKey<>(FlightSqlConnectionProperties.SESSION_OPTIONS, String.class), + "{}")); + assertThat(ex.getStatus()).isEqualTo(AdbcStatusCode.NOT_IMPLEMENTED); + } + + // ----- Server-side producer ----- + + static class SessionProducer implements FlightSqlProducer { + private final Map sessionOptions = new HashMap<>(); + final AtomicBoolean closeSessionCalled = new AtomicBoolean(false); + + void reset() { + sessionOptions.clear(); + closeSessionCalled.set(false); + } + + @Override + public void setSessionOptions( + SetSessionOptionsRequest request, + CallContext context, + StreamListener listener) { + for (Map.Entry e : request.getSessionOptions().entrySet()) { + if (e.getValue().isEmpty()) { + sessionOptions.remove(e.getKey()); + } else { + sessionOptions.put(e.getKey(), e.getValue()); + } + } + listener.onNext(new SetSessionOptionsResult(Collections.emptyMap())); + listener.onCompleted(); + } + + @Override + public void getSessionOptions( + GetSessionOptionsRequest request, + CallContext context, + StreamListener listener) { + listener.onNext(new GetSessionOptionsResult(new HashMap<>(sessionOptions))); + listener.onCompleted(); + } + + @Override + public void closeSession( + CloseSessionRequest request, + CallContext context, + StreamListener listener) { + closeSessionCalled.set(true); + listener.onNext(new CloseSessionResult(CloseSessionResult.Status.CLOSED)); + listener.onCompleted(); + } + + // Required abstract method implementations (no-op) + + @Override + public void createPreparedStatement( + FlightSql.ActionCreatePreparedStatementRequest req, + CallContext ctx, + StreamListener listener) { + listener.onError(CallStatus.UNIMPLEMENTED.toRuntimeException()); + } + + @Override + public void closePreparedStatement( + FlightSql.ActionClosePreparedStatementRequest req, + CallContext ctx, + StreamListener listener) { + listener.onError(CallStatus.UNIMPLEMENTED.toRuntimeException()); + } + + @Override + public FlightInfo getFlightInfoStatement( + FlightSql.CommandStatementQuery cmd, CallContext ctx, FlightDescriptor descriptor) { + throw CallStatus.UNIMPLEMENTED.toRuntimeException(); + } + + @Override + public FlightInfo getFlightInfoPreparedStatement( + FlightSql.CommandPreparedStatementQuery cmd, CallContext ctx, FlightDescriptor descriptor) { + throw CallStatus.UNIMPLEMENTED.toRuntimeException(); + } + + @Override + public SchemaResult getSchemaStatement( + FlightSql.CommandStatementQuery cmd, CallContext ctx, FlightDescriptor descriptor) { + throw CallStatus.UNIMPLEMENTED.toRuntimeException(); + } + + @Override + public void getStreamStatement( + FlightSql.TicketStatementQuery ticket, CallContext ctx, ServerStreamListener listener) { + listener.error(CallStatus.UNIMPLEMENTED.toRuntimeException()); + } + + @Override + public void getStreamPreparedStatement( + FlightSql.CommandPreparedStatementQuery cmd, CallContext ctx, ServerStreamListener listener) { + listener.error(CallStatus.UNIMPLEMENTED.toRuntimeException()); + } + + @Override + public Runnable acceptPutStatement( + FlightSql.CommandStatementUpdate cmd, + CallContext ctx, + FlightStream stream, + StreamListener listener) { + return null; + } + + @Override + public Runnable acceptPutPreparedStatementUpdate( + FlightSql.CommandPreparedStatementUpdate cmd, + CallContext ctx, + FlightStream stream, + StreamListener listener) { + return null; + } + + @Override + public Runnable acceptPutPreparedStatementQuery( + FlightSql.CommandPreparedStatementQuery cmd, + CallContext ctx, + FlightStream stream, + StreamListener listener) { + return null; + } + + @Override + public FlightInfo getFlightInfoSqlInfo( + FlightSql.CommandGetSqlInfo cmd, CallContext ctx, FlightDescriptor descriptor) { + throw CallStatus.UNIMPLEMENTED.toRuntimeException(); + } + + @Override + public void getStreamSqlInfo( + FlightSql.CommandGetSqlInfo cmd, CallContext ctx, ServerStreamListener listener) { + listener.error(CallStatus.UNIMPLEMENTED.toRuntimeException()); + } + + @Override + public FlightInfo getFlightInfoTypeInfo( + FlightSql.CommandGetXdbcTypeInfo cmd, CallContext ctx, FlightDescriptor descriptor) { + throw CallStatus.UNIMPLEMENTED.toRuntimeException(); + } + + @Override + public void getStreamTypeInfo( + FlightSql.CommandGetXdbcTypeInfo cmd, CallContext ctx, ServerStreamListener listener) { + listener.error(CallStatus.UNIMPLEMENTED.toRuntimeException()); + } + + @Override + public FlightInfo getFlightInfoCatalogs( + FlightSql.CommandGetCatalogs cmd, CallContext ctx, FlightDescriptor descriptor) { + throw CallStatus.UNIMPLEMENTED.toRuntimeException(); + } + + @Override + public void getStreamCatalogs(CallContext ctx, ServerStreamListener listener) { + listener.error(CallStatus.UNIMPLEMENTED.toRuntimeException()); + } + + @Override + public FlightInfo getFlightInfoSchemas( + FlightSql.CommandGetDbSchemas cmd, CallContext ctx, FlightDescriptor descriptor) { + throw CallStatus.UNIMPLEMENTED.toRuntimeException(); + } + + @Override + public void getStreamSchemas( + FlightSql.CommandGetDbSchemas cmd, CallContext ctx, ServerStreamListener listener) { + listener.error(CallStatus.UNIMPLEMENTED.toRuntimeException()); + } + + @Override + public FlightInfo getFlightInfoTables( + FlightSql.CommandGetTables cmd, CallContext ctx, FlightDescriptor descriptor) { + throw CallStatus.UNIMPLEMENTED.toRuntimeException(); + } + + @Override + public void getStreamTables( + FlightSql.CommandGetTables cmd, CallContext ctx, ServerStreamListener listener) { + listener.error(CallStatus.UNIMPLEMENTED.toRuntimeException()); + } + + @Override + public FlightInfo getFlightInfoTableTypes( + FlightSql.CommandGetTableTypes cmd, CallContext ctx, FlightDescriptor descriptor) { + throw CallStatus.UNIMPLEMENTED.toRuntimeException(); + } + + @Override + public void getStreamTableTypes(CallContext ctx, ServerStreamListener listener) { + listener.error(CallStatus.UNIMPLEMENTED.toRuntimeException()); + } + + @Override + public FlightInfo getFlightInfoPrimaryKeys( + FlightSql.CommandGetPrimaryKeys cmd, CallContext ctx, FlightDescriptor descriptor) { + throw CallStatus.UNIMPLEMENTED.toRuntimeException(); + } + + @Override + public void getStreamPrimaryKeys( + FlightSql.CommandGetPrimaryKeys cmd, CallContext ctx, ServerStreamListener listener) { + listener.error(CallStatus.UNIMPLEMENTED.toRuntimeException()); + } + + @Override + public FlightInfo getFlightInfoExportedKeys( + FlightSql.CommandGetExportedKeys cmd, CallContext ctx, FlightDescriptor descriptor) { + throw CallStatus.UNIMPLEMENTED.toRuntimeException(); + } + + @Override + public FlightInfo getFlightInfoImportedKeys( + FlightSql.CommandGetImportedKeys cmd, CallContext ctx, FlightDescriptor descriptor) { + throw CallStatus.UNIMPLEMENTED.toRuntimeException(); + } + + @Override + public FlightInfo getFlightInfoCrossReference( + FlightSql.CommandGetCrossReference cmd, CallContext ctx, FlightDescriptor descriptor) { + throw CallStatus.UNIMPLEMENTED.toRuntimeException(); + } + + @Override + public void getStreamExportedKeys( + FlightSql.CommandGetExportedKeys cmd, CallContext ctx, ServerStreamListener listener) { + listener.error(CallStatus.UNIMPLEMENTED.toRuntimeException()); + } + + @Override + public void getStreamImportedKeys( + FlightSql.CommandGetImportedKeys cmd, CallContext ctx, ServerStreamListener listener) { + listener.error(CallStatus.UNIMPLEMENTED.toRuntimeException()); + } + + @Override + public void getStreamCrossReference( + FlightSql.CommandGetCrossReference cmd, CallContext ctx, ServerStreamListener listener) { + listener.error(CallStatus.UNIMPLEMENTED.toRuntimeException()); + } + + @Override + public void listFlights( + CallContext ctx, Criteria criteria, StreamListener listener) {} + + @Override + public void close() {} + } +} From b69cb837428d00b3c90afadcfe2d6082adee6daf Mon Sep 17 00:00:00 2001 From: unikdahal Date: Sat, 4 Jul 2026 01:54:24 +0530 Subject: [PATCH 02/10] Added Jackson Library, Fixed Issues --- java/driver/flight-sql/pom.xml | 5 + .../driver/flightsql/FlightSqlConnection.java | 114 +++---- .../flightsql/FlightSqlSessionUtil.java | 295 ++++++++---------- .../flightsql/FlightSqlSessionTest.java | 53 +++- java/pom.xml | 1 + 5 files changed, 227 insertions(+), 241 deletions(-) diff --git a/java/driver/flight-sql/pom.xml b/java/driver/flight-sql/pom.xml index da97becb76..3d1cc854eb 100644 --- a/java/driver/flight-sql/pom.xml +++ b/java/driver/flight-sql/pom.xml @@ -37,6 +37,11 @@ caffeine 3.2.4 + + com.fasterxml.jackson.core + jackson-databind + ${dep.jackson.version} + com.google.protobuf protobuf-java diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java index e6b0b39750..02fe64d076 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java @@ -48,7 +48,6 @@ import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.flight.FlightStatusCode; import org.apache.arrow.flight.GetSessionOptionsRequest; -import org.apache.arrow.flight.GetSessionOptionsResult; import org.apache.arrow.flight.HeaderCallOption; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.SessionOptionValue; @@ -214,74 +213,44 @@ public void setAutoCommit(boolean enableAutoCommit) throws AdbcException { } @Override - @SuppressWarnings("unchecked") public T getOption(TypedKey key) throws AdbcException { final String k = key.getKey(); if (k.equals(FlightSqlConnectionProperties.SESSION_OPTIONS)) { - return key.cast(FlightSqlSessionUtil.sessionOptionsToJson(fetchSessionOptions())); - - } else if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_BOOL_PREFIX)) { - final String name = - k.substring(FlightSqlConnectionProperties.SESSION_OPTION_BOOL_PREFIX.length()); - final SessionOptionValue val = fetchSessionOptions().get(name); - if (val == null) { - throw new AdbcException( - "[Flight SQL] Session option not found: " + name, - null, - AdbcStatusCode.NOT_FOUND, - null, - 0); - } - if (key.getType() == Boolean.class) { - return key.cast(val.acceptVisitor(FlightSqlSessionUtil.BOOL_VISITOR)); - } - return key.cast(val.acceptVisitor(FlightSqlSessionUtil.STRING_VISITOR)); - - } else if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_STRING_LIST_PREFIX)) { - final String name = - k.substring(FlightSqlConnectionProperties.SESSION_OPTION_STRING_LIST_PREFIX.length()); - final SessionOptionValue val = fetchSessionOptions().get(name); - if (val == null) { - throw new AdbcException( - "[Flight SQL] Session option not found: " + name, - null, - AdbcStatusCode.NOT_FOUND, - null, - 0); - } - if (key.getType() != String[].class) { + if (key.getType() != String.class) { return AdbcConnection.super.getOption(key); } - return key.cast(val.acceptVisitor(FlightSqlSessionUtil.STRING_ARRAY_VISITOR)); + return key.cast(FlightSqlSessionUtil.toJson(fetchSessionOptionsOrEmpty())); + } + final String prefix; + if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_BOOL_PREFIX)) { + prefix = FlightSqlConnectionProperties.SESSION_OPTION_BOOL_PREFIX; + } else if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_STRING_LIST_PREFIX)) { + prefix = FlightSqlConnectionProperties.SESSION_OPTION_STRING_LIST_PREFIX; } else if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_PREFIX)) { - final String name = - k.substring(FlightSqlConnectionProperties.SESSION_OPTION_PREFIX.length()); - final SessionOptionValue val = fetchSessionOptions().get(name); - if (val == null) { - throw new AdbcException( - "[Flight SQL] Session option not found: " + name, - null, - AdbcStatusCode.NOT_FOUND, - null, - 0); - } - try { - if (key.getType() == Long.class) { - return key.cast(val.acceptVisitor(FlightSqlSessionUtil.LONG_VISITOR)); - } else if (key.getType() == Double.class) { - return key.cast(val.acceptVisitor(FlightSqlSessionUtil.DOUBLE_VISITOR)); - } - } catch (NumberFormatException e) { - throw AdbcException.invalidArgument( - "[Flight SQL] Session option '" + name + "' value cannot be parsed as " - + key.getType().getSimpleName() + ": " + e.getMessage()); - } - return key.cast(val.acceptVisitor(FlightSqlSessionUtil.STRING_VISITOR)); + prefix = FlightSqlConnectionProperties.SESSION_OPTION_PREFIX; + } else { + return AdbcConnection.super.getOption(key); } - return AdbcConnection.super.getOption(key); + final String name = k.substring(prefix.length()); + if (name.isEmpty()) { + throw AdbcException.invalidArgument("[Flight SQL] Session option name must not be empty"); + } + final Object raw = + FlightSqlSessionUtil.require(fetchSessionOptions(), name) + .acceptVisitor(FlightSqlSessionUtil.TO_JAVA); + if (raw == null) { + throw new AdbcException( + "[Flight SQL] Session option not found: " + name, + null, + AdbcStatusCode.NOT_FOUND, + null, + 0); + } + final T result = FlightSqlSessionUtil.cast(key, raw, name); + return result != null ? result : AdbcConnection.super.getOption(key); } @Override @@ -290,7 +259,8 @@ public void setOption(TypedKey key, T value) throws AdbcException { if (value == null) { throw AdbcException.invalidArgument( - "[Flight SQL] null value not allowed for key: " + k + "[Flight SQL] null value not allowed for key: " + + k + " — use adbc.flight.sql.session.optionerase. to erase an option"); } @@ -322,8 +292,7 @@ public void setOption(TypedKey key, T value) throws AdbcException { doSetSessionOption(name, SessionOptionValueFactory.makeSessionOptionValue(arr)); } else if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_PREFIX)) { - final String name = - k.substring(FlightSqlConnectionProperties.SESSION_OPTION_PREFIX.length()); + final String name = k.substring(FlightSqlConnectionProperties.SESSION_OPTION_PREFIX.length()); final SessionOptionValue sv; if (value instanceof Long) { sv = SessionOptionValueFactory.makeSessionOptionValue((Long) value); @@ -347,8 +316,10 @@ public void setOption(TypedKey key, T value) throws AdbcException { public void close() throws Exception { try { client.closeSession(new CloseSessionRequest()); - } catch (FlightRuntimeException ignored) { - // best-effort: UNIMPLEMENTED is normal; other errors are silently swallowed at close time + } catch (FlightRuntimeException e) { + if (e.status().code() != FlightStatusCode.UNIMPLEMENTED) { + throw FlightSqlDriverUtil.fromFlightException(e); + } } clientCache.invalidateAll(); AutoCloseables.close(client, allocator); @@ -361,9 +332,15 @@ public String toString() { private Map fetchSessionOptions() throws AdbcException { try { - GetSessionOptionsResult result = - client.getSessionOptions(new GetSessionOptionsRequest()); - return result.getSessionOptions(); + return client.getSessionOptions(new GetSessionOptionsRequest()).getSessionOptions(); + } catch (FlightRuntimeException e) { + throw FlightSqlDriverUtil.fromFlightException(e); + } + } + + private Map fetchSessionOptionsOrEmpty() throws AdbcException { + try { + return client.getSessionOptions(new GetSessionOptionsRequest()).getSessionOptions(); } catch (FlightRuntimeException e) { if (e.status().code() == FlightStatusCode.UNIMPLEMENTED) { return Collections.emptyMap(); @@ -374,8 +351,7 @@ private Map fetchSessionOptions() throws AdbcExcepti private void doSetSessionOption(String name, SessionOptionValue value) throws AdbcException { if (name.isEmpty()) { - throw AdbcException.invalidArgument( - "[Flight SQL] Session option name must not be empty"); + throw AdbcException.invalidArgument("[Flight SQL] Session option name must not be empty"); } final SetSessionOptionsResult result; try { diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java index df3d97c18c..9a533afa0e 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java @@ -14,198 +14,159 @@ * See the License for the specific language governing permissions and * limitations under the License. */ + package org.apache.arrow.adbc.driver.flightsql; -import java.util.ArrayList; -import java.util.List; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.util.LinkedHashMap; import java.util.Map; import org.apache.arrow.adbc.core.AdbcException; +import org.apache.arrow.adbc.core.AdbcStatusCode; +import org.apache.arrow.adbc.core.TypedKey; +import org.apache.arrow.flight.NoOpSessionOptionValueVisitor; import org.apache.arrow.flight.SessionOptionValue; -import org.apache.arrow.flight.SessionOptionValueVisitor; -/** Helpers for Flight SQL session option serialization and type conversion. */ +/** Package-private helpers for Flight SQL session option serialization and type conversion. */ final class FlightSqlSessionUtil { - private FlightSqlSessionUtil() {} + static final ObjectMapper MAPPER = new ObjectMapper(); + + /** + * Extracts a native Java value from a {@link SessionOptionValue}. NaN/Infinity doubles are + * returned as-is and filtered in {@link #toJson}; String[] is defensively cloned; Void returns + * {@code null} (callers must handle null before calling {@link #cast}). + */ + static final NoOpSessionOptionValueVisitor TO_JAVA = + new NoOpSessionOptionValueVisitor() { + @Override + public Object visit(String v) { + return v; + } + + @Override + public Object visit(boolean v) { + return v; + } - // -- JSON helpers -- + @Override + public Object visit(long v) { + return v; + } + + @Override + public Object visit(double v) { + return v; + } + + @Override + public Object visit(String[] v) { + return v.clone(); + } + }; - static String sessionOptionsToJson(Map opts) { - final StringBuilder sb = new StringBuilder("{"); - boolean first = true; + /** Serializes all session options to a JSON object string. */ + static String toJson(Map opts) throws AdbcException { + Map map = new LinkedHashMap<>(); for (Map.Entry e : opts.entrySet()) { - if (!first) sb.append(','); - sb.append('"').append(escapeJson(e.getKey())).append("\":").append(sessionValueToJson(e.getValue())); - first = false; + Object v = e.getValue().acceptVisitor(TO_JAVA); + // JSON does not support NaN or Infinity; represent them as null + if (v instanceof Double && !Double.isFinite((Double) v)) { + v = null; + } + map.put(e.getKey(), v); + } + try { + return MAPPER.writeValueAsString(map); + } catch (JsonProcessingException e) { + throw new AdbcException( + "[Flight SQL] Failed to serialize session options", e, AdbcStatusCode.INTERNAL, null, 0); } - return sb.append('}').toString(); - } - - static String sessionValueToJson(SessionOptionValue val) { - return val.acceptVisitor(JSON_VALUE_VISITOR); } - static String toJsonArray(String[] values) { - final StringBuilder sb = new StringBuilder("["); - for (int i = 0; i < values.length; i++) { - if (i > 0) sb.append(','); - sb.append('"').append(escapeJson(values[i])).append('"'); + /** Parses a JSON string array (used when a string-list option is supplied as JSON). */ + static String[] parseJsonArray(String json) throws AdbcException { + try { + return MAPPER.readValue(json, String[].class); + } catch (JsonProcessingException e) { + throw new AdbcException( + "[Flight SQL] Expected JSON array for string list option, got: " + json, + e, + AdbcStatusCode.INVALID_ARGUMENT, + null, + 0); } - return sb.append(']').toString(); } - static String[] parseJsonArray(String json) throws AdbcException { - final String trimmed = json.trim(); - if (!trimmed.startsWith("[") || !trimmed.endsWith("]")) { - throw AdbcException.invalidArgument( - "[Flight SQL] Expected JSON array for string list option, got: " + json); + /** + * Casts a raw Java value extracted via {@link #TO_JAVA} to the type requested by a {@link + * TypedKey}. {@code raw} must not be {@code null} (Void options must be rejected before calling + * this). Returns {@code null} for unsupported types so the caller can delegate to the default + * {@code AdbcConnection} implementation. + */ + @SuppressWarnings("unchecked") + static T cast(TypedKey key, Object raw, String optionName) throws AdbcException { + final Class type = key.getType(); + if (type == String.class) { + if (raw instanceof String[]) { + try { + return (T) MAPPER.writeValueAsString(raw); + } catch (JsonProcessingException e) { + throw new AdbcException( + "[Flight SQL] Failed to serialize string list option as JSON", + e, + AdbcStatusCode.INTERNAL, + null, + 0); + } + } + return (T) String.valueOf(raw); } - final String inner = trimmed.substring(1, trimmed.length() - 1).trim(); - if (inner.isEmpty()) { - return new String[0]; + if (type == Boolean.class) { + if (raw instanceof Boolean) return (T) raw; + return (T) Boolean.valueOf(Boolean.parseBoolean(String.valueOf(raw))); } - final List result = new ArrayList<>(); - int i = 0; - while (i < inner.length()) { - while (i < inner.length() && Character.isWhitespace(inner.charAt(i))) i++; - if (i >= inner.length()) break; - if (inner.charAt(i) != '"') { - throw AdbcException.invalidArgument( - "[Flight SQL] Invalid JSON string array element at position " + i + " in: " + json); - } - i++; // skip opening quote - final StringBuilder elem = new StringBuilder(); - while (i < inner.length() && inner.charAt(i) != '"') { - if (inner.charAt(i) == '\\' && i + 1 < inner.length()) { - i++; - switch (inner.charAt(i)) { - case '"': elem.append('"'); break; - case '\\': elem.append('\\'); break; - case '/': elem.append('/'); break; - case 'n': elem.append('\n'); break; - case 'r': elem.append('\r'); break; - case 't': elem.append('\t'); break; - case 'b': elem.append('\b'); break; - case 'f': elem.append('\f'); break; - case 'u': - if (i + 4 < inner.length()) { - final String hex = inner.substring(i + 1, i + 5); - try { - elem.append((char) Integer.parseInt(hex, 16)); - i += 4; - } catch (NumberFormatException e) { - elem.append('u'); - } - } else { - elem.append('u'); - } - break; - default: elem.append(inner.charAt(i)); - } - } else { - elem.append(inner.charAt(i)); - } - i++; + if (type == String[].class) { + if (raw instanceof String[]) return (T) raw; + return (T) new String[] {String.valueOf(raw)}; + } + try { + if (type == Long.class) { + if (raw instanceof Long) return (T) raw; + if (raw instanceof Number) return (T) Long.valueOf(((Number) raw).longValue()); + return (T) Long.valueOf(Long.parseLong(String.valueOf(raw))); } - if (i >= inner.length()) { - throw AdbcException.invalidArgument( - "[Flight SQL] Unterminated string in JSON array: " + json); + if (type == Double.class) { + if (raw instanceof Double) return (T) raw; + if (raw instanceof Number) return (T) Double.valueOf(((Number) raw).doubleValue()); + return (T) Double.valueOf(Double.parseDouble(String.valueOf(raw))); } - i++; // skip closing quote - result.add(elem.toString()); - while (i < inner.length() && Character.isWhitespace(inner.charAt(i))) i++; - if (i < inner.length() && inner.charAt(i) == ',') i++; + } catch (NumberFormatException e) { + throw AdbcException.invalidArgument( + "[Flight SQL] Session option '" + + optionName + + "' cannot be parsed as " + + type.getSimpleName() + + ": " + + raw); } - return result.toArray(new String[0]); + return null; } - static String escapeJson(String s) { - final StringBuilder sb = new StringBuilder(s.length()); - for (int i = 0; i < s.length(); i++) { - final char c = s.charAt(i); - switch (c) { - case '"': sb.append("\\\""); break; - case '\\': sb.append("\\\\"); break; - case '\n': sb.append("\\n"); break; - case '\r': sb.append("\\r"); break; - case '\t': sb.append("\\t"); break; - case '\b': sb.append("\\b"); break; - case '\f': sb.append("\\f"); break; - default: - if (c < 0x20) { - sb.append(String.format("\\u%04x", (int) c)); - } else { - sb.append(c); - } - } + /** Looks up a session option by name, throwing {@code NOT_FOUND} if absent. */ + static SessionOptionValue require(Map opts, String name) + throws AdbcException { + SessionOptionValue val = opts.get(name); + if (val == null) { + throw new AdbcException( + "[Flight SQL] Session option not found: " + name, + null, + AdbcStatusCode.NOT_FOUND, + null, + 0); } - return sb.toString(); + return val; } - // -- SessionOptionValue visitors -- - - static final SessionOptionValueVisitor BOOL_VISITOR = - new SessionOptionValueVisitor() { - @Override public Boolean visit(String v) { return Boolean.parseBoolean(v); } - @Override public Boolean visit(boolean v) { return v; } - @Override public Boolean visit(long v) { return v != 0; } - @Override public Boolean visit(double v) { return v != 0; } - @Override public Boolean visit(String[] v) { return false; } - @Override public Boolean visit(Void v) { return false; } - }; - - static final SessionOptionValueVisitor STRING_ARRAY_VISITOR = - new SessionOptionValueVisitor() { - @Override public String[] visit(String v) { return new String[]{v}; } - @Override public String[] visit(boolean v) { return new String[]{String.valueOf(v)}; } - @Override public String[] visit(long v) { return new String[]{String.valueOf(v)}; } - @Override public String[] visit(double v) { return new String[]{String.valueOf(v)}; } - @Override public String[] visit(String[] v) { return v.clone(); } - @Override public String[] visit(Void v) { return new String[0]; } - }; - - static final SessionOptionValueVisitor STRING_VISITOR = - new SessionOptionValueVisitor() { - @Override public String visit(String v) { return v; } - @Override public String visit(boolean v) { return String.valueOf(v); } - @Override public String visit(long v) { return String.valueOf(v); } - @Override public String visit(double v) { return String.valueOf(v); } - @Override public String visit(String[] v) { return toJsonArray(v); } - @Override public String visit(Void v) { return ""; } - }; - - static final SessionOptionValueVisitor LONG_VISITOR = - new SessionOptionValueVisitor() { - @Override public Long visit(String v) { return Long.parseLong(v); } - @Override public Long visit(boolean v) { return v ? 1L : 0L; } - @Override public Long visit(long v) { return v; } - @Override public Long visit(double v) { return (long) v; } - @Override public Long visit(String[] v) { return 0L; } - @Override public Long visit(Void v) { return 0L; } - }; - - static final SessionOptionValueVisitor DOUBLE_VISITOR = - new SessionOptionValueVisitor() { - @Override public Double visit(String v) { return Double.parseDouble(v); } - @Override public Double visit(boolean v) { return v ? 1.0 : 0.0; } - @Override public Double visit(long v) { return (double) v; } - @Override public Double visit(double v) { return v; } - @Override public Double visit(String[] v) { return 0.0; } - @Override public Double visit(Void v) { return 0.0; } - }; - - static final SessionOptionValueVisitor JSON_VALUE_VISITOR = - new SessionOptionValueVisitor() { - @Override public String visit(String v) { return '"' + escapeJson(v) + '"'; } - @Override public String visit(boolean v) { return String.valueOf(v); } - @Override public String visit(long v) { return String.valueOf(v); } - @Override public String visit(double v) { - if (Double.isNaN(v) || Double.isInfinite(v)) { - return "null"; - } - return String.valueOf(v); - } - @Override public String visit(String[] v) { return toJsonArray(v); } - @Override public String visit(Void v) { return "null"; } - }; + private FlightSqlSessionUtil() {} } diff --git a/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionTest.java b/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionTest.java index 571964a1db..115e7b91aa 100644 --- a/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionTest.java +++ b/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionTest.java @@ -14,6 +14,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ + package org.apache.arrow.adbc.driver.flightsql; import static org.assertj.core.api.Assertions.assertThat; @@ -31,8 +32,8 @@ import org.apache.arrow.adbc.core.TypedKey; import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.CloseSessionRequest; -import org.apache.arrow.flight.Criteria; import org.apache.arrow.flight.CloseSessionResult; +import org.apache.arrow.flight.Criteria; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightServer; @@ -44,7 +45,6 @@ import org.apache.arrow.flight.Result; import org.apache.arrow.flight.SchemaResult; import org.apache.arrow.flight.SessionOptionValue; -import org.apache.arrow.flight.SessionOptionValueFactory; import org.apache.arrow.flight.SetSessionOptionsRequest; import org.apache.arrow.flight.SetSessionOptionsResult; import org.apache.arrow.flight.sql.FlightSqlProducer; @@ -186,8 +186,7 @@ void testEraseOption() throws Exception { @Test void testGetSessionOptionsBlob() throws Exception { connection.setOption( - new TypedKey<>( - FlightSqlConnectionProperties.SESSION_OPTION_PREFIX + "k1", String.class), + new TypedKey<>(FlightSqlConnectionProperties.SESSION_OPTION_PREFIX + "k1", String.class), "v1"); String blob = @@ -205,6 +204,36 @@ void testCloseSessionCalledOnClose() throws Exception { assertThat(producer.closeSessionCalled.get()).isTrue(); } + @Test + void testCloseDoesNotThrowWhenServerReturnsUnimplemented() throws Exception { + producer.rejectClose.set(true); + connection.close(); // must not throw + connection = null; + } + + @Test + void testGetOptionNotImplementedWhenServerUnsupported() throws Exception { + producer.rejectGetSession.set(true); + AdbcException ex = + assertThrows( + AdbcException.class, + () -> + connection.getOption( + new TypedKey<>( + FlightSqlConnectionProperties.SESSION_OPTION_PREFIX + "foo", + String.class))); + assertThat(ex.getStatus()).isEqualTo(AdbcStatusCode.NOT_IMPLEMENTED); + } + + @Test + void testGetSessionOptionsBlobEmptyWhenServerUnsupported() throws Exception { + producer.rejectGetSession.set(true); + String blob = + connection.getOption( + new TypedKey<>(FlightSqlConnectionProperties.SESSION_OPTIONS, String.class)); + assertThat(blob).isEqualTo("{}"); + } + @Test void testSetOptionReadOnlyBlobThrows() { AdbcException ex = @@ -222,10 +251,14 @@ void testSetOptionReadOnlyBlobThrows() { static class SessionProducer implements FlightSqlProducer { private final Map sessionOptions = new HashMap<>(); final AtomicBoolean closeSessionCalled = new AtomicBoolean(false); + final AtomicBoolean rejectClose = new AtomicBoolean(false); + final AtomicBoolean rejectGetSession = new AtomicBoolean(false); void reset() { sessionOptions.clear(); closeSessionCalled.set(false); + rejectClose.set(false); + rejectGetSession.set(false); } @Override @@ -249,6 +282,10 @@ public void getSessionOptions( GetSessionOptionsRequest request, CallContext context, StreamListener listener) { + if (rejectGetSession.get()) { + listener.onError(CallStatus.UNIMPLEMENTED.toRuntimeException()); + return; + } listener.onNext(new GetSessionOptionsResult(new HashMap<>(sessionOptions))); listener.onCompleted(); } @@ -258,6 +295,10 @@ public void closeSession( CloseSessionRequest request, CallContext context, StreamListener listener) { + if (rejectClose.get()) { + listener.onError(CallStatus.UNIMPLEMENTED.toRuntimeException()); + return; + } closeSessionCalled.set(true); listener.onNext(new CloseSessionResult(CloseSessionResult.Status.CLOSED)); listener.onCompleted(); @@ -307,7 +348,9 @@ public void getStreamStatement( @Override public void getStreamPreparedStatement( - FlightSql.CommandPreparedStatementQuery cmd, CallContext ctx, ServerStreamListener listener) { + FlightSql.CommandPreparedStatementQuery cmd, + CallContext ctx, + ServerStreamListener listener) { listener.error(CallStatus.UNIMPLEMENTED.toRuntimeException()); } diff --git a/java/pom.xml b/java/pom.xml index abd9139f51..bb3aa9c35c 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -97,6 +97,7 @@ 19.0.0 + 2.21.0 4.2.0 11 From ee9b8907dcb83242380535092d03d3a93ac7413e Mon Sep 17 00:00:00 2001 From: unikdahal Date: Sat, 4 Jul 2026 16:53:37 +0530 Subject: [PATCH 03/10] fix for if closeSession throws any other exceptions then client and allocator leaking memory --- .../arrow/adbc/driver/flightsql/FlightSqlConnection.java | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java index 02fe64d076..0ffa82e5fa 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java @@ -320,9 +320,13 @@ public void close() throws Exception { if (e.status().code() != FlightStatusCode.UNIMPLEMENTED) { throw FlightSqlDriverUtil.fromFlightException(e); } + } finally { + try { + clientCache.invalidateAll(); + } finally { + AutoCloseables.close(client, allocator); + } } - clientCache.invalidateAll(); - AutoCloseables.close(client, allocator); } @Override From 52eb8a296a0ca8bd0807691ca296b5518bf4de98 Mon Sep 17 00:00:00 2001 From: unikdahal Date: Sat, 4 Jul 2026 17:05:52 +0530 Subject: [PATCH 04/10] fix for not throwing exception in case of non boolean values when parsing session options --- .../adbc/driver/flightsql/FlightSqlSessionUtil.java | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java index 9a533afa0e..3bf574b7f7 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java @@ -124,7 +124,14 @@ static T cast(TypedKey key, Object raw, String optionName) throws AdbcExc } if (type == Boolean.class) { if (raw instanceof Boolean) return (T) raw; - return (T) Boolean.valueOf(Boolean.parseBoolean(String.valueOf(raw))); + String s = String.valueOf(raw); + if ("true".equalsIgnoreCase(s)) return (T) Boolean.TRUE; + if ("false".equalsIgnoreCase(s)) return (T) Boolean.FALSE; + throw AdbcException.invalidArgument( + "[Flight SQL] Session option '" + + optionName + + "' cannot be parsed as Boolean: " + + raw); } if (type == String[].class) { if (raw instanceof String[]) return (T) raw; From cfcf3f5c30f064248c78d9412a4d450cc308777d Mon Sep 17 00:00:00 2001 From: unikdahal Date: Sun, 5 Jul 2026 05:22:27 +0530 Subject: [PATCH 05/10] Fixed formatting --- .../arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java index 3bf574b7f7..3b6ebe6ed2 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java @@ -128,10 +128,7 @@ static T cast(TypedKey key, Object raw, String optionName) throws AdbcExc if ("true".equalsIgnoreCase(s)) return (T) Boolean.TRUE; if ("false".equalsIgnoreCase(s)) return (T) Boolean.FALSE; throw AdbcException.invalidArgument( - "[Flight SQL] Session option '" - + optionName - + "' cannot be parsed as Boolean: " - + raw); + "[Flight SQL] Session option '" + optionName + "' cannot be parsed as Boolean: " + raw); } if (type == String[].class) { if (raw instanceof String[]) return (T) raw; From 7580eaec3558f6610b0b4c85c32a297924ae6361 Mon Sep 17 00:00:00 2001 From: unikdahal Date: Mon, 13 Jul 2026 00:44:01 +0530 Subject: [PATCH 06/10] refactored as per review comments --- .../adbc/driver/flightsql/FlightSqlConnection.java | 12 ++---------- .../adbc/driver/flightsql/FlightSqlSessionUtil.java | 9 +++------ 2 files changed, 5 insertions(+), 16 deletions(-) diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java index 0ffa82e5fa..a3bb033c73 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java @@ -239,7 +239,7 @@ public T getOption(TypedKey key) throws AdbcException { throw AdbcException.invalidArgument("[Flight SQL] Session option name must not be empty"); } final Object raw = - FlightSqlSessionUtil.require(fetchSessionOptions(), name) + FlightSqlSessionUtil.require(fetchSessionOptionsOrEmpty(), name) .acceptVisitor(FlightSqlSessionUtil.TO_JAVA); if (raw == null) { throw new AdbcException( @@ -261,7 +261,7 @@ public void setOption(TypedKey key, T value) throws AdbcException { throw AdbcException.invalidArgument( "[Flight SQL] null value not allowed for key: " + k - + " — use adbc.flight.sql.session.optionerase. to erase an option"); + + " - use adbc.flight.sql.session.optionerase. to erase an option"); } if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_ERASE_PREFIX)) { @@ -334,14 +334,6 @@ public String toString() { return "FlightSqlConnection{" + "client=" + client + '}'; } - private Map fetchSessionOptions() throws AdbcException { - try { - return client.getSessionOptions(new GetSessionOptionsRequest()).getSessionOptions(); - } catch (FlightRuntimeException e) { - throw FlightSqlDriverUtil.fromFlightException(e); - } - } - private Map fetchSessionOptionsOrEmpty() throws AdbcException { try { return client.getSessionOptions(new GetSessionOptionsRequest()).getSessionOptions(); diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java index 3b6ebe6ed2..ad9eca0500 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java @@ -89,12 +89,9 @@ static String[] parseJsonArray(String json) throws AdbcException { try { return MAPPER.readValue(json, String[].class); } catch (JsonProcessingException e) { - throw new AdbcException( - "[Flight SQL] Expected JSON array for string list option, got: " + json, - e, - AdbcStatusCode.INVALID_ARGUMENT, - null, - 0); + throw AdbcException.invalidArgument( + "[Flight SQL] Expected JSON array for string list option, got: " + json) + .withCause(e); } } From f2341e4bcc8197d129e976e565046b3e96aa83ec Mon Sep 17 00:00:00 2001 From: unikdahal Date: Mon, 13 Jul 2026 01:02:30 +0530 Subject: [PATCH 07/10] refactored as per review comments --- .../apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java index 45d14f9292..669386494a 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java @@ -239,7 +239,7 @@ public T getOption(TypedKey key) throws AdbcException { throw AdbcException.invalidArgument("[Flight SQL] Session option name must not be empty"); } final Object raw = - FlightSqlSessionUtil.require(fetchSessionOptions(), name) + FlightSqlSessionUtil.require(fetchSessionOptionsOrEmpty(), name) .acceptVisitor(FlightSqlSessionUtil.TO_JAVA); if (raw == null) { throw new AdbcException( From 72bfccd13c5583d5a9a8f0477fbc636738467b55 Mon Sep 17 00:00:00 2001 From: unikdahal Date: Mon, 13 Jul 2026 01:05:49 +0530 Subject: [PATCH 08/10] removed unused fetchSessionOptions method --- .../arrow/adbc/driver/flightsql/FlightSqlConnection.java | 8 -------- 1 file changed, 8 deletions(-) diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java index 669386494a..688cf99104 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java @@ -338,14 +338,6 @@ public String toString() { return "FlightSqlConnection{" + "client=" + client + '}'; } - private Map fetchSessionOptions() throws AdbcException { - try { - return client.getSessionOptions(new GetSessionOptionsRequest()).getSessionOptions(); - } catch (FlightRuntimeException e) { - throw FlightSqlDriverUtil.fromFlightException(e); - } - } - private Map fetchSessionOptionsOrEmpty() throws AdbcException { try { return client.getSessionOptions(new GetSessionOptionsRequest()).getSessionOptions(); From d9f766f4356c62f2d9ac358ff7dc71d0846873a7 Mon Sep 17 00:00:00 2001 From: unikdahal Date: Mon, 13 Jul 2026 01:05:49 +0530 Subject: [PATCH 09/10] removed unused fetchSessionOptions method --- .../adbc/driver/flightsql/FlightSqlConnection.java | 8 -------- .../driver/flightsql/FlightSqlSessionTest.java | 14 -------------- 2 files changed, 22 deletions(-) diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java index 669386494a..688cf99104 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java @@ -338,14 +338,6 @@ public String toString() { return "FlightSqlConnection{" + "client=" + client + '}'; } - private Map fetchSessionOptions() throws AdbcException { - try { - return client.getSessionOptions(new GetSessionOptionsRequest()).getSessionOptions(); - } catch (FlightRuntimeException e) { - throw FlightSqlDriverUtil.fromFlightException(e); - } - } - private Map fetchSessionOptionsOrEmpty() throws AdbcException { try { return client.getSessionOptions(new GetSessionOptionsRequest()).getSessionOptions(); diff --git a/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionTest.java b/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionTest.java index 115e7b91aa..5ceb757b60 100644 --- a/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionTest.java +++ b/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionTest.java @@ -211,20 +211,6 @@ void testCloseDoesNotThrowWhenServerReturnsUnimplemented() throws Exception { connection = null; } - @Test - void testGetOptionNotImplementedWhenServerUnsupported() throws Exception { - producer.rejectGetSession.set(true); - AdbcException ex = - assertThrows( - AdbcException.class, - () -> - connection.getOption( - new TypedKey<>( - FlightSqlConnectionProperties.SESSION_OPTION_PREFIX + "foo", - String.class))); - assertThat(ex.getStatus()).isEqualTo(AdbcStatusCode.NOT_IMPLEMENTED); - } - @Test void testGetSessionOptionsBlobEmptyWhenServerUnsupported() throws Exception { producer.rejectGetSession.set(true); From 0d475c8a3bf34d23cdb68212121977903fda63cc Mon Sep 17 00:00:00 2001 From: unikdahal Date: Wed, 15 Jul 2026 00:40:34 +0530 Subject: [PATCH 10/10] some refactoring and pipeline failure fix --- .../driver/flightsql/FlightSqlConnection.java | 25 +++++++++----- .../flightsql/FlightSqlSessionUtil.java | 34 +++++++++++-------- 2 files changed, 37 insertions(+), 22 deletions(-) diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java index 688cf99104..e633991cf7 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java @@ -257,30 +257,29 @@ public T getOption(TypedKey key) throws AdbcException { public void setOption(TypedKey key, T value) throws AdbcException { final String k = key.getKey(); - if (value == null) { - throw AdbcException.invalidArgument( - "[Flight SQL] null value not allowed for key: " - + k - + " - use adbc.flight.sql.session.optionerase. to erase an option"); - } - if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_ERASE_PREFIX)) { final String name = k.substring(FlightSqlConnectionProperties.SESSION_OPTION_ERASE_PREFIX.length()); doSetSessionOption(name, SessionOptionValueFactory.makeEmptySessionOptionValue()); } else if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_BOOL_PREFIX)) { + if (value == null) { + throw invalidNullValue(k); + } final String name = k.substring(FlightSqlConnectionProperties.SESSION_OPTION_BOOL_PREFIX.length()); final boolean b; if (value instanceof Boolean) { b = (Boolean) value; } else { - b = Boolean.parseBoolean(value.toString()); + b = FlightSqlSessionUtil.parseStrictBoolean(value.toString(), name); } doSetSessionOption(name, SessionOptionValueFactory.makeSessionOptionValue(b)); } else if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_STRING_LIST_PREFIX)) { + if (value == null) { + throw invalidNullValue(k); + } final String name = k.substring(FlightSqlConnectionProperties.SESSION_OPTION_STRING_LIST_PREFIX.length()); final String[] arr; @@ -292,6 +291,9 @@ public void setOption(TypedKey key, T value) throws AdbcException { doSetSessionOption(name, SessionOptionValueFactory.makeSessionOptionValue(arr)); } else if (k.startsWith(FlightSqlConnectionProperties.SESSION_OPTION_PREFIX)) { + if (value == null) { + throw invalidNullValue(k); + } final String name = k.substring(FlightSqlConnectionProperties.SESSION_OPTION_PREFIX.length()); final SessionOptionValue sv; if (value instanceof Long) { @@ -312,6 +314,13 @@ public void setOption(TypedKey key, T value) throws AdbcException { } } + private static AdbcException invalidNullValue(String key) { + return AdbcException.invalidArgument( + "[Flight SQL] null value not allowed for key: " + + key + + " - use adbc.flight.sql.session.optionerase. to erase an option"); + } + @Override public void close() throws AdbcException { try { diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java index ad9eca0500..c16b63a751 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java @@ -26,6 +26,7 @@ import org.apache.arrow.adbc.core.TypedKey; import org.apache.arrow.flight.NoOpSessionOptionValueVisitor; import org.apache.arrow.flight.SessionOptionValue; +import org.checkerframework.checker.nullness.qual.Nullable; /** Package-private helpers for Flight SQL session option serialization and type conversion. */ final class FlightSqlSessionUtil { @@ -67,9 +68,9 @@ public Object visit(String[] v) { /** Serializes all session options to a JSON object string. */ static String toJson(Map opts) throws AdbcException { - Map map = new LinkedHashMap<>(); + Map map = new LinkedHashMap<>(); for (Map.Entry e : opts.entrySet()) { - Object v = e.getValue().acceptVisitor(TO_JAVA); + @Nullable Object v = e.getValue().acceptVisitor(TO_JAVA); // JSON does not support NaN or Infinity; represent them as null if (v instanceof Double && !Double.isFinite((Double) v)) { v = null; @@ -102,7 +103,7 @@ static String[] parseJsonArray(String json) throws AdbcException { * {@code AdbcConnection} implementation. */ @SuppressWarnings("unchecked") - static T cast(TypedKey key, Object raw, String optionName) throws AdbcException { + static @Nullable T cast(TypedKey key, Object raw, String optionName) throws AdbcException { final Class type = key.getType(); if (type == String.class) { if (raw instanceof String[]) { @@ -121,11 +122,7 @@ static T cast(TypedKey key, Object raw, String optionName) throws AdbcExc } if (type == Boolean.class) { if (raw instanceof Boolean) return (T) raw; - String s = String.valueOf(raw); - if ("true".equalsIgnoreCase(s)) return (T) Boolean.TRUE; - if ("false".equalsIgnoreCase(s)) return (T) Boolean.FALSE; - throw AdbcException.invalidArgument( - "[Flight SQL] Session option '" + optionName + "' cannot be parsed as Boolean: " + raw); + return (T) Boolean.valueOf(parseStrictBoolean(String.valueOf(raw), optionName)); } if (type == String[].class) { if (raw instanceof String[]) return (T) raw; @@ -144,12 +141,13 @@ static T cast(TypedKey key, Object raw, String optionName) throws AdbcExc } } catch (NumberFormatException e) { throw AdbcException.invalidArgument( - "[Flight SQL] Session option '" - + optionName - + "' cannot be parsed as " - + type.getSimpleName() - + ": " - + raw); + "[Flight SQL] Session option '" + + optionName + + "' cannot be parsed as " + + type.getSimpleName() + + ": " + + raw) + .withCause(e); } return null; } @@ -169,5 +167,13 @@ static SessionOptionValue require(Map opts, String n return val; } + /** Strictly parses "true"/"false" (case-insensitive); rejects anything else. */ + static boolean parseStrictBoolean(String s, String optionName) throws AdbcException { + if ("true".equalsIgnoreCase(s)) return true; + if ("false".equalsIgnoreCase(s)) return false; + throw AdbcException.invalidArgument( + "[Flight SQL] Session option '" + optionName + "' cannot be parsed as Boolean: " + s); + } + private FlightSqlSessionUtil() {} }