Copilot commented on code in PR #4444:
URL: https://github.com/apache/arrow-adbc/pull/4444#discussion_r3488658020
##########
java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java:
##########
@@ -203,8 +213,143 @@ public void setAutoCommit(boolean enableAutoCommit)
throws AdbcException {
}
}
+ @Override
+ @SuppressWarnings("unchecked")
+ public <T> T getOption(TypedKey<T> 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));
Review Comment:
For the string-list prefix, unsupported key types should be rejected/handled
explicitly. As written, callers using String.class (to get a JSON
representation) or any other non-String[] type will hit
AdbcConnection.super.getOption(key), while callers using an unexpected type
could still end up with a ClassCastException elsewhere. Consider supporting
String.class here (consistent with cross-language JSON support) and delegating
all other types to the default implementation.
##########
java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java:
##########
@@ -203,8 +213,143 @@ public void setAutoCommit(boolean enableAutoCommit)
throws AdbcException {
}
}
+ @Override
+ @SuppressWarnings("unchecked")
+ public <T> T getOption(TypedKey<T> 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));
+
Review Comment:
For the session bool prefix, any non-Boolean key type falls through to
STRING_VISITOR and is then cast via TypedKey.cast(...). If a caller uses an
unexpected type (e.g., Integer.class), this will throw ClassCastException
instead of the usual AdbcException.notImplemented/invalidArgument behavior.
##########
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<String, SessionOptionValue> opts) {
+ final StringBuilder sb = new StringBuilder("{");
+ boolean first = true;
+ for (Map.Entry<String, SessionOptionValue> 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<String> 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;
Review Comment:
parseJsonArray silently accepts invalid \\uXXXX escapes: on malformed/short
hex it appends a literal 'u' instead of failing. That can corrupt option values
while still appearing to parse successfully. Since this helper is meant for
cross-language JSON compatibility, it should reject invalid escapes with
INVALID_ARGUMENT.
##########
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<String, SessionOptionValue> opts) {
+ final StringBuilder sb = new StringBuilder("{");
+ boolean first = true;
+ for (Map.Entry<String, SessionOptionValue> 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<String> 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++;
+ }
Review Comment:
parseJsonArray currently allows invalid JSON like missing commas between
elements and trailing commas (because the comma is optional after each
element). This permissiveness can mask typos and lead to inconsistent behavior
across languages/parsers.
##########
java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java:
##########
@@ -203,8 +213,143 @@ public void setAutoCommit(boolean enableAutoCommit)
throws AdbcException {
}
}
+ @Override
+ @SuppressWarnings("unchecked")
+ public <T> T getOption(TypedKey<T> 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));
+ }
Review Comment:
For the generic session option prefix, unsupported TypedKey types will
currently be cast from STRING_VISITOR output, which can throw
ClassCastException. It’s safer to only return values for the explicitly
supported types (String/Long/Double) and delegate everything else to
AdbcConnection.super.getOption(key).
##########
java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java:
##########
@@ -203,8 +213,143 @@ public void setAutoCommit(boolean enableAutoCommit)
throws AdbcException {
}
}
+ @Override
+ @SuppressWarnings("unchecked")
+ public <T> T getOption(TypedKey<T> 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 <T> void setOption(TypedKey<T> 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.<name> 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
+ }
Review Comment:
close() currently swallows *all* FlightRuntimeException from CloseSession.
The linked issue/PR description calls out best-effort handling specifically for
UNIMPLEMENTED; swallowing other statuses (e.g., UNAUTHENTICATED/UNAVAILABLE)
can hide real shutdown failures and make debugging harder.
##########
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<String, Object> 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();
Review Comment:
PR description says the new unit tests cover graceful close when the server
returns UNIMPLEMENTED for CloseSession, but this test class only exercises the
successful CloseSession path. Either the description should be updated, or a
test should be added that makes closeSession return UNIMPLEMENTED and asserts
connection.close() remains best-effort.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
