feat(java/driver/flight-sql): implement Flight SQL session management #4444
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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> 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)); | ||
|
|
||
|
unikdahal marked this conversation as resolved.
Outdated
|
||
| } 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)); | ||
|
unikdahal marked this conversation as resolved.
Outdated
|
||
|
|
||
| } 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)); | ||
| } | ||
|
unikdahal marked this conversation as resolved.
Outdated
|
||
|
|
||
| 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"); | ||
|
lidavidm marked this conversation as resolved.
Outdated
|
||
| } | ||
|
|
||
| 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()); | ||
| } | ||
|
unikdahal marked this conversation as resolved.
|
||
| 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()); | ||
|
Contributor
There was a problem hiding this comment. Some servers use tokens, so this would need to be something like
Contributor
There was a problem hiding this comment. Would need to be passed in from FlightSqlDatabase. I do that in another open PR.
Contributor
Author
There was a problem hiding this comment. Thanks hadn't thought of this took a pull and refactored it
Contributor
There was a problem hiding this comment. Nice. Thanks for doing all of this :) |
||
| } catch (FlightRuntimeException ignored) { | ||
| // best-effort: UNIMPLEMENTED is normal; other errors are silently swallowed at close time | ||
| } | ||
|
unikdahal marked this conversation as resolved.
|
||
| clientCache.invalidateAll(); | ||
| AutoCloseables.close(client, allocator); | ||
| } | ||
|
unikdahal marked this conversation as resolved.
|
||
|
|
@@ -214,6 +359,40 @@ public String toString() { | |
| return "FlightSqlConnection{" + "client=" + client + '}'; | ||
| } | ||
|
|
||
| private Map<String, SessionOptionValue> fetchSessionOptions() throws AdbcException { | ||
| try { | ||
| GetSessionOptionsResult result = | ||
| client.getSessionOptions(new GetSessionOptionsRequest()); | ||
| return result.getSessionOptions(); | ||
| } catch (FlightRuntimeException e) { | ||
| if (e.status().code() == FlightStatusCode.UNIMPLEMENTED) { | ||
|
unikdahal marked this conversation as resolved.
Outdated
|
||
| 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. | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.