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/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 99c9fb7350..fd43185eba 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,22 @@ 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.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; @@ -208,11 +217,125 @@ public void setAutoCommit(boolean enableAutoCommit) throws AdbcException { } } + @Override + public T getOption(TypedKey key) throws AdbcException { + final String k = key.getKey(); + + if (k.equals(FlightSqlConnectionProperties.SESSION_OPTIONS)) { + if (key.getType() != String.class) { + return AdbcConnection.super.getOption(key); + } + 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)) { + prefix = FlightSqlConnectionProperties.SESSION_OPTION_PREFIX; + } else { + 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(fetchSessionOptionsOrEmpty(), 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 + public void setOption(TypedKey key, T value) throws AdbcException { + final String k = key.getKey(); + + 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 = 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; + 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)) { + if (value == null) { + throw invalidNullValue(k); + } + 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); + } + } + + 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 { - clientCache.invalidateAll(); try { - AutoCloseables.close(client, allocator); + // Best-effort: the Go driver also ignores all errors closing the session. + client.closeSession(new CloseSessionRequest(), callOptions); + } catch (FlightRuntimeException e) { + // ignore + } + try { + AutoCloseables.close(clientCache::invalidateAll, client, allocator); } catch (Exception e) { throw AdbcException.internal("[Flight SQL] Failed to close connection").withCause(e); } @@ -223,6 +346,39 @@ public String toString() { return "FlightSqlConnection{" + "client=" + client + '}'; } + private Map fetchSessionOptionsOrEmpty() throws AdbcException { + try { + return client.getSessionOptions(new GetSessionOptionsRequest()).getSessionOptions(); + } catch (FlightRuntimeException e) { + // Go also treats INVALID_ARGUMENT as "server doesn't support sessions" here. + if (e.status().code() == FlightStatusCode.UNIMPLEMENTED + || e.status().code() == FlightStatusCode.INVALID_ARGUMENT) { + 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..d500c5de33 --- /dev/null +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionUtil.java @@ -0,0 +1,174 @@ +/* + * 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 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.checkerframework.checker.nullness.qual.Nullable; + +/** Package-private helpers for Flight SQL session option serialization and type conversion. */ +final class FlightSqlSessionUtil { + + static final ObjectMapper MAPPER = new ObjectMapper(); + + /** + * Extracts a native Java value from a {@link SessionOptionValue}. JSON has no representation for + * NaN/Infinity, so non-finite doubles are converted to their {@link String} form (e.g. {@code + * "NaN"}) here rather than special-cased later; 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; + } + + @Override + public Object visit(long v) { + return v; + } + + @Override + public Object visit(double v) { + return Double.isFinite(v) ? (Object) v : Double.toString(v); + } + + @Override + public Object visit(String[] v) { + return v.clone(); + } + }; + + /** 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()) { + map.put(e.getKey(), e.getValue().acceptVisitor(TO_JAVA)); + } + try { + return MAPPER.writeValueAsString(map); + } catch (JsonProcessingException e) { + throw AdbcException.internal("[Flight SQL] Failed to serialize session options").withCause(e); + } + } + + /** 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 AdbcException.invalidArgument( + "[Flight SQL] Expected JSON array for string list option, got: " + json) + .withCause(e); + } + } + + /** + * 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 @Nullable 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); + } + if (type == Boolean.class) { + if (raw instanceof Boolean) return (T) raw; + return (T) Boolean.valueOf(parseStrictBoolean(String.valueOf(raw), optionName)); + } + 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 (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))); + } + } catch (NumberFormatException e) { + throw AdbcException.invalidArgument( + "[Flight SQL] Session option '" + + optionName + + "' cannot be parsed as " + + type.getSimpleName() + + ": " + + raw) + .withCause(e); + } + return null; + } + + /** 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 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() {} +} 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..5ceb757b60 --- /dev/null +++ b/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlSessionTest.java @@ -0,0 +1,495 @@ +/* + * 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.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; +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.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 testCloseDoesNotThrowWhenServerReturnsUnimplemented() throws Exception { + producer.rejectClose.set(true); + connection.close(); // must not throw + connection = null; + } + + @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 = + 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); + 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 + 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) { + if (rejectGetSession.get()) { + listener.onError(CallStatus.UNIMPLEMENTED.toRuntimeException()); + return; + } + listener.onNext(new GetSessionOptionsResult(new HashMap<>(sessionOptions))); + listener.onCompleted(); + } + + @Override + 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(); + } + + // 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() {} + } +} diff --git a/java/pom.xml b/java/pom.xml index c53859d5f0..2f67fdf504 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -97,6 +97,7 @@ 19.0.0 + 2.21.0 4.2.1 11