From 064b7bf26c743e82b6fd82f2fd9a9a4db67a5e33 Mon Sep 17 00:00:00 2001 From: Devin Smith Date: Wed, 20 Nov 2024 18:06:08 -0800 Subject: [PATCH] feat: Add Flight SQL server support (#6023) This adds [Arrow Flight SQL](https://arrow.apache.org/docs/format/FlightSql.html) support to the Deephaven server, with the goal of supporting various SQL drivers built on top of Flight SQL (JDBC, ADBC, ODBC). It is limited to query statements (ie, no UPDATEs). The implementation supports ad-hoc query statements and ad-hoc looking prepared query statements, but not parameterized prepared statements. Queries are able to reference tables from the global scope by name; we might be able to expand support to other resolvers in the future, potentially with catalogs and namespaces. The scope of supported queries is guided by our Sql engine `io.deephaven.engine.sql.Sql` which is based on a Calcite query parser. It is important to note that while the Flight SQL implementation respects `io.deephaven.server.session.TicketResolver.Authorization.transform` it is **not** hooked up to `io.deephaven.auth.ServiceAuthWiring` nor `io.deephaven.server.table.validation.ColumnExpressionValidator`. So while Flight SQL does not allow users script execution, it does not limit the table operations a user may perform, nor does it restrict the calling of arbitrary functions from a formula. As such, the security posture of Flight SQL sits somewhere between "full access" and "limited access". A cookie-based authentication extension has been added to support some Flight SQL clients which don't operate via the normal Flight authentication "Authorization" header (yes, it's a misnomer), and instead expect the server to send "Set-Cookie" to the clients, and for the clients to echo back the cookie(s) via the "Cookie" header. The server will only send the authentication token as a cookie when the client sends the header and value "x-deephaven-auth-cookie-request=true". To support this, the `io.grpc.Context` has been captured and preserved during the export submission logic. The full Flight SQL action and command spectrum has been skeleton-ed out with appropriate "unimplemented" messages in anticipation that we will want to expand the scope of the Flight SQL implementation in the future. It also serves as a more explicit guide to readers of the code for what is and is not supported. Fixes #5339 --------- Co-authored-by: jianfengmao --- .../java/io/deephaven/engine/sql/Sql.java | 32 +- .../sql/TableCreatorTicketInterceptor.java | 44 +- .../io/deephaven/engine/util/TableTools.java | 31 +- .../deephaven/engine/liveness/Liveness.java | 4 +- .../extensions/barrage/util/BarrageUtil.java | 48 +- extensions/flight-sql/README.md | 24 + extensions/flight-sql/build.gradle | 78 + extensions/flight-sql/gradle.properties | 1 + .../server/DeephavenServerTestBase.java | 62 + .../flightsql/FlightSqlJdbcTestBase.java | 175 ++ .../FlightSqlJdbcUnauthenticatedTestBase.java | 89 + .../server/flightsql/FlightSqlTestModule.java | 124 ++ .../server/flightsql/JettyTestComponent.java | 38 + .../jetty/FlightSqlJdbcTestJetty.java | 15 + ...FlightSqlJdbcUnauthenticatedTestJetty.java | 15 + .../flightsql/FlightSqlActionHelper.java | 191 ++ .../flightsql/FlightSqlCommandHelper.java | 225 +++ .../flightsql/FlightSqlErrorHelper.java | 25 + .../server/flightsql/FlightSqlModule.java | 25 + .../server/flightsql/FlightSqlResolver.java | 1748 +++++++++++++++++ .../flightsql/FlightSqlSharedConstants.java | 45 + .../flightsql/FlightSqlTicketHelper.java | 178 ++ .../flightsql/TableCreatorScopeTickets.java | 33 + .../FlightSqlFilterPredicateTest.java | 179 ++ .../server/flightsql/FlightSqlTest.java | 1095 +++++++++++ .../FlightSqlTicketResolverTest.java | 170 ++ .../FlightSqlUnauthenticatedTest.java | 350 ++++ gradle/libs.versions.toml | 2 + .../apache/arrow/flight/ProtocolExposer.java | 54 + .../io/deephaven/proto/util/ByteHelper.java | 6 +- .../io/deephaven/proto/util/Exceptions.java | 15 + py/embedded-server/java-runtime/build.gradle | 2 + .../python/server/EmbeddedServer.java | 2 + .../deephaven/qst/TableCreatorDelegate.java | 57 + server/jetty-app/build.gradle | 5 + .../jetty/CommunityComponentFactory.java | 4 + .../JettyClientChannelFactoryModule.java | 2 +- .../jetty-app/src/main/resources/logback.xml | 3 +- server/jetty/build.gradle | 5 +- .../deephaven/server/arrow/ArrowModule.java | 16 + .../server/arrow/FlightServiceGrpcImpl.java | 37 +- .../ServerCallStreamObserverAdapter.java | 75 + .../server/console/ScopeTicketResolver.java | 3 + .../server/session/ActionResolver.java | 57 + .../server/session/ActionRouter.java | 99 + .../deephaven/server/session/AuthCookie.java | 73 + .../server/session/CommandResolver.java | 43 + .../server/session/PathResolver.java | 23 + .../session/PathResolverPrefixedBase.java | 48 + .../session/SessionServiceGrpcImpl.java | 54 +- .../server/session/SessionState.java | 10 +- .../server/session/TicketResolver.java | 27 +- .../server/session/TicketResolverBase.java | 19 +- .../server/session/TicketRouter.java | 151 +- server/test-utils/build.gradle | 6 +- .../runner/DeephavenApiServerTestBase.java | 6 +- .../test/TestAuthorizationProvider.java | 8 + settings.gradle | 3 + .../deephaven/sql/RelNodeVisitorAdapter.java | 18 +- .../java/io/deephaven/sql/RexVisitorBase.java | 33 +- .../sql/UnsupportedSqlOperation.java | 23 + 61 files changed, 5855 insertions(+), 178 deletions(-) create mode 100644 extensions/flight-sql/README.md create mode 100644 extensions/flight-sql/build.gradle create mode 100644 extensions/flight-sql/gradle.properties create mode 100644 extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/DeephavenServerTestBase.java create mode 100644 extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/FlightSqlJdbcTestBase.java create mode 100644 extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/FlightSqlJdbcUnauthenticatedTestBase.java create mode 100644 extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/FlightSqlTestModule.java create mode 100644 extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/JettyTestComponent.java create mode 100644 extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/jetty/FlightSqlJdbcTestJetty.java create mode 100644 extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/jetty/FlightSqlJdbcUnauthenticatedTestJetty.java create mode 100644 extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlActionHelper.java create mode 100644 extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlCommandHelper.java create mode 100644 extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlErrorHelper.java create mode 100644 extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlModule.java create mode 100644 extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlResolver.java create mode 100644 extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlSharedConstants.java create mode 100644 extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlTicketHelper.java create mode 100644 extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/TableCreatorScopeTickets.java create mode 100644 extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlFilterPredicateTest.java create mode 100644 extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlTest.java create mode 100644 extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlTicketResolverTest.java create mode 100644 extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlUnauthenticatedTest.java create mode 100644 java-client/flight/src/main/java/org/apache/arrow/flight/ProtocolExposer.java create mode 100644 qst/src/main/java/io/deephaven/qst/TableCreatorDelegate.java rename server/{jetty => jetty-app}/src/main/java/io/deephaven/server/jetty/CommunityComponentFactory.java (94%) rename server/{jetty => jetty-app}/src/main/java/io/deephaven/server/jetty/JettyClientChannelFactoryModule.java (96%) create mode 100644 server/src/main/java/io/deephaven/server/arrow/ServerCallStreamObserverAdapter.java create mode 100644 server/src/main/java/io/deephaven/server/session/ActionResolver.java create mode 100644 server/src/main/java/io/deephaven/server/session/ActionRouter.java create mode 100644 server/src/main/java/io/deephaven/server/session/AuthCookie.java create mode 100644 server/src/main/java/io/deephaven/server/session/CommandResolver.java create mode 100644 server/src/main/java/io/deephaven/server/session/PathResolver.java create mode 100644 server/src/main/java/io/deephaven/server/session/PathResolverPrefixedBase.java create mode 100644 sql/src/main/java/io/deephaven/sql/UnsupportedSqlOperation.java diff --git a/engine/sql/src/main/java/io/deephaven/engine/sql/Sql.java b/engine/sql/src/main/java/io/deephaven/engine/sql/Sql.java index 665a4378ed8..fb87d365147 100644 --- a/engine/sql/src/main/java/io/deephaven/engine/sql/Sql.java +++ b/engine/sql/src/main/java/io/deephaven/engine/sql/Sql.java @@ -19,10 +19,12 @@ import io.deephaven.qst.table.TableHeader.Builder; import io.deephaven.qst.table.TableSpec; import io.deephaven.qst.table.TicketTable; +import io.deephaven.qst.type.Type; import io.deephaven.sql.Scope; import io.deephaven.sql.ScopeStaticImpl; import io.deephaven.sql.SqlAdapter; import io.deephaven.sql.TableInformation; +import io.deephaven.util.annotations.InternalUseOnly; import io.deephaven.util.annotations.ScriptApi; import java.nio.charset.StandardCharsets; @@ -30,6 +32,7 @@ import java.util.List; import java.util.Map; import java.util.Map.Entry; +import java.util.function.Function; /** * Experimental SQL execution. Subject to change. @@ -47,35 +50,38 @@ public static TableSpec dryRun(String sql) { return dryRun(sql, currentScriptSessionNamedTables()); } + @InternalUseOnly + public static TableSpec parseSql(String sql, Map scope, Function ticketFunction, + Map out) { + return SqlAdapter.parseSql(sql, scope(scope, out, ticketFunction)); + } + private static Table evaluate(String sql, Map scope) { final Map map = new HashMap<>(scope.size()); - final TableSpec tableSpec = parseSql(sql, scope, map); + final TableSpec tableSpec = parseSql(sql, scope, Sql::sqlref, map); log.debug().append("Executing. Graphviz representation:").nl().append(ToGraphvizDot.INSTANCE, tableSpec).endl(); return tableSpec.logic().create(new TableCreatorTicketInterceptor(TableCreatorImpl.INSTANCE, map)); } private static TableSpec dryRun(String sql, Map scope) { - final TableSpec tableSpec = parseSql(sql, scope, null); + final TableSpec tableSpec = parseSql(sql, scope, Sql::sqlref, null); log.info().append("Dry run. Graphviz representation:").nl().append(ToGraphvizDot.INSTANCE, tableSpec).endl(); return tableSpec; } - private static TableSpec parseSql(String sql, Map scope, Map out) { - return SqlAdapter.parseSql(sql, scope(scope, out)); - } - private static TicketTable sqlref(String tableName) { + // The TicketTable can technically be anything unique (incrementing number, random, ...), but for + // visualization purposes it makes sense to use the (already unique) table name. return TicketTable.of(("sqlref/" + tableName).getBytes(StandardCharsets.UTF_8)); } - private static Scope scope(Map scope, Map out) { + private static Scope scope(Map scope, Map out, + Function ticketFunction) { final ScopeStaticImpl.Builder builder = ScopeStaticImpl.builder(); for (Entry e : scope.entrySet()) { final String tableName = e.getKey(); final Table table = e.getValue(); - // The TicketTable can technically be anything unique (incrementing number, random, ...), but for - // visualization purposes it makes sense to use the (already unique) table name. - final TicketTable spec = sqlref(tableName); + final TicketTable spec = ticketFunction.apply(tableName); final List qualifiedName = List.of(tableName); final TableHeader header = adapt(table.getDefinition()); builder.addTables(TableInformation.of(qualifiedName, header, spec)); @@ -103,11 +109,7 @@ private static TableHeader adapt(TableDefinition tableDef) { } private static ColumnHeader adapt(ColumnDefinition columnDef) { - if (columnDef.getComponentType() == null) { - return ColumnHeader.of(columnDef.getName(), columnDef.getDataType()); - } - // SQLTODO(array-type) - throw new UnsupportedOperationException("SQLTODO(array-type)"); + return ColumnHeader.of(columnDef.getName(), Type.find(columnDef.getDataType())); } private enum ToGraphvizDot implements ObjFormatter { diff --git a/engine/sql/src/main/java/io/deephaven/engine/sql/TableCreatorTicketInterceptor.java b/engine/sql/src/main/java/io/deephaven/engine/sql/TableCreatorTicketInterceptor.java index c3a81c9590b..a98d5da40b2 100644 --- a/engine/sql/src/main/java/io/deephaven/engine/sql/TableCreatorTicketInterceptor.java +++ b/engine/sql/src/main/java/io/deephaven/engine/sql/TableCreatorTicketInterceptor.java @@ -5,23 +5,17 @@ import io.deephaven.engine.table.Table; import io.deephaven.qst.TableCreator; -import io.deephaven.qst.table.EmptyTable; -import io.deephaven.qst.table.InputTable; -import io.deephaven.qst.table.MultiJoinInput; -import io.deephaven.qst.table.NewTable; +import io.deephaven.qst.TableCreatorDelegate; import io.deephaven.qst.table.TicketTable; -import io.deephaven.qst.table.TimeTable; -import java.util.List; import java.util.Map; import java.util.Objects; -class TableCreatorTicketInterceptor implements TableCreator { - private final TableCreator
delegate; +class TableCreatorTicketInterceptor extends TableCreatorDelegate
{ private final Map map; public TableCreatorTicketInterceptor(TableCreator
delegate, Map map) { - this.delegate = Objects.requireNonNull(delegate); + super(delegate); this.map = Objects.requireNonNull(map); } @@ -31,36 +25,6 @@ public Table of(TicketTable ticketTable) { if (table != null) { return table; } - return delegate.of(ticketTable); - } - - @Override - public Table of(NewTable newTable) { - return delegate.of(newTable); - } - - @Override - public Table of(EmptyTable emptyTable) { - return delegate.of(emptyTable); - } - - @Override - public Table of(TimeTable timeTable) { - return delegate.of(timeTable); - } - - @Override - public Table of(InputTable inputTable) { - return delegate.of(inputTable); - } - - @Override - public Table multiJoin(List> multiJoinInputs) { - return delegate.multiJoin(multiJoinInputs); - } - - @Override - public Table merge(Iterable
tables) { - return delegate.merge(tables); + return super.of(ticketTable); } } diff --git a/engine/table/src/main/java/io/deephaven/engine/util/TableTools.java b/engine/table/src/main/java/io/deephaven/engine/util/TableTools.java index 924137c6947..2c4176012ad 100644 --- a/engine/table/src/main/java/io/deephaven/engine/util/TableTools.java +++ b/engine/table/src/main/java/io/deephaven/engine/util/TableTools.java @@ -32,6 +32,7 @@ import io.deephaven.util.annotations.ScriptApi; import io.deephaven.util.type.ArrayTypeUtils; import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; import java.io.*; import java.nio.charset.StandardCharsets; @@ -67,13 +68,13 @@ private static BinaryOperator throwingMerger() { }; } - private static Collector> toLinkedMap( + private static Collector> toLinkedMap( Function keyMapper, Function valueMapper) { return Collectors.toMap(keyMapper, valueMapper, throwingMerger(), LinkedHashMap::new); } - private static final Collector, ?, Map>> COLUMN_HOLDER_LINKEDMAP_COLLECTOR = + private static final Collector, ?, LinkedHashMap>> COLUMN_HOLDER_LINKEDMAP_COLLECTOR = toLinkedMap(ColumnHolder::getName, ColumnHolder::getColumnSource); /////////// Utilities To Display Tables ///////////////// @@ -752,22 +753,24 @@ public static Table newTable(ColumnHolder... columnHolders) { checkSizes(columnHolders); WritableRowSet rowSet = getRowSet(columnHolders); Map> columns = Arrays.stream(columnHolders).collect(COLUMN_HOLDER_LINKEDMAP_COLLECTOR); - return new QueryTable(rowSet.toTracking(), columns) { - { - setFlat(); - } - }; + QueryTable queryTable = new QueryTable(rowSet.toTracking(), columns); + queryTable.setFlat(); + return queryTable; } public static Table newTable(TableDefinition definition, ColumnHolder... columnHolders) { + return newTable(definition, null, columnHolders); + } + + public static Table newTable(TableDefinition definition, @Nullable Map attributes, + ColumnHolder... columnHolders) { checkSizes(columnHolders); - WritableRowSet rowSet = getRowSet(columnHolders); - Map> columns = Arrays.stream(columnHolders).collect(COLUMN_HOLDER_LINKEDMAP_COLLECTOR); - return new QueryTable(definition, rowSet.toTracking(), columns) { - { - setFlat(); - } - }; + final WritableRowSet rowSet = getRowSet(columnHolders); + final LinkedHashMap> columns = + Arrays.stream(columnHolders).collect(COLUMN_HOLDER_LINKEDMAP_COLLECTOR); + final QueryTable queryTable = new QueryTable(definition, rowSet.toTracking(), columns, null, attributes); + queryTable.setFlat(); + return queryTable; } /** diff --git a/engine/updategraph/src/main/java/io/deephaven/engine/liveness/Liveness.java b/engine/updategraph/src/main/java/io/deephaven/engine/liveness/Liveness.java index 3d3543d3dd9..320e98efcee 100644 --- a/engine/updategraph/src/main/java/io/deephaven/engine/liveness/Liveness.java +++ b/engine/updategraph/src/main/java/io/deephaven/engine/liveness/Liveness.java @@ -97,8 +97,8 @@ private Liveness() {} /** *

* Determine whether a cached object should be reused, w.r.t. liveness. Null inputs are never safe for reuse. If the - * object is a {@link LivenessReferent} and not a non-refreshing {@link DynamicNode}, this method will return the - * result of trying to manage object with the top of the current thread's {@link LivenessScopeStack}. + * object is a {@link LivenessReferent} and is a refreshing {@link DynamicNode}, this method will return the result + * of trying to manage object with the top of the current thread's {@link LivenessScopeStack}. * * @param object The object * @return True if the object did not need management, or if it was successfully managed, false otherwise diff --git a/extensions/barrage/src/main/java/io/deephaven/extensions/barrage/util/BarrageUtil.java b/extensions/barrage/src/main/java/io/deephaven/extensions/barrage/util/BarrageUtil.java index 3d83db05833..362bdc67353 100755 --- a/extensions/barrage/src/main/java/io/deephaven/extensions/barrage/util/BarrageUtil.java +++ b/extensions/barrage/src/main/java/io/deephaven/extensions/barrage/util/BarrageUtil.java @@ -9,10 +9,12 @@ import com.google.protobuf.ByteStringAccess; import com.google.rpc.Code; import io.deephaven.UncheckedDeephavenException; +import io.deephaven.api.util.NameValidator; import io.deephaven.barrage.flatbuf.BarrageMessageWrapper; import io.deephaven.base.ArrayUtil; import io.deephaven.base.ClassUtil; import io.deephaven.base.verify.Assert; +import io.deephaven.chunk.ChunkType; import io.deephaven.configuration.Configuration; import io.deephaven.engine.rowset.RowSequence; import io.deephaven.engine.rowset.RowSet; @@ -27,6 +29,8 @@ import io.deephaven.engine.table.impl.sources.ReinterpretUtils; import io.deephaven.engine.table.impl.util.BarrageMessage; import io.deephaven.engine.updategraph.impl.PeriodicUpdateGraph; +import io.deephaven.engine.util.ColumnFormatting; +import io.deephaven.engine.util.input.InputTableUpdater; import io.deephaven.extensions.barrage.BarragePerformanceLog; import io.deephaven.extensions.barrage.BarrageSnapshotOptions; import io.deephaven.extensions.barrage.BarrageStreamGenerator; @@ -35,15 +39,10 @@ import io.deephaven.extensions.barrage.chunk.vector.VectorExpansionKernel; import io.deephaven.internal.log.LoggerFactory; import io.deephaven.io.logger.Logger; +import io.deephaven.proto.backplane.grpc.ExportedTableCreationResponse; import io.deephaven.proto.flight.util.MessageHelper; import io.deephaven.proto.flight.util.SchemaHelper; import io.deephaven.proto.util.Exceptions; -import io.deephaven.api.util.NameValidator; -import io.deephaven.engine.util.ColumnFormatting; -import io.deephaven.engine.util.input.InputTableUpdater; -import io.deephaven.chunk.ChunkType; -import io.deephaven.proto.backplane.grpc.ExportedTableCreationResponse; -import io.deephaven.qst.column.Column; import io.deephaven.util.type.TypeUtils; import io.deephaven.vector.Vector; import io.grpc.stub.StreamObserver; @@ -69,8 +68,21 @@ import java.time.LocalDate; import java.time.LocalTime; import java.time.ZonedDateTime; -import java.util.*; -import java.util.function.*; +import java.util.Arrays; +import java.util.BitSet; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.function.BiConsumer; +import java.util.function.Consumer; +import java.util.function.Function; +import java.util.function.IntFunction; +import java.util.function.ToIntFunction; import java.util.regex.Matcher; import java.util.regex.Pattern; import java.util.stream.Collectors; @@ -207,6 +219,14 @@ public static ByteString schemaBytesFromTableDefinition( fbb, DEFAULT_SNAPSHOT_DESER_OPTIONS, tableDefinition, attributes, isFlat)); } + public static Schema schemaFromTable(@NotNull final Table table) { + return makeSchema(DEFAULT_SNAPSHOT_DESER_OPTIONS, table.getDefinition(), table.getAttributes(), table.isFlat()); + } + + public static Schema toSchema(final TableDefinition definition, Map attributes, boolean isFlat) { + return makeSchema(DEFAULT_SNAPSHOT_DESER_OPTIONS, definition, attributes, isFlat); + } + public static ByteString schemaBytes(@NotNull final ToIntFunction schemaPayloadWriter) { // note that flight expects the Schema to be wrapped in a Message prefixed by a 4-byte identifier @@ -226,8 +246,15 @@ public static int makeTableSchemaPayload( @NotNull final TableDefinition tableDefinition, @NotNull final Map attributes, final boolean isFlat) { - final Map schemaMetadata = attributesToMetadata(attributes, isFlat); + return makeSchema(options, tableDefinition, attributes, isFlat).getSchema(builder); + } + public static Schema makeSchema( + @NotNull final StreamReaderOptions options, + @NotNull final TableDefinition tableDefinition, + @NotNull final Map attributes, + final boolean isFlat) { + final Map schemaMetadata = attributesToMetadata(attributes, isFlat); final Map descriptions = GridAttributes.getColumnDescriptions(attributes); final InputTableUpdater inputTableUpdater = (InputTableUpdater) attributes.get(Table.INPUT_TABLE_ATTRIBUTE); final List fields = columnDefinitionsToFields( @@ -235,8 +262,7 @@ public static int makeTableSchemaPayload( ignored -> new HashMap<>(), attributes, options.columnsAsList()) .collect(Collectors.toList()); - - return new Schema(fields, schemaMetadata).getSchema(builder); + return new Schema(fields, schemaMetadata); } @NotNull diff --git a/extensions/flight-sql/README.md b/extensions/flight-sql/README.md new file mode 100644 index 00000000000..9f5c3e54907 --- /dev/null +++ b/extensions/flight-sql/README.md @@ -0,0 +1,24 @@ +# Flight SQL + +See [FlightSqlResolver](src/main/java/io/deephaven/server/flightsql/FlightSqlResolver.java) for documentation on +Deephaven's Flight SQL service. + +## Client + +The Flight SQL client is simple constructed based on the underlying Flight client. + +```java +FlightClient flightClient = ...; +FlightSqlClient flightSqlClient = new FlightSqlClient(flightClient); +``` + +## JDBC + +The default Flight SQL JDBC driver uses cookie authorization; by default, this is not enabled on the Deephaven server. +To enable this, the request header "x-deephaven-auth-cookie-request" must be set to "true". + +Example JDBC connection string to self-signed TLS: + +``` +jdbc:arrow-flight-sql://localhost:8443/?Authorization=Anonymous&useEncryption=1&disableCertificateVerification=1&x-deephaven-auth-cookie-request=true +``` diff --git a/extensions/flight-sql/build.gradle b/extensions/flight-sql/build.gradle new file mode 100644 index 00000000000..3c4e155e411 --- /dev/null +++ b/extensions/flight-sql/build.gradle @@ -0,0 +1,78 @@ +plugins { + id 'java-library' + id 'io.deephaven.project.register' +} + +description = 'The Deephaven Flight SQL library' + +sourceSets { + jdbcTest { + compileClasspath += sourceSets.main.output + runtimeClasspath += sourceSets.main.output + } +} + +configurations { + jdbcTestImplementation.extendsFrom implementation + jdbcTestRuntimeOnly.extendsFrom runtimeOnly +} + +dependencies { + api project(':server') + implementation project(':sql') + implementation project(':engine-sql') + // :sql does not expose calcite as a dependency (maybe it should?); in the meantime, we want to make sure we can + // provide reasonable error messages to the client + implementation libs.calcite.core + + implementation libs.dagger + implementation libs.arrow.flight.sql + + // FlightSqlClient testing does not depend on a public port being bound (ie, does not require server-jetty) because + // it can use io.grpc.inprocess.InProcessChannelBuilder (via io.deephaven.server.runner.ServerBuilderInProcessModule). + + testImplementation project(':authorization') + testImplementation project(':server-test-utils') + testAnnotationProcessor libs.dagger.compiler + testImplementation libs.assertj + testImplementation platform(libs.junit.bom) + testImplementation libs.junit.jupiter + testRuntimeOnly libs.junit.platform.launcher + testRuntimeOnly libs.junit.vintage.engine + testRuntimeOnly project(':log-to-slf4j') + testRuntimeOnly libs.slf4j.simple + + // JDBC testing needs an actually server instance bound to a port because it can only connect over JDBC URIs like + // jdbc:arrow-flight-sql://localhost:1000. + jdbcTestImplementation project(':server-jetty') + jdbcTestRuntimeOnly libs.arrow.flight.sql.jdbc + + jdbcTestImplementation project(':server-test-utils') + jdbcTestAnnotationProcessor libs.dagger.compiler + jdbcTestImplementation libs.assertj + jdbcTestImplementation platform(libs.junit.bom) + jdbcTestImplementation libs.junit.jupiter + jdbcTestRuntimeOnly libs.junit.platform.launcher + jdbcTestRuntimeOnly libs.junit.vintage.engine + jdbcTestRuntimeOnly project(':log-to-slf4j') + jdbcTestRuntimeOnly libs.slf4j.simple +} + +test { + useJUnitPlatform() +} + +def jdbcTest = tasks.register('jdbcTest', Test) { + description = 'Runs JDBC tests.' + group = 'verification' + + testClassesDirs = sourceSets.jdbcTest.output.classesDirs + classpath = sourceSets.jdbcTest.runtimeClasspath + shouldRunAfter test + + useJUnitPlatform() +} + +check.dependsOn jdbcTest + +apply plugin: 'io.deephaven.java-open-nio' diff --git a/extensions/flight-sql/gradle.properties b/extensions/flight-sql/gradle.properties new file mode 100644 index 00000000000..c186bbfdde1 --- /dev/null +++ b/extensions/flight-sql/gradle.properties @@ -0,0 +1 @@ +io.deephaven.project.ProjectType=JAVA_PUBLIC diff --git a/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/DeephavenServerTestBase.java b/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/DeephavenServerTestBase.java new file mode 100644 index 00000000000..76458d91a57 --- /dev/null +++ b/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/DeephavenServerTestBase.java @@ -0,0 +1,62 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server; + +import io.deephaven.engine.context.ExecutionContext; +import io.deephaven.io.logger.LogBuffer; +import io.deephaven.io.logger.LogBufferGlobal; +import io.deephaven.server.runner.GrpcServer; +import io.deephaven.server.runner.MainHelper; +import io.deephaven.util.SafeCloseable; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Timeout; + +import java.io.IOException; +import java.util.concurrent.TimeUnit; + +@Timeout(30) +public abstract class DeephavenServerTestBase { + + public interface TestComponent { + + GrpcServer server(); + + ExecutionContext executionContext(); + } + + protected TestComponent component; + + private LogBuffer logBuffer; + private SafeCloseable executionContext; + private GrpcServer server; + protected int localPort; + + protected abstract TestComponent component(); + + @BeforeAll + static void setupOnce() throws IOException { + MainHelper.bootstrapProjectDirectories(); + } + + @BeforeEach + void setup() throws IOException { + logBuffer = new LogBuffer(128); + LogBufferGlobal.setInstance(logBuffer); + component = component(); + executionContext = component.executionContext().open(); + server = component.server(); + server.start(); + localPort = server.getPort(); + } + + @AfterEach + void tearDown() throws InterruptedException { + server.stopWithTimeout(10, TimeUnit.SECONDS); + server.join(); + executionContext.close(); + LogBufferGlobal.clear(logBuffer); + } +} diff --git a/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/FlightSqlJdbcTestBase.java b/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/FlightSqlJdbcTestBase.java new file mode 100644 index 00000000000..0c5c4993edb --- /dev/null +++ b/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/FlightSqlJdbcTestBase.java @@ -0,0 +1,175 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql; + +import io.deephaven.server.DeephavenServerTestBase; +import org.junit.jupiter.api.Test; + +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Statement; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.failBecauseExceptionWasNotThrown; + +public abstract class FlightSqlJdbcTestBase extends DeephavenServerTestBase { + + private String jdbcUrl(boolean requestCookie) { + return String.format( + "jdbc:arrow-flight-sql://localhost:%d/?Authorization=Anonymous&useEncryption=false%s", + localPort, + (requestCookie ? "&x-deephaven-auth-cookie-request=true" : "")); + } + + private Connection connect(boolean requestCookie) throws SQLException { + return DriverManager.getConnection(jdbcUrl(requestCookie)); + } + + @Test + void execute() throws SQLException { + try ( + final Connection connection = connect(true); + final Statement statement = connection.createStatement()) { + if (statement.execute("SELECT 1 as Foo, 2 as Bar")) { + consume(statement.getResultSet(), 2, 1); + } + } + } + + @Test + void executeQuery() throws SQLException { + try ( + final Connection connection = connect(true); + final Statement statement = connection.createStatement()) { + consume(statement.executeQuery("SELECT 1 as Foo, 2 as Bar"), 2, 1); + } + } + + @Test + void executeUpdate() throws SQLException { + try ( + final Connection connection = connect(true); + final Statement statement = connection.createStatement()) { + try { + statement.executeUpdate("INSERT INTO fake(name) VALUES('Smith')"); + failBecauseExceptionWasNotThrown(SQLException.class); + } catch (SQLException e) { + assertThat((Throwable) e).getRootCause() + .hasMessageContaining("Object 'fake' not found"); + } + } + } + + @Test + void preparedExecute() throws SQLException { + try ( + final Connection connection = connect(true); + final PreparedStatement preparedStatement = connection.prepareStatement("SELECT 1 as Foo, 2 as Bar")) { + if (preparedStatement.execute()) { + consume(preparedStatement.getResultSet(), 2, 1); + } + consume(preparedStatement.executeQuery(), 2, 1); + try { + preparedStatement.executeUpdate(); + failBecauseExceptionWasNotThrown(SQLException.class); + } catch (SQLException e) { + assertThat((Throwable) e).getRootCause() + .hasMessageContaining("Flight SQL descriptors cannot be published to"); + } + } + } + + @Test + void preparedExecuteQuery() throws SQLException { + try ( + final Connection connection = connect(true); + final PreparedStatement preparedStatement = connection.prepareStatement("SELECT 1 as Foo, 2 as Bar")) { + consume(preparedStatement.executeQuery(), 2, 1); + } + } + + @Test + void preparedUpdate() throws SQLException { + try ( + final Connection connection = connect(true); + final PreparedStatement preparedStatement = + connection.prepareStatement("INSERT INTO fake(name) VALUES('Smith')")) { + try { + preparedStatement.executeUpdate(); + failBecauseExceptionWasNotThrown(SQLException.class); + } catch (SQLException e) { + assertThat((Throwable) e).getRootCause() + .hasMessageContaining("Flight SQL descriptors cannot be published to"); + } + } + } + + @Test + void executeQueryNoCookie() throws SQLException { + try (final Connection connection = connect(false)) { + final Statement statement = connection.createStatement(); + try { + statement.executeQuery("SELECT 1 as Foo, 2 as Bar"); + failBecauseExceptionWasNotThrown(SQLException.class); + } catch (SQLException e) { + assertThat((Throwable) e).getRootCause() + .hasMessageContaining( + "Flight SQL: Must use the original session; is the client echoing the authentication token properly?"); + } + try { + statement.close(); + failBecauseExceptionWasNotThrown(SQLException.class); + } catch (SQLException e) { + assertThat((Throwable) e).getRootCause() + .hasMessageContaining( + "Flight SQL: Must use the original session; is the client echoing the authentication token properly?"); + } + } + } + + @Test + void preparedExecuteQueryNoCookie() throws SQLException { + try (final Connection connection = connect(false)) { + final PreparedStatement preparedStatement = connection.prepareStatement("SELECT 1 as Foo, 2 as Bar"); + try { + preparedStatement.executeQuery(); + failBecauseExceptionWasNotThrown(SQLException.class); + } catch (RuntimeException e) { + assertHelpfulClientErrorMessage(e); + } + // If our authentication is bad, we won't be able to close the prepared statement either. If we want to + // solve for this scenario, we would probably need to use randomized handles for the prepared statements + // (instead of incrementing handle ids). + try { + preparedStatement.close(); + failBecauseExceptionWasNotThrown(RuntimeException.class); + } catch (RuntimeException e) { + assertHelpfulClientErrorMessage(e); + } + } + } + + private static void assertHelpfulClientErrorMessage(RuntimeException e) { + // Note: this is arguably a JDBC implementation bug; it should be throwing java.sql.SQLException, but it's + // exposing shadowed internal error from Flight. + assertThat(e.getClass().getName()).isEqualTo( + "org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.FlightRuntimeException"); + assertThat(e).hasMessageContaining( + "Flight SQL: Must use the original session; is the client echoing the authentication token properly?"); + } + + private static void consume(ResultSet rs, int numCols, int numRows) throws SQLException { + final ResultSetMetaData rsmd = rs.getMetaData(); + assertThat(rsmd.getColumnCount()).isEqualTo(numCols); + int rows = 0; + while (rs.next()) { + ++rows; + } + assertThat(rows).isEqualTo(numRows); + } +} diff --git a/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/FlightSqlJdbcUnauthenticatedTestBase.java b/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/FlightSqlJdbcUnauthenticatedTestBase.java new file mode 100644 index 00000000000..746949cd4cd --- /dev/null +++ b/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/FlightSqlJdbcUnauthenticatedTestBase.java @@ -0,0 +1,89 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql; + +import io.deephaven.server.DeephavenServerTestBase; +import org.junit.jupiter.api.Test; + +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.SQLException; +import java.sql.Statement; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.failBecauseExceptionWasNotThrown; + +public abstract class FlightSqlJdbcUnauthenticatedTestBase extends DeephavenServerTestBase { + private String jdbcUrl() { + return String.format( + "jdbc:arrow-flight-sql://localhost:%d/?useEncryption=false", + localPort); + } + + private Connection connect() throws SQLException { + return DriverManager.getConnection(jdbcUrl()); + } + + @Test + void executeQuery() throws SQLException { + // uses prepared statement internally + try ( + final Connection connection = connect(); + final Statement statement = connection.createStatement()) { + try { + statement.executeQuery("SELECT 1 as Foo, 2 as Bar"); + failBecauseExceptionWasNotThrown(SQLException.class); + } catch (SQLException e) { + unauthenticated(e); + } + } + } + + @Test + void execute() throws SQLException { + // uses prepared statement internally + try ( + final Connection connection = connect(); + final Statement statement = connection.createStatement()) { + try { + statement.execute("SELECT 1 as Foo, 2 as Bar"); + failBecauseExceptionWasNotThrown(SQLException.class); + } catch (SQLException e) { + unauthenticated(e); + } + } + } + + @Test + void executeUpdate() throws SQLException { + // uses prepared statement internally + try ( + final Connection connection = connect(); + final Statement statement = connection.createStatement()) { + try { + statement.executeUpdate("INSERT INTO fake(name) VALUES('Smith')"); + failBecauseExceptionWasNotThrown(SQLException.class); + } catch (SQLException e) { + unauthenticated(e); + } + } + } + + @Test + void prepareStatement() throws SQLException { + try ( + final Connection connection = connect()) { + try { + connection.prepareStatement("SELECT 1"); + } catch (SQLException e) { + unauthenticated(e); + } + } + } + + private static void unauthenticated(SQLException e) { + assertThat((Throwable) e).getRootCause() + .hasMessageContaining("Flight SQL: Must be authenticated"); + } +} diff --git a/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/FlightSqlTestModule.java b/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/FlightSqlTestModule.java new file mode 100644 index 00000000000..203b2863784 --- /dev/null +++ b/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/FlightSqlTestModule.java @@ -0,0 +1,124 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql; + +import dagger.Module; +import dagger.Provides; +import dagger.multibindings.IntoSet; +import io.deephaven.base.clock.Clock; +import io.deephaven.engine.context.ExecutionContext; +import io.deephaven.engine.updategraph.OperationInitializer; +import io.deephaven.engine.updategraph.UpdateGraph; +import io.deephaven.engine.util.AbstractScriptSession; +import io.deephaven.engine.util.NoLanguageDeephavenSession; +import io.deephaven.engine.util.ScriptSession; +import io.deephaven.server.arrow.ArrowModule; +import io.deephaven.server.auth.AuthorizationProvider; +import io.deephaven.server.config.ConfigServiceModule; +import io.deephaven.server.console.ConsoleModule; +import io.deephaven.server.log.LogModule; +import io.deephaven.server.plugin.PluginsModule; +import io.deephaven.server.session.ExportTicketResolver; +import io.deephaven.server.session.ObfuscatingErrorTransformerModule; +import io.deephaven.server.session.SessionModule; +import io.deephaven.server.session.TicketResolver; +import io.deephaven.server.table.TableModule; +import io.deephaven.server.test.TestAuthModule; +import io.deephaven.server.test.TestAuthorizationProvider; +import io.deephaven.server.util.Scheduler; +import org.jetbrains.annotations.Nullable; + +import javax.inject.Named; +import javax.inject.Singleton; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; + +@Module(includes = { + ArrowModule.class, + ConfigServiceModule.class, + ConsoleModule.class, + LogModule.class, + SessionModule.class, + TableModule.class, + TestAuthModule.class, + ObfuscatingErrorTransformerModule.class, + PluginsModule.class, + FlightSqlModule.class +}) +public class FlightSqlTestModule { + @IntoSet + @Provides + TicketResolver ticketResolver(ExportTicketResolver resolver) { + return resolver; + } + + @Singleton + @Provides + AbstractScriptSession provideAbstractScriptSession( + final UpdateGraph updateGraph, + final OperationInitializer operationInitializer) { + return new NoLanguageDeephavenSession( + updateGraph, operationInitializer, "non-script-session"); + } + + @Provides + ScriptSession provideScriptSession(AbstractScriptSession scriptSession) { + return scriptSession; + } + + @Provides + @Singleton + ScheduledExecutorService provideExecutorService() { + return Executors.newScheduledThreadPool(1); + } + + @Provides + Scheduler provideScheduler(ScheduledExecutorService concurrentExecutor) { + return new Scheduler.DelegatingImpl( + Executors.newSingleThreadExecutor(), + concurrentExecutor, + Clock.system()); + } + + @Provides + @Named("session.tokenExpireMs") + long provideTokenExpireMs() { + return 60_000_000; + } + + @Provides + @Named("http.port") + int provideHttpPort() { + return 0;// 'select first available' + } + + @Provides + @Named("grpc.maxInboundMessageSize") + int provideMaxInboundMessageSize() { + return 1024 * 1024; + } + + @Provides + AuthorizationProvider provideAuthorizationProvider(TestAuthorizationProvider provider) { + return provider; + } + + @Provides + @Singleton + TestAuthorizationProvider provideTestAuthorizationProvider() { + return new TestAuthorizationProvider(); + } + + @Provides + @Singleton + static UpdateGraph provideUpdateGraph() { + return ExecutionContext.getContext().getUpdateGraph(); + } + + @Provides + @Singleton + static OperationInitializer provideOperationInitializer() { + return ExecutionContext.getContext().getOperationInitializer(); + } +} diff --git a/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/JettyTestComponent.java b/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/JettyTestComponent.java new file mode 100644 index 00000000000..add6799ad28 --- /dev/null +++ b/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/JettyTestComponent.java @@ -0,0 +1,38 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql; + +import dagger.Component; +import dagger.Module; +import dagger.Provides; +import io.deephaven.server.DeephavenServerTestBase.TestComponent; +import io.deephaven.server.flightsql.JettyTestComponent.JettyTestConfig; +import io.deephaven.server.jetty.JettyConfig; +import io.deephaven.server.jetty.JettyServerModule; +import io.deephaven.server.runner.ExecutionContextUnitTestModule; + +import javax.inject.Singleton; +import java.time.Duration; +import java.time.temporal.ChronoUnit; + +@Singleton +@Component(modules = { + ExecutionContextUnitTestModule.class, + JettyServerModule.class, + JettyTestConfig.class, + FlightSqlTestModule.class, +}) +public interface JettyTestComponent extends TestComponent { + + @Module + interface JettyTestConfig { + @Provides + static JettyConfig providesJettyConfig() { + return JettyConfig.builder() + .port(0) + .tokenExpire(Duration.of(5, ChronoUnit.MINUTES)) + .build(); + } + } +} diff --git a/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/jetty/FlightSqlJdbcTestJetty.java b/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/jetty/FlightSqlJdbcTestJetty.java new file mode 100644 index 00000000000..892c2cd1a81 --- /dev/null +++ b/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/jetty/FlightSqlJdbcTestJetty.java @@ -0,0 +1,15 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql.jetty; + +import io.deephaven.server.flightsql.DaggerJettyTestComponent; +import io.deephaven.server.flightsql.FlightSqlJdbcTestBase; + +public class FlightSqlJdbcTestJetty extends FlightSqlJdbcTestBase { + + @Override + protected TestComponent component() { + return DaggerJettyTestComponent.create(); + } +} diff --git a/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/jetty/FlightSqlJdbcUnauthenticatedTestJetty.java b/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/jetty/FlightSqlJdbcUnauthenticatedTestJetty.java new file mode 100644 index 00000000000..2941df3ee15 --- /dev/null +++ b/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/jetty/FlightSqlJdbcUnauthenticatedTestJetty.java @@ -0,0 +1,15 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql.jetty; + +import io.deephaven.server.flightsql.DaggerJettyTestComponent; +import io.deephaven.server.flightsql.FlightSqlJdbcUnauthenticatedTestBase; + +public class FlightSqlJdbcUnauthenticatedTestJetty extends FlightSqlJdbcUnauthenticatedTestBase { + + @Override + protected TestComponent component() { + return DaggerJettyTestComponent.create(); + } +} diff --git a/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlActionHelper.java b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlActionHelper.java new file mode 100644 index 00000000000..289423ecfd0 --- /dev/null +++ b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlActionHelper.java @@ -0,0 +1,191 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql; + +import com.google.protobuf.Any; +import com.google.protobuf.InvalidProtocolBufferException; +import com.google.protobuf.Message; +import io.deephaven.base.verify.Assert; +import io.deephaven.util.annotations.VisibleForTesting; +import io.grpc.Status; +import org.apache.arrow.flight.Action; +import org.apache.arrow.flight.ActionType; +import org.apache.arrow.flight.sql.FlightSqlUtils; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionBeginSavepointRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionBeginTransactionRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionCancelQueryRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionClosePreparedStatementRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedSubstraitPlanRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionEndSavepointRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionEndTransactionRequest; + +import java.util.Set; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +final class FlightSqlActionHelper { + + @VisibleForTesting + static final String CREATE_PREPARED_STATEMENT_ACTION_TYPE = "CreatePreparedStatement"; + + @VisibleForTesting + static final String CLOSE_PREPARED_STATEMENT_ACTION_TYPE = "ClosePreparedStatement"; + + @VisibleForTesting + static final String BEGIN_SAVEPOINT_ACTION_TYPE = "BeginSavepoint"; + + @VisibleForTesting + static final String END_SAVEPOINT_ACTION_TYPE = "EndSavepoint"; + + @VisibleForTesting + static final String BEGIN_TRANSACTION_ACTION_TYPE = "BeginTransaction"; + + @VisibleForTesting + static final String END_TRANSACTION_ACTION_TYPE = "EndTransaction"; + + @VisibleForTesting + static final String CANCEL_QUERY_ACTION_TYPE = "CancelQuery"; + + @VisibleForTesting + static final String CREATE_PREPARED_SUBSTRAIT_PLAN_ACTION_TYPE = "CreatePreparedSubstraitPlan"; + + /** + * Note: FlightSqlUtils.FLIGHT_SQL_ACTIONS is not all the actions, see + * Add all ActionTypes to FlightSqlUtils.FLIGHT_SQL_ACTIONS + * + *

+ * It is unfortunate that there is no proper prefix or namespace for these action types, which would make it much + * easier to route correctly. + */ + private static final Set FLIGHT_SQL_ACTION_TYPES = Stream.of( + FlightSqlUtils.FLIGHT_SQL_BEGIN_SAVEPOINT, + FlightSqlUtils.FLIGHT_SQL_BEGIN_TRANSACTION, + FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_STATEMENT, + FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT, + FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_SUBSTRAIT_PLAN, + FlightSqlUtils.FLIGHT_SQL_CANCEL_QUERY, + FlightSqlUtils.FLIGHT_SQL_END_SAVEPOINT, + FlightSqlUtils.FLIGHT_SQL_END_TRANSACTION) + .map(ActionType::getType) + .collect(Collectors.toSet()); + + interface ActionVisitor { + + T visit(ActionCreatePreparedStatementRequest action); + + T visit(ActionClosePreparedStatementRequest action); + + T visit(ActionBeginSavepointRequest action); + + T visit(ActionEndSavepointRequest action); + + T visit(ActionBeginTransactionRequest action); + + T visit(ActionEndTransactionRequest action); + + T visit(@SuppressWarnings("deprecation") ActionCancelQueryRequest action); + + T visit(ActionCreatePreparedSubstraitPlanRequest action); + } + + public static boolean handlesAction(String type) { + // There is no prefix for Flight SQL action types, so the best we can do is a set-based lookup. This also means + // that this resolver will not be able to respond with an appropriately scoped error message for new Flight SQL + // action types (io.deephaven.server.flightsql.FlightSqlResolver.UnsupportedAction). + return FLIGHT_SQL_ACTION_TYPES.contains(type); + } + + public static T visit(Action action, ActionVisitor visitor) { + final String type = action.getType(); + switch (type) { + case CREATE_PREPARED_STATEMENT_ACTION_TYPE: + return visitor.visit(unpack(action.getBody(), ActionCreatePreparedStatementRequest.class)); + case CLOSE_PREPARED_STATEMENT_ACTION_TYPE: + return visitor.visit(unpack(action.getBody(), ActionClosePreparedStatementRequest.class)); + case BEGIN_SAVEPOINT_ACTION_TYPE: + return visitor.visit(unpack(action.getBody(), ActionBeginSavepointRequest.class)); + case END_SAVEPOINT_ACTION_TYPE: + return visitor.visit(unpack(action.getBody(), ActionEndSavepointRequest.class)); + case BEGIN_TRANSACTION_ACTION_TYPE: + return visitor.visit(unpack(action.getBody(), ActionBeginTransactionRequest.class)); + case END_TRANSACTION_ACTION_TYPE: + return visitor.visit(unpack(action.getBody(), ActionEndTransactionRequest.class)); + case CANCEL_QUERY_ACTION_TYPE: + // noinspection deprecation + return visitor.visit(unpack(action.getBody(), ActionCancelQueryRequest.class)); + case CREATE_PREPARED_SUBSTRAIT_PLAN_ACTION_TYPE: + return visitor.visit(unpack(action.getBody(), ActionCreatePreparedSubstraitPlanRequest.class)); + } + // noinspection DataFlowIssue + throw Assert.statementNeverExecuted(); + } + + private static T unpack(byte[] body, Class clazz) { + final Any any = parseActionOrThrow(body); + return unpackActionOrThrow(any, clazz); + } + + private static Any parseActionOrThrow(byte[] data) { + try { + return Any.parseFrom(data); + } catch (final InvalidProtocolBufferException e) { + throw FlightSqlErrorHelper.error(Status.Code.INVALID_ARGUMENT, "Invalid action"); + } + } + + private static T unpackActionOrThrow(Any source, Class clazz) { + try { + return source.unpack(clazz); + } catch (final InvalidProtocolBufferException e) { + throw FlightSqlErrorHelper.error(Status.Code.INVALID_ARGUMENT, + "Invalid action, provided message cannot be unpacked as " + clazz.getName(), e); + } + } + + public static abstract class ActionVisitorBase implements ActionVisitor { + + public abstract T visitDefault(ActionType actionType, Object action); + + @Override + public T visit(ActionCreatePreparedStatementRequest action) { + return visitDefault(FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_STATEMENT, action); + } + + @Override + public T visit(ActionClosePreparedStatementRequest action) { + return visitDefault(FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT, action); + } + + @Override + public T visit(ActionBeginSavepointRequest action) { + return visitDefault(FlightSqlUtils.FLIGHT_SQL_BEGIN_SAVEPOINT, action); + } + + @Override + public T visit(ActionEndSavepointRequest action) { + return visitDefault(FlightSqlUtils.FLIGHT_SQL_END_SAVEPOINT, action); + } + + @Override + public T visit(ActionBeginTransactionRequest action) { + return visitDefault(FlightSqlUtils.FLIGHT_SQL_BEGIN_TRANSACTION, action); + } + + @Override + public T visit(ActionEndTransactionRequest action) { + return visitDefault(FlightSqlUtils.FLIGHT_SQL_END_TRANSACTION, action); + } + + @Override + public T visit(@SuppressWarnings("deprecation") ActionCancelQueryRequest action) { + return visitDefault(FlightSqlUtils.FLIGHT_SQL_CANCEL_QUERY, action); + } + + @Override + public T visit(ActionCreatePreparedSubstraitPlanRequest action) { + return visitDefault(FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_SUBSTRAIT_PLAN, action); + } + } +} diff --git a/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlCommandHelper.java b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlCommandHelper.java new file mode 100644 index 00000000000..09c9417d88f --- /dev/null +++ b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlCommandHelper.java @@ -0,0 +1,225 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql; + +import com.google.protobuf.Any; +import com.google.protobuf.ByteString; +import com.google.protobuf.Descriptors.Descriptor; +import com.google.protobuf.InvalidProtocolBufferException; +import com.google.protobuf.Message; +import io.deephaven.base.verify.Assert; +import io.grpc.Status; +import org.apache.arrow.flight.impl.Flight.FlightDescriptor; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCatalogs; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCrossReference; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetDbSchemas; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetExportedKeys; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetImportedKeys; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetPrimaryKeys; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetSqlInfo; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTableTypes; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTables; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetXdbcTypeInfo; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementQuery; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementUpdate; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementQuery; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementSubstraitPlan; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementUpdate; + +final class FlightSqlCommandHelper { + + interface CommandVisitor { + + T visit(CommandGetCatalogs command); + + T visit(CommandGetDbSchemas command); + + T visit(CommandGetTableTypes command); + + T visit(CommandGetImportedKeys command); + + T visit(CommandGetExportedKeys command); + + T visit(CommandGetPrimaryKeys command); + + T visit(CommandGetTables command); + + T visit(CommandStatementQuery command); + + T visit(CommandPreparedStatementQuery command); + + T visit(CommandGetSqlInfo command); + + T visit(CommandStatementUpdate command); + + T visit(CommandGetCrossReference command); + + T visit(CommandStatementSubstraitPlan command); + + T visit(CommandPreparedStatementUpdate command); + + T visit(CommandGetXdbcTypeInfo command); + + T visit(CommandStatementIngest command); + } + + public static boolean handlesCommand(FlightDescriptor descriptor) { + // If not CMD, there is an error with io.deephaven.server.session.TicketRouter.getPathResolver / handlesPath + Assert.eq(descriptor.getType(), "descriptor.getType()", FlightDescriptor.DescriptorType.CMD, "CMD"); + // No good way to check if this is a valid command without parsing to Any first. + final Any command = parseOrNull(descriptor.getCmd()); + return command != null + && command.getTypeUrl().startsWith(FlightSqlSharedConstants.FLIGHT_SQL_COMMAND_TYPE_PREFIX); + } + + public static T visit(FlightDescriptor descriptor, CommandVisitor visitor, String logId) { + // If not CMD, there is an error with io.deephaven.server.session.TicketRouter.getPathResolver / handlesPath + Assert.eq(descriptor.getType(), "descriptor.getType()", FlightDescriptor.DescriptorType.CMD, "CMD"); + final Any command = parseOrNull(descriptor.getCmd()); + // If null, there is an error with io.deephaven.server.session.TicketRouter.getCommandResolver / handlesCommand + Assert.neqNull(command, "command"); + final String typeUrl = command.getTypeUrl(); + // If not true, there is an error with io.deephaven.server.session.TicketRouter.getCommandResolver / + // handlesCommand + Assert.eqTrue(typeUrl.startsWith(FlightSqlSharedConstants.FLIGHT_SQL_COMMAND_TYPE_PREFIX), + "typeUrl.startsWith"); + switch (typeUrl) { + case FlightSqlSharedConstants.COMMAND_STATEMENT_QUERY_TYPE_URL: + return visitor.visit(unpack(command, CommandStatementQuery.class, logId)); + case FlightSqlSharedConstants.COMMAND_PREPARED_STATEMENT_QUERY_TYPE_URL: + return visitor.visit(unpack(command, CommandPreparedStatementQuery.class, logId)); + case FlightSqlSharedConstants.COMMAND_GET_TABLES_TYPE_URL: + return visitor.visit(unpack(command, CommandGetTables.class, logId)); + case FlightSqlSharedConstants.COMMAND_GET_TABLE_TYPES_TYPE_URL: + return visitor.visit(unpack(command, CommandGetTableTypes.class, logId)); + case FlightSqlSharedConstants.COMMAND_GET_CATALOGS_TYPE_URL: + return visitor.visit(unpack(command, CommandGetCatalogs.class, logId)); + case FlightSqlSharedConstants.COMMAND_GET_DB_SCHEMAS_TYPE_URL: + return visitor.visit(unpack(command, CommandGetDbSchemas.class, logId)); + case FlightSqlSharedConstants.COMMAND_GET_PRIMARY_KEYS_TYPE_URL: + return visitor.visit(unpack(command, CommandGetPrimaryKeys.class, logId)); + case FlightSqlSharedConstants.COMMAND_GET_IMPORTED_KEYS_TYPE_URL: + return visitor.visit(unpack(command, CommandGetImportedKeys.class, logId)); + case FlightSqlSharedConstants.COMMAND_GET_EXPORTED_KEYS_TYPE_URL: + return visitor.visit(unpack(command, CommandGetExportedKeys.class, logId)); + case FlightSqlSharedConstants.COMMAND_GET_SQL_INFO_TYPE_URL: + return visitor.visit(unpack(command, CommandGetSqlInfo.class, logId)); + case FlightSqlSharedConstants.COMMAND_STATEMENT_UPDATE_TYPE_URL: + return visitor.visit(unpack(command, CommandStatementUpdate.class, logId)); + case FlightSqlSharedConstants.COMMAND_GET_CROSS_REFERENCE_TYPE_URL: + return visitor.visit(unpack(command, CommandGetCrossReference.class, logId)); + case FlightSqlSharedConstants.COMMAND_STATEMENT_SUBSTRAIT_PLAN_TYPE_URL: + return visitor.visit(unpack(command, CommandStatementSubstraitPlan.class, logId)); + case FlightSqlSharedConstants.COMMAND_PREPARED_STATEMENT_UPDATE_TYPE_URL: + return visitor.visit(unpack(command, CommandPreparedStatementUpdate.class, logId)); + case FlightSqlSharedConstants.COMMAND_GET_XDBC_TYPE_INFO_TYPE_URL: + return visitor.visit(unpack(command, CommandGetXdbcTypeInfo.class, logId)); + case FlightSqlSharedConstants.COMMAND_STATEMENT_INGEST_TYPE_URL: + return visitor.visit(unpack(command, CommandStatementIngest.class, logId)); + } + throw FlightSqlErrorHelper.error(Status.Code.UNIMPLEMENTED, String.format("command '%s' is unknown", typeUrl)); + } + + private static Any parseOrNull(ByteString data) { + try { + return Any.parseFrom(data); + } catch (final InvalidProtocolBufferException e) { + return null; + } + } + + private static T unpack(Any command, Class clazz, String logId) { + try { + return command.unpack(clazz); + } catch (InvalidProtocolBufferException e) { + throw FlightSqlErrorHelper.error(Status.Code.FAILED_PRECONDITION, String + .format("Invalid command, provided message cannot be unpacked as %s, %s", clazz.getName(), logId)); + } + } + + public static abstract class CommandVisitorBase implements CommandVisitor { + public abstract T visitDefault(Descriptor descriptor, Object command); + + @Override + public T visit(CommandGetCatalogs command) { + return visitDefault(CommandGetCatalogs.getDescriptor(), command); + } + + @Override + public T visit(CommandGetDbSchemas command) { + return visitDefault(CommandGetDbSchemas.getDescriptor(), command); + } + + @Override + public T visit(CommandGetTableTypes command) { + return visitDefault(CommandGetTableTypes.getDescriptor(), command); + } + + @Override + public T visit(CommandGetImportedKeys command) { + return visitDefault(CommandGetImportedKeys.getDescriptor(), command); + } + + @Override + public T visit(CommandGetExportedKeys command) { + return visitDefault(CommandGetExportedKeys.getDescriptor(), command); + } + + @Override + public T visit(CommandGetPrimaryKeys command) { + return visitDefault(CommandGetPrimaryKeys.getDescriptor(), command); + } + + @Override + public T visit(CommandGetTables command) { + return visitDefault(CommandGetTables.getDescriptor(), command); + } + + @Override + public T visit(CommandStatementQuery command) { + return visitDefault(CommandStatementQuery.getDescriptor(), command); + } + + @Override + public T visit(CommandPreparedStatementQuery command) { + return visitDefault(CommandPreparedStatementQuery.getDescriptor(), command); + } + + @Override + public T visit(CommandGetSqlInfo command) { + return visitDefault(CommandGetSqlInfo.getDescriptor(), command); + } + + @Override + public T visit(CommandStatementUpdate command) { + return visitDefault(CommandStatementUpdate.getDescriptor(), command); + } + + @Override + public T visit(CommandGetCrossReference command) { + return visitDefault(CommandGetCrossReference.getDescriptor(), command); + } + + @Override + public T visit(CommandStatementSubstraitPlan command) { + return visitDefault(CommandStatementSubstraitPlan.getDescriptor(), command); + } + + @Override + public T visit(CommandPreparedStatementUpdate command) { + return visitDefault(CommandPreparedStatementUpdate.getDescriptor(), command); + } + + @Override + public T visit(CommandGetXdbcTypeInfo command) { + return visitDefault(CommandGetXdbcTypeInfo.getDescriptor(), command); + } + + @Override + public T visit(CommandStatementIngest command) { + return visitDefault(CommandStatementIngest.getDescriptor(), command); + } + } +} diff --git a/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlErrorHelper.java b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlErrorHelper.java new file mode 100644 index 00000000000..3b9a49334e7 --- /dev/null +++ b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlErrorHelper.java @@ -0,0 +1,25 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql; + +import io.grpc.Status; +import io.grpc.StatusRuntimeException; + +final class FlightSqlErrorHelper { + + static StatusRuntimeException error(Status.Code code, String message) { + return code + .toStatus() + .withDescription("Flight SQL: " + message) + .asRuntimeException(); + } + + static StatusRuntimeException error(Status.Code code, String message, Throwable cause) { + return code + .toStatus() + .withDescription("Flight SQL: " + message) + .withCause(cause) + .asRuntimeException(); + } +} diff --git a/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlModule.java b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlModule.java new file mode 100644 index 00000000000..4e3c4b2b812 --- /dev/null +++ b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlModule.java @@ -0,0 +1,25 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql; + +import dagger.Binds; +import dagger.Module; +import dagger.multibindings.IntoSet; +import io.deephaven.server.session.ActionResolver; +import io.deephaven.server.session.TicketResolver; + +/** + * Binds {@link FlightSqlResolver} as a {@link TicketResolver} and an {@link ActionResolver}. + */ +@Module +public interface FlightSqlModule { + + @Binds + @IntoSet + TicketResolver bindFlightSqlAsTicketResolver(FlightSqlResolver resolver); + + @Binds + @IntoSet + ActionResolver bindFlightSqlAsActionResolver(FlightSqlResolver resolver); +} diff --git a/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlResolver.java b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlResolver.java new file mode 100644 index 00000000000..2f103832a0a --- /dev/null +++ b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlResolver.java @@ -0,0 +1,1748 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql; + +import com.google.protobuf.Any; +import com.google.protobuf.ByteString; +import com.google.protobuf.ByteStringAccess; +import com.google.protobuf.Descriptors.Descriptor; +import com.google.protobuf.Message; +import com.google.protobuf.Timestamp; +import io.deephaven.base.log.LogOutput; +import io.deephaven.base.verify.Assert; +import io.deephaven.configuration.Configuration; +import io.deephaven.engine.context.ExecutionContext; +import io.deephaven.engine.context.QueryScope; +import io.deephaven.engine.liveness.LivenessScopeStack; +import io.deephaven.engine.sql.Sql; +import io.deephaven.engine.table.ColumnDefinition; +import io.deephaven.engine.table.Table; +import io.deephaven.engine.table.TableDefinition; +import io.deephaven.engine.table.impl.TableCreatorImpl; +import io.deephaven.engine.table.impl.perf.QueryPerformanceNugget; +import io.deephaven.engine.table.impl.perf.QueryPerformanceRecorder; +import io.deephaven.engine.table.impl.util.ColumnHolder; +import io.deephaven.engine.util.TableTools; +import io.deephaven.extensions.barrage.util.ArrowIpcUtil; +import io.deephaven.extensions.barrage.util.BarrageUtil; +import io.deephaven.extensions.barrage.util.GrpcUtil; +import io.deephaven.hash.KeyedObjectHashMap; +import io.deephaven.hash.KeyedObjectKey; +import io.deephaven.internal.log.LoggerFactory; +import io.deephaven.io.logger.Logger; +import io.deephaven.proto.backplane.grpc.ExportNotification; +import io.deephaven.proto.util.ByteHelper; +import io.deephaven.qst.table.TableSpec; +import io.deephaven.qst.table.TicketTable; +import io.deephaven.server.auth.AuthorizationProvider; +import io.deephaven.server.console.ScopeTicketResolver; +import io.deephaven.server.session.ActionResolver; +import io.deephaven.server.session.CommandResolver; +import io.deephaven.server.session.SessionState; +import io.deephaven.server.session.SessionState.ExportObject; +import io.deephaven.server.session.TicketRouter; +import io.deephaven.server.util.Scheduler; +import io.deephaven.sql.SqlParseException; +import io.deephaven.sql.UnsupportedSqlOperation; +import io.deephaven.util.SafeCloseable; +import io.deephaven.util.annotations.VisibleForTesting; +import io.grpc.Status.Code; +import io.grpc.StatusRuntimeException; +import io.grpc.stub.StreamObserver; +import org.apache.arrow.flight.Action; +import org.apache.arrow.flight.ActionType; +import org.apache.arrow.flight.Result; +import org.apache.arrow.flight.impl.Flight; +import org.apache.arrow.flight.impl.Flight.Empty; +import org.apache.arrow.flight.impl.Flight.FlightDescriptor; +import org.apache.arrow.flight.impl.Flight.FlightEndpoint; +import org.apache.arrow.flight.impl.Flight.FlightInfo; +import org.apache.arrow.flight.impl.Flight.Ticket; +import org.apache.arrow.flight.sql.FlightSqlProducer; +import org.apache.arrow.flight.sql.FlightSqlUtils; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionClosePreparedStatementRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementResult; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCatalogs; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetDbSchemas; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetExportedKeys; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetImportedKeys; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetPrimaryKeys; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTableTypes; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTables; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementQuery; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementQuery; +import org.apache.arrow.flight.sql.impl.FlightSql.TicketStatementQuery; +import org.apache.arrow.vector.types.pojo.ArrowType.Utf8; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.calcite.rex.RexDynamicParam; +import org.apache.calcite.runtime.CalciteContextException; +import org.apache.calcite.sql.validate.SqlValidatorException; +import org.jetbrains.annotations.Nullable; + +import javax.inject.Inject; +import javax.inject.Singleton; +import java.io.ByteArrayOutputStream; +import java.io.Closeable; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.nio.ByteBuffer; +import java.security.SecureRandom; +import java.time.Duration; +import java.time.Instant; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Map.Entry; +import java.util.Objects; +import java.util.PrimitiveIterator; +import java.util.Set; +import java.util.function.Consumer; +import java.util.function.Function; +import java.util.function.Predicate; +import java.util.regex.Pattern; +import java.util.stream.IntStream; + +import static io.deephaven.server.flightsql.FlightSqlErrorHelper.error; + +/** + * A Flight SQL resolver. This supports the read-only + * querying of the global query scope, which is presented simply with the query scope variables names as the table names + * without a catalog and schema name. + * + *

+ * This implementation does not currently follow the Flight SQL protocol to exact specification. Namely, all the + * returned {@link Schema Flight schemas} have nullable {@link Field fields}, and some of the fields on specific + * commands have different types (see {@link #flightInfoFor(SessionState, FlightDescriptor, String)} for specifics). + * + *

+ * All commands, actions, and resolution must be called by authenticated users. + */ +@Singleton +public final class FlightSqlResolver implements ActionResolver, CommandResolver { + + private static final String CATALOG_NAME = "catalog_name"; + private static final String PK_CATALOG_NAME = "pk_catalog_name"; + private static final String FK_CATALOG_NAME = "fk_catalog_name"; + + private static final String DB_SCHEMA_NAME = "db_schema_name"; + private static final String PK_DB_SCHEMA_NAME = "pk_db_schema_name"; + private static final String FK_DB_SCHEMA_NAME = "fk_db_schema_name"; + + private static final String TABLE_NAME = "table_name"; + private static final String PK_TABLE_NAME = "pk_table_name"; + private static final String FK_TABLE_NAME = "fk_table_name"; + + private static final String COLUMN_NAME = "column_name"; + private static final String PK_COLUMN_NAME = "pk_column_name"; + private static final String FK_COLUMN_NAME = "fk_column_name"; + + private static final String KEY_NAME = "key_name"; + private static final String PK_KEY_NAME = "pk_key_name"; + private static final String FK_KEY_NAME = "fk_key_name"; + + private static final String TABLE_TYPE = "table_type"; + private static final String KEY_SEQUENCE = "key_sequence"; + private static final String TABLE_SCHEMA = "table_schema"; + private static final String UPDATE_RULE = "update_rule"; + private static final String DELETE_RULE = "delete_rule"; + + private static final String TABLE_TYPE_TABLE = "TABLE"; + private static final Duration FIXED_TICKET_EXPIRE_DURATION = Duration.ofMinutes(1); + private static final long QUERY_WATCHDOG_TIMEOUT_MILLIS = Duration + .parse(Configuration.getInstance().getStringWithDefault("FlightSQL.queryTimeout", "PT5s")).toMillis(); + + private static final Logger log = LoggerFactory.getLogger(FlightSqlResolver.class); + + private static final KeyedObjectKey QUERY_KEY = + new KeyedObjectKey.BasicAdapter<>(QueryBase::handleId); + + private static final KeyedObjectKey PREPARED_STATEMENT_KEY = + new KeyedObjectKey.BasicAdapter<>(PreparedStatement::handleId); + + @VisibleForTesting + static final Schema DATASET_SCHEMA_SENTINEL = new Schema(List.of(Field.nullable("DO_NOT_USE", Utf8.INSTANCE))); + + // Unable to depends on TicketRouter, would be a circular dependency atm (since TicketRouter depends on all the + // TicketResolvers). + // private final TicketRouter router; + private final ScopeTicketResolver scopeTicketResolver; + private final Scheduler scheduler; + private final Authorization authorization; + private final KeyedObjectHashMap queries; + private final KeyedObjectHashMap preparedStatements; + + @Inject + public FlightSqlResolver( + final AuthorizationProvider authProvider, + final ScopeTicketResolver scopeTicketResolver, + final Scheduler scheduler) { + this.authorization = Objects.requireNonNull(authProvider.getTicketResolverAuthorization()); + this.scopeTicketResolver = Objects.requireNonNull(scopeTicketResolver); + this.scheduler = Objects.requireNonNull(scheduler); + this.queries = new KeyedObjectHashMap<>(QUERY_KEY); + this.preparedStatements = new KeyedObjectHashMap<>(PREPARED_STATEMENT_KEY); + } + + /** + * The Flight SQL ticket route, equal to {@value FlightSqlTicketHelper#TICKET_PREFIX}. + * + * @return the Flight SQL ticket route + */ + @Override + public byte ticketRoute() { + return FlightSqlTicketHelper.TICKET_PREFIX; + } + + // --------------------------------------------------------------------------------------------------------------- + + /** + * Returns {@code true} if the given command {@code descriptor} appears to be a valid Flight SQL command; that is, + * it is parsable as an {@code Any} protobuf message with the type URL prefixed with + * {@value FlightSqlSharedConstants#FLIGHT_SQL_COMMAND_TYPE_PREFIX}. + * + * @param descriptor the descriptor + * @return {@code true} if the given command appears to be a valid Flight SQL command + */ + @Override + public boolean handlesCommand(Flight.FlightDescriptor descriptor) { + return FlightSqlCommandHelper.handlesCommand(descriptor); + } + + /** + * Executes the given {@code descriptor} command. Only supports authenticated access. + * + *

+ * {@link CommandStatementQuery}: Executes the given SQL query. The returned Flight info should be promptly + * resolved, and resolved at most once. Transactions are not currently supported. + * + *

+ * {@link CommandPreparedStatementQuery}: Executes the prepared SQL query (must be executed within the scope of a + * {@link FlightSqlUtils#FLIGHT_SQL_CREATE_PREPARED_STATEMENT} / + * {@link FlightSqlUtils#FLIGHT_SQL_CLOSE_PREPARED_STATEMENT}). The returned Flight info should be promptly + * resolved, and resolved at most once. + * + *

+ * {@link CommandGetTables}: Retrieve the tables authorized for the user. The {@value TABLE_NAME}, + * {@value TABLE_TYPE}, and (optional) {@value TABLE_SCHEMA} fields will be out-of-spec as nullable columns (the + * returned data for these columns will never be {@code null}). + * + *

+ * {@link CommandGetCatalogs}: Retrieves the catalogs authorized for the user. The {@value CATALOG_NAME} field will + * be out-of-spec as a nullable column (the returned data for this column will never be {@code null}). Currently, + * always an empty table. + * + *

+ * {@link CommandGetDbSchemas}: Retrieves the catalogs and schemas authorized for the user. The + * {@value DB_SCHEMA_NAME} field will be out-of-spec as a nullable (the returned data for this column will never be + * {@code null}). Currently, always an empty table. + * + *

+ * {@link CommandGetTableTypes}: Retrieves the table types authorized for the user. The {@value TABLE_TYPE} field + * will be out-of-spec as a nullable (the returned data for this column will never be {@code null}). Currently, + * always a table with a single row with value {@value TABLE_TYPE_TABLE}. + * + *

+ * {@link CommandGetPrimaryKeys}: Retrieves the primary keys for a table if the user is authorized. If the table + * does not exist (or the user is not authorized), a {@link Code#NOT_FOUND} exception will be thrown. The + * {@value TABLE_NAME}, {@value COLUMN_NAME}, and {@value KEY_SEQUENCE} will be out-of-spec as nullable columns (the + * returned data for these columns will never be {@code null}). Currently, always an empty table. + * + *

+ * {@link CommandGetImportedKeys}: Retrieves the imported keys for a table if the user is authorized. If the table + * does not exist (or the user is not authorized), a {@link Code#NOT_FOUND} exception will be thrown. The + * {@value PK_TABLE_NAME}, {@value PK_COLUMN_NAME}, {@value FK_TABLE_NAME}, {@value FK_COLUMN_NAME}, and + * {@value KEY_SEQUENCE} will be out-of-spec as nullable columns (the returned data for these columns will never be + * {@code null}). The {@value UPDATE_RULE} and {@value DELETE_RULE} will be out-of-spec as nullable {@code int8} + * types instead of {@code uint8} (the returned data for these columns will never be {@code null}). Currently, + * always an empty table. + * + *

+ * {@link CommandGetExportedKeys}: Retrieves the exported keys for a table if the user is authorized. If the table + * does not exist (or the user is not authorized), a {@link Code#NOT_FOUND} exception will be thrown. The + * {@value PK_TABLE_NAME}, {@value PK_COLUMN_NAME}, {@value FK_TABLE_NAME}, {@value FK_COLUMN_NAME}, and + * {@value KEY_SEQUENCE} will be out-of-spec as nullable columns (the returned data for these columns will never be + * {@code null}). The {@value UPDATE_RULE} and {@value DELETE_RULE} will be out-of-spec as nullable {@code int8} + * types instead of {@code uint8} (the returned data for these columns will never be {@code null}). Currently, + * always an empty table. + * + *

+ * All other commands will throw an {@link Code#UNIMPLEMENTED} exception. + * + * @param session the session + * @param descriptor the flight descriptor to retrieve a ticket for + * @param logId an end-user friendly identification of the ticket should an error occur + * @return the flight info for the given {@code descriptor} command + */ + @Override + public ExportObject flightInfoFor( + @Nullable final SessionState session, final Flight.FlightDescriptor descriptor, final String logId) { + if (session == null) { + throw unauthenticatedError(); + } + return FlightSqlCommandHelper.visit(descriptor, new GetFlightInfoImpl(session, descriptor), logId); + } + + private class GetFlightInfoImpl extends FlightSqlCommandHelper.CommandVisitorBase> { + private final SessionState session; + private final FlightDescriptor descriptor; + + public GetFlightInfoImpl(SessionState session, FlightDescriptor descriptor) { + this.session = Objects.requireNonNull(session); + this.descriptor = Objects.requireNonNull(descriptor); + } + + @Override + public ExportObject visitDefault(Descriptor descriptor, Object command) { + return submit(new UnsupportedCommand<>(descriptor), command); + } + + @Override + public ExportObject visit(CommandGetCatalogs command) { + return submit(CommandGetCatalogsConstants.HANDLER, command); + } + + @Override + public ExportObject visit(CommandGetDbSchemas command) { + return submit(CommandGetDbSchemasConstants.HANDLER, command); + } + + @Override + public ExportObject visit(CommandGetTableTypes command) { + return submit(CommandGetTableTypesConstants.HANDLER, command); + } + + @Override + public ExportObject visit(CommandGetImportedKeys command) { + return submit(commandGetImportedKeysHandler, command); + } + + @Override + public ExportObject visit(CommandGetExportedKeys command) { + return submit(commandGetExportedKeysHandler, command); + } + + @Override + public ExportObject visit(CommandGetPrimaryKeys command) { + return submit(commandGetPrimaryKeysHandler, command); + } + + @Override + public ExportObject visit(CommandGetTables command) { + return submit(new CommandGetTablesImpl(), command); + } + + @Override + public ExportObject visit(CommandStatementQuery command) { + return submit(new CommandStatementQueryImpl(session), command); + } + + @Override + public ExportObject visit(CommandPreparedStatementQuery command) { + return submit(new CommandPreparedStatementQueryImpl(session), command); + } + + private ExportObject submit(CommandHandler handler, T command) { + return session.nonExport().submit(() -> getInfo(handler, command)); + } + + private FlightInfo getInfo(CommandHandler handler, T command) { + final QueryPerformanceRecorder qpr = QueryPerformanceRecorder.getInstance(); + try (final QueryPerformanceNugget ignore = + qpr.getNugget(String.format("FlightSQL.getInfo/%s", command.getClass().getSimpleName()))) { + return flightInfo(handler, command); + } + } + + private FlightInfo flightInfo(CommandHandler handler, T command) { + final TicketHandler ticketHandler = handler.execute(command); + try { + return ticketHandler.getInfo(descriptor); + } catch (Throwable t) { + if (ticketHandler instanceof TicketHandlerReleasable) { + ((TicketHandlerReleasable) ticketHandler).release(); + } + throw t; + } + } + } + + // --------------------------------------------------------------------------------------------------------------- + + /** + * Only supports authenticated access. + * + * @param session the user session context + * @param ticket (as ByteByffer) the ticket to resolve + * @param logId an end-user friendly identification of the ticket should an error occur + * @return the exported table + * @param the type, must be Table + */ + @Override + public SessionState.ExportObject resolve( + @Nullable final SessionState session, final ByteBuffer ticket, final String logId) { + if (session == null) { + throw unauthenticatedError(); + } + final ExportObject

tableExport = FlightSqlTicketHelper.visit(ticket, new ResolveImpl(session), logId); + // noinspection unchecked + return (ExportObject) tableExport; + } + + private class ResolveImpl implements FlightSqlTicketHelper.TicketVisitor> { + private final SessionState session; + + public ResolveImpl(SessionState session) { + this.session = Objects.requireNonNull(session); + } + + @Override + public ExportObject
visit(CommandGetCatalogs ticket) { + return submit(CommandGetCatalogsConstants.HANDLER, ticket); + } + + @Override + public ExportObject
visit(CommandGetDbSchemas ticket) { + return submit(CommandGetDbSchemasConstants.HANDLER, ticket); + } + + @Override + public ExportObject
visit(CommandGetTableTypes ticket) { + return submit(CommandGetTableTypesConstants.HANDLER, ticket); + } + + @Override + public ExportObject
visit(CommandGetImportedKeys ticket) { + return submit(commandGetImportedKeysHandler, ticket); + } + + @Override + public ExportObject
visit(CommandGetExportedKeys ticket) { + return submit(commandGetExportedKeysHandler, ticket); + } + + @Override + public ExportObject
visit(CommandGetPrimaryKeys ticket) { + return submit(commandGetPrimaryKeysHandler, ticket); + } + + @Override + public ExportObject
visit(CommandGetTables ticket) { + return submit(commandGetTables, ticket); + } + + private ExportObject
submit(CommandHandlerFixedBase fixed, C command) { + // We know this is a trivial execute, okay to do on RPC thread + return submit(fixed.execute(command)); + } + + @Override + public ExportObject
visit(TicketStatementQuery ticket) { + final TicketHandler ticketHandler = queries.get(ticket.getStatementHandle()); + if (ticketHandler == null) { + throw error(Code.NOT_FOUND, + "Unable to find Flight SQL query. Flight SQL tickets should be resolved promptly and resolved at most once."); + } + if (!ticketHandler.isOwner(session)) { + // We should not be concerned about returning "NOT_FOUND" here; the handleId is sufficiently random that + // it is much more likely an authentication setup issue. + throw permissionDeniedWithHelpfulMessage(); + } + return submit(ticketHandler); + } + + // Note: we could be more efficient and do the static table resolution on thread instead of submitting. For + // simplicity purposes for now, we will submit all of them for resolution. + private ExportObject
submit(TicketHandler handler) { + return new TableResolver(session, handler).submit(); + } + } + + private static class TableResolver implements SessionState.ExportErrorHandler { + private final SessionState session; + private final TicketHandler handler; + + public TableResolver(SessionState session, TicketHandler handler) { + this.handler = Objects.requireNonNull(handler); + this.session = Objects.requireNonNull(session); + } + + public ExportObject
submit() { + // We need to provide clean handoff of the Table for Liveness management between the resolver and the + // export; as such, we _can't_ unmanage the Table during a call to TicketHandler.resolve, so we must rely + // on onSuccess / onError callbacks (after export has started managing the Table). + return session.
nonExport() + .onSuccess(this::onSuccess) + .onError(this) + .submit(handler::resolve); + } + + private void onSuccess() { + release(); + } + + @Override + public void onError(ExportNotification.State resultState, String errorContext, @Nullable Exception cause, + @Nullable String dependentExportId) { + release(); + } + + private void release() { + if (!(handler instanceof TicketHandlerReleasable)) { + return; + } + ((TicketHandlerReleasable) handler).release(); + } + } + + @Override + public SessionState.ExportObject resolve( + @Nullable final SessionState session, final Flight.FlightDescriptor descriptor, final String logId) { + // This general interface does not make sense; resolution should always be done against a _ticket_. Nothing + // calls io.deephaven.server.session.TicketRouter.resolve(SessionState, FlightDescriptor, String) + // noinspection DataFlowIssue + throw Assert.statementNeverExecuted(); + } + + // --------------------------------------------------------------------------------------------------------------- + + /** + * Supports unauthenticated access. When unauthenticated, will not return any actions types. When authenticated, + * will return the action types the user is authorized to access. Currently, supports + * {@link FlightSqlUtils#FLIGHT_SQL_CREATE_PREPARED_STATEMENT} and + * {@link FlightSqlUtils#FLIGHT_SQL_CLOSE_PREPARED_STATEMENT}. + * + * @param session the session + * @param visitor the visitor + */ + @Override + public void listActions(@Nullable SessionState session, Consumer visitor) { + if (session == null) { + return; + } + visitor.accept(FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_STATEMENT); + visitor.accept(FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT); + } + + /** + * Returns {@code true} if {@code type} is a known Flight SQL action type (even if this implementation does not + * implement it). + * + * @param type the action type + * @return if {@code type} is a known Flight SQL action type + */ + @Override + public boolean handlesActionType(String type) { + return FlightSqlActionHelper.handlesAction(type); + } + + /** + * Executes the given {@code action}. Only supports authenticated access. Currently, supports + * {@link FlightSqlUtils#FLIGHT_SQL_CREATE_PREPARED_STATEMENT} and + * {@link FlightSqlUtils#FLIGHT_SQL_CLOSE_PREPARED_STATEMENT}; all other action types will throw an + * {@link Code#UNIMPLEMENTED} exception. Transactions are not currently supported. + * + * @param session the session + * @param action the action + * @param observer the observer + */ + @Override + public void doAction( + @Nullable final SessionState session, + final Action action, + final StreamObserver observer) { + // If false, there is an error with io.deephaven.server.session.ActionRouter.doAction / handlesActionType + Assert.eqTrue(handlesActionType(action.getType()), "handlesActionType(action.getType())"); + if (session == null) { + throw unauthenticatedError(); + } + executeAction(session, FlightSqlActionHelper.visit(action, new ActionHandlerVisitor()), observer); + } + + // --------------------------------------------------------------------------------------------------------------- + + /** + * Supports unauthenticated access. When unauthenticated, will not return any Flight info. When authenticated, this + * may return Flight info the user is authorized to access. Currently, no Flight info is returned. + * + * @param session optional session that the resolver can use to filter which flights a visitor sees + * @param visitor the callback to invoke per descriptor path + */ + @Override + public void forAllFlightInfo(@Nullable final SessionState session, final Consumer visitor) { + if (session == null) { + return; + } + // Potential support for listing here in the future + } + + // --------------------------------------------------------------------------------------------------------------- + + /** + * Publishing to Flight SQL descriptors is not currently supported. Throws a {@link Code#FAILED_PRECONDITION} error. + */ + @Override + public SessionState.ExportBuilder publish( + final SessionState session, + final Flight.FlightDescriptor descriptor, + final String logId, + @Nullable final Runnable onPublish) { + if (session == null) { + throw unauthenticatedError(); + } + throw error(Code.FAILED_PRECONDITION, + "Could not publish '" + logId + "': Flight SQL descriptors cannot be published to"); + } + + /** + * Publishing to Flight SQL tickets is not currently supported. Throws a {@link Code#FAILED_PRECONDITION} error. + */ + @Override + public SessionState.ExportBuilder publish( + final SessionState session, + final ByteBuffer ticket, + final String logId, + @Nullable final Runnable onPublish) { + if (session == null) { + throw unauthenticatedError(); + } + throw error(Code.FAILED_PRECONDITION, + "Could not publish '" + logId + "': Flight SQL tickets cannot be published to"); + } + + // --------------------------------------------------------------------------------------------------------------- + + @Override + public String getLogNameFor(final ByteBuffer ticket, final String logId) { + // This is a bit different from the other resolvers; a ticket may be a very long byte string here since it + // may represent a command. + return FlightSqlTicketHelper.toReadableString(ticket, logId); + } + + // --------------------------------------------------------------------------------------------------------------- + + interface CommandHandler { + + TicketHandler execute(C command); + } + + interface TicketHandler { + + boolean isOwner(SessionState session); + + FlightInfo getInfo(FlightDescriptor descriptor); + + Table resolve(); + } + + interface TicketHandlerReleasable extends TicketHandler { + + void release(); + } + + private Table executeSqlQuery(SessionState session, String sql) { + // See SQLTODO(catalog-reader-implementation) + final QueryScope queryScope = ExecutionContext.getContext().getQueryScope(); + // noinspection unchecked,rawtypes + final Map queryScopeTables = + (Map) (Map) queryScope.toMap(queryScope::unwrapObject, (n, t) -> t instanceof Table); + final TableSpec tableSpec = Sql.parseSql(sql, queryScopeTables, TicketTable::fromQueryScopeField, null); + // Note: this is doing io.deephaven.server.session.TicketResolver.Authorization.transform, but not + // io.deephaven.auth.ServiceAuthWiring + // TODO(deephaven-core#6307): Declarative server-side table execution logic that preserves authorization logic + try (final SafeCloseable ignored = LivenessScopeStack.open()) { + final Table table = tableSpec.logic() + .create(new TableCreatorScopeTickets(TableCreatorImpl.INSTANCE, scopeTicketResolver, session)); + if (table.isRefreshing()) { + table.retainReference(); + } + return table; + } + } + + /** + * This is the base class for "easy" commands; that is, commands that have a fixed schema and are cheap to + * initialize. + */ + static abstract class CommandHandlerFixedBase implements CommandHandler { + + /** + * This is called as the first part of {@link TicketHandler#getInfo(FlightDescriptor)} for the handler returned + * from {@link #execute(T)}. It can be used as an early signal to let clients know that the command is not + * supported, or one of the arguments is not valid. + */ + void checkForGetInfo(T command) { + + } + + /** + * This is called as the first part of {@link TicketHandler#resolve()} for the handler returned from + * {@link #execute(T)}. + */ + void checkForResolve(T command) { + // This is provided for completeness, but the current implementations don't use it. + // + // The callers that override checkForGetInfo, for example, all involve table names; if that table exists and + // they are authorized, they will get back a ticket that upon resolve will be an (empty) table. Otherwise, + // they will get a NOT_FOUND exception at getFlightInfo time. + // + // In this context, it is incorrect to do the same check at resolve time because we need to ensure that + // getFlightInfo / doGet (/ doExchange) appears stateful - it would be incorrect to return getFlightInfo + // with the semantics "this table exists" and then potentially throw a NOT_FOUND at resolve time. + // + // If Deephaven Flight SQL implements CommandGetExportedKeys, CommandGetImportedKeys, or + // CommandGetPrimaryKeys, we'll likely need to "upgrade" the implementation to a properly stateful one like + // QueryBase with handle-based tickets. + } + + long totalRecords() { + return -1; + } + + abstract Ticket ticket(T command); + + abstract ByteString schemaBytes(T command); + + abstract Table table(T command); + + /** + * The handler. Will invoke {@link #checkForGetInfo(Message)} as the first part of + * {@link TicketHandler#getInfo(FlightDescriptor)}. Will invoke {@link #checkForResolve(Message)} as the first + * part of {@link TicketHandler#resolve()}. + */ + @Override + public final TicketHandler execute(T command) { + return new TicketHandlerFixed(command); + } + + private class TicketHandlerFixed implements TicketHandler { + private final T command; + + private TicketHandlerFixed(T command) { + this.command = Objects.requireNonNull(command); + } + + @Override + public boolean isOwner(SessionState session) { + return true; + } + + @Override + public FlightInfo getInfo(FlightDescriptor descriptor) { + checkForGetInfo(command); + // Note: the presence of expirationTime is mainly a way to let clients know that they can retry a DoGet + // / DoExchange with an existing ticket without needing to do a new getFlightInfo first (for at least + // the amount of time as specified by expirationTime). Given that the tables resolved via this code-path + // are "easy" (in a lot of the cases, they are static empty tables or "easily" computable)... + // We are not setting an expiration timestamp for the SQL queries - they are only meant to be resolvable + // once; this is different from the watchdog concept. + return FlightInfo.newBuilder() + .setFlightDescriptor(descriptor) + .setSchema(schemaBytes(command)) + .addEndpoint(FlightEndpoint.newBuilder() + .setTicket(ticket(command)) + .setExpirationTime(timestamp(Instant.now().plus(FIXED_TICKET_EXPIRE_DURATION))) + .build()) + .setTotalRecords(totalRecords()) + .setTotalBytes(-1) + .build(); + } + + @Override + public Table resolve() { + checkForResolve(command); + final Table table = CommandHandlerFixedBase.this.table(command); + final long totalRecords = totalRecords(); + if (totalRecords != -1) { + // If true, TicketHandler implementation error; should only override totalRecords for + // non-refreshing tables + Assert.eqFalse(table.isRefreshing(), "table.isRefreshing()"); + // If false, Ticket handler implementation error; totalRecords does not match the table size + Assert.eq(table.size(), "table.size()", totalRecords, "totalRecords"); + } + return table; + } + } + } + + private static final class UnsupportedCommand implements CommandHandler, TicketHandler { + private final Descriptor descriptor; + + UnsupportedCommand(Descriptor descriptor) { + this.descriptor = Objects.requireNonNull(descriptor); + } + + @Override + public TicketHandler execute(T command) { + return this; + } + + @Override + public boolean isOwner(SessionState session) { + return true; + } + + @Override + public FlightInfo getInfo(FlightDescriptor descriptor) { + throw error(Code.UNIMPLEMENTED, + String.format("command '%s' is unimplemented", this.descriptor.getFullName())); + } + + @Override + public Table resolve() { + throw error(Code.INVALID_ARGUMENT, String.format( + "client is misbehaving, should use getInfo for command '%s'", this.descriptor.getFullName())); + } + } + + abstract class QueryBase implements CommandHandler, TicketHandlerReleasable { + private final ByteString handleId; + protected final SessionState session; + + private boolean initialized; + private boolean resolved; + private Table table; + + QueryBase(SessionState session) { + this.handleId = randomHandleId(); + this.session = Objects.requireNonNull(session); + queries.put(handleId, this); + } + + public ByteString handleId() { + return handleId; + } + + @Override + public final TicketHandlerReleasable execute(C command) { + try { + return executeImpl(command); + } catch (Throwable t) { + release(); + throw t; + } + } + + private synchronized QueryBase executeImpl(C command) { + Assert.eqFalse(initialized, "initialized"); + initialized = true; + executeSql(command); + Assert.neqNull(table, "table"); + // Note: we aren't currently providing a way to proactively cleanup query watchdogs - given their + // short-lived nature, they will execute "quick enough" for most use cases. + scheduler.runAfterDelay(QUERY_WATCHDOG_TIMEOUT_MILLIS, this::onWatchdog); + return this; + } + + // responsible for setting table and schemaBytes + protected abstract void executeSql(C command); + + protected void executeSql(String sql) { + try { + table = executeSqlQuery(session, sql); + } catch (SqlParseException e) { + throw error(Code.INVALID_ARGUMENT, "query can't be parsed", e); + } catch (UnsupportedSqlOperation e) { + if (e.clazz() == RexDynamicParam.class) { + throw queryParametersNotSupported(e); + } + throw error(Code.INVALID_ARGUMENT, + String.format("Unsupported calcite type '%s'", e.clazz().getName()), + e); + } catch (CalciteContextException e) { + // See org.apache.calcite.runtime.CalciteResource for the various messages we might encounter + final Throwable cause = e.getCause(); + if (cause instanceof SqlValidatorException) { + if (cause.getMessage().contains("not found")) { + throw error(Code.NOT_FOUND, cause.getMessage(), cause); + } + throw error(Code.INVALID_ARGUMENT, cause.getMessage(), cause); + } + throw e; + } + } + + // ---------------------------------------------------------------------------------------------------------- + + @Override + public final boolean isOwner(SessionState session) { + return this.session.equals(session); + } + + @Override + public final synchronized FlightInfo getInfo(FlightDescriptor descriptor) { + return TicketRouter.getFlightInfo(table, descriptor, ticket()); + } + + @Override + public final synchronized Table resolve() { + if (resolved) { + throw error(Code.FAILED_PRECONDITION, "Should only resolve once"); + } + resolved = true; + if (table == null) { + throw error(Code.FAILED_PRECONDITION, "Should resolve table quicker"); + } + return table; + } + + @Override + public synchronized void release() { + if (!queries.remove(handleId, this)) { + return; + } + doRelease(); + } + + private void doRelease() { + if (table != null) { + if (table.isRefreshing()) { + table.dropReference(); + } + table = null; + } + } + + // ---------------------------------------------------------------------------------------------------------- + + private synchronized void onWatchdog() { + if (!queries.remove(handleId, this)) { + return; + } + log.debug().append("Watchdog cleaning up query handleId=") + .append(ByteStringAsHex.INSTANCE, handleId) + .endl(); + doRelease(); + } + + private Ticket ticket() { + return FlightSqlTicketHelper.ticketCreator().visit(TicketStatementQuery.newBuilder() + .setStatementHandle(handleId) + .build()); + } + } + + final class CommandStatementQueryImpl extends QueryBase { + + CommandStatementQueryImpl(SessionState session) { + super(session); + } + + @Override + public void executeSql(CommandStatementQuery command) { + if (command.hasTransactionId()) { + throw transactionIdsNotSupported(); + } + executeSql(command.getQuery()); + } + } + + final class CommandPreparedStatementQueryImpl extends QueryBase { + + private PreparedStatement prepared; + + CommandPreparedStatementQueryImpl(SessionState session) { + super(session); + } + + @Override + public void executeSql(CommandPreparedStatementQuery command) { + prepared = getPreparedStatement(session, command.getPreparedStatementHandle()); + // Assumed this is not actually parameterized. + final String sql = prepared.parameterizedQuery(); + executeSql(sql); + prepared.attach(this); + } + + @Override + public void release() { + releaseImpl(true); + } + + private void releaseImpl(boolean detach) { + if (detach && prepared != null) { + prepared.detach(this); + } + super.release(); + } + } + + private static class CommandStaticTable extends CommandHandlerFixedBase { + private final Table table; + private final Function f; + private final ByteString schemaBytes; + + CommandStaticTable(Table table, Function f) { + super(); + Assert.eqFalse(table.isRefreshing(), "table.isRefreshing()"); + this.table = Objects.requireNonNull(table); + this.f = Objects.requireNonNull(f); + this.schemaBytes = BarrageUtil.schemaBytesFromTable(table); + } + + @Override + Ticket ticket(T command) { + return f.apply(command); + } + + @Override + ByteString schemaBytes(T command) { + return schemaBytes; + } + + @Override + Table table(T command) { + return table; + } + + @Override + long totalRecords() { + return table.size(); + } + } + + @VisibleForTesting + static final class CommandGetTableTypesConstants { + + /** + * Models return type for {@link CommandGetTableTypes}, + * {@link FlightSqlProducer.Schemas#GET_TABLE_TYPES_SCHEMA}. + * + * 
+         * table_type: utf8 not null
+         *
+ */ + @VisibleForTesting + static final TableDefinition DEFINITION = TableDefinition.of( + ColumnDefinition.ofString(TABLE_TYPE) // out-of-spec + ); + private static final Map ATTRIBUTES = Map.of(); + private static final Table TABLE = + TableTools.newTable(DEFINITION, ATTRIBUTES, TableTools.stringCol(TABLE_TYPE, TABLE_TYPE_TABLE)); + + public static final CommandHandlerFixedBase HANDLER = new CommandStaticTable<>( + TABLE, FlightSqlTicketHelper.ticketCreator()::visit); + } + + @VisibleForTesting + static final class CommandGetCatalogsConstants { + + /** + * Models return type for {@link CommandGetCatalogs}, {@link FlightSqlProducer.Schemas#GET_CATALOGS_SCHEMA}. + * + * 
+         * catalog_name: utf8 not null
+         *
+ */ + @VisibleForTesting + static final TableDefinition DEFINITION = TableDefinition.of( + ColumnDefinition.ofString(CATALOG_NAME) // out-of-spec + ); + private static final Map ATTRIBUTES = Map.of(); + private static final Table TABLE = TableTools.newTable(DEFINITION, ATTRIBUTES); + + public static final CommandHandlerFixedBase HANDLER = + new CommandStaticTable<>(TABLE, FlightSqlTicketHelper.ticketCreator()::visit); + } + + @VisibleForTesting + static final class CommandGetDbSchemasConstants { + + /** + * Models return type for {@link CommandGetDbSchemas}, {@link FlightSqlProducer.Schemas#GET_SCHEMAS_SCHEMA}. + * + * 
+         * catalog_name: utf8,
+         * db_schema_name: utf8 not null
+         *
+ */ + @VisibleForTesting + static final TableDefinition DEFINITION = TableDefinition.of( + ColumnDefinition.ofString(CATALOG_NAME), + ColumnDefinition.ofString(DB_SCHEMA_NAME) // out-of-spec + ); + private static final Map ATTRIBUTES = Map.of(); + private static final Table TABLE = TableTools.newTable(DEFINITION, ATTRIBUTES); + public static final CommandHandlerFixedBase HANDLER = new CommandStaticTable<>( + TABLE, FlightSqlTicketHelper.ticketCreator()::visit); + } + + @VisibleForTesting + static final class CommandGetKeysConstants { + + /** + * Models return type for {@link CommandGetImportedKeys} / {@link CommandGetExportedKeys}, + * {@link FlightSqlProducer.Schemas#GET_IMPORTED_KEYS_SCHEMA}, + * {@link FlightSqlProducer.Schemas#GET_EXPORTED_KEYS_SCHEMA}. + * + * 
+         * pk_catalog_name: utf8,
+         * pk_db_schema_name: utf8,
+         * pk_table_name: utf8 not null,
+         * pk_column_name: utf8 not null,
+         * fk_catalog_name: utf8,
+         * fk_db_schema_name: utf8,
+         * fk_table_name: utf8 not null,
+         * fk_column_name: utf8 not null,
+         * key_sequence: int32 not null,
+         * fk_key_name: utf8,
+         * pk_key_name: utf8,
+         * update_rule: uint8 not null,
+         * delete_rule: uint8 not null
+         *
+ */ + @VisibleForTesting + static final TableDefinition DEFINITION = TableDefinition.of( + ColumnDefinition.ofString(PK_CATALOG_NAME), + ColumnDefinition.ofString(PK_DB_SCHEMA_NAME), + ColumnDefinition.ofString(PK_TABLE_NAME), // out-of-spec + ColumnDefinition.ofString(PK_COLUMN_NAME), // out-of-spec + ColumnDefinition.ofString(FK_CATALOG_NAME), + ColumnDefinition.ofString(FK_DB_SCHEMA_NAME), + ColumnDefinition.ofString(FK_TABLE_NAME), // out-of-spec + ColumnDefinition.ofString(FK_COLUMN_NAME), // out-of-spec + ColumnDefinition.ofInt(KEY_SEQUENCE), // out-of-spec + ColumnDefinition.ofString(FK_KEY_NAME), // yes, this does come _before_ the PK version + ColumnDefinition.ofString(PK_KEY_NAME), + ColumnDefinition.ofByte(UPDATE_RULE), // out-of-spec + ColumnDefinition.ofByte(DELETE_RULE) // out-of-spec + ); + + private static final Map ATTRIBUTES = Map.of(); + private static final Table TABLE = TableTools.newTable(DEFINITION, ATTRIBUTES); + } + + @VisibleForTesting + static final class CommandGetPrimaryKeysConstants { + + /** + * Models return type for {@link CommandGetPrimaryKeys} / + * {@link FlightSqlProducer.Schemas#GET_PRIMARY_KEYS_SCHEMA}. + * + * 
+         * catalog_name: utf8,
+         * db_schema_name: utf8,
+         * table_name: utf8 not null,
+         * column_name: utf8 not null,
+         * key_name: utf8,
+         * key_sequence: int32 not null
+         *
+ */ + @VisibleForTesting + static final TableDefinition DEFINITION = TableDefinition.of( + ColumnDefinition.ofString(CATALOG_NAME), + ColumnDefinition.ofString(DB_SCHEMA_NAME), + ColumnDefinition.ofString(TABLE_NAME), // out-of-spec + ColumnDefinition.ofString(COLUMN_NAME), // out-of-spec + ColumnDefinition.ofString(KEY_NAME), + ColumnDefinition.ofInt(KEY_SEQUENCE) // out-of-spec + ); + + private static final Map ATTRIBUTES = Map.of(); + private static final Table TABLE = TableTools.newTable(DEFINITION, ATTRIBUTES); + } + + private boolean hasTable(String catalog, String dbSchema, String table) { + if (catalog != null && !catalog.isEmpty()) { + return false; + } + if (dbSchema != null && !dbSchema.isEmpty()) { + return false; + } + final Object obj; + final QueryScope scope = ExecutionContext.getContext().getQueryScope(); + try { + obj = scope.readParamValue(table); + } catch (QueryScope.MissingVariableException e) { + return false; + } + if (!(obj instanceof Table)) { + return false; + } + return !authorization.isDeniedAccess(obj); + } + + private final CommandHandlerFixedBase commandGetPrimaryKeysHandler = + new CommandStaticTable<>(CommandGetPrimaryKeysConstants.TABLE, + FlightSqlTicketHelper.ticketCreator()::visit) { + @Override + void checkForGetInfo(CommandGetPrimaryKeys command) { + if (CommandGetPrimaryKeys.getDefaultInstance().equals(command)) { + // We need to pretend that CommandGetPrimaryKeys.getDefaultInstance() is a valid command until + // we can plumb getSchema through to the resolvers. + // TODO(deephaven-core#6218): feat: expose getSchema to TicketResolvers + return; + } + if (!hasTable( + command.hasCatalog() ? command.getCatalog() : null, + command.hasDbSchema() ? command.getDbSchema() : null, + command.getTable())) { + throw tableNotFound(); + } + } + }; + + private final CommandHandlerFixedBase commandGetImportedKeysHandler = + new CommandStaticTable<>(CommandGetKeysConstants.TABLE, FlightSqlTicketHelper.ticketCreator()::visit) { + @Override + void checkForGetInfo(CommandGetImportedKeys command) { + if (CommandGetImportedKeys.getDefaultInstance().equals(command)) { + // We need to pretend that CommandGetImportedKeys.getDefaultInstance() is a valid command until + // we can plumb getSchema through to the resolvers. + // TODO(deephaven-core#6218): feat: expose getSchema to TicketResolvers + return; + } + if (!hasTable( + command.hasCatalog() ? command.getCatalog() : null, + command.hasDbSchema() ? command.getDbSchema() : null, + command.getTable())) { + throw tableNotFound(); + } + } + }; + + private final CommandHandlerFixedBase commandGetExportedKeysHandler = + new CommandStaticTable<>(CommandGetKeysConstants.TABLE, FlightSqlTicketHelper.ticketCreator()::visit) { + @Override + void checkForGetInfo(CommandGetExportedKeys command) { + if (CommandGetExportedKeys.getDefaultInstance().equals(command)) { + // We need to pretend that CommandGetExportedKeys.getDefaultInstance() is a valid command until + // we can plumb getSchema through to the resolvers. + // TODO(deephaven-core#6218): feat: expose getSchema to TicketResolvers + return; + } + if (!hasTable( + command.hasCatalog() ? command.getCatalog() : null, + command.hasDbSchema() ? command.getDbSchema() : null, + command.getTable())) { + throw tableNotFound(); + } + } + }; + + private final CommandHandlerFixedBase commandGetTables = new CommandGetTablesImpl(); + + @VisibleForTesting + static final class CommandGetTablesConstants { + + /** + * Models return type for {@link CommandGetTables} / {@link FlightSqlProducer.Schemas#GET_TABLES_SCHEMA}. + * + * 
+         * catalog_name: utf8,
+         * db_schema_name: utf8,
+         * table_name: utf8 not null,
+         * table_type: utf8 not null,
+         * table_schema: bytes not null
+         *
+ */ + @VisibleForTesting + static final TableDefinition DEFINITION = TableDefinition.of( + ColumnDefinition.ofString(CATALOG_NAME), + ColumnDefinition.ofString(DB_SCHEMA_NAME), + ColumnDefinition.ofString(TABLE_NAME), // out-of-spec + ColumnDefinition.ofString(TABLE_TYPE), // out-of-spec + ColumnDefinition.fromGenericType(TABLE_SCHEMA, Schema.class) // out-of-spec + ); + + /** + * Models return type for {@link CommandGetTables} / + * {@link FlightSqlProducer.Schemas#GET_TABLES_SCHEMA_NO_SCHEMA}. + * + * 
+         * catalog_name: utf8,
+         * db_schema_name: utf8,
+         * table_name: utf8 not null,
+         * table_type: utf8 not null,
+         *
+ */ + @VisibleForTesting + static final TableDefinition DEFINITION_NO_SCHEMA = TableDefinition.of( + ColumnDefinition.ofString(CATALOG_NAME), + ColumnDefinition.ofString(DB_SCHEMA_NAME), // out-of-spec + ColumnDefinition.ofString(TABLE_NAME), // out-of-spec + ColumnDefinition.ofString(TABLE_TYPE)); + + private static final Map ATTRIBUTES = Map.of(); + + private static final ByteString SCHEMA_BYTES_NO_SCHEMA = + BarrageUtil.schemaBytesFromTableDefinition(DEFINITION_NO_SCHEMA, ATTRIBUTES, true); + + private static final ByteString SCHEMA_BYTES = + BarrageUtil.schemaBytesFromTableDefinition(DEFINITION, ATTRIBUTES, true); + } + + private class CommandGetTablesImpl extends CommandHandlerFixedBase { + + @Override + Ticket ticket(CommandGetTables command) { + return FlightSqlTicketHelper.ticketCreator().visit(command); + } + + @Override + ByteString schemaBytes(CommandGetTables command) { + return command.getIncludeSchema() + ? CommandGetTablesConstants.SCHEMA_BYTES + : CommandGetTablesConstants.SCHEMA_BYTES_NO_SCHEMA; + } + + @Override + public Table table(CommandGetTables request) { + // A not present `catalog` means "don't filter based on catalog". + // An empty `catalog` string explicitly means "only return tables that don't have a catalog". + // In our case (since we don't expose catalogs ATM), we can combine them. + final boolean hasCatalog = request.hasCatalog() && !request.getCatalog().isEmpty(); + + // `table_types` is a set that the user wants to include, empty means "include all". + final boolean hasTableTypeTable = + request.getTableTypesCount() == 0 || request.getTableTypesList().contains(TABLE_TYPE_TABLE); + + final boolean includeSchema = request.getIncludeSchema(); + if (hasCatalog || !hasTableTypeTable || request.hasDbSchemaFilterPattern()) { + return getTablesEmpty(includeSchema, CommandGetTablesConstants.ATTRIBUTES); + } + final Predicate tableNameFilter = request.hasTableNameFilterPattern() + ? flightSqlFilterPredicate(request.getTableNameFilterPattern()) + : x -> true; + return getTables(includeSchema, ExecutionContext.getContext().getQueryScope(), + CommandGetTablesConstants.ATTRIBUTES, tableNameFilter); + } + + private Table getTablesEmpty(boolean includeSchema, Map attributes) { + return includeSchema + ? TableTools.newTable(CommandGetTablesConstants.DEFINITION, attributes) + : TableTools.newTable(CommandGetTablesConstants.DEFINITION_NO_SCHEMA, attributes); + } + + private Table getTables(boolean includeSchema, QueryScope queryScope, Map attributes, + Predicate tableNameFilter) { + Objects.requireNonNull(attributes); + final Map queryScopeTables = + (Map) (Map) queryScope.toMap(queryScope::unwrapObject, (n, t) -> t instanceof Table); + final int size = queryScopeTables.size(); + final String[] catalogNames = new String[size]; + final String[] dbSchemaNames = new String[size]; + final String[] tableNames = new String[size]; + final String[] tableTypes = new String[size]; + final Schema[] tableSchemas = includeSchema ? new Schema[size] : null; + int count = 0; + for (Entry e : queryScopeTables.entrySet()) { + final String tableName = e.getKey(); + if (!tableNameFilter.test(tableName)) { + continue; + } + final Schema schema; + if (includeSchema) { + final Table table = authorization.transform(e.getValue()); + if (table == null) { + continue; + } + schema = BarrageUtil.schemaFromTable(table); + } else { + if (authorization.isDeniedAccess(e.getValue())) { + continue; + } + schema = null; + } + catalogNames[count] = null; + dbSchemaNames[count] = null; + tableNames[count] = tableName; + tableTypes[count] = TABLE_TYPE_TABLE; + if (includeSchema) { + tableSchemas[count] = schema; + } + ++count; + } + final ColumnHolder c1 = TableTools.stringCol(CATALOG_NAME, catalogNames); + final ColumnHolder c2 = TableTools.stringCol(DB_SCHEMA_NAME, dbSchemaNames); + final ColumnHolder c3 = TableTools.stringCol(TABLE_NAME, tableNames); + final ColumnHolder c4 = TableTools.stringCol(TABLE_TYPE, tableTypes); + final ColumnHolder c5 = includeSchema + ? new ColumnHolder<>(TABLE_SCHEMA, Schema.class, null, false, tableSchemas) + : null; + final Table newTable = includeSchema + ? TableTools.newTable(CommandGetTablesConstants.DEFINITION, attributes, c1, c2, c3, c4, c5) + : TableTools.newTable(CommandGetTablesConstants.DEFINITION_NO_SCHEMA, attributes, c1, c2, c3, c4); + return count == size + ? newTable + : newTable.head(count); + } + } + + // --------------------------------------------------------------------------------------------------------------- + + private void executeAction( + final SessionState session, + final ActionHandler handler, + final StreamObserver observer) { + // If there was complicated logic going on (actual building of tables), or needed to block, we would instead use + // exports or some other mechanism of doing this work off-thread. For now, it's simple enough that we can do it + // all on this RPC thread. + handler.execute(session, new SafelyOnNextConsumer<>(observer)); + GrpcUtil.safelyComplete(observer); + } + + private class ActionHandlerVisitor + extends FlightSqlActionHelper.ActionVisitorBase> { + @Override + public ActionHandler visit(ActionCreatePreparedStatementRequest action) { + return new CreatePreparedStatementImpl(action); + } + + @Override + public ActionHandler visit(ActionClosePreparedStatementRequest action) { + return new ClosePreparedStatementImpl(action); + } + + @Override + public ActionHandler visitDefault(ActionType actionType, Object action) { + return new UnsupportedAction<>(actionType); + } + } + + private static org.apache.arrow.flight.Result pack(com.google.protobuf.Message message) { + return new org.apache.arrow.flight.Result(Any.pack(message).toByteArray()); + } + + private PreparedStatement getPreparedStatement(SessionState session, ByteString handle) { + Objects.requireNonNull(session); + final PreparedStatement preparedStatement = preparedStatements.get(handle); + if (preparedStatement == null) { + throw error(Code.NOT_FOUND, "Unknown Prepared Statement"); + } + preparedStatement.verifyOwner(session); + return preparedStatement; + } + + interface ActionHandler { + + void execute(SessionState session, Consumer visitor); + } + + static abstract class ActionBase + implements ActionHandler { + + final ActionType type; + final Request request; + + public ActionBase(Request request, ActionType type) { + this.type = Objects.requireNonNull(type); + this.request = Objects.requireNonNull(request); + } + } + + final class CreatePreparedStatementImpl + extends ActionBase { + public CreatePreparedStatementImpl(ActionCreatePreparedStatementRequest request) { + super(request, FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_STATEMENT); + } + + @Override + public void execute( + final SessionState session, + final Consumer visitor) { + if (request.hasTransactionId()) { + throw transactionIdsNotSupported(); + } + // It could be good to parse the query at this point in time to ensure it's valid and _not_ parameterized; + // we will need to dig into Calcite further to explore this possibility. For now, we will error out either + // when the client tries to do a DoPut for the parameter value, or during the Ticket execution, if the query + // is invalid. + final PreparedStatement prepared = new PreparedStatement(session, request.getQuery()); + + // Note: we are providing a fake dataset schema here since the Flight SQL JDBC driver uses the results as an + // indication of whether the query is a SELECT or UPDATE, see + // org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler.PreparedStatement.getType. There should + // likely be some better way the driver could be implemented... + // + // Regardless, the client is not allowed to assume correctness of the returned schema. For example, the + // parameterized query `SELECT ?` is undefined at this point in time. We may need to re-examine this if we + // eventually support non-trivial parameterized queries (it may be necessary to set setParameterSchema, even + // if we don't know exactly what they will be). + // + // There does seem to be some conflicting guidance on whether dataset_schema is actually required or not. + // + // From the FlightSql.proto: + // + // > If a result set generating query was provided, dataset_schema contains the schema of the result set. + // It should be an IPC-encapsulated Schema, as described in Schema.fbs. For some queries, the schema of the + // results may depend on the schema of the parameters. The server should provide its best guess as to the + // schema at this point. Clients must not assume that this schema, if provided, will be accurate. + // + // From https://arrow.apache.org/docs/format/FlightSql.html#query-execution + // + // > The response will contain an opaque handle used to identify the prepared statement. It may also contain + // two optional schemas: the Arrow schema of the result set, and the Arrow schema of the bind parameters (if + // any). Because the schema of the result set may depend on the bind parameters, the schemas may not + // necessarily be provided here as a result, or if provided, they may not be accurate. Clients should not + // assume the schema provided here will be the schema of any data actually returned by executing the + // prepared statement. + // + // > Some statements may have bind parameters without any specific type. (As a trivial example for SQL, + // consider SELECT ?.) It is not currently specified how this should be handled in the bind parameter schema + // above. We suggest either using a union type to enumerate the possible types, or using the NA (null) type + // as a wildcard/placeholder. + final ByteString datasetSchemaBytes; + try { + datasetSchemaBytes = serializeToByteString(DATASET_SCHEMA_SENTINEL); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + final ActionCreatePreparedStatementResult response = ActionCreatePreparedStatementResult.newBuilder() + .setPreparedStatementHandle(prepared.handleId()) + .setDatasetSchema(datasetSchemaBytes) + // .setParameterSchema(...) + .build(); + visitor.accept(response); + } + } + + // Faking it as Empty message so it types check + final class ClosePreparedStatementImpl extends ActionBase { + public ClosePreparedStatementImpl(ActionClosePreparedStatementRequest request) { + super(request, FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT); + } + + @Override + public void execute( + final SessionState session, + final Consumer visitor) { + final PreparedStatement prepared = getPreparedStatement(session, request.getPreparedStatementHandle()); + prepared.close(); + // no responses + } + } + + static final class UnsupportedAction implements ActionHandler { + private final ActionType type; + + public UnsupportedAction(ActionType type) { + this.type = Objects.requireNonNull(type); + } + + @Override + public void execute(SessionState session, Consumer visitor) { + throw error(Code.UNIMPLEMENTED, + String.format("Action type '%s' is unimplemented", type.getType())); + } + } + + private static class SafelyOnNextConsumer implements Consumer { + private final StreamObserver delegate; + + public SafelyOnNextConsumer(StreamObserver delegate) { + this.delegate = Objects.requireNonNull(delegate); + } + + @Override + public void accept(Response response) { + GrpcUtil.safelyOnNext(delegate, pack(response)); + } + } + + // --------------------------------------------------------------------------------------------------------------- + + private static StatusRuntimeException unauthenticatedError() { + return error(Code.UNAUTHENTICATED, "Must be authenticated"); + } + + private static StatusRuntimeException permissionDeniedWithHelpfulMessage() { + return error(Code.PERMISSION_DENIED, + "Must use the original session; is the client echoing the authentication token properly? Some clients may need to explicitly enable cookie-based authentication with the header x-deephaven-auth-cookie-request=true (namely, Java Flight SQL JDBC drivers, and maybe others)."); + } + + private static StatusRuntimeException tableNotFound() { + return error(Code.NOT_FOUND, "table not found"); + } + + private static StatusRuntimeException transactionIdsNotSupported() { + return error(Code.INVALID_ARGUMENT, "transaction ids are not supported"); + } + + private static StatusRuntimeException queryParametersNotSupported(RuntimeException cause) { + return error(Code.INVALID_ARGUMENT, "query parameters are not supported", cause); + } + + /* + * The random number generator used by this class to create random based UUIDs. In a holder class to defer + * initialization until needed. + */ + private static class Holder { + static final SecureRandom SECURE_RANDOM = new SecureRandom(); + } + + private static ByteString randomHandleId() { + // While we don't _rely_ on security through obscurity, we don't want to have a simple incrementing counter + // since it would be trivial to deduce other users' handleIds. + final byte[] handleIdBytes = new byte[16]; + Holder.SECURE_RANDOM.nextBytes(handleIdBytes); + return ByteStringAccess.wrap(handleIdBytes); + } + + private enum ByteStringAsHex implements LogOutput.ObjFormatter { + INSTANCE; + + @Override + public void format(LogOutput logOutput, ByteString bytes) { + logOutput.append("0x").append(ByteHelper.byteBufToHex(bytes.asReadOnlyByteBuffer())); + } + } + + private class PreparedStatement { + private final ByteString handleId; + private final SessionState session; + private final String parameterizedQuery; + private final Set queries; + private final Closeable onSessionClosedCallback; + + PreparedStatement(SessionState session, String parameterizedQuery) { + this.session = Objects.requireNonNull(session); + this.parameterizedQuery = Objects.requireNonNull(parameterizedQuery); + this.handleId = randomHandleId(); + this.queries = new HashSet<>(); + preparedStatements.put(handleId, this); + this.session.addOnCloseCallback(onSessionClosedCallback = this::onSessionClosed); + } + + public ByteString handleId() { + return handleId; + } + + public String parameterizedQuery() { + return parameterizedQuery; + } + + public void verifyOwner(SessionState session) { + if (!this.session.equals(session)) { + // We should not be concerned about returning "NOT_FOUND" here; the handleId is sufficiently random that + // it is much more likely an authentication setup issue. + throw permissionDeniedWithHelpfulMessage(); + } + } + + public synchronized void attach(CommandPreparedStatementQueryImpl query) { + queries.add(query); + } + + public synchronized void detach(CommandPreparedStatementQueryImpl query) { + queries.remove(query); + } + + public void close() { + closeImpl(); + session.removeOnCloseCallback(onSessionClosedCallback); + } + + private void onSessionClosed() { + log.debug() + .append("onSessionClosed: removing prepared statement handleId=") + .append(ByteStringAsHex.INSTANCE, handleId) + .endl(); + closeImpl(); + } + + private synchronized void closeImpl() { + if (!preparedStatements.remove(handleId, this)) { + return; + } + for (CommandPreparedStatementQueryImpl query : queries) { + query.releaseImpl(false); + } + queries.clear(); + } + } + + /** + * The Arrow "specification" for filter pattern leaves a lot to be desired. In totality: + * + * 
+     * In the pattern string, two special characters can be used to denote matching rules:
+     *    - "%" means to match any substring with 0 or more characters.
+     *    - "_" means to match any one character.
+     *
+ * + * There does not seem to be any potential for escaping, which means that underscores can't explicitly be matched + * against, which is a common pattern used in Deephaven table names. As mentioned below, it also follows that an + * empty string should only explicitly match against an empty string. + * + *

+ * The flight-sql-jdbc-core + * implement of sqlToRegexLike uses a similar approach, but appears more fragile as it is doing manual escaping + * of regex as opposed to {@link Pattern#quote(String)}. + */ + @VisibleForTesting + static Predicate flightSqlFilterPredicate(String flightSqlPattern) { + // TODO(deephaven-core#6403): Flight SQL filter pattern improvements + // This is the technically correct, although likely represents a Flight SQL client mis-use, as the results will + // be empty (unless an empty db_schema_name is allowed). + // + // Unlike the "catalog" field in CommandGetDbSchemas (/ CommandGetTables) where an empty string means + // "retrieves those without a catalog", an empty filter pattern does not seem to be meant to match the + // respective field where the value is not present. + // + // The Arrow schema for CommandGetDbSchemas explicitly points out that the returned db_schema_name is not null, + // which implies that filter patterns are not meant to match against fields where the value is not present + // (null). + if (flightSqlPattern.isEmpty()) { + // If Deephaven supports catalog / db_schema_name in the future and db_schema_name can be empty, we'd need + // to match on that. + // return String::isEmpty; + return x -> false; + } + if ("%".equals(flightSqlPattern)) { + return x -> true; + } + if (flightSqlPattern.indexOf('%') == -1 && flightSqlPattern.indexOf('_') == -1) { + // If there are no special characters, search for an exact match; this case was explicitly seen via the + // Flight SQL JDBC driver. + return flightSqlPattern::equals; + } + final StringBuilder pattern = new StringBuilder(); + final StringBuilder quoted = new StringBuilder(); + final Runnable appendQuoted = () -> { + if (quoted.length() != 0) { + pattern.append(Pattern.quote(quoted.toString())); + quoted.setLength(0); + } + }; + try (final IntStream codePoints = flightSqlPattern.codePoints()) { + final PrimitiveIterator.OfInt it = codePoints.iterator(); + while (it.hasNext()) { + final int codePoint = it.nextInt(); + if (Character.isBmpCodePoint(codePoint)) { + final char c = (char) codePoint; + if (c == '%') { + appendQuoted.run(); + pattern.append(".*"); + } else if (c == '_') { + appendQuoted.run(); + pattern.append('.'); + } else { + quoted.append(c); + } + } else { + quoted.appendCodePoint(codePoint); + } + } + } + appendQuoted.run(); + final Pattern p = Pattern.compile(pattern.toString()); + return x -> p.matcher(x).matches(); + } + + private static ByteString serializeToByteString(Schema schema) throws IOException { + final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); + ArrowIpcUtil.serialize(outputStream, schema); + return ByteStringAccess.wrap(outputStream.toByteArray()); + } + + private static Timestamp timestamp(Instant instant) { + return Timestamp.newBuilder() + .setSeconds(instant.getEpochSecond()) + .setNanos(instant.getNano()) + .build(); + } +} diff --git a/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlSharedConstants.java b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlSharedConstants.java new file mode 100644 index 00000000000..f605d068cde --- /dev/null +++ b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlSharedConstants.java @@ -0,0 +1,45 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql; + +final class FlightSqlSharedConstants { + static final String FLIGHT_SQL_TYPE_PREFIX = "type.googleapis.com/arrow.flight.protocol.sql."; + + static final String FLIGHT_SQL_COMMAND_TYPE_PREFIX = FLIGHT_SQL_TYPE_PREFIX + "Command"; + + static final String COMMAND_GET_XDBC_TYPE_INFO_TYPE_URL = FLIGHT_SQL_COMMAND_TYPE_PREFIX + "GetXdbcTypeInfo"; + + static final String COMMAND_GET_PRIMARY_KEYS_TYPE_URL = FLIGHT_SQL_COMMAND_TYPE_PREFIX + "GetPrimaryKeys"; + + static final String COMMAND_GET_IMPORTED_KEYS_TYPE_URL = FLIGHT_SQL_COMMAND_TYPE_PREFIX + "GetImportedKeys"; + + static final String COMMAND_GET_EXPORTED_KEYS_TYPE_URL = FLIGHT_SQL_COMMAND_TYPE_PREFIX + "GetExportedKeys"; + + static final String COMMAND_GET_CROSS_REFERENCE_TYPE_URL = FLIGHT_SQL_COMMAND_TYPE_PREFIX + "GetCrossReference"; + + static final String COMMAND_GET_SQL_INFO_TYPE_URL = FLIGHT_SQL_COMMAND_TYPE_PREFIX + "GetSqlInfo"; + + static final String COMMAND_GET_TABLES_TYPE_URL = FLIGHT_SQL_COMMAND_TYPE_PREFIX + "GetTables"; + + static final String COMMAND_GET_DB_SCHEMAS_TYPE_URL = FLIGHT_SQL_COMMAND_TYPE_PREFIX + "GetDbSchemas"; + + static final String COMMAND_GET_CATALOGS_TYPE_URL = FLIGHT_SQL_COMMAND_TYPE_PREFIX + "GetCatalogs"; + + static final String COMMAND_GET_TABLE_TYPES_TYPE_URL = FLIGHT_SQL_COMMAND_TYPE_PREFIX + "GetTableTypes"; + + static final String COMMAND_PREPARED_STATEMENT_UPDATE_TYPE_URL = + FLIGHT_SQL_COMMAND_TYPE_PREFIX + "PreparedStatementUpdate"; + + static final String COMMAND_PREPARED_STATEMENT_QUERY_TYPE_URL = + FLIGHT_SQL_COMMAND_TYPE_PREFIX + "PreparedStatementQuery"; + + static final String COMMAND_STATEMENT_SUBSTRAIT_PLAN_TYPE_URL = + FLIGHT_SQL_COMMAND_TYPE_PREFIX + "StatementSubstraitPlan"; + + static final String COMMAND_STATEMENT_UPDATE_TYPE_URL = FLIGHT_SQL_COMMAND_TYPE_PREFIX + "StatementUpdate"; + + static final String COMMAND_STATEMENT_QUERY_TYPE_URL = FLIGHT_SQL_COMMAND_TYPE_PREFIX + "StatementQuery"; + + static final String COMMAND_STATEMENT_INGEST_TYPE_URL = FLIGHT_SQL_COMMAND_TYPE_PREFIX + "StatementIngest"; +} diff --git a/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlTicketHelper.java b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlTicketHelper.java new file mode 100644 index 00000000000..77c4681cf39 --- /dev/null +++ b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlTicketHelper.java @@ -0,0 +1,178 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql; + +import com.google.protobuf.Any; +import com.google.protobuf.ByteString; +import com.google.protobuf.InvalidProtocolBufferException; +import com.google.protobuf.Message; +import io.deephaven.base.verify.Assert; +import io.deephaven.util.annotations.VisibleForTesting; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import org.apache.arrow.flight.impl.Flight.Ticket; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCatalogs; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetDbSchemas; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetExportedKeys; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetImportedKeys; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetPrimaryKeys; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTableTypes; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTables; +import org.apache.arrow.flight.sql.impl.FlightSql.TicketStatementQuery; + +import java.nio.ByteBuffer; + +import static io.deephaven.server.flightsql.FlightSqlSharedConstants.COMMAND_GET_CATALOGS_TYPE_URL; +import static io.deephaven.server.flightsql.FlightSqlSharedConstants.COMMAND_GET_DB_SCHEMAS_TYPE_URL; +import static io.deephaven.server.flightsql.FlightSqlSharedConstants.COMMAND_GET_EXPORTED_KEYS_TYPE_URL; +import static io.deephaven.server.flightsql.FlightSqlSharedConstants.COMMAND_GET_IMPORTED_KEYS_TYPE_URL; +import static io.deephaven.server.flightsql.FlightSqlSharedConstants.COMMAND_GET_PRIMARY_KEYS_TYPE_URL; +import static io.deephaven.server.flightsql.FlightSqlSharedConstants.COMMAND_GET_TABLES_TYPE_URL; +import static io.deephaven.server.flightsql.FlightSqlSharedConstants.COMMAND_GET_TABLE_TYPES_TYPE_URL; + +final class FlightSqlTicketHelper { + + public static final char TICKET_PREFIX = 'q'; + + // This is a server-implementation detail, but happens to be the same scheme that Flight SQL + // org.apache.arrow.flight.sql.FlightSqlProducer uses + @VisibleForTesting + static final String TICKET_STATEMENT_QUERY_TYPE_URL = + FlightSqlSharedConstants.FLIGHT_SQL_TYPE_PREFIX + "TicketStatementQuery"; + + private static final ByteString PREFIX = ByteString.copyFrom(new byte[] {(byte) TICKET_PREFIX}); + + interface TicketVisitor { + + // These ticket objects could be anything we want; they don't _have_ to be the protobuf objects. But, they are + // convenient as they already contain the information needed to act on the ticket. + + T visit(CommandGetCatalogs ticket); + + T visit(CommandGetDbSchemas ticket); + + T visit(CommandGetTableTypes ticket); + + T visit(CommandGetImportedKeys ticket); + + T visit(CommandGetExportedKeys ticket); + + T visit(CommandGetPrimaryKeys ticket); + + T visit(CommandGetTables ticket); + + T visit(TicketStatementQuery ticket); + } + + public static TicketVisitor ticketCreator() { + return TicketCreator.INSTANCE; + } + + public static String toReadableString(final ByteBuffer ticket, final String logId) { + final Any any = partialUnpackTicket(ticket, logId); + // We don't necessarily want to print out the full protobuf; this will at least give some more logging info on + // the type of the ticket. + return any.getTypeUrl(); + } + + public static T visit(ByteBuffer ticket, TicketVisitor visitor, String logId) { + return visit(partialUnpackTicket(ticket, logId), visitor, logId); + } + + private static Any partialUnpackTicket(ByteBuffer ticket, final String logId) { + ticket = ticket.slice(); + // If false, it means there is an error with FlightSqlResolver.ticketRoute / TicketRouter.getResolver + Assert.eq(ticket.get(), "ticket.get()", TICKET_PREFIX, "TICKET_PREFIX"); + try { + return Any.parseFrom(ticket); + } catch (InvalidProtocolBufferException e) { + throw invalidTicket(logId); + } + } + + private static T visit(Any ticket, TicketVisitor visitor, String logId) { + switch (ticket.getTypeUrl()) { + case TICKET_STATEMENT_QUERY_TYPE_URL: + return visitor.visit(unpack(ticket, TicketStatementQuery.class, logId)); + case COMMAND_GET_TABLES_TYPE_URL: + return visitor.visit(unpack(ticket, CommandGetTables.class, logId)); + case COMMAND_GET_TABLE_TYPES_TYPE_URL: + return visitor.visit(unpack(ticket, CommandGetTableTypes.class, logId)); + case COMMAND_GET_CATALOGS_TYPE_URL: + return visitor.visit(unpack(ticket, CommandGetCatalogs.class, logId)); + case COMMAND_GET_DB_SCHEMAS_TYPE_URL: + return visitor.visit(unpack(ticket, CommandGetDbSchemas.class, logId)); + case COMMAND_GET_PRIMARY_KEYS_TYPE_URL: + return visitor.visit(unpack(ticket, CommandGetPrimaryKeys.class, logId)); + case COMMAND_GET_IMPORTED_KEYS_TYPE_URL: + return visitor.visit(unpack(ticket, CommandGetImportedKeys.class, logId)); + case COMMAND_GET_EXPORTED_KEYS_TYPE_URL: + return visitor.visit(unpack(ticket, CommandGetExportedKeys.class, logId)); + } + throw invalidTicket(logId); + } + + private enum TicketCreator implements TicketVisitor { + INSTANCE; + + @Override + public Ticket visit(CommandGetCatalogs ticket) { + return packedTicket(ticket); + } + + @Override + public Ticket visit(CommandGetDbSchemas ticket) { + return packedTicket(ticket); + } + + @Override + public Ticket visit(CommandGetTableTypes ticket) { + return packedTicket(ticket); + } + + @Override + public Ticket visit(CommandGetImportedKeys ticket) { + return packedTicket(ticket); + } + + @Override + public Ticket visit(CommandGetExportedKeys ticket) { + return packedTicket(ticket); + } + + @Override + public Ticket visit(CommandGetPrimaryKeys ticket) { + return packedTicket(ticket); + } + + @Override + public Ticket visit(CommandGetTables ticket) { + return packedTicket(ticket); + } + + @Override + public Ticket visit(TicketStatementQuery ticket) { + return packedTicket(ticket); + } + + private static Ticket packedTicket(Message message) { + // Note: this is _similar_ to how the Flight SQL example server implementation constructs tickets using + // Any#pack; the big difference is that all DH tickets must (currently) be uniquely route-able based on the + // first byte of the ticket. + return Ticket.newBuilder().setTicket(PREFIX.concat(Any.pack(message).toByteString())).build(); + } + } + + private static StatusRuntimeException invalidTicket(String logId) { + return FlightSqlErrorHelper.error(Status.Code.FAILED_PRECONDITION, String.format("Invalid ticket, %s", logId)); + } + + private static T unpack(Any ticket, Class clazz, String logId) { + try { + return ticket.unpack(clazz); + } catch (InvalidProtocolBufferException e) { + throw invalidTicket(logId); + } + } +} diff --git a/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/TableCreatorScopeTickets.java b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/TableCreatorScopeTickets.java new file mode 100644 index 00000000000..c254d5ad99a --- /dev/null +++ b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/TableCreatorScopeTickets.java @@ -0,0 +1,33 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql; + +import io.deephaven.engine.table.Table; +import io.deephaven.qst.TableCreator; +import io.deephaven.qst.TableCreatorDelegate; +import io.deephaven.qst.table.TicketTable; +import io.deephaven.server.console.ScopeTicketResolver; +import io.deephaven.server.session.SessionState; + +import java.nio.ByteBuffer; +import java.util.Objects; + +final class TableCreatorScopeTickets extends TableCreatorDelegate

{ + + private final ScopeTicketResolver scopeTicketResolver; + private final SessionState session; + + TableCreatorScopeTickets(TableCreator
delegate, ScopeTicketResolver scopeTicketResolver, + SessionState session) { + super(delegate); + this.scopeTicketResolver = Objects.requireNonNull(scopeTicketResolver); + this.session = session; + } + + @Override + public Table of(TicketTable ticketTable) { + return scopeTicketResolver.
resolve(session, ByteBuffer.wrap(ticketTable.ticket()), + TableCreatorScopeTickets.class.getSimpleName()).get(); + } +} diff --git a/extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlFilterPredicateTest.java b/extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlFilterPredicateTest.java new file mode 100644 index 00000000000..9ac5c4221a0 --- /dev/null +++ b/extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlFilterPredicateTest.java @@ -0,0 +1,179 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql; + +import org.assertj.core.api.PredicateAssert; +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +class FlightSqlFilterPredicateTest { + + private static final List EMPTY_STRING = List.of(""); + + private static final List ONE_CHARS = List.of(" ", "X", "%", "_", ".", "*", "\uD83D\uDCA9"); + + private static final List TWO_CHARS = + List.of("ab", "Cd", "F ", " f", "_-", " ", "\uD83D\uDCA9\uD83D\uDCA9"); + + private static final List THREE_CHARS = + List.of("abc", "Cde", "F ", " f", " v ", "_-_", " ", "\uD83D\uDCA9\uD83D\uDCA9\uD83D\uDCA9"); + + @Test + void rejectAll() { + predicate("") + .rejectsAll(EMPTY_STRING) + .rejectsAll(ONE_CHARS) + .rejectsAll(TWO_CHARS) + .rejectsAll(THREE_CHARS); + } + + @Test + void acceptAll() { + for (String flightSqlPattern : new String[] {"%", "%%", "%%%", "%%%%"}) { + predicate(flightSqlPattern) + .acceptsAll(EMPTY_STRING) + .acceptsAll(ONE_CHARS) + .acceptsAll(TWO_CHARS) + .acceptsAll(THREE_CHARS); + } + } + + @Test + void acceptsAnyOneChar() { + predicate("_") + .rejectsAll(EMPTY_STRING) + .acceptsAll(ONE_CHARS) + .rejectsAll(TWO_CHARS) + .rejectsAll(THREE_CHARS); + } + + @Test + void acceptsOnePlusChar() { + for (String flightSqlPattern : new String[] {"_%", "%_"}) { + predicate(flightSqlPattern) + .rejectsAll(EMPTY_STRING) + .acceptsAll(ONE_CHARS) + .acceptsAll(TWO_CHARS) + .acceptsAll(THREE_CHARS); + } + } + + @Test + void acceptsTwoPlusChar() { + for (String flightSqlPattern : new String[] {"__%", "%__", "_%_"}) { + predicate(flightSqlPattern) + .rejectsAll(EMPTY_STRING) + .rejectsAll(ONE_CHARS) + .acceptsAll(TWO_CHARS) + .acceptsAll(THREE_CHARS); + } + } + + @Test + void acceptLiteralString() { + predicate("Foo") + .accepts("Foo") + .rejects("Bar") + .rejectsAll(EMPTY_STRING) + .rejectsAll(ONE_CHARS) + .rejectsAll(TWO_CHARS) + .rejectsAll(THREE_CHARS); + + } + + @Test + void acceptUndescoreAsAnyOne() { + predicate("foo_ball") + .accepts("foo_ball", "foosball", "foodball", "foo\uD83D\uDCA9ball") + .rejects("foo__ball", "Foo_ball", "foo_all", "foo\uD83D\uDCA9\uD83D\uDCA9ball") + .rejectsAll(EMPTY_STRING) + .rejectsAll(ONE_CHARS) + .rejectsAll(TWO_CHARS) + .rejectsAll(THREE_CHARS); + } + + @Test + void acceptUndescoreAsOnePlus() { + predicate("foo%ball") + .accepts("foo_ball", "foosball", "foodball", "foo\uD83D\uDCA9ball", "foo__ball", + "foo\uD83D\uDCA9\uD83D\uDCA9ball") + .rejects("Foo_ball", "foo_all") + .rejectsAll(EMPTY_STRING) + .rejectsAll(ONE_CHARS) + .rejectsAll(TWO_CHARS) + .rejectsAll(THREE_CHARS); + } + + @Disabled("No way to match literal underscore") + @Test + void matchLiteralUnderscore() { + + } + + @Disabled("No way to match literal percentage") + @Test + void matchLiteralPercentage() { + + } + + @Test + void plusIsNotSpecial() { + predicate("A+") + .accepts("A+") + .rejects("A", "AA", "A ") + .rejectsAll(EMPTY_STRING) + .rejectsAll(ONE_CHARS) + .rejectsAll(TWO_CHARS) + .rejectsAll(THREE_CHARS); + } + + @Test + void starIsNotSpecial() { + predicate("A*") + .accepts("A*") + .rejects("A", "AA", "A ", "AAA") + .rejectsAll(EMPTY_STRING) + .rejectsAll(ONE_CHARS) + .rejectsAll(TWO_CHARS) + .rejectsAll(THREE_CHARS); + } + + @Test + void dotstarIsNotSpecial() { + predicate(".*") + .accepts(".*") + .rejects("A", "AA", "A ", "AAA") + .rejectsAll(EMPTY_STRING) + .rejectsAll(ONE_CHARS) + .rejectsAll(TWO_CHARS) + .rejectsAll(THREE_CHARS); + } + + @Test + void predicateContainsUnicode() { + // A better test would be to include a Unicode character that contains '_' or '%' encoded as part of the low + // surrogate (if this is possible), that way we could ensure that it is not treated as a special character + predicate("𞸷") + .accepts("𞸷") + .rejectsAll(EMPTY_STRING) + .rejectsAll(ONE_CHARS) + .rejectsAll(TWO_CHARS) + .rejectsAll(THREE_CHARS); + predicate("_𞸷_") + .accepts("𞸷𞸷𞸷", "x𞸷X", "T𞸷a", " 𞸷 ") + .rejects("𞸷", "𞸷𞸷") + .rejectsAll(EMPTY_STRING) + .rejectsAll(ONE_CHARS) + .rejectsAll(TWO_CHARS) + .rejectsAll(THREE_CHARS); + } + + private static PredicateAssert predicate(String flightSqlPattern) { + return assertThat(FlightSqlResolver.flightSqlFilterPredicate(flightSqlPattern)); + } +} diff --git a/extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlTest.java b/extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlTest.java new file mode 100644 index 00000000000..d20d8d50c9e --- /dev/null +++ b/extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlTest.java @@ -0,0 +1,1095 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql; + +import com.google.protobuf.Any; +import com.google.protobuf.ByteString; +import com.google.protobuf.Descriptors.Descriptor; +import com.google.protobuf.Message; +import dagger.BindsInstance; +import dagger.Component; +import dagger.Module; +import io.deephaven.engine.context.ExecutionContext; +import io.deephaven.engine.table.ColumnDefinition; +import io.deephaven.engine.table.Table; +import io.deephaven.engine.table.TableDefinition; +import io.deephaven.engine.util.TableTools; +import io.deephaven.proto.backplane.grpc.WrappedAuthenticationRequest; +import io.deephaven.server.auth.AuthorizationProvider; +import io.deephaven.server.config.ServerConfig; +import io.deephaven.server.runner.DeephavenApiServerTestBase; +import io.deephaven.server.runner.DeephavenApiServerTestBase.TestComponent.Builder; +import io.grpc.ManagedChannel; +import org.apache.arrow.flight.Action; +import org.apache.arrow.flight.ActionType; +import org.apache.arrow.flight.CancelFlightInfoRequest; +import org.apache.arrow.flight.Criteria; +import org.apache.arrow.flight.FlightClient; +import org.apache.arrow.flight.FlightConstants; +import org.apache.arrow.flight.FlightDescriptor; +import org.apache.arrow.flight.FlightEndpoint; +import org.apache.arrow.flight.FlightGrpcUtilsExtension; +import org.apache.arrow.flight.FlightInfo; +import org.apache.arrow.flight.FlightRuntimeException; +import org.apache.arrow.flight.FlightStatusCode; +import org.apache.arrow.flight.FlightStream; +import org.apache.arrow.flight.ProtocolExposer; +import org.apache.arrow.flight.Result; +import org.apache.arrow.flight.SchemaResult; +import org.apache.arrow.flight.Ticket; +import org.apache.arrow.flight.auth.ClientAuthHandler; +import org.apache.arrow.flight.impl.Flight; +import org.apache.arrow.flight.sql.FlightSqlClient; +import org.apache.arrow.flight.sql.FlightSqlClient.PreparedStatement; +import org.apache.arrow.flight.sql.FlightSqlClient.Savepoint; +import org.apache.arrow.flight.sql.FlightSqlClient.SubstraitPlan; +import org.apache.arrow.flight.sql.FlightSqlClient.Transaction; +import org.apache.arrow.flight.sql.FlightSqlUtils; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionBeginSavepointRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionBeginTransactionRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionCancelQueryRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionClosePreparedStatementRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedSubstraitPlanRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionEndSavepointRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.ActionEndTransactionRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCatalogs; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCrossReference; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetDbSchemas; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetExportedKeys; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetImportedKeys; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetPrimaryKeys; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetSqlInfo; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTableTypes; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTables; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetXdbcTypeInfo; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementQuery; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementQuery; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementSubstraitPlan; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementUpdate; +import org.apache.arrow.flight.sql.util.TableRef; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.types.Types.MinorType; +import org.apache.arrow.vector.types.pojo.ArrowType.Utf8; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +import javax.inject.Named; +import javax.inject.Singleton; +import java.io.PrintStream; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.Iterator; +import java.util.List; +import java.util.Map; + +import static io.deephaven.server.flightsql.FlightSqlTicketHelper.TICKET_PREFIX; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.failBecauseExceptionWasNotThrown; + +// using JUnit4 so we can inherit properly from DeephavenApiServerTestBase +@RunWith(JUnit4.class) +public class FlightSqlTest extends DeephavenApiServerTestBase { + + private static final Map DEEPHAVEN_STRING = Map.of( + "deephaven:isSortable", "true", + "deephaven:isRowStyle", "false", + "deephaven:isPartitioning", "false", + "deephaven:type", "java.lang.String", + "deephaven:isNumberFormat", "false", + "deephaven:isStyle", "false", + "deephaven:isDateFormat", "false"); + + private static final Map DEEPHAVEN_INT = Map.of( + "deephaven:isSortable", "true", + "deephaven:isRowStyle", "false", + "deephaven:isPartitioning", "false", + "deephaven:type", "int", + "deephaven:isNumberFormat", "false", + "deephaven:isStyle", "false", + "deephaven:isDateFormat", "false"); + + private static final Map DEEPHAVEN_BYTE = Map.of( + "deephaven:isSortable", "true", + "deephaven:isRowStyle", "false", + "deephaven:isPartitioning", "false", + "deephaven:type", "byte", + "deephaven:isNumberFormat", "false", + "deephaven:isStyle", "false", + "deephaven:isDateFormat", "false"); + + private static final Map DEEPHAVEN_SCHEMA = Map.of( + "deephaven:isSortable", "false", + "deephaven:isRowStyle", "false", + "deephaven:isPartitioning", "false", + "deephaven:type", "org.apache.arrow.vector.types.pojo.Schema", + "deephaven:isNumberFormat", "false", + "deephaven:isStyle", "false", + "deephaven:isDateFormat", "false"); + + private static final Map FLAT_ATTRIBUTES = Map.of( + "deephaven:attribute_type.IsFlat", "java.lang.Boolean", + "deephaven:attribute.IsFlat", "true"); + + private static final Field CATALOG_NAME = + new Field("catalog_name", new FieldType(true, Utf8.INSTANCE, null, DEEPHAVEN_STRING), null); + + private static final Field PK_CATALOG_NAME = + new Field("pk_catalog_name", new FieldType(true, Utf8.INSTANCE, null, DEEPHAVEN_STRING), null); + + private static final Field FK_CATALOG_NAME = + new Field("fk_catalog_name", new FieldType(true, Utf8.INSTANCE, null, DEEPHAVEN_STRING), null); + + private static final Field DB_SCHEMA_NAME = + new Field("db_schema_name", new FieldType(true, Utf8.INSTANCE, null, DEEPHAVEN_STRING), null); + + private static final Field PK_DB_SCHEMA_NAME = + new Field("pk_db_schema_name", new FieldType(true, Utf8.INSTANCE, null, DEEPHAVEN_STRING), null); + + private static final Field FK_DB_SCHEMA_NAME = + new Field("fk_db_schema_name", new FieldType(true, Utf8.INSTANCE, null, DEEPHAVEN_STRING), null); + + private static final Field TABLE_NAME = + new Field("table_name", new FieldType(true, Utf8.INSTANCE, null, DEEPHAVEN_STRING), null); + + private static final Field COLUMN_NAME = + new Field("column_name", new FieldType(true, Utf8.INSTANCE, null, DEEPHAVEN_STRING), null); + + private static final Field KEY_NAME = + new Field("key_name", new FieldType(true, Utf8.INSTANCE, null, DEEPHAVEN_STRING), null); + + private static final Field PK_TABLE_NAME = + new Field("pk_table_name", new FieldType(true, Utf8.INSTANCE, null, DEEPHAVEN_STRING), null); + + private static final Field FK_TABLE_NAME = + new Field("fk_table_name", new FieldType(true, Utf8.INSTANCE, null, DEEPHAVEN_STRING), null); + + private static final Field TABLE_TYPE = + new Field("table_type", new FieldType(true, Utf8.INSTANCE, null, DEEPHAVEN_STRING), null); + + private static final Field PK_COLUMN_NAME = + new Field("pk_column_name", new FieldType(true, Utf8.INSTANCE, null, DEEPHAVEN_STRING), null); + + private static final Field FK_COLUMN_NAME = + new Field("fk_column_name", new FieldType(true, Utf8.INSTANCE, null, DEEPHAVEN_STRING), null); + + private static final Field KEY_SEQUENCE = + new Field("key_sequence", new FieldType(true, MinorType.INT.getType(), null, DEEPHAVEN_INT), null); + + private static final Field PK_KEY_NAME = + new Field("pk_key_name", new FieldType(true, Utf8.INSTANCE, null, DEEPHAVEN_STRING), null); + + private static final Field FK_KEY_NAME = + new Field("fk_key_name", new FieldType(true, Utf8.INSTANCE, null, DEEPHAVEN_STRING), null); + + private static final Field UPDATE_RULE = + new Field("update_rule", new FieldType(true, MinorType.TINYINT.getType(), null, DEEPHAVEN_BYTE), null); + + private static final Field DELETE_RULE = + new Field("delete_rule", new FieldType(true, MinorType.TINYINT.getType(), null, DEEPHAVEN_BYTE), null); + + private static final Field TABLE_SCHEMA = + new Field("table_schema", new FieldType(true, MinorType.VARBINARY.getType(), null, DEEPHAVEN_SCHEMA), null); + + private static final TableRef FOO_TABLE_REF = TableRef.of(null, null, "foo_table"); + public static final TableRef BAR_TABLE_REF = TableRef.of(null, null, "barTable"); + + @Module(includes = { + TestModule.class, + FlightSqlModule.class, + }) + public interface MyModule { + + } + + @Singleton + @Component(modules = MyModule.class) + public interface MyComponent extends TestComponent { + + @Component.Builder + interface Builder extends TestComponent.Builder { + + @BindsInstance + Builder withServerConfig(ServerConfig serverConfig); + + @BindsInstance + Builder withOut(@Named("out") PrintStream out); + + @BindsInstance + Builder withErr(@Named("err") PrintStream err); + + @BindsInstance + Builder withAuthorizationProvider(AuthorizationProvider authorizationProvider); + + MyComponent build(); + } + } + + BufferAllocator bufferAllocator; + FlightClient flightClient; + FlightSqlClient flightSqlClient; + + @Override + protected Builder testComponentBuilder() { + return DaggerFlightSqlTest_MyComponent.builder(); + } + + @Before + @Override + public void setUp() throws Exception { + super.setUp(); + ManagedChannel channel = channelBuilder().build(); + register(channel); + bufferAllocator = new RootAllocator(); + // Note: this pattern of FlightClient owning the ManagedChannel does not mesh well with the idea that some + // other entity may be managing the authentication lifecycle. We'd prefer to pass in the stubs or "intercepted" + // channel directly, but that's not supported. So, we need to create the specific middleware interfaces so + // flight can do its own shims. + flightClient = FlightGrpcUtilsExtension.createFlightClientWithSharedChannel(bufferAllocator, channel, + new ArrayList<>()); + // Note: this is not extensible, at least not with Auth v2 / JDBC. + flightClient.authenticate(new ClientAuthHandler() { + private byte[] callToken = new byte[0]; + + @Override + public void authenticate(ClientAuthSender outgoing, Iterator incoming) { + WrappedAuthenticationRequest request = WrappedAuthenticationRequest.newBuilder() + .setType("Anonymous") + .setPayload(ByteString.EMPTY) + .build(); + outgoing.send(request.toByteArray()); + callToken = incoming.next(); + } + + @Override + public byte[] getCallToken() { + return callToken; + } + }); + flightSqlClient = new FlightSqlClient(flightClient); + } + + @Override + public void tearDown() throws Exception { + // this also closes flightClient + flightSqlClient.close(); + bufferAllocator.close(); + super.tearDown(); + } + + @Test + public void listFlights() { + assertThat(flightClient.listFlights(Criteria.ALL)).isEmpty(); + } + + @Test + public void listActions() { + assertThat(flightClient.listActions()) + .usingElementComparator(Comparator.comparing(ActionType::getType)) + .containsExactlyInAnyOrder( + FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_STATEMENT, + FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT); + } + + @Test + public void getCatalogs() throws Exception { + final Schema expectedSchema = flatTableSchema(CATALOG_NAME); + { + final SchemaResult schemaResult = flightSqlClient.getCatalogsSchema(); + assertThat(schemaResult.getSchema()).isEqualTo(expectedSchema); + } + { + final FlightInfo info = flightSqlClient.getCatalogs(); + assertThat(info.getSchema()).isEqualTo(expectedSchema); + consume(info, 0, 0, true); + } + unpackable(CommandGetCatalogs.getDescriptor(), CommandGetCatalogs.class); + } + + @Test + public void getSchemas() throws Exception { + final Schema expectedSchema = flatTableSchema(CATALOG_NAME, DB_SCHEMA_NAME); + { + final SchemaResult schemasSchema = flightSqlClient.getSchemasSchema(); + assertThat(schemasSchema.getSchema()).isEqualTo(expectedSchema); + } + for (final FlightInfo info : new FlightInfo[] { + flightSqlClient.getSchemas(null, null), + flightSqlClient.getSchemas("DoesNotExist", null), + flightSqlClient.getSchemas(null, ""), + flightSqlClient.getSchemas(null, "%"), + flightSqlClient.getSchemas(null, "SomeSchema"), + }) { + assertThat(info.getSchema()).isEqualTo(expectedSchema); + consume(info, 0, 0, true); + } + unpackable(CommandGetDbSchemas.getDescriptor(), CommandGetDbSchemas.class); + } + + @Test + public void getTables() throws Exception { + setFooTable(); + setFoodTable(); + setBarTable(); + for (final boolean includeSchema : new boolean[] {false, true}) { + final Schema expectedSchema = includeSchema + ? flatTableSchema(CATALOG_NAME, DB_SCHEMA_NAME, TABLE_NAME, TABLE_TYPE, TABLE_SCHEMA) + : flatTableSchema(CATALOG_NAME, DB_SCHEMA_NAME, TABLE_NAME, TABLE_TYPE); + { + final SchemaResult schema = flightSqlClient.getTablesSchema(includeSchema); + assertThat(schema.getSchema()).isEqualTo(expectedSchema); + } + // Any of these queries will fetch everything from query scope + for (final FlightInfo info : new FlightInfo[] { + flightSqlClient.getTables(null, null, null, null, includeSchema), + flightSqlClient.getTables("", null, null, null, includeSchema), + flightSqlClient.getTables(null, null, null, List.of("TABLE"), includeSchema), + flightSqlClient.getTables(null, null, null, List.of("IRRELEVANT_TYPE", "TABLE"), includeSchema), + flightSqlClient.getTables(null, null, "%", null, includeSchema), + flightSqlClient.getTables(null, null, "%able", null, includeSchema), + }) { + assertThat(info.getSchema()).isEqualTo(expectedSchema); + consume(info, 1, 3, true); + } + + // Any of these queries will fetch foo_table and foodtable; there is no way to uniquely filter based on the + // `_` literal + for (final FlightInfo info : new FlightInfo[] { + flightSqlClient.getTables(null, null, "foo_table", null, includeSchema), + flightSqlClient.getTables(null, null, "foo_%", null, includeSchema), + flightSqlClient.getTables(null, null, "f%", null, includeSchema), + flightSqlClient.getTables(null, null, "%table", null, includeSchema), + }) { + assertThat(info.getSchema()).isEqualTo(expectedSchema); + consume(info, 1, 2, true); + } + + // Any of these queries will fetch foodtable + for (final FlightInfo info : new FlightInfo[] { + flightSqlClient.getTables(null, null, "foodtable", null, includeSchema), + }) { + assertThat(info.getSchema()).isEqualTo(expectedSchema); + consume(info, 1, 1, true); + } + + // Any of these queries will fetch barTable + for (final FlightInfo info : new FlightInfo[] { + flightSqlClient.getTables(null, null, "barTable", null, includeSchema), + flightSqlClient.getTables(null, null, "bar%", null, includeSchema), + flightSqlClient.getTables(null, null, "b%", null, includeSchema), + flightSqlClient.getTables(null, null, "%Table", null, includeSchema), + }) { + assertThat(info.getSchema()).isEqualTo(expectedSchema); + consume(info, 1, 1, true); + } + + // Any of these queries will fetch an empty table + for (final FlightInfo info : new FlightInfo[] { + flightSqlClient.getTables("DoesNotExistCatalog", null, null, null, includeSchema), + flightSqlClient.getTables(null, null, null, List.of("IRRELEVANT_TYPE"), includeSchema), + flightSqlClient.getTables(null, "", null, null, includeSchema), + flightSqlClient.getTables(null, "%", null, null, includeSchema), + flightSqlClient.getTables(null, null, "", null, includeSchema), + flightSqlClient.getTables(null, null, "doesNotExist", null, includeSchema), + flightSqlClient.getTables(null, null, "%_table2", null, includeSchema), + }) { + assertThat(info.getSchema()).isEqualTo(expectedSchema); + consume(info, 0, 0, true); + } + } + unpackable(CommandGetTables.getDescriptor(), CommandGetTables.class); + } + + @Test + public void getTableTypes() throws Exception { + final Schema expectedSchema = flatTableSchema(TABLE_TYPE); + { + final SchemaResult schema = flightSqlClient.getTableTypesSchema(); + assertThat(schema.getSchema()).isEqualTo(expectedSchema); + } + { + final FlightInfo info = flightSqlClient.getTableTypes(); + assertThat(info.getSchema()).isEqualTo(expectedSchema); + consume(info, 1, 1, true); + } + unpackable(CommandGetTableTypes.getDescriptor(), CommandGetTableTypes.class); + } + + @Test + public void select1() throws Exception { + final Schema expectedSchema = new Schema( + List.of(new Field("Foo", new FieldType(true, MinorType.INT.getType(), null, DEEPHAVEN_INT), null))); + { + final SchemaResult schema = flightSqlClient.getExecuteSchema("SELECT 1 as Foo"); + assertThat(schema.getSchema()).isEqualTo(expectedSchema); + } + { + final FlightInfo info = flightSqlClient.execute("SELECT 1 as Foo"); + assertThat(info.getSchema()).isEqualTo(expectedSchema); + consume(info, 1, 1, false); + } + unpackable(CommandStatementQuery.getDescriptor(), CommandStatementQuery.class); + } + + @Test + public void select1Prepared() throws Exception { + final Schema expectedSchema = new Schema( + List.of(new Field("Foo", new FieldType(true, MinorType.INT.getType(), null, DEEPHAVEN_INT), null))); + try (final PreparedStatement prepared = flightSqlClient.prepare("SELECT 1 as Foo")) { + assertThat(prepared.getResultSetSchema()).isEqualTo(FlightSqlResolver.DATASET_SCHEMA_SENTINEL); + { + final SchemaResult schema = prepared.fetchSchema(); + assertThat(schema.getSchema()).isEqualTo(expectedSchema); + } + { + final FlightInfo info = prepared.execute(); + assertThat(info.getSchema()).isEqualTo(expectedSchema); + consume(info, 1, 1, false); + } + unpackable(CommandPreparedStatementQuery.getDescriptor(), CommandPreparedStatementQuery.class); + } + } + + @Test + public void selectStarFromQueryScopeTable() throws Exception { + setFooTable(); + + final Schema expectedSchema = flatTableSchema( + new Field("Foo", new FieldType(true, MinorType.INT.getType(), null, DEEPHAVEN_INT), null)); + { + final SchemaResult schema = flightSqlClient.getExecuteSchema("SELECT * FROM foo_table"); + assertThat(schema.getSchema()).isEqualTo(expectedSchema); + } + { + final FlightInfo info = flightSqlClient.execute("SELECT * FROM foo_table"); + assertThat(info.getSchema()).isEqualTo(expectedSchema); + consume(info, 1, 3, false); + } + // The Flight SQL resolver will maintain state to ensure results are resolvable, even if the underlying table + // goes away between flightInfo and doGet. + { + final FlightInfo info = flightSqlClient.execute("SELECT * FROM foo_table"); + assertThat(info.getSchema()).isEqualTo(expectedSchema); + removeFooTable(); + consume(info, 1, 3, false); + } + unpackable(CommandStatementQuery.getDescriptor(), CommandStatementQuery.class); + + } + + @Test + public void selectStarPreparedFromQueryScopeTable() throws Exception { + setFooTable(); + { + final Schema expectedSchema = flatTableSchema( + new Field("Foo", new FieldType(true, MinorType.INT.getType(), null, DEEPHAVEN_INT), null)); + try (final PreparedStatement prepared = flightSqlClient.prepare("SELECT * FROM foo_table")) { + assertThat(prepared.getResultSetSchema()).isEqualTo(FlightSqlResolver.DATASET_SCHEMA_SENTINEL); + { + final SchemaResult schema = prepared.fetchSchema(); + assertThat(schema.getSchema()).isEqualTo(expectedSchema); + } + { + final FlightInfo info = prepared.execute(); + assertThat(info.getSchema()).isEqualTo(expectedSchema); + consume(info, 1, 3, false); + } + // The Flight SQL resolver will maintain state to ensure results are resolvable, even if the underlying + // table goes away between flightInfo and doGet. + { + final FlightInfo info = prepared.execute(); + assertThat(info.getSchema()).isEqualTo(expectedSchema); + removeFooTable(); + consume(info, 1, 3, false); + } + // The states in _not_ maintained by the PreparedStatement state though, and will not be available for + // the next execute + expectException(prepared::execute, FlightStatusCode.NOT_FOUND, "Object 'foo_table' not found"); + unpackable(CommandPreparedStatementQuery.getDescriptor(), CommandPreparedStatementQuery.class); + } + unpackable(FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_STATEMENT, ActionCreatePreparedStatementRequest.class); + unpackable(FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT, ActionClosePreparedStatementRequest.class); + } + } + + @Test + public void preparedStatementIsLazy() throws Exception { + try (final PreparedStatement prepared = flightSqlClient.prepare("SELECT * FROM foo_table")) { + assertThat(prepared.getResultSetSchema()).isEqualTo(FlightSqlResolver.DATASET_SCHEMA_SENTINEL); + expectException(prepared::fetchSchema, FlightStatusCode.NOT_FOUND, "Object 'foo_table' not found"); + expectException(prepared::execute, FlightStatusCode.NOT_FOUND, "Object 'foo_table' not found"); + // If the state-of-the-world changes, this will be reflected in new calls against the prepared statement. + // This also implies that we won't error out at the start of prepare call if the table doesn't exist. + // + // We could introduce some sort of reference-based Transactional model (orthogonal to a snapshot-based + // Transactional model) which would ensure schema consistency, but that would be an effort outside of a + // PreparedStatement. + setFooTable(); + final Schema expectedSchema = flatTableSchema( + new Field("Foo", new FieldType(true, MinorType.INT.getType(), null, DEEPHAVEN_INT), null)); + { + final SchemaResult schema = prepared.fetchSchema(); + assertThat(schema.getSchema()).isEqualTo(expectedSchema); + } + { + final FlightInfo info = prepared.execute(); + assertThat(info.getSchema()).isEqualTo(expectedSchema); + consume(info, 1, 3, false); + } + } + } + + @Test + public void selectQuestionMark() { + queryError("SELECT ?", FlightStatusCode.INVALID_ARGUMENT, "Illegal use of dynamic parameter"); + } + + @Test + public void selectFooParam() { + setFooTable(); + queryError("SELECT Foo FROM foo_table WHERE Foo = ?", FlightStatusCode.INVALID_ARGUMENT, + "Flight SQL: query parameters are not supported"); + } + + @Test + public void selectTableDoesNotExist() { + queryError("SELECT * FROM my_table", FlightStatusCode.NOT_FOUND, "Object 'my_table' not found"); + } + + @Test + public void selectColumnDoesNotExist() { + setFooTable(); + queryError("SELECT BadColumn FROM foo_table", FlightStatusCode.NOT_FOUND, + "Column 'BadColumn' not found in any table"); + } + + @Test + public void selectFunctionDoesNotExist() { + setFooTable(); + queryError("SELECT my_function(Foo) FROM foo_table", FlightStatusCode.INVALID_ARGUMENT, + "No match found for function signature"); + } + + @Test + public void badSqlQuery() { + queryError("this is not SQL", FlightStatusCode.INVALID_ARGUMENT, "Flight SQL: query can't be parsed"); + } + + @Test + public void executeSubstrait() { + getSchemaUnimplemented(() -> flightSqlClient.getExecuteSubstraitSchema(fakePlan()), + CommandStatementSubstraitPlan.getDescriptor()); + commandUnimplemented(() -> flightSqlClient.executeSubstrait(fakePlan()), + CommandStatementSubstraitPlan.getDescriptor()); + misbehave(CommandStatementSubstraitPlan.getDefaultInstance(), CommandStatementSubstraitPlan.getDescriptor()); + unpackable(CommandStatementSubstraitPlan.getDescriptor(), CommandStatementSubstraitPlan.class); + } + + @Test + public void executeSubstraitUpdate() { + // Note: this is the same descriptor as the executeSubstrait + getSchemaUnimplemented(() -> flightSqlClient.getExecuteSubstraitSchema(fakePlan()), + CommandStatementSubstraitPlan.getDescriptor()); + expectUnpublishable(() -> flightSqlClient.executeSubstraitUpdate(fakePlan())); + unpackable(CommandStatementSubstraitPlan.getDescriptor(), CommandStatementSubstraitPlan.class); + } + + @Test + public void insert1() { + expectUnpublishable(() -> flightSqlClient.executeUpdate("INSERT INTO fake(name) VALUES('Smith')")); + unpackable(CommandStatementUpdate.getDescriptor(), CommandStatementUpdate.class); + } + + private void queryError(String query, FlightStatusCode expectedCode, String expectedMessage) { + expectException(() -> flightSqlClient.getExecuteSchema(query), expectedCode, expectedMessage); + expectException(() -> flightSqlClient.execute(query), expectedCode, expectedMessage); + try (final PreparedStatement prepared = flightSqlClient.prepare(query)) { + assertThat(prepared.getResultSetSchema()).isEqualTo(FlightSqlResolver.DATASET_SCHEMA_SENTINEL); + expectException(prepared::fetchSchema, expectedCode, expectedMessage); + expectException(prepared::execute, expectedCode, expectedMessage); + } + } + + @Test + public void insertPrepared() { + setFooTable(); + try (final PreparedStatement prepared = flightSqlClient.prepare("INSERT INTO foo_table(Foo) VALUES(42)")) { + expectException(prepared::fetchSchema, FlightStatusCode.INVALID_ARGUMENT, + "Flight SQL: Unsupported calcite type 'org.apache.calcite.rel.logical.LogicalTableModify'"); + expectException(prepared::execute, FlightStatusCode.INVALID_ARGUMENT, + "Flight SQL: Unsupported calcite type 'org.apache.calcite.rel.logical.LogicalTableModify'"); + } + try (final PreparedStatement prepared = flightSqlClient.prepare("INSERT INTO foo_table(MyArg) VALUES(42)")) { + expectException(prepared::fetchSchema, FlightStatusCode.INVALID_ARGUMENT, + "Flight SQL: Unknown target column 'MyArg'"); + expectException(prepared::execute, FlightStatusCode.INVALID_ARGUMENT, + "Flight SQL: Unknown target column 'MyArg'"); + } + try (final PreparedStatement prepared = flightSqlClient.prepare("INSERT INTO x(Foo) VALUES(42)")) { + expectException(prepared::fetchSchema, FlightStatusCode.NOT_FOUND, "Flight SQL: Object 'x' not found"); + expectException(prepared::execute, FlightStatusCode.NOT_FOUND, "Flight SQL: Object 'x' not found"); + } + } + + @Test + public void getSqlInfo() { + getSchemaUnimplemented(() -> flightSqlClient.getSqlInfoSchema(), CommandGetSqlInfo.getDescriptor()); + commandUnimplemented(() -> flightSqlClient.getSqlInfo(), CommandGetSqlInfo.getDescriptor()); + misbehave(CommandGetSqlInfo.getDefaultInstance(), CommandGetSqlInfo.getDescriptor()); + unpackable(CommandGetSqlInfo.getDescriptor(), CommandGetSqlInfo.class); + } + + @Test + public void getXdbcTypeInfo() { + getSchemaUnimplemented(() -> flightSqlClient.getXdbcTypeInfoSchema(), CommandGetXdbcTypeInfo.getDescriptor()); + commandUnimplemented(() -> flightSqlClient.getXdbcTypeInfo(), CommandGetXdbcTypeInfo.getDescriptor()); + misbehave(CommandGetXdbcTypeInfo.getDefaultInstance(), CommandGetXdbcTypeInfo.getDescriptor()); + unpackable(CommandGetXdbcTypeInfo.getDescriptor(), CommandGetXdbcTypeInfo.class); + } + + @Test + public void getCrossReference() { + setFooTable(); + setBarTable(); + getSchemaUnimplemented(() -> flightSqlClient.getCrossReferenceSchema(), + CommandGetCrossReference.getDescriptor()); + commandUnimplemented(() -> flightSqlClient.getCrossReference(FOO_TABLE_REF, BAR_TABLE_REF), + CommandGetCrossReference.getDescriptor()); + misbehave(CommandGetCrossReference.getDefaultInstance(), CommandGetCrossReference.getDescriptor()); + unpackable(CommandGetCrossReference.getDescriptor(), CommandGetCrossReference.class); + } + + @Test + public void getPrimaryKeys() throws Exception { + setFooTable(); + final Schema expectedSchema = + flatTableSchema(CATALOG_NAME, DB_SCHEMA_NAME, TABLE_NAME, COLUMN_NAME, KEY_NAME, KEY_SEQUENCE); + { + final SchemaResult exportedKeysSchema = flightSqlClient.getPrimaryKeysSchema(); + assertThat(exportedKeysSchema.getSchema()).isEqualTo(expectedSchema); + } + { + final FlightInfo info = flightSqlClient.getPrimaryKeys(FOO_TABLE_REF); + assertThat(info.getSchema()).isEqualTo(expectedSchema); + consume(info, 0, 0, true); + } + // Note: the info must remain valid even if the server state goes away. + { + final FlightInfo info = flightSqlClient.getPrimaryKeys(FOO_TABLE_REF); + assertThat(info.getSchema()).isEqualTo(expectedSchema); + removeFooTable(); + // resolve should still be OK + consume(info, 0, 0, true); + } + expectException(() -> flightSqlClient.getPrimaryKeys(BAR_TABLE_REF), FlightStatusCode.NOT_FOUND, + "Flight SQL: table not found"); + + // Note: misbehaving clients who fudge tickets directly will not get errors; but they will also not learn any + // information on whether the tables actually exist or not since the returned table is always empty. + for (final CommandGetPrimaryKeys command : new CommandGetPrimaryKeys[] { + CommandGetPrimaryKeys.newBuilder().setTable("DoesNotExist").build(), + CommandGetPrimaryKeys.newBuilder().setCatalog("Catalog").setDbSchema("DbSchema") + .setTable("DoesNotExist").build() + }) { + final Ticket ticket = + ProtocolExposer.fromProtocol(FlightSqlTicketHelper.ticketCreator().visit(command)); + try (final FlightStream stream = flightSqlClient.getStream(ticket)) { + consume(stream, 0, 0); + } + } + unpackable(CommandGetPrimaryKeys.getDescriptor(), CommandGetPrimaryKeys.class); + } + + @Test + public void getExportedKeys() throws Exception { + setFooTable(); + final Schema expectedSchema = flatTableSchema(PK_CATALOG_NAME, PK_DB_SCHEMA_NAME, PK_TABLE_NAME, PK_COLUMN_NAME, + FK_CATALOG_NAME, FK_DB_SCHEMA_NAME, FK_TABLE_NAME, FK_COLUMN_NAME, KEY_SEQUENCE, FK_KEY_NAME, + PK_KEY_NAME, UPDATE_RULE, DELETE_RULE); + { + final SchemaResult exportedKeysSchema = flightSqlClient.getExportedKeysSchema(); + assertThat(exportedKeysSchema.getSchema()).isEqualTo(expectedSchema); + } + { + final FlightInfo info = flightSqlClient.getExportedKeys(FOO_TABLE_REF); + assertThat(info.getSchema()).isEqualTo(expectedSchema); + consume(info, 0, 0, true); + } + // Note: the info must remain valid even if the server state goes away. + { + final FlightInfo info = flightSqlClient.getExportedKeys(FOO_TABLE_REF); + assertThat(info.getSchema()).isEqualTo(expectedSchema); + removeFooTable(); + // resolve should still be OK + consume(info, 0, 0, true); + } + expectException(() -> flightSqlClient.getExportedKeys(BAR_TABLE_REF), FlightStatusCode.NOT_FOUND, + "Flight SQL: table not found"); + + // Note: misbehaving clients who fudge tickets directly will not get errors; but they will also not learn any + // information on whether the tables actually exist or not since the returned table is always empty. + for (final CommandGetExportedKeys command : new CommandGetExportedKeys[] { + CommandGetExportedKeys.newBuilder().setTable("DoesNotExist").build(), + CommandGetExportedKeys.newBuilder().setCatalog("Catalog").setDbSchema("DbSchema") + .setTable("DoesNotExist").build() + }) { + final Ticket ticket = + ProtocolExposer.fromProtocol(FlightSqlTicketHelper.ticketCreator().visit(command)); + try (final FlightStream stream = flightSqlClient.getStream(ticket)) { + consume(stream, 0, 0); + } + } + unpackable(CommandGetExportedKeys.getDescriptor(), CommandGetExportedKeys.class); + } + + @Test + public void getImportedKeys() throws Exception { + setFooTable(); + final Schema expectedSchema = flatTableSchema(PK_CATALOG_NAME, PK_DB_SCHEMA_NAME, PK_TABLE_NAME, PK_COLUMN_NAME, + FK_CATALOG_NAME, FK_DB_SCHEMA_NAME, FK_TABLE_NAME, FK_COLUMN_NAME, KEY_SEQUENCE, FK_KEY_NAME, + PK_KEY_NAME, UPDATE_RULE, DELETE_RULE); + { + final SchemaResult importedKeysSchema = flightSqlClient.getImportedKeysSchema(); + assertThat(importedKeysSchema.getSchema()).isEqualTo(expectedSchema); + } + { + final FlightInfo info = flightSqlClient.getImportedKeys(FOO_TABLE_REF); + assertThat(info.getSchema()).isEqualTo(expectedSchema); + consume(info, 0, 0, true); + } + // Note: the info must remain valid even if the server state goes away. + { + final FlightInfo info = flightSqlClient.getImportedKeys(FOO_TABLE_REF); + assertThat(info.getSchema()).isEqualTo(expectedSchema); + removeFooTable(); + // resolve should still be OK + consume(info, 0, 0, true); + } + + expectException(() -> flightSqlClient.getImportedKeys(BAR_TABLE_REF), FlightStatusCode.NOT_FOUND, + "Flight SQL: table not found"); + + // Note: misbehaving clients who fudge tickets directly will not get errors; but they will also not learn any + // information on whether the tables actually exist or not since the returned table is always empty. + for (final CommandGetImportedKeys command : new CommandGetImportedKeys[] { + CommandGetImportedKeys.newBuilder().setTable("DoesNotExist").build(), + CommandGetImportedKeys.newBuilder().setCatalog("Catalog").setDbSchema("DbSchema") + .setTable("DoesNotExist").build() + }) { + final Ticket ticket = + ProtocolExposer.fromProtocol(FlightSqlTicketHelper.ticketCreator().visit(command)); + try (final FlightStream stream = flightSqlClient.getStream(ticket)) { + consume(stream, 0, 0); + } + } + unpackable(CommandGetImportedKeys.getDescriptor(), CommandGetImportedKeys.class); + } + + @Test + public void commandStatementIngest() { + // Note: if we actually want to test out FlightSqlClient behavior more directly, we can set up scaffolding to + // call FlightSqlClient#executeIngest. + final FlightDescriptor ingestCommand = + FlightDescriptor.command(Any.pack(CommandStatementIngest.getDefaultInstance()).toByteArray()); + getSchemaUnimplemented(() -> flightClient.getSchema(ingestCommand), CommandStatementIngest.getDescriptor()); + commandUnimplemented(() -> flightClient.getInfo(ingestCommand), CommandStatementIngest.getDescriptor()); + misbehave(CommandStatementIngest.getDefaultInstance(), CommandStatementIngest.getDescriptor()); + unpackable(CommandStatementIngest.getDescriptor(), CommandStatementIngest.class); + } + + @Test + public void unknownCommandLooksLikeFlightSql() { + final String typeUrl = "type.googleapis.com/arrow.flight.protocol.sql.CommandLooksRealButDoesNotExist"; + final FlightDescriptor descriptor = unpackableCommand(typeUrl); + getSchemaUnknown(() -> flightClient.getSchema(descriptor), typeUrl); + commandUnknown(() -> flightClient.getInfo(descriptor), typeUrl); + } + + @Test + public void unknownCommand() { + // Note: this should likely be tested in the context of Flight, not Flight SQL + final String typeUrl = "type.googleapis.com/com.example.SomeRandomCommand"; + final FlightDescriptor descriptor = unpackableCommand(typeUrl); + expectException(() -> flightClient.getSchema(descriptor), FlightStatusCode.INVALID_ARGUMENT, + "no resolver for command"); + expectException(() -> flightClient.getInfo(descriptor), FlightStatusCode.INVALID_ARGUMENT, + "no resolver for command"); + } + + @Test + public void prepareSubstrait() { + actionUnimplemented(() -> flightSqlClient.prepare(fakePlan()), + FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_SUBSTRAIT_PLAN); + unpackable(FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_SUBSTRAIT_PLAN, + ActionCreatePreparedSubstraitPlanRequest.class); + } + + @Test + public void beginTransaction() { + actionUnimplemented(() -> flightSqlClient.beginTransaction(), FlightSqlUtils.FLIGHT_SQL_BEGIN_TRANSACTION); + unpackable(FlightSqlUtils.FLIGHT_SQL_BEGIN_TRANSACTION, ActionBeginTransactionRequest.class); + } + + @Test + public void commit() { + actionUnimplemented(() -> flightSqlClient.commit(fakeTxn()), FlightSqlUtils.FLIGHT_SQL_END_TRANSACTION); + unpackable(FlightSqlUtils.FLIGHT_SQL_END_TRANSACTION, ActionEndTransactionRequest.class); + } + + @Test + public void rollbackTxn() { + actionUnimplemented(() -> flightSqlClient.rollback(fakeTxn()), FlightSqlUtils.FLIGHT_SQL_END_TRANSACTION); + unpackable(FlightSqlUtils.FLIGHT_SQL_END_TRANSACTION, ActionEndTransactionRequest.class); + } + + @Test + public void beginSavepoint() { + actionUnimplemented(() -> flightSqlClient.beginSavepoint(fakeTxn(), "fakeName"), + FlightSqlUtils.FLIGHT_SQL_BEGIN_SAVEPOINT); + unpackable(FlightSqlUtils.FLIGHT_SQL_BEGIN_SAVEPOINT, ActionBeginSavepointRequest.class); + } + + @Test + public void release() { + actionUnimplemented(() -> flightSqlClient.release(fakeSavepoint()), FlightSqlUtils.FLIGHT_SQL_END_SAVEPOINT); + unpackable(FlightSqlUtils.FLIGHT_SQL_END_SAVEPOINT, ActionEndSavepointRequest.class); + } + + @Test + public void rollbackSavepoint() { + actionUnimplemented(() -> flightSqlClient.rollback(fakeSavepoint()), FlightSqlUtils.FLIGHT_SQL_END_SAVEPOINT); + unpackable(FlightSqlUtils.FLIGHT_SQL_END_SAVEPOINT, ActionEndSavepointRequest.class); + } + + @Test + public void cancelQuery() { + final FlightInfo info = flightSqlClient.execute("SELECT 1"); + actionUnimplemented(() -> flightSqlClient.cancelQuery(info), FlightSqlUtils.FLIGHT_SQL_CANCEL_QUERY); + unpackable(FlightSqlUtils.FLIGHT_SQL_CANCEL_QUERY, ActionCancelQueryRequest.class); + } + + @Test + public void cancelFlightInfo() { + // Note: this should likely be tested in the context of Flight, not Flight SQL + final FlightInfo info = flightSqlClient.execute("SELECT 1"); + actionNoResolver(() -> flightClient.cancelFlightInfo(new CancelFlightInfoRequest(info)), + FlightConstants.CANCEL_FLIGHT_INFO.getType()); + } + + @Test + public void unknownAction() { + // Note: this should likely be tested in the context of Flight, not Flight SQL + final String type = "SomeFakeAction"; + final Action action = new Action(type, new byte[0]); + actionNoResolver(() -> doAction(action), type); + } + + private Result doAction(Action action) { + final Iterator it = flightClient.doAction(action); + if (!it.hasNext()) { + throw new IllegalStateException(); + } + final Result result = it.next(); + if (it.hasNext()) { + throw new IllegalStateException(); + } + return result; + } + + private void misbehave(Message message, Descriptor descriptor) { + final Ticket ticket = ProtocolExposer.fromProtocol(Flight.Ticket.newBuilder() + .setTicket( + ByteString.copyFrom(new byte[] {(byte) TICKET_PREFIX}).concat(Any.pack(message).toByteString())) + .build()); + expectException(() -> flightSqlClient.getStream(ticket).next(), FlightStatusCode.INVALID_ARGUMENT, + "Flight SQL: Invalid ticket"); + } + + private static FlightDescriptor unpackableCommand(Descriptor descriptor) { + return unpackableCommand("type.googleapis.com/" + descriptor.getFullName()); + } + + private static FlightDescriptor unpackableCommand(String typeUrl) { + return FlightDescriptor.command( + Any.newBuilder().setTypeUrl(typeUrl).setValue(ByteString.copyFrom(new byte[1])).build().toByteArray()); + } + + private void getSchemaUnimplemented(Runnable r, Descriptor command) { + // right now our server impl routes all getSchema through their respective commands + commandUnimplemented(r, command); + } + + private void commandUnimplemented(Runnable r, Descriptor command) { + expectException(r, FlightStatusCode.UNIMPLEMENTED, + String.format("Flight SQL: command '%s' is unimplemented", command.getFullName())); + } + + private void getSchemaUnknown(Runnable r, String command) { + // right now our server impl routes all getSchema through their respective commands + commandUnknown(r, command); + } + + private void commandUnknown(Runnable r, String command) { + expectException(r, FlightStatusCode.UNIMPLEMENTED, + String.format("Flight SQL: command '%s' is unknown", command)); + } + + private void unpackable(Descriptor descriptor, Class clazz) { + final FlightDescriptor flightDescriptor = unpackableCommand(descriptor); + getSchemaUnpackable(() -> flightClient.getSchema(flightDescriptor), clazz); + commandUnpackable(() -> flightClient.getInfo(flightDescriptor), clazz); + } + + + private void unpackable(ActionType type, Class actionProto) { + { + final Action action = new Action(type.getType(), Any.getDefaultInstance().toByteArray()); + expectException(() -> doAction(action), FlightStatusCode.INVALID_ARGUMENT, String.format( + "Flight SQL: Invalid action, provided message cannot be unpacked as %s", actionProto.getName())); + } + { + final Action action = new Action(type.getType(), new byte[] {-1}); + expectException(() -> doAction(action), FlightStatusCode.INVALID_ARGUMENT, "Flight SQL: Invalid action"); + } + } + + private void getSchemaUnpackable(Runnable r, Class clazz) { + commandUnpackable(r, clazz); + } + + private void commandUnpackable(Runnable r, Class clazz) { + expectUnpackableCommand(r, clazz); + } + + private void expectUnpackableCommand(Runnable r, Class clazz) { + expectException(r, FlightStatusCode.INVALID_ARGUMENT, + String.format("Flight SQL: Invalid command, provided message cannot be unpacked as %s", + clazz.getName())); + } + + private void expectUnpublishable(Runnable r) { + expectException(r, FlightStatusCode.INVALID_ARGUMENT, "Flight SQL descriptors cannot be published to"); + } + + private void actionUnimplemented(Runnable r, ActionType actionType) { + expectException(r, FlightStatusCode.UNIMPLEMENTED, + String.format("Flight SQL: Action type '%s' is unimplemented", actionType.getType())); + } + + private void actionNoResolver(Runnable r, String actionType) { + expectException(r, FlightStatusCode.UNIMPLEMENTED, + String.format("No action resolver found for action type '%s'", actionType)); + } + + private static void expectException(Runnable r, FlightStatusCode code, String messagePart) { + try { + r.run(); + failBecauseExceptionWasNotThrown(FlightRuntimeException.class); + } catch (FlightRuntimeException e) { + assertThat(e.status().code()).isEqualTo(code); + assertThat(e).hasMessageContaining(messagePart); + } + } + + private static FlightEndpoint endpoint(FlightInfo info) { + assertThat(info.getEndpoints()).hasSize(1); + return info.getEndpoints().get(0); + } + + private static Schema flatTableSchema(Field... fields) { + return new Schema(List.of(fields), FLAT_ATTRIBUTES); + } + + private static void setFooTable() { + setSimpleTable("foo_table", "Foo"); + } + + private static void setFoodTable() { + setSimpleTable("foodtable", "Food"); + } + + private static void setBarTable() { + setSimpleTable("barTable", "Bar"); + } + + private static void removeFooTable() { + removeTable("foo_table"); + } + + private static void removeFoodTable() { + removeTable("foodtable"); + } + + private static void removeBarTable() { + removeTable("barTable"); + } + + private static void setSimpleTable(String tableName, String columnName) { + final TableDefinition td = TableDefinition.of(ColumnDefinition.ofInt(columnName)); + final Table table = TableTools.newTable(td, TableTools.intCol(columnName, 1, 2, 3)); + ExecutionContext.getContext().getQueryScope().putParam(tableName, table); + } + + private static void removeTable(String tableName) { + ExecutionContext.getContext().getQueryScope().putParam(tableName, null); + } + + private void consume(FlightInfo info, int expectedFlightCount, int expectedNumRows, boolean expectReusable) + throws Exception { + final FlightEndpoint endpoint = endpoint(info); + if (expectReusable) { + assertThat(endpoint.getExpirationTime()).isPresent(); + } else { + assertThat(endpoint.getExpirationTime()).isEmpty(); + } + try (final FlightStream stream = flightSqlClient.getStream(endpoint.getTicket())) { + consume(stream, expectedFlightCount, expectedNumRows); + } + if (expectReusable) { + try (final FlightStream stream = flightSqlClient.getStream(endpoint.getTicket())) { + consume(stream, expectedFlightCount, expectedNumRows); + } + } else { + try (final FlightStream stream = flightSqlClient.getStream(endpoint.getTicket())) { + consumeNotFound(stream); + } + } + } + + private static void consume(FlightStream stream, int expectedFlightCount, int expectedNumRows) { + int numRows = 0; + int flightCount = 0; + while (stream.next()) { + ++flightCount; + numRows += stream.getRoot().getRowCount(); + } + assertThat(flightCount).isEqualTo(expectedFlightCount); + assertThat(numRows).isEqualTo(expectedNumRows); + } + + private static void consumeNotFound(FlightStream stream) { + expectException(stream::next, FlightStatusCode.NOT_FOUND, + "Unable to find Flight SQL query. Flight SQL tickets should be resolved promptly and resolved at most once."); + } + + private static SubstraitPlan fakePlan() { + return new SubstraitPlan("fake".getBytes(StandardCharsets.UTF_8), "1"); + } + + private static Transaction fakeTxn() { + return new Transaction("fake".getBytes(StandardCharsets.UTF_8)); + } + + private static Savepoint fakeSavepoint() { + return new Savepoint("fake".getBytes(StandardCharsets.UTF_8)); + } +} diff --git a/extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlTicketResolverTest.java b/extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlTicketResolverTest.java new file mode 100644 index 00000000000..abb7dca0d5a --- /dev/null +++ b/extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlTicketResolverTest.java @@ -0,0 +1,170 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql; + +import com.google.protobuf.Any; +import com.google.protobuf.Message; +import io.deephaven.engine.table.TableDefinition; +import io.deephaven.extensions.barrage.util.BarrageUtil; +import io.deephaven.server.flightsql.FlightSqlResolver.CommandGetCatalogsConstants; +import io.deephaven.server.flightsql.FlightSqlResolver.CommandGetDbSchemasConstants; +import io.deephaven.server.flightsql.FlightSqlResolver.CommandGetKeysConstants; +import io.deephaven.server.flightsql.FlightSqlResolver.CommandGetPrimaryKeysConstants; +import io.deephaven.server.flightsql.FlightSqlResolver.CommandGetTableTypesConstants; +import io.deephaven.server.flightsql.FlightSqlResolver.CommandGetTablesConstants; +import org.apache.arrow.flight.ActionType; +import org.apache.arrow.flight.sql.FlightSqlProducer.Schemas; +import org.apache.arrow.flight.sql.FlightSqlUtils; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCatalogs; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetCrossReference; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetDbSchemas; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetExportedKeys; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetImportedKeys; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetPrimaryKeys; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetSqlInfo; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTableTypes; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTables; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetXdbcTypeInfo; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementQuery; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementUpdate; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementQuery; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementSubstraitPlan; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementUpdate; +import org.apache.arrow.flight.sql.impl.FlightSql.TicketStatementQuery; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.Test; + +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +public class FlightSqlTicketResolverTest { + @Test + public void actionTypes() { + checkActionType(FlightSqlActionHelper.CREATE_PREPARED_STATEMENT_ACTION_TYPE, + FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_STATEMENT); + checkActionType(FlightSqlActionHelper.CLOSE_PREPARED_STATEMENT_ACTION_TYPE, + FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT); + checkActionType(FlightSqlActionHelper.BEGIN_SAVEPOINT_ACTION_TYPE, FlightSqlUtils.FLIGHT_SQL_BEGIN_SAVEPOINT); + checkActionType(FlightSqlActionHelper.END_SAVEPOINT_ACTION_TYPE, FlightSqlUtils.FLIGHT_SQL_END_SAVEPOINT); + checkActionType(FlightSqlActionHelper.BEGIN_TRANSACTION_ACTION_TYPE, + FlightSqlUtils.FLIGHT_SQL_BEGIN_TRANSACTION); + checkActionType(FlightSqlActionHelper.END_TRANSACTION_ACTION_TYPE, FlightSqlUtils.FLIGHT_SQL_END_TRANSACTION); + checkActionType(FlightSqlActionHelper.CANCEL_QUERY_ACTION_TYPE, FlightSqlUtils.FLIGHT_SQL_CANCEL_QUERY); + checkActionType(FlightSqlActionHelper.CREATE_PREPARED_SUBSTRAIT_PLAN_ACTION_TYPE, + FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_SUBSTRAIT_PLAN); + } + + @Test + public void packedTypeUrls() { + checkPackedType(FlightSqlSharedConstants.COMMAND_STATEMENT_QUERY_TYPE_URL, + CommandStatementQuery.getDefaultInstance()); + checkPackedType(FlightSqlSharedConstants.COMMAND_STATEMENT_UPDATE_TYPE_URL, + CommandStatementUpdate.getDefaultInstance()); + // Need to update to newer FlightSql version for this + // checkPackedType(FlightSqlTicketResolver.COMMAND_STATEMENT_INGEST_TYPE_URL, + // CommandStatementIngest.getDefaultInstance()); + checkPackedType(FlightSqlSharedConstants.COMMAND_STATEMENT_SUBSTRAIT_PLAN_TYPE_URL, + CommandStatementSubstraitPlan.getDefaultInstance()); + checkPackedType(FlightSqlSharedConstants.COMMAND_PREPARED_STATEMENT_QUERY_TYPE_URL, + CommandPreparedStatementQuery.getDefaultInstance()); + checkPackedType(FlightSqlSharedConstants.COMMAND_PREPARED_STATEMENT_UPDATE_TYPE_URL, + CommandPreparedStatementUpdate.getDefaultInstance()); + checkPackedType(FlightSqlSharedConstants.COMMAND_GET_TABLE_TYPES_TYPE_URL, + CommandGetTableTypes.getDefaultInstance()); + checkPackedType(FlightSqlSharedConstants.COMMAND_GET_CATALOGS_TYPE_URL, + CommandGetCatalogs.getDefaultInstance()); + checkPackedType(FlightSqlSharedConstants.COMMAND_GET_DB_SCHEMAS_TYPE_URL, + CommandGetDbSchemas.getDefaultInstance()); + checkPackedType(FlightSqlSharedConstants.COMMAND_GET_TABLES_TYPE_URL, + CommandGetTables.getDefaultInstance()); + checkPackedType(FlightSqlSharedConstants.COMMAND_GET_SQL_INFO_TYPE_URL, + CommandGetSqlInfo.getDefaultInstance()); + checkPackedType(FlightSqlSharedConstants.COMMAND_GET_CROSS_REFERENCE_TYPE_URL, + CommandGetCrossReference.getDefaultInstance()); + checkPackedType(FlightSqlSharedConstants.COMMAND_GET_EXPORTED_KEYS_TYPE_URL, + CommandGetExportedKeys.getDefaultInstance()); + checkPackedType(FlightSqlSharedConstants.COMMAND_GET_IMPORTED_KEYS_TYPE_URL, + CommandGetImportedKeys.getDefaultInstance()); + checkPackedType(FlightSqlSharedConstants.COMMAND_GET_PRIMARY_KEYS_TYPE_URL, + CommandGetPrimaryKeys.getDefaultInstance()); + checkPackedType(FlightSqlSharedConstants.COMMAND_GET_XDBC_TYPE_INFO_TYPE_URL, + CommandGetXdbcTypeInfo.getDefaultInstance()); + checkPackedType(FlightSqlTicketHelper.TICKET_STATEMENT_QUERY_TYPE_URL, + TicketStatementQuery.getDefaultInstance()); + } + + @Test + void getTableTypesSchema() { + isSimilar(CommandGetTableTypesConstants.DEFINITION, Schemas.GET_TABLE_TYPES_SCHEMA); + } + + @Test + void getCatalogsSchema() { + isSimilar(CommandGetCatalogsConstants.DEFINITION, Schemas.GET_CATALOGS_SCHEMA); + } + + @Test + void getDbSchemasSchema() { + isSimilar(CommandGetDbSchemasConstants.DEFINITION, Schemas.GET_SCHEMAS_SCHEMA); + } + + @Disabled("Deephaven is unable to serialize byte as uint8") + @Test + void getImportedKeysSchema() { + isSimilar(CommandGetKeysConstants.DEFINITION, Schemas.GET_IMPORTED_KEYS_SCHEMA); + } + + @Disabled("Deephaven is unable to serialize byte as uint8") + @Test + void getExportedKeysSchema() { + isSimilar(CommandGetKeysConstants.DEFINITION, Schemas.GET_EXPORTED_KEYS_SCHEMA); + } + + @Disabled("Arrow Java Flight SQL has a bug in ordering, not the same as documented in the protobuf spec, see https://github.com/apache/arrow/issues/44521") + @Test + void getPrimaryKeysSchema() { + isSimilar(CommandGetPrimaryKeysConstants.DEFINITION, Schemas.GET_PRIMARY_KEYS_SCHEMA); + } + + @Test + void getTablesSchema() { + isSimilar(CommandGetTablesConstants.DEFINITION, Schemas.GET_TABLES_SCHEMA); + isSimilar(CommandGetTablesConstants.DEFINITION_NO_SCHEMA, Schemas.GET_TABLES_SCHEMA_NO_SCHEMA); + } + + private static void checkActionType(String actionType, ActionType expected) { + assertThat(actionType).isEqualTo(expected.getType()); + } + + private static void checkPackedType(String typeUrl, Message expected) { + assertThat(typeUrl).isEqualTo(Any.pack(expected).getTypeUrl()); + } + + private static void isSimilar(TableDefinition definition, Schema expected) { + isSimilar(BarrageUtil.toSchema(definition, Map.of(), true), expected); + } + + private static void isSimilar(Schema actual, Schema expected) { + assertThat(actual.getFields()).hasSameSizeAs(expected.getFields()); + int L = actual.getFields().size(); + for (int i = 0; i < L; ++i) { + isSimilar(actual.getFields().get(i), expected.getFields().get(i)); + } + } + + private static void isSimilar(Field actual, Field expected) { + assertThat(actual.getName()).isEqualTo(expected.getName()); + assertThat(actual.getChildren()).isEqualTo(expected.getChildren()); + isSimilar(actual.getFieldType(), expected.getFieldType()); + } + + private static void isSimilar(FieldType actual, FieldType expected) { + assertThat(actual.getType()).isEqualTo(expected.getType()); + assertThat(actual.getDictionary()).isEqualTo(expected.getDictionary()); + } +} diff --git a/extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlUnauthenticatedTest.java b/extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlUnauthenticatedTest.java new file mode 100644 index 00000000000..35ddf3e4ce3 --- /dev/null +++ b/extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlUnauthenticatedTest.java @@ -0,0 +1,350 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql; + +import com.google.protobuf.Any; +import com.google.protobuf.ByteString; +import dagger.BindsInstance; +import dagger.Component; +import dagger.Module; +import io.deephaven.server.auth.AuthorizationProvider; +import io.deephaven.server.config.ServerConfig; +import io.deephaven.server.runner.DeephavenApiServerTestBase; +import io.deephaven.server.runner.DeephavenApiServerTestBase.TestComponent.Builder; +import io.grpc.ManagedChannel; +import org.apache.arrow.flight.*; +import org.apache.arrow.flight.sql.FlightSqlClient; +import org.apache.arrow.flight.sql.FlightSqlClient.Savepoint; +import org.apache.arrow.flight.sql.FlightSqlClient.SubstraitPlan; +import org.apache.arrow.flight.sql.FlightSqlClient.Transaction; +import org.apache.arrow.flight.sql.util.TableRef; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +import javax.inject.Named; +import javax.inject.Singleton; +import java.io.PrintStream; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Iterator; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.failBecauseExceptionWasNotThrown; + +// using JUnit4 so we can inherit properly from DeephavenApiServerTestBase +@RunWith(JUnit4.class) +public class FlightSqlUnauthenticatedTest extends DeephavenApiServerTestBase { + + private static final TableRef FOO_TABLE_REF = TableRef.of(null, null, "foo_table"); + public static final TableRef BAR_TABLE_REF = TableRef.of(null, null, "bar_table"); + + @Module(includes = { + TestModule.class, + FlightSqlModule.class, + }) + public interface MyModule { + + } + + @Singleton + @Component(modules = MyModule.class) + public interface MyComponent extends TestComponent { + + @Component.Builder + interface Builder extends TestComponent.Builder { + + @BindsInstance + Builder withServerConfig(ServerConfig serverConfig); + + @BindsInstance + Builder withOut(@Named("out") PrintStream out); + + @BindsInstance + Builder withErr(@Named("err") PrintStream err); + + @BindsInstance + Builder withAuthorizationProvider(AuthorizationProvider authorizationProvider); + + MyComponent build(); + } + } + + private BufferAllocator bufferAllocator; + private FlightClient flightClient; + private FlightSqlClient flightSqlClient; + + @Override + protected Builder testComponentBuilder() { + return DaggerFlightSqlTest_MyComponent.builder(); + } + + @Before + @Override + public void setUp() throws Exception { + super.setUp(); + ManagedChannel channel = channelBuilder().build(); + register(channel); + bufferAllocator = new RootAllocator(); + // Note: this pattern of FlightClient owning the ManagedChannel does not mesh well with the idea that some + // other entity may be managing the authentication lifecycle. We'd prefer to pass in the stubs or "intercepted" + // channel directly, but that's not supported. So, we need to create the specific middleware interfaces so + // flight can do its own shims. + flightClient = FlightGrpcUtilsExtension.createFlightClientWithSharedChannel(bufferAllocator, channel, + new ArrayList<>()); + flightSqlClient = new FlightSqlClient(flightClient); + } + + @Override + public void tearDown() throws Exception { + // this also closes flightClient + flightSqlClient.close(); + bufferAllocator.close(); + super.tearDown(); + } + + @Test + public void listActions() { + // Note: this should likely be tested in the context of Flight, not Flight SQL + assertThat(flightClient.listActions()).isEmpty(); + } + + @Test + public void listFlights() { + // Note: this should likely be tested in the context of Flight, not Flight SQL + assertThat(flightClient.listFlights(Criteria.ALL)).isEmpty(); + } + + @Test + public void getCatalogs() { + unauthenticated(() -> flightSqlClient.getCatalogsSchema()); + unauthenticated(() -> flightSqlClient.getCatalogs()); + } + + @Test + public void getSchemas() { + unauthenticated(() -> flightSqlClient.getSchemasSchema()); + unauthenticated(() -> flightSqlClient.getSchemas(null, null)); + } + + @Test + public void getTables() throws Exception { + unauthenticated(() -> flightSqlClient.getTablesSchema(false)); + unauthenticated(() -> flightSqlClient.getTablesSchema(true)); + unauthenticated(() -> flightSqlClient.getTables(null, null, null, null, false)); + unauthenticated(() -> flightSqlClient.getTables(null, null, null, null, true)); + } + + @Test + public void getTableTypes() throws Exception { + unauthenticated(() -> flightSqlClient.getTableTypesSchema()); + unauthenticated(() -> flightSqlClient.getTableTypes()); + } + + @Test + public void select1() throws Exception { + unauthenticated(() -> flightSqlClient.getExecuteSchema("SELECT 1 as Foo")); + unauthenticated(() -> flightSqlClient.execute("SELECT 1 as Foo")); + } + + @Test + public void select1Prepared() throws Exception { + unauthenticated(() -> flightSqlClient.prepare("SELECT 1 as Foo")); + } + + @Test + public void executeSubstrait() { + unauthenticated(() -> flightSqlClient.getExecuteSubstraitSchema(fakePlan())); + unauthenticated(() -> flightSqlClient.executeSubstrait(fakePlan())); + } + + @Test + public void executeUpdate() { + // We are unable to hook in earlier atm than + // io.deephaven.server.arrow.ArrowFlightUtil.DoPutObserver.DoPutObserver + // so we are unable to provide Flight SQL-specific error message. This could be remedied in the future with an + // update to TicketResolver. + try { + flightSqlClient.executeUpdate("INSERT INTO fake(name) VALUES('Smith')"); + failBecauseExceptionWasNotThrown(FlightRuntimeException.class); + } catch (FlightRuntimeException e) { + assertThat(e.status().code()).isEqualTo(FlightStatusCode.UNAUTHENTICATED); + assertThat(e).hasMessage(""); + } + } + + @Test + public void executeSubstraitUpdate() { + // We are unable to hook in earlier atm than + // io.deephaven.server.arrow.ArrowFlightUtil.DoPutObserver.DoPutObserver + // so we are unable to provide Flight SQL-specific error message. This could be remedied in the future with an + // update to TicketResolver. + try { + flightSqlClient.executeSubstraitUpdate(fakePlan()); + failBecauseExceptionWasNotThrown(FlightRuntimeException.class); + } catch (FlightRuntimeException e) { + assertThat(e.status().code()).isEqualTo(FlightStatusCode.UNAUTHENTICATED); + assertThat(e).hasMessage(""); + } + } + + @Test + public void getSqlInfo() { + unauthenticated(() -> flightSqlClient.getSqlInfoSchema()); + unauthenticated(() -> flightSqlClient.getSqlInfo()); + } + + @Test + public void getXdbcTypeInfo() { + unauthenticated(() -> flightSqlClient.getXdbcTypeInfoSchema()); + unauthenticated(() -> flightSqlClient.getXdbcTypeInfo()); + } + + @Test + public void getCrossReference() { + unauthenticated(() -> flightSqlClient.getCrossReferenceSchema()); + unauthenticated(() -> flightSqlClient.getCrossReference(FOO_TABLE_REF, BAR_TABLE_REF)); + } + + @Test + public void getPrimaryKeys() { + unauthenticated(() -> flightSqlClient.getPrimaryKeysSchema()); + unauthenticated(() -> flightSqlClient.getPrimaryKeys(FOO_TABLE_REF)); + } + + @Test + public void getExportedKeys() { + unauthenticated(() -> flightSqlClient.getExportedKeysSchema()); + unauthenticated(() -> flightSqlClient.getExportedKeys(FOO_TABLE_REF)); + } + + @Test + public void getImportedKeys() { + unauthenticated(() -> flightSqlClient.getImportedKeysSchema()); + unauthenticated(() -> flightSqlClient.getImportedKeys(FOO_TABLE_REF)); + } + + @Test + public void commandStatementIngest() { + // This is a real newer Flight SQL command. + // Once we upgrade to newer Flight SQL, we can change this to Unimplemented and use the proper APIs. + final String typeUrl = "type.googleapis.com/arrow.flight.protocol.sql.CommandStatementIngest"; + final FlightDescriptor descriptor = unpackableCommand(typeUrl); + unauthenticated(() -> flightClient.getSchema(descriptor)); + unauthenticated(() -> flightClient.getInfo(descriptor)); + } + + @Test + public void unknownCommandLooksLikeFlightSql() { + final String typeUrl = "type.googleapis.com/arrow.flight.protocol.sql.CommandLooksRealButDoesNotExist"; + final FlightDescriptor descriptor = unpackableCommand(typeUrl); + unauthenticated(() -> flightClient.getSchema(descriptor)); + unauthenticated(() -> flightClient.getInfo(descriptor)); + } + + @Test + public void unknownCommand() { + // Note: this should likely be tested in the context of Flight, not Flight SQL + final String typeUrl = "type.googleapis.com/com.example.SomeRandomCommand"; + final FlightDescriptor descriptor = unpackableCommand(typeUrl); + expectException(() -> flightClient.getSchema(descriptor), FlightStatusCode.INVALID_ARGUMENT, + "no resolver for command"); + expectException(() -> flightClient.getInfo(descriptor), FlightStatusCode.INVALID_ARGUMENT, + "no resolver for command"); + } + + @Test + public void prepareSubstrait() { + unauthenticated(() -> flightSqlClient.prepare(fakePlan())); + } + + @Test + public void beginTransaction() { + unauthenticated(() -> flightSqlClient.beginTransaction()); + } + + @Test + public void commit() { + unauthenticated(() -> flightSqlClient.commit(fakeTxn())); + } + + @Test + public void rollbackTxn() { + unauthenticated(() -> flightSqlClient.rollback(fakeTxn())); + } + + @Test + public void beginSavepoint() { + unauthenticated(() -> flightSqlClient.beginSavepoint(fakeTxn(), "fakeName")); + } + + @Test + public void release() { + unauthenticated(() -> flightSqlClient.release(fakeSavepoint())); + } + + @Test + public void rollbackSavepoint() { + unauthenticated(() -> flightSqlClient.rollback(fakeSavepoint())); + } + + @Test + public void unknownAction() { + // Note: this should likely be tested in the context of Flight, not Flight SQL + final String type = "SomeFakeAction"; + final Action action = new Action(type, new byte[0]); + actionNoResolver(() -> doAction(action), type); + } + + private Result doAction(Action action) { + final Iterator it = flightClient.doAction(action); + if (!it.hasNext()) { + throw new IllegalStateException(); + } + final Result result = it.next(); + if (it.hasNext()) { + throw new IllegalStateException(); + } + return result; + } + + private static FlightDescriptor unpackableCommand(String typeUrl) { + return FlightDescriptor.command( + Any.newBuilder().setTypeUrl(typeUrl).setValue(ByteString.copyFrom(new byte[1])).build().toByteArray()); + } + + private void unauthenticated(Runnable r) { + expectException(r, FlightStatusCode.UNAUTHENTICATED, "Flight SQL: Must be authenticated"); + } + + private void actionNoResolver(Runnable r, String actionType) { + expectException(r, FlightStatusCode.UNIMPLEMENTED, + String.format("No action resolver found for action type '%s'", actionType)); + } + + private static void expectException(Runnable r, FlightStatusCode code, String messagePart) { + try { + r.run(); + failBecauseExceptionWasNotThrown(FlightRuntimeException.class); + } catch (FlightRuntimeException e) { + assertThat(e.status().code()).isEqualTo(code); + assertThat(e).hasMessageContaining(messagePart); + } + } + + private static SubstraitPlan fakePlan() { + return new SubstraitPlan("fake".getBytes(StandardCharsets.UTF_8), "1"); + } + + private static Transaction fakeTxn() { + return new Transaction("fake".getBytes(StandardCharsets.UTF_8)); + } + + private static Savepoint fakeSavepoint() { + return new Savepoint("fake".getBytes(StandardCharsets.UTF_8)); + } +} diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index 059bcd9b053..c688410b7ee 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -98,6 +98,8 @@ arrow-compression = { module = "org.apache.arrow:arrow-compression", version.ref arrow-format = { module = "org.apache.arrow:arrow-format", version.ref = "arrow" } arrow-vector = { module = "org.apache.arrow:arrow-vector", version.ref = "arrow" } arrow-flight-core = { module = "org.apache.arrow:flight-core", version.ref = "arrow" } +arrow-flight-sql = { module = "org.apache.arrow:flight-sql", version.ref = "arrow" } +arrow-flight-sql-jdbc = { module = "org.apache.arrow:flight-sql-jdbc-driver", version.ref = "arrow" } autoservice = { module = "com.google.auto.service:auto-service-annotations", version.ref = "autoservice" } autoservice-compiler = { module = "com.google.auto.service:auto-service", version.ref = "autoservice" } diff --git a/java-client/flight/src/main/java/org/apache/arrow/flight/ProtocolExposer.java b/java-client/flight/src/main/java/org/apache/arrow/flight/ProtocolExposer.java new file mode 100644 index 00000000000..cf8b9308442 --- /dev/null +++ b/java-client/flight/src/main/java/org/apache/arrow/flight/ProtocolExposer.java @@ -0,0 +1,54 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package org.apache.arrow.flight; + +import org.apache.arrow.flight.impl.Flight; + + +public final class ProtocolExposer { + + /** + * Workaround for [Java][Flight] Add ActionType description + * getter + */ + public static Flight.ActionType toProtocol(ActionType actionType) { + return actionType.toProtocol(); + } + + public static ActionType fromProtocol(Flight.ActionType actionType) { + return new ActionType(actionType); + } + + public static Flight.Action toProtocol(Action action) { + return action.toProtocol(); + } + + public static Action fromProtocol(Flight.Action action) { + return new Action(action); + } + + public static Flight.Result toProtocol(Result action) { + return action.toProtocol(); + } + + public static Result fromProtocol(Flight.Result result) { + return new Result(result); + } + + public static Flight.FlightDescriptor toProtocol(FlightDescriptor descriptor) { + return descriptor.toProtocol(); + } + + public static FlightDescriptor fromProtocol(Flight.FlightDescriptor descriptor) { + return new FlightDescriptor(descriptor); + } + + public static Flight.Ticket toProtocol(Ticket ticket) { + return ticket.toProtocol(); + } + + public static Ticket fromProtocol(Flight.Ticket ticket) { + return new Ticket(ticket); + } +} diff --git a/proto/proto-backplane-grpc/src/main/java/io/deephaven/proto/util/ByteHelper.java b/proto/proto-backplane-grpc/src/main/java/io/deephaven/proto/util/ByteHelper.java index 8136cb31801..a995990534b 100644 --- a/proto/proto-backplane-grpc/src/main/java/io/deephaven/proto/util/ByteHelper.java +++ b/proto/proto-backplane-grpc/src/main/java/io/deephaven/proto/util/ByteHelper.java @@ -6,10 +6,10 @@ import java.nio.ByteBuffer; public class ByteHelper { - public static String byteBufToHex(final ByteBuffer ticket) { + public static String byteBufToHex(final ByteBuffer buffer) { StringBuilder sb = new StringBuilder(); - for (int i = ticket.position(); i < ticket.limit(); ++i) { - sb.append(String.format("%02x", ticket.get(i))); + for (int i = buffer.position(); i < buffer.limit(); ++i) { + sb.append(String.format("%02x", buffer.get(i))); } return sb.toString(); } diff --git a/proto/proto-backplane-grpc/src/main/java/io/deephaven/proto/util/Exceptions.java b/proto/proto-backplane-grpc/src/main/java/io/deephaven/proto/util/Exceptions.java index 3a1d98610b9..cdeaa94ce08 100644 --- a/proto/proto-backplane-grpc/src/main/java/io/deephaven/proto/util/Exceptions.java +++ b/proto/proto-backplane-grpc/src/main/java/io/deephaven/proto/util/Exceptions.java @@ -14,4 +14,19 @@ public static StatusRuntimeException statusRuntimeException(final Code statusCod return StatusProto.toStatusRuntimeException( Status.newBuilder().setCode(statusCode.getNumber()).setMessage(details).build()); } + + static StatusRuntimeException error(io.grpc.Status.Code code, String message) { + return code + .toStatus() + .withDescription("Flight SQL: " + message) + .asRuntimeException(); + } + + static StatusRuntimeException error(io.grpc.Status.Code code, String message, Throwable cause) { + return code + .toStatus() + .withDescription("Flight SQL: " + message) + .withCause(cause) + .asRuntimeException(); + } } diff --git a/py/embedded-server/java-runtime/build.gradle b/py/embedded-server/java-runtime/build.gradle index a2b8bc35d30..04db6fa4ac1 100644 --- a/py/embedded-server/java-runtime/build.gradle +++ b/py/embedded-server/java-runtime/build.gradle @@ -13,6 +13,8 @@ dependencies { implementation project(":util-processenvironment") implementation project(":util-thread") + implementation project(':extensions-flight-sql') + implementation libs.dagger annotationProcessor libs.dagger.compiler diff --git a/py/embedded-server/java-runtime/src/main/java/io/deephaven/python/server/EmbeddedServer.java b/py/embedded-server/java-runtime/src/main/java/io/deephaven/python/server/EmbeddedServer.java index 3bc427f8904..3c7133d347e 100644 --- a/py/embedded-server/java-runtime/src/main/java/io/deephaven/python/server/EmbeddedServer.java +++ b/py/embedded-server/java-runtime/src/main/java/io/deephaven/python/server/EmbeddedServer.java @@ -21,6 +21,7 @@ import io.deephaven.server.console.groovy.GroovyConsoleSessionModule; import io.deephaven.server.console.python.PythonConsoleSessionModule; import io.deephaven.server.console.python.PythonGlobalScopeModule; +import io.deephaven.server.flightsql.FlightSqlModule; import io.deephaven.server.healthcheck.HealthCheckModule; import io.deephaven.server.jetty.JettyConfig; import io.deephaven.server.jetty.JettyConfig.Builder; @@ -73,6 +74,7 @@ static String providesUserAgent() { HealthCheckModule.class, PythonPluginsRegistration.Module.class, JettyServerModule.class, + FlightSqlModule.class, HealthCheckModule.class, PythonConsoleSessionModule.class, GroovyConsoleSessionModule.class, diff --git a/qst/src/main/java/io/deephaven/qst/TableCreatorDelegate.java b/qst/src/main/java/io/deephaven/qst/TableCreatorDelegate.java new file mode 100644 index 00000000000..be0e59c9f2e --- /dev/null +++ b/qst/src/main/java/io/deephaven/qst/TableCreatorDelegate.java @@ -0,0 +1,57 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.qst; + +import io.deephaven.qst.table.EmptyTable; +import io.deephaven.qst.table.InputTable; +import io.deephaven.qst.table.MultiJoinInput; +import io.deephaven.qst.table.NewTable; +import io.deephaven.qst.table.TicketTable; +import io.deephaven.qst.table.TimeTable; + +import java.util.List; +import java.util.Objects; + +public abstract class TableCreatorDelegate
implements TableCreator
{ + private final TableCreator
delegate; + + public TableCreatorDelegate(TableCreator
delegate) { + this.delegate = Objects.requireNonNull(delegate); + } + + @Override + public TABLE of(NewTable newTable) { + return delegate.of(newTable); + } + + @Override + public TABLE of(EmptyTable emptyTable) { + return delegate.of(emptyTable); + } + + @Override + public TABLE of(TimeTable timeTable) { + return delegate.of(timeTable); + } + + @Override + public TABLE of(TicketTable ticketTable) { + return delegate.of(ticketTable); + } + + @Override + public TABLE of(InputTable inputTable) { + return delegate.of(inputTable); + } + + @Override + public TABLE multiJoin(List> multiJoinInputs) { + return delegate.multiJoin(multiJoinInputs); + } + + @Override + public TABLE merge(Iterable
tables) { + return delegate.merge(tables); + } +} diff --git a/server/jetty-app/build.gradle b/server/jetty-app/build.gradle index acce172c114..fae715c29fa 100644 --- a/server/jetty-app/build.gradle +++ b/server/jetty-app/build.gradle @@ -11,6 +11,11 @@ configurations { dependencies { implementation project(':server-jetty') + implementation project(':extensions-flight-sql') + + implementation libs.dagger + annotationProcessor libs.dagger.compiler + runtimeOnly project(':log-to-slf4j') runtimeOnly project(':logback-print-stream-globals') runtimeOnly project(':logback-logbuffer') diff --git a/server/jetty/src/main/java/io/deephaven/server/jetty/CommunityComponentFactory.java b/server/jetty-app/src/main/java/io/deephaven/server/jetty/CommunityComponentFactory.java similarity index 94% rename from server/jetty/src/main/java/io/deephaven/server/jetty/CommunityComponentFactory.java rename to server/jetty-app/src/main/java/io/deephaven/server/jetty/CommunityComponentFactory.java index 014eaae59eb..c67c03f0d19 100644 --- a/server/jetty/src/main/java/io/deephaven/server/jetty/CommunityComponentFactory.java +++ b/server/jetty-app/src/main/java/io/deephaven/server/jetty/CommunityComponentFactory.java @@ -7,6 +7,7 @@ import dagger.Module; import io.deephaven.configuration.Configuration; import io.deephaven.server.auth.CommunityAuthorizationModule; +import io.deephaven.server.flightsql.FlightSqlModule; import io.deephaven.server.runner.CommunityDefaultsModule; import io.deephaven.server.runner.ComponentFactoryBase; @@ -64,11 +65,14 @@ interface Builder extends JettyServerComponent.Builder - - + diff --git a/server/jetty/build.gradle b/server/jetty/build.gradle index 0004397cf48..419f4475e80 100644 --- a/server/jetty/build.gradle +++ b/server/jetty/build.gradle @@ -11,7 +11,10 @@ dependencies { because 'downstream dagger compile' } - implementation project(":util-thread") + implementation project(':util-thread') + compileOnlyApi(project(':util-thread')) { + because 'downstream dagger compile' + } runtimeOnly(project(':web')) diff --git a/server/src/main/java/io/deephaven/server/arrow/ArrowModule.java b/server/src/main/java/io/deephaven/server/arrow/ArrowModule.java index 7f2b22aa464..13805733381 100644 --- a/server/src/main/java/io/deephaven/server/arrow/ArrowModule.java +++ b/server/src/main/java/io/deephaven/server/arrow/ArrowModule.java @@ -6,6 +6,7 @@ import dagger.Binds; import dagger.Module; import dagger.Provides; +import dagger.multibindings.ElementsIntoSet; import dagger.multibindings.IntoSet; import io.deephaven.barrage.flatbuf.BarrageSnapshotRequest; import io.deephaven.barrage.flatbuf.BarrageSubscriptionRequest; @@ -14,9 +15,12 @@ import io.deephaven.extensions.barrage.BarrageSubscriptionOptions; import io.deephaven.server.barrage.BarrageMessageProducer; import io.deephaven.extensions.barrage.BarrageStreamGeneratorImpl; +import io.deephaven.server.session.ActionResolver; +import io.deephaven.server.session.TicketResolver; import io.grpc.BindableService; import javax.inject.Singleton; +import java.util.Set; @Module public abstract class ArrowModule { @@ -43,4 +47,16 @@ static BarrageMessageProducer.Adapter snapshotOptAdapter() { return BarrageSnapshotOptions::of; } + + @Provides + @ElementsIntoSet + static Set primesEmptyTicketResolvers() { + return Set.of(); + } + + @Provides + @ElementsIntoSet + static Set primesEmptyActionResolvers() { + return Set.of(); + } } diff --git a/server/src/main/java/io/deephaven/server/arrow/FlightServiceGrpcImpl.java b/server/src/main/java/io/deephaven/server/arrow/FlightServiceGrpcImpl.java index dc7fc6f1ce0..49925075e23 100644 --- a/server/src/main/java/io/deephaven/server/arrow/FlightServiceGrpcImpl.java +++ b/server/src/main/java/io/deephaven/server/arrow/FlightServiceGrpcImpl.java @@ -9,6 +9,7 @@ import com.google.protobuf.ByteStringAccess; import com.google.protobuf.InvalidProtocolBufferException; import com.google.rpc.Code; +import io.deephaven.auth.AuthContext; import io.deephaven.auth.AuthenticationException; import io.deephaven.auth.AuthenticationRequestHandler; import io.deephaven.auth.BasicAuthMarshaller; @@ -22,15 +23,19 @@ import io.deephaven.proto.backplane.grpc.ExportNotification; import io.deephaven.proto.backplane.grpc.WrappedAuthenticationRequest; import io.deephaven.proto.util.Exceptions; +import io.deephaven.server.session.ActionRouter; import io.deephaven.server.session.SessionService; import io.deephaven.server.session.SessionState; import io.deephaven.server.session.TicketRouter; -import io.deephaven.auth.AuthContext; import io.deephaven.util.SafeCloseable; import io.grpc.StatusRuntimeException; +import io.grpc.stub.ServerCallStreamObserver; import io.grpc.stub.StreamObserver; +import org.apache.arrow.flight.ProtocolExposer; import org.apache.arrow.flight.auth2.Auth2Constants; import org.apache.arrow.flight.impl.Flight; +import org.apache.arrow.flight.impl.Flight.ActionType; +import org.apache.arrow.flight.impl.Flight.Empty; import org.apache.arrow.flight.impl.FlightServiceGrpc; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -43,6 +48,8 @@ import java.util.Optional; import java.util.UUID; import java.util.concurrent.ScheduledExecutorService; +import java.util.function.Consumer; +import java.util.function.Function; import static io.deephaven.extensions.barrage.util.GrpcUtil.safelyComplete; import static io.deephaven.extensions.barrage.util.GrpcUtil.safelyError; @@ -57,6 +64,7 @@ public class FlightServiceGrpcImpl extends FlightServiceGrpc.FlightServiceImplBa private final SessionService sessionService; private final SessionService.ErrorTransformer errorTransformer; private final TicketRouter ticketRouter; + private final ActionRouter actionRouter; private final ArrowFlightUtil.DoExchangeMarshaller.Factory doExchangeFactory; private final Map authRequestHandlers; @@ -68,6 +76,7 @@ public FlightServiceGrpcImpl( final SessionService sessionService, final SessionService.ErrorTransformer errorTransformer, final TicketRouter ticketRouter, + final ActionRouter actionRouter, final ArrowFlightUtil.DoExchangeMarshaller.Factory doExchangeFactory, Map authRequestHandlers) { this.executorService = executorService; @@ -75,6 +84,7 @@ public FlightServiceGrpcImpl( this.sessionService = sessionService; this.errorTransformer = errorTransformer; this.ticketRouter = ticketRouter; + this.actionRouter = actionRouter; this.doExchangeFactory = doExchangeFactory; this.authRequestHandlers = authRequestHandlers; } @@ -203,14 +213,35 @@ public void onCompleted() { } } + @Override + public void doAction(Flight.Action request, StreamObserver responseObserver) { + actionRouter.doAction( + sessionService.getOptionalSession(), + ProtocolExposer.fromProtocol(request), + new ServerCallStreamObserverAdapter<>( + (ServerCallStreamObserver) responseObserver, ProtocolExposer::toProtocol)); + } + @Override public void listFlights( @NotNull final Flight.Criteria request, @NotNull final StreamObserver responseObserver) { + if (!request.getExpression().isEmpty()) { + responseObserver.onError( + Exceptions.statusRuntimeException(Code.INVALID_ARGUMENT, "Criteria expressions are not supported")); + return; + } ticketRouter.visitFlightInfo(sessionService.getOptionalSession(), responseObserver::onNext); responseObserver.onCompleted(); } + @Override + public void listActions(Empty request, StreamObserver responseObserver) { + actionRouter.listActions(sessionService.getOptionalSession(), + adapt(responseObserver::onNext, ProtocolExposer::toProtocol)); + responseObserver.onCompleted(); + } + @Override public void getFlightInfo( @NotNull final Flight.FlightDescriptor request, @@ -333,4 +364,8 @@ public StreamObserver doPutCustom(final StreamObserver doExchangeCustom(final StreamObserver responseObserver) { return doExchangeFactory.openExchange(sessionService.getCurrentSession(), responseObserver); } + + private static Consumer adapt(Consumer consumer, Function function) { + return t -> consumer.accept(function.apply(t)); + } } diff --git a/server/src/main/java/io/deephaven/server/arrow/ServerCallStreamObserverAdapter.java b/server/src/main/java/io/deephaven/server/arrow/ServerCallStreamObserverAdapter.java new file mode 100644 index 00000000000..77f15c136d1 --- /dev/null +++ b/server/src/main/java/io/deephaven/server/arrow/ServerCallStreamObserverAdapter.java @@ -0,0 +1,75 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.arrow; + +import io.grpc.stub.ServerCallStreamObserver; + +import java.util.Objects; +import java.util.function.Function; + +final class ServerCallStreamObserverAdapter extends ServerCallStreamObserver { + + private final ServerCallStreamObserver delegate; + private final Function f; + + ServerCallStreamObserverAdapter(ServerCallStreamObserver delegate, Function f) { + this.delegate = Objects.requireNonNull(delegate); + this.f = Objects.requireNonNull(f); + } + + @Override + public void onNext(T value) { + delegate.onNext(f.apply(value)); + } + + @Override + public void onError(Throwable t) { + delegate.onError(t); + } + + @Override + public void onCompleted() { + delegate.onCompleted(); + } + + @Override + public boolean isCancelled() { + return delegate.isCancelled(); + } + + @Override + public void setOnCancelHandler(Runnable onCancelHandler) { + delegate.setOnCancelHandler(onCancelHandler); + } + + @Override + public void setCompression(String compression) { + delegate.setCompression(compression); + } + + @Override + public boolean isReady() { + return delegate.isReady(); + } + + @Override + public void setOnReadyHandler(Runnable onReadyHandler) { + delegate.setOnReadyHandler(onReadyHandler); + } + + @Override + public void request(int count) { + delegate.request(count); + } + + @Override + public void setMessageCompression(boolean enable) { + delegate.setMessageCompression(enable); + } + + @Override + public void disableAutoInboundFlowControl() { + delegate.disableAutoInboundFlowControl(); + } +} diff --git a/server/src/main/java/io/deephaven/server/console/ScopeTicketResolver.java b/server/src/main/java/io/deephaven/server/console/ScopeTicketResolver.java index 2ee01446d39..42af427c4c0 100644 --- a/server/src/main/java/io/deephaven/server/console/ScopeTicketResolver.java +++ b/server/src/main/java/io/deephaven/server/console/ScopeTicketResolver.java @@ -74,6 +74,9 @@ public SessionState.ExportObject flightInfoFor( @Override public void forAllFlightInfo(@Nullable final SessionState session, final Consumer visitor) { + if (session == null) { + return; + } final QueryScope queryScope = ExecutionContext.getContext().getQueryScope(); queryScope.toMap(queryScope::unwrapObject, (n, t) -> t instanceof Table).forEach((name, table) -> { final Table transformedTable = authorization.transform((Table) table); diff --git a/server/src/main/java/io/deephaven/server/session/ActionResolver.java b/server/src/main/java/io/deephaven/server/session/ActionResolver.java new file mode 100644 index 00000000000..c99a2225a56 --- /dev/null +++ b/server/src/main/java/io/deephaven/server/session/ActionResolver.java @@ -0,0 +1,57 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.session; + +import io.grpc.stub.StreamObserver; +import org.apache.arrow.flight.Action; +import org.apache.arrow.flight.ActionType; +import org.apache.arrow.flight.Result; +import org.jetbrains.annotations.Nullable; + +import java.util.function.Consumer; + +public interface ActionResolver { + + /** + * Invokes the {@code visitor} for the specific action types that this implementation supports for the given + * {@code session}; it should be the case that all consumed action types return {@code true} against + * {@link #handlesActionType(String)}. Unlike {@link #handlesActionType(String)}, the implementations should + * not invoke the visitor for action types in their domain that they do not implement. + * + *

+ * This is called in the context of {@link ActionRouter#listActions(SessionState, Consumer)} to allow flight + * consumers to understand the capabilities of this flight service. + * + * @param session the session + * @param visitor the visitor + */ + void listActions(@Nullable SessionState session, Consumer visitor); + + /** + * Returns {@code true} if this resolver is responsible for handling the action {@code type}. Implementations should + * prefer to return {@code true} if they know the action type is in their domain even if they don't implement it; + * this allows them to provide a more specific error message for unimplemented action types. + * + *

+ * This is used in support of routing in {@link ActionRouter#doAction(SessionState, Action, StreamObserver)} calls. + * + * @param type the action type + * @return {@code true} if this resolver handles the action type + */ + boolean handlesActionType(String type); + + /** + * Executes the given {@code action}. Should only be called if {@link #handlesActionType(String)} is {@code true} + * for the given {@code action}. + * + *

+ * This is called in the context of {@link ActionRouter#doAction(SessionState, Action, StreamObserver)} to allow + * flight consumers to execute an action against this flight service. + * + * @param session the session + * @param action the action + * @param observer the observer + */ + void doAction(@Nullable final SessionState session, final Action action, final StreamObserver observer); +} diff --git a/server/src/main/java/io/deephaven/server/session/ActionRouter.java b/server/src/main/java/io/deephaven/server/session/ActionRouter.java new file mode 100644 index 00000000000..33ebcf901da --- /dev/null +++ b/server/src/main/java/io/deephaven/server/session/ActionRouter.java @@ -0,0 +1,99 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.session; + +import com.google.rpc.Code; +import io.deephaven.configuration.Configuration; +import io.deephaven.engine.table.impl.perf.QueryPerformanceNugget; +import io.deephaven.engine.table.impl.perf.QueryPerformanceRecorder; +import io.deephaven.proto.util.Exceptions; +import io.grpc.stub.StreamObserver; +import org.apache.arrow.flight.Action; +import org.apache.arrow.flight.ActionType; +import org.apache.arrow.flight.Result; +import org.jetbrains.annotations.Nullable; + +import javax.inject.Inject; +import java.util.Set; +import java.util.function.Consumer; +import java.util.stream.Collectors; + +public final class ActionRouter { + + private static boolean enabled(ActionResolver resolver) { + final String property = + ActionResolver.class.getSimpleName() + "." + resolver.getClass().getSimpleName() + ".enabled"; + return Configuration.getInstance().getBooleanWithDefault(property, true); + } + + private final Set resolvers; + + @Inject + public ActionRouter(Set resolvers) { + this.resolvers = resolvers.stream() + .filter(ActionRouter::enabled) + .collect(Collectors.toSet()); + } + + /** + * Invokes {@code visitor} for all of the resolvers. Used as the basis for implementing FlightService ListActions. + * + * @param session the session + * @param visitor the visitor + */ + public void listActions(@Nullable final SessionState session, final Consumer visitor) { + final QueryPerformanceRecorder qpr = QueryPerformanceRecorder.getInstance(); + try (final QueryPerformanceNugget ignored = qpr.getNugget("listActions")) { + for (ActionResolver resolver : resolvers) { + resolver.listActions(session, visitor); + } + } + } + + /** + * Routes {@code action} to the appropriate {@link ActionResolver}. Used as the basis for implementing FlightService + * DoAction. + * + * @param session the session + * @param action the action + * @param observer the observer + * @throws io.grpc.StatusRuntimeException if zero or more than one resolver is found + */ + public void doAction(@Nullable final SessionState session, final Action action, + final StreamObserver observer) { + final QueryPerformanceRecorder qpr = QueryPerformanceRecorder.getInstance(); + try (final QueryPerformanceNugget ignored = qpr.getNugget(String.format("doAction:%s", action.getType()))) { + getResolver(action.getType()).doAction(session, action, observer); + } + } + + private ActionResolver getResolver(final String type) { + ActionResolver actionResolver = null; + // This is the most "naive" resolution logic; it scales linearly with the number of resolvers, but it is the + // most general and may be the best we can do for certain types of action protocols built on top of Flight. If + // we find the number of action resolvers scaling up, we could devise a more efficient strategy in some cases + // either based on a prefix model and/or a fixed set model (which could be communicated either through new + // method(s) on ActionResolver, or through subclasses). + // ` + // Regardless, even with a moderate amount of action resolvers, the linear nature of this should not be a + // bottleneck. + for (ActionResolver resolver : resolvers) { + if (!resolver.handlesActionType(type)) { + continue; + } + if (actionResolver != null) { + throw Exceptions.statusRuntimeException(Code.INTERNAL, + String.format("Found multiple doAction resolvers for action type '%s'", type)); + } + actionResolver = resolver; + } + if (actionResolver == null) { + // Similar to the default unimplemented message from + // org.apache.arrow.flight.impl.FlightServiceGrpc.AsyncService.doAction + throw Exceptions.statusRuntimeException(Code.UNIMPLEMENTED, + String.format("No action resolver found for action type '%s'", type)); + } + return actionResolver; + } +} diff --git a/server/src/main/java/io/deephaven/server/session/AuthCookie.java b/server/src/main/java/io/deephaven/server/session/AuthCookie.java new file mode 100644 index 00000000000..5e06973ac06 --- /dev/null +++ b/server/src/main/java/io/deephaven/server/session/AuthCookie.java @@ -0,0 +1,73 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.session; + +import com.github.f4b6a3.uuid.UuidCreator; +import com.github.f4b6a3.uuid.exception.InvalidUuidException; +import io.grpc.Context; +import io.grpc.Contexts; +import io.grpc.Metadata; +import io.grpc.ServerCall; +import io.grpc.ServerCall.Listener; +import io.grpc.ServerCallHandler; +import io.grpc.ServerInterceptor; + +import java.util.Optional; +import java.util.UUID; + +/** + * This exists to work around how the Flight SQL JDBC driver works out-of-the-box. + */ +final class AuthCookie { + + private static final String HEADER = "x-deephaven-auth-cookie-request"; + + private static final String DEEPHAVEN_AUTH_COOKIE = "deephaven-auth-cookie"; + + private static final Metadata.Key REQUEST_AUTH_COOKIE_HEADER_KEY = + Metadata.Key.of(HEADER, Metadata.ASCII_STRING_MARSHALLER); + + private static final Metadata.Key SET_COOKIE = + Metadata.Key.of("set-cookie", Metadata.ASCII_STRING_MARSHALLER); + + private static final Metadata.Key COOKIE = + Metadata.Key.of("cookie", Metadata.ASCII_STRING_MARSHALLER); + + /** + * Returns {@code true} if the metadata contains the header {@value HEADER} with value "true". + */ + public static boolean hasDeephavenAuthCookieRequest(Metadata md) { + return Boolean.parseBoolean(md.get(REQUEST_AUTH_COOKIE_HEADER_KEY)); + } + + /** + * Sets the auth cookie {@value DEEPHAVEN_AUTH_COOKIE} to {@code token}. + */ + public static void setDeephavenAuthCookie(Metadata md, UUID token) { + md.put(SET_COOKIE, DEEPHAVEN_AUTH_COOKIE + "=" + token.toString()); + } + + /** + * Parses the "cookie" header for the Deephaven auth cookie if it is of the form "deephaven-auth-cookie=". + */ + public static Optional parseAuthCookie(Metadata md) { + final String cookie = md.get(COOKIE); + if (cookie == null) { + return Optional.empty(); + } + // DH will only ever set one cookie of the form "deephaven-auth-cookie="; anything that doesn't match this + // is invalid. + final String[] split = cookie.split("="); + if (split.length != 2 || !DEEPHAVEN_AUTH_COOKIE.equals(split[0])) { + return Optional.empty(); + } + final UUID uuid; + try { + uuid = UuidCreator.fromString(split[1]); + } catch (InvalidUuidException e) { + return Optional.empty(); + } + return Optional.of(uuid); + } +} diff --git a/server/src/main/java/io/deephaven/server/session/CommandResolver.java b/server/src/main/java/io/deephaven/server/session/CommandResolver.java new file mode 100644 index 00000000000..1aa111c2dbf --- /dev/null +++ b/server/src/main/java/io/deephaven/server/session/CommandResolver.java @@ -0,0 +1,43 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.session; + + +import org.apache.arrow.flight.impl.Flight.FlightDescriptor; + +/** + * A specialization of {@link TicketResolver} that signifies this resolver supports Flight descriptor commands. + * + *

+ * Unfortunately, there is no universal way to know whether a command belongs to a given Flight protocol or not; at + * best, we can assume (or mandate) that all the supportable command bytes are sufficiently unique such that there is no + * potential for overlap amongst the installed Flight protocols. + * + *

+ * For example there could be command protocols built on top of Flight that simply use integer ordinals as their command + * serialization format. In such a case, only one such protocol could safely be installed; otherwise, there would be no + * reliable way of differentiating between them from the command bytes. (It's possible that other means of + * differentiating could be established, like header values.) + * + *

+ * If Deephaven is in a position to create a protocol that uses Flight commands, or advise on their creation, it would + * probably be wise to use a command serialization format that has a "unique" magic value as its prefix. + * + *

+ * The Flight SQL approach is to use the protobuf message Any to wrap up the respective protobuf Flight SQL command + * message. While this approach is very likely to produce a sufficiently unique selection criteria, it requires + * "non-trivial" parsing to determine whether the command is supported or not. + */ +public interface CommandResolver extends TicketResolver { + + /** + * Returns {@code true} if this resolver is responsible for handling the {@code descriptor} command. Implementations + * should prefer to return {@code true} here if they know the command is in their domain even if they don't + * implement it; this allows them to provide a more specific error message for unsupported commands. + * + * @param descriptor the descriptor + * @return {@code true} if this resolver handles the descriptor command + */ + boolean handlesCommand(FlightDescriptor descriptor); +} diff --git a/server/src/main/java/io/deephaven/server/session/PathResolver.java b/server/src/main/java/io/deephaven/server/session/PathResolver.java new file mode 100644 index 00000000000..b844e346d62 --- /dev/null +++ b/server/src/main/java/io/deephaven/server/session/PathResolver.java @@ -0,0 +1,23 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.session; + + +import org.apache.arrow.flight.impl.Flight.FlightDescriptor; + +/** + * A specialization of {@link TicketResolver} that signifies this resolver supports Flight descriptor paths. + */ +public interface PathResolver extends TicketResolver { + + /** + * Returns {@code true} if this resolver is responsible for handling the {@code descriptor} path. Implementations + * should prefer to return {@code true} here if they know the path is in their domain even if they don't implement + * it; this allows them to provide a more specific error message for unsupported paths. + * + * @param descriptor the descriptor + * @return {@code true} if this resolver handles the descriptor path + */ + boolean handlesPath(FlightDescriptor descriptor); +} diff --git a/server/src/main/java/io/deephaven/server/session/PathResolverPrefixedBase.java b/server/src/main/java/io/deephaven/server/session/PathResolverPrefixedBase.java new file mode 100644 index 00000000000..f048781d6cb --- /dev/null +++ b/server/src/main/java/io/deephaven/server/session/PathResolverPrefixedBase.java @@ -0,0 +1,48 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.session; + +import org.apache.arrow.flight.impl.Flight; + +import java.util.Objects; + +/** + * A specialization of {@link PathResolver} whose path {@link Flight.FlightDescriptor} resolution is based on the first + * path in the list. + */ +public abstract class PathResolverPrefixedBase implements PathResolver { + + private final String flightDescriptorRoute; + + public PathResolverPrefixedBase(String flightDescriptorRoute) { + this.flightDescriptorRoute = Objects.requireNonNull(flightDescriptorRoute); + } + + /** + * The first path entry on a route indicates which resolver to use. The remaining path elements are used to resolve + * the descriptor. + * + * @return the string that will route from flight descriptor to this resolver + */ + public final String flightDescriptorRoute() { + return flightDescriptorRoute; + } + + /** + * Returns {@code true} if the first path in {@code descriptor} is equal to {@link #flightDescriptorRoute()}. + * + * @param descriptor the descriptor + * @return {@code true} if this resolver handles the descriptor path + */ + @Override + public final boolean handlesPath(Flight.FlightDescriptor descriptor) { + if (descriptor.getType() != Flight.FlightDescriptor.DescriptorType.PATH) { + throw new IllegalStateException("descriptor is not a path"); + } + if (descriptor.getPathCount() == 0) { + return false; + } + return flightDescriptorRoute.equals(descriptor.getPath(0)); + } +} diff --git a/server/src/main/java/io/deephaven/server/session/SessionServiceGrpcImpl.java b/server/src/main/java/io/deephaven/server/session/SessionServiceGrpcImpl.java index ea9e1b39f23..b74ee4d0fbf 100644 --- a/server/src/main/java/io/deephaven/server/session/SessionServiceGrpcImpl.java +++ b/server/src/main/java/io/deephaven/server/session/SessionServiceGrpcImpl.java @@ -43,7 +43,9 @@ import java.nio.charset.StandardCharsets; import java.util.LinkedHashMap; import java.util.Map; +import java.util.Objects; import java.util.Set; +import java.util.UUID; public class SessionServiceGrpcImpl extends SessionServiceGrpc.SessionServiceImplBase { /** @@ -53,6 +55,7 @@ public class SessionServiceGrpcImpl extends SessionServiceGrpc.SessionServiceImp public static final String DEEPHAVEN_SESSION_ID = Auth2Constants.AUTHORIZATION_HEADER; public static final Metadata.Key SESSION_HEADER_KEY = Metadata.Key.of(Auth2Constants.AUTHORIZATION_HEADER, Metadata.ASCII_STRING_MARSHALLER); + public static final Context.Key SESSION_CONTEXT_KEY = Context.key(Auth2Constants.AUTHORIZATION_HEADER); @@ -264,12 +267,17 @@ public static class InterceptedCall extends SimpleForwardingServerC private final SessionService service; private final SessionState session; private final Map, String> extraHeaders = new LinkedHashMap<>(); + private final boolean setDeephavenAuthCookie; - private InterceptedCall(final SessionService service, final ServerCall call, - @Nullable final SessionState session) { - super(call); - this.service = service; + private InterceptedCall( + final SessionService service, + final ServerCall call, + @Nullable final SessionState session, + boolean setDeephavenAuthCookie) { + super(Objects.requireNonNull(call)); + this.service = Objects.requireNonNull(service); this.session = session; + this.setDeephavenAuthCookie = setDeephavenAuthCookie; } @Override @@ -305,6 +313,9 @@ private void addHeaders(final Metadata md) { final SessionService.TokenExpiration exp = service.refreshToken(session); if (exp != null) { md.put(SESSION_HEADER_KEY, Auth2Constants.BEARER_PREFIX + exp.token.toString()); + if (setDeephavenAuthCookie) { + AuthCookie.setDeephavenAuthCookie(md, exp.token); + } } } } @@ -330,8 +341,8 @@ public static class SessionServiceInterceptor implements ServerInterceptor { public SessionServiceInterceptor( final SessionService service, final SessionService.ErrorTransformer errorTransformer) { - this.service = service; - this.errorTransformer = errorTransformer; + this.service = Objects.requireNonNull(service); + this.errorTransformer = Objects.requireNonNull(errorTransformer); } @Override @@ -350,20 +361,31 @@ public ServerCall.Listener interceptCall(final ServerCall() {}; + if (session == null) { + // Lookup the session using the auth cookie + final UUID uuid = AuthCookie.parseAuthCookie(metadata).orElse(null); + if (uuid != null) { + session = service.getSessionForToken(uuid); + } + } + + if (session == null) { + // Lookup the session using Flight Auth 2.0 token. + final String token = metadata.get(SESSION_HEADER_KEY); + if (token != null) { + try { + session = service.getSessionForAuthToken(token); + } catch (AuthenticationException e) { + // As an interceptor, we can't throw, so ignoring this and just returning the no-op listener. + safeClose(call, AUTHENTICATION_DETAILS_INVALID, new Metadata(), false); + return new ServerCall.Listener<>() {}; + } } } // On the outer half of the call we'll install the context that includes our session. - final InterceptedCall serverCall = new InterceptedCall<>(service, call, session); + final InterceptedCall serverCall = new InterceptedCall<>(service, call, session, + AuthCookie.hasDeephavenAuthCookieRequest(metadata)); final Context context = Context.current().withValues( SESSION_CONTEXT_KEY, session, SESSION_CALL_KEY, serverCall); diff --git a/server/src/main/java/io/deephaven/server/session/SessionState.java b/server/src/main/java/io/deephaven/server/session/SessionState.java index 526726c78a1..581c7b4abc7 100644 --- a/server/src/main/java/io/deephaven/server/session/SessionState.java +++ b/server/src/main/java/io/deephaven/server/session/SessionState.java @@ -717,7 +717,7 @@ private synchronized void setWork( return; } - this.exportMain = exportMain; + this.exportMain = Objects.requireNonNull(exportMain); this.errorHandler = errorHandler; this.successHandler = successHandler; @@ -799,6 +799,13 @@ public Ticket getExportId() { return ExportTicketHelper.wrapExportIdInTicket(exportId); } + /** + * @return the export id for this export + */ + public int getExportIdInt() { + return exportId; + } + /** * Add dependency if object export has not yet completed. * @@ -1369,7 +1376,6 @@ public class ExportBuilder { ExportBuilder(final int exportId) { this.exportId = exportId; - if (exportId == NON_EXPORT_ID) { this.export = new ExportObject<>(SessionState.this.errorTransformer, SessionState.this, NON_EXPORT_ID); } else { diff --git a/server/src/main/java/io/deephaven/server/session/TicketResolver.java b/server/src/main/java/io/deephaven/server/session/TicketResolver.java index cb74ffcc49f..8a7555050b3 100644 --- a/server/src/main/java/io/deephaven/server/session/TicketResolver.java +++ b/server/src/main/java/io/deephaven/server/session/TicketResolver.java @@ -14,6 +14,23 @@ public interface TicketResolver { interface Authorization { + + /** + * Check if the caller is denied access to {@code source}; semantically equivalent to + * {@code transform(source) == null}. A {@code false} result does not mean that the caller may use + * {@code source} untransformed; they must still call {@link #transform(Object)} as needed. + * + *

+ * The default implementation is equivalent to {@code transform(source) == null}. Implementations that perform + * expensive transformations may want to override this method to provide a more efficient check. + * + * @param source the source object + * @return if the transform of {@code source} will result in {@code null}. + */ + default boolean isDeniedAccess(Object source) { + return transform(source) == null; + } + /** * Implementations must type check the provided source as any type of object can be stored in an export. *

@@ -61,14 +78,6 @@ interface Authorization { */ byte ticketRoute(); - /** - * The first path entry on a route indicates which resolver to use. The remaining path elements are used to resolve - * the descriptor. - * - * @return the string that will route from flight descriptor to this resolver - */ - String flightDescriptorRoute(); - /** * Resolve a flight ticket to an export object future. * @@ -175,4 +184,6 @@ SessionState.ExportObject flightInfoFor(@Nullable SessionStat * @param visitor the callback to invoke per descriptor path */ void forAllFlightInfo(@Nullable SessionState session, Consumer visitor); + + // TODO(deephaven-core#6295): Consider use of Flight POJOs instead of protobufs } diff --git a/server/src/main/java/io/deephaven/server/session/TicketResolverBase.java b/server/src/main/java/io/deephaven/server/session/TicketResolverBase.java index b206000f334..7d025ac2959 100644 --- a/server/src/main/java/io/deephaven/server/session/TicketResolverBase.java +++ b/server/src/main/java/io/deephaven/server/session/TicketResolverBase.java @@ -5,27 +5,24 @@ import io.deephaven.server.auth.AuthorizationProvider; -public abstract class TicketResolverBase implements TicketResolver { +import java.util.Objects; + +public abstract class TicketResolverBase extends PathResolverPrefixedBase { protected final Authorization authorization; private final byte ticketPrefix; - private final String flightDescriptorRoute; public TicketResolverBase( final AuthorizationProvider authProvider, - final byte ticketPrefix, final String flightDescriptorRoute) { - this.authorization = authProvider.getTicketResolverAuthorization(); + final byte ticketPrefix, + final String flightDescriptorRoute) { + super(flightDescriptorRoute); + this.authorization = Objects.requireNonNull(authProvider.getTicketResolverAuthorization()); this.ticketPrefix = ticketPrefix; - this.flightDescriptorRoute = flightDescriptorRoute; } @Override - public byte ticketRoute() { + public final byte ticketRoute() { return ticketPrefix; } - - @Override - public String flightDescriptorRoute() { - return flightDescriptorRoute; - } } diff --git a/server/src/main/java/io/deephaven/server/session/TicketRouter.java b/server/src/main/java/io/deephaven/server/session/TicketRouter.java index f5741cb0e01..e7ea4ec57f8 100644 --- a/server/src/main/java/io/deephaven/server/session/TicketRouter.java +++ b/server/src/main/java/io/deephaven/server/session/TicketRouter.java @@ -4,7 +4,9 @@ package io.deephaven.server.session; import com.google.rpc.Code; +import io.deephaven.configuration.Configuration; import io.deephaven.engine.table.Table; +import io.deephaven.engine.table.impl.perf.QueryPerformanceNugget; import io.deephaven.engine.table.impl.perf.QueryPerformanceRecorder; import io.deephaven.extensions.barrage.util.BarrageUtil; import io.deephaven.hash.KeyedIntObjectHashMap; @@ -16,38 +18,67 @@ import io.deephaven.server.auth.AuthorizationProvider; import io.deephaven.util.SafeCloseable; import org.apache.arrow.flight.impl.Flight; +import org.apache.arrow.flight.impl.Flight.FlightDescriptor; +import org.apache.arrow.flight.impl.Flight.FlightDescriptor.DescriptorType; import org.jetbrains.annotations.Nullable; import javax.inject.Inject; import javax.inject.Singleton; import java.nio.ByteBuffer; +import java.util.Objects; +import java.util.Optional; import java.util.Set; import java.util.function.Consumer; +import java.util.function.Predicate; +import java.util.stream.Collectors; @Singleton public class TicketRouter { + + private static boolean enabled(TicketResolver resolver) { + final String property = + TicketResolver.class.getSimpleName() + "." + resolver.getClass().getSimpleName() + ".enabled"; + return Configuration.getInstance().getBooleanWithDefault(property, true); + } + private final KeyedIntObjectHashMap byteResolverMap = new KeyedIntObjectHashMap<>(RESOLVER_OBJECT_TICKET_ID); - private final KeyedObjectHashMap descriptorResolverMap = + private final KeyedObjectHashMap prefixedPathResolverMap = new KeyedObjectHashMap<>(RESOLVER_OBJECT_DESCRIPTOR_ID); private final TicketResolver.Authorization authorization; + private final Set commandResolvers; + private final Set genericPathResolvers; @Inject public TicketRouter( final AuthorizationProvider authorizationProvider, - final Set resolvers) { + Set resolvers) { + resolvers = resolvers.stream().filter(TicketRouter::enabled).collect(Collectors.toSet()); this.authorization = authorizationProvider.getTicketResolverAuthorization(); - resolvers.forEach(resolver -> { + this.commandResolvers = resolvers.stream() + .filter(CommandResolver.class::isInstance) + .map(CommandResolver.class::cast) + .collect(Collectors.toSet()); + this.genericPathResolvers = resolvers.stream() + .filter(PathResolver.class::isInstance) + .filter(Predicate.not(PathResolverPrefixedBase.class::isInstance)) + .map(PathResolver.class::cast) + .collect(Collectors.toSet()); + for (TicketResolver resolver : resolvers) { if (!byteResolverMap.add(resolver)) { throw new IllegalArgumentException("Duplicate ticket resolver for ticket route " + resolver.ticketRoute()); } - if (!descriptorResolverMap.add(resolver)) { + if (!(resolver instanceof PathResolverPrefixedBase)) { + continue; + } + final PathResolverPrefixedBase prefixedPathResolver = (PathResolverPrefixedBase) resolver; + if (!prefixedPathResolverMap.add(prefixedPathResolver)) { throw new IllegalArgumentException("Duplicate ticket resolver for descriptor route " - + resolver.flightDescriptorRoute()); + + prefixedPathResolver.flightDescriptorRoute()); } - }); + } } /** @@ -244,10 +275,14 @@ public void publish( @Nullable final Runnable onPublish, final SessionState.ExportErrorHandler errorHandler, final SessionState.ExportObject source) { - final ByteBuffer ticketBuffer = ticket.getTicket().asReadOnlyByteBuffer(); - final TicketResolver resolver = getResolver(ticketBuffer.get(ticketBuffer.position()), logId); - authorization.authorizePublishRequest(resolver, ticketBuffer); - resolver.publish(session, ticketBuffer, logId, onPublish, errorHandler, source); + final String ticketName = getLogNameFor(ticket, logId); + try (final SafeCloseable ignored = + QueryPerformanceRecorder.getInstance().getNugget("publishTicket:" + ticketName)) { + final ByteBuffer ticketBuffer = ticket.getTicket().asReadOnlyByteBuffer(); + final TicketResolver resolver = getResolver(ticketBuffer.get(ticketBuffer.position()), logId); + authorization.authorizePublishRequest(resolver, ticketBuffer); + resolver.publish(session, ticketBuffer, logId, onPublish, errorHandler, source); + } } /** @@ -264,11 +299,17 @@ public SessionState.ExportObject flightInfoFor( @Nullable final SessionState session, final Flight.FlightDescriptor descriptor, final String logId) { + // noinspection CaughtExceptionImmediatelyRethrown try (final SafeCloseable ignored = QueryPerformanceRecorder.getInstance().getNugget( "flightInfoForDescriptor:" + descriptor)) { return getResolver(descriptor, logId).flightInfoFor(session, descriptor, logId); } catch (RuntimeException e) { - return SessionState.wrapAsFailedExport(e); + // io.deephaven.server.flightsql.FlightSqlUnauthenticatedTest RPC never finishes when this path is used + // return SessionState.wrapAsFailedExport(e); + // This is a partial workaround for + // TODO(deephaven-core#6374): FlightServiceGrpcImpl getFlightInfo / getSchema unauthenticated path + // misimplemented + throw e; } } @@ -312,7 +353,10 @@ public String getLogNameFor(final ByteBuffer ticket, final String logId) { * @param visitor the callback to invoke per descriptor path */ public void visitFlightInfo(@Nullable final SessionState session, final Consumer visitor) { - byteResolverMap.iterator().forEachRemaining(resolver -> resolver.forAllFlightInfo(session, visitor)); + final QueryPerformanceRecorder qpr = QueryPerformanceRecorder.getInstance(); + try (final QueryPerformanceNugget ignored = qpr.getNugget("visitFlightInfo")) { + byteResolverMap.iterator().forEachRemaining(resolver -> resolver.forAllFlightInfo(session, visitor)); + } } public static Flight.FlightInfo getFlightInfo(final Table table, @@ -339,37 +383,96 @@ private TicketResolver getResolver(final byte route, final String logId) { } private TicketResolver getResolver(final Flight.FlightDescriptor descriptor, final String logId) { - if (descriptor.getType() != Flight.FlightDescriptor.DescriptorType.PATH) { - throw Exceptions.statusRuntimeException(Code.FAILED_PRECONDITION, - "Could not resolve '" + logId + "': flight descriptor is not a path"); + if (descriptor.getType() == Flight.FlightDescriptor.DescriptorType.PATH) { + return getPathResolver(descriptor, logId); + } + if (descriptor.getType() == DescriptorType.CMD) { + return getCommandResolver(descriptor, logId); + } + throw Exceptions.statusRuntimeException(Code.FAILED_PRECONDITION, + "Could not resolve '" + logId + "': unexpected type"); + } + + private PathResolver getPathResolver(FlightDescriptor descriptor, String logId) { + if (descriptor.getType() != DescriptorType.PATH) { + throw new IllegalStateException("descriptor is not a path"); } if (descriptor.getPathCount() <= 0) { throw Exceptions.statusRuntimeException(Code.FAILED_PRECONDITION, "Could not resolve '" + logId + "': flight descriptor does not have route path"); } - final String route = descriptor.getPath(0); - final TicketResolver resolver = descriptorResolverMap.get(route); - if (resolver == null) { + final PathResolverPrefixedBase prefixedResolver = prefixedPathResolverMap.get(route); + final PathResolver genericResolver = getGenericPathResolver(descriptor, logId, route).orElse(null); + if (prefixedResolver == null && genericResolver == null) { throw Exceptions.statusRuntimeException(Code.FAILED_PRECONDITION, - "Could not resolve '" + logId + "': no resolver for route '" + route + "'"); + "Could not resolve '" + logId + "': no resolver for path route '" + route + "'"); + } + if (prefixedResolver != null && genericResolver != null) { + throw Exceptions.statusRuntimeException(Code.INTERNAL, + "Could not resolve '" + logId + "': multiple resolvers for path route '" + route + "'"); } + return prefixedResolver != null ? prefixedResolver : Objects.requireNonNull(genericResolver); + } - return resolver; + private Optional getGenericPathResolver(FlightDescriptor descriptor, String logId, String route) { + PathResolver genericResolver = null; + for (PathResolver resolver : genericPathResolvers) { + if (!resolver.handlesPath(descriptor)) { + continue; + } + if (genericResolver != null) { + throw Exceptions.statusRuntimeException(Code.INTERNAL, + "Could not resolve '" + logId + "': multiple resolvers for path route '" + route + "'"); + } + genericResolver = resolver; + } + return Optional.ofNullable(genericResolver); + } + + private CommandResolver getCommandResolver(FlightDescriptor descriptor, String logId) { + if (descriptor.getType() != DescriptorType.CMD) { + throw new IllegalStateException("descriptor is not a command"); + } + // This is the most "naive" resolution logic; it scales linearly with the number of command resolvers, but it is + // the most general and may be the best we can do for certain types of command protocols built on top of Flight. + // If we find the number of command resolvers scaling up, we could devise a more efficient strategy in some + // cases either based on a prefix model and/or a fixed set model (which could be communicated either through new + // method(s) on CommandResolver, or through subclasses). + // + // Regardless, even with a moderate amount of command resolvers, the linear nature of this should not be a + // bottleneck. + CommandResolver commandResolver = null; + for (CommandResolver resolver : commandResolvers) { + if (!resolver.handlesCommand(descriptor)) { + continue; + } + if (commandResolver != null) { + // Is there any good way to give a friendly string for unknown command bytes? Probably not. + throw Exceptions.statusRuntimeException(Code.INTERNAL, + "Could not resolve '" + logId + "': multiple resolvers for command"); + } + commandResolver = resolver; + } + if (commandResolver == null) { + throw Exceptions.statusRuntimeException(Code.FAILED_PRECONDITION, + "Could not resolve '" + logId + "': no resolver for command"); + } + return commandResolver; } private static final KeyedIntObjectKey RESOLVER_OBJECT_TICKET_ID = - new KeyedIntObjectKey.BasicStrict() { + new KeyedIntObjectKey.BasicStrict<>() { @Override public int getIntKey(final TicketResolver ticketResolver) { return ticketResolver.ticketRoute(); } }; - private static final KeyedObjectKey RESOLVER_OBJECT_DESCRIPTOR_ID = - new KeyedObjectKey.Basic() { + private static final KeyedObjectKey RESOLVER_OBJECT_DESCRIPTOR_ID = + new KeyedObjectKey.Basic<>() { @Override - public String getKey(TicketResolver ticketResolver) { + public String getKey(PathResolverPrefixedBase ticketResolver) { return ticketResolver.flightDescriptorRoute(); } }; diff --git a/server/test-utils/build.gradle b/server/test-utils/build.gradle index 6f3c56e1af0..f3eab96e8d6 100644 --- a/server/test-utils/build.gradle +++ b/server/test-utils/build.gradle @@ -5,7 +5,11 @@ plugins { } dependencies { - implementation project(":util-thread") + implementation project(':util-thread') + compileOnlyApi(project(':util-thread')) { + because 'downstream dagger compile' + } + implementation project(':Base') implementation project(':authentication') implementation project(':authorization') diff --git a/server/test-utils/src/main/java/io/deephaven/server/runner/DeephavenApiServerTestBase.java b/server/test-utils/src/main/java/io/deephaven/server/runner/DeephavenApiServerTestBase.java index 1d39ac9bbd3..168a4330467 100644 --- a/server/test-utils/src/main/java/io/deephaven/server/runner/DeephavenApiServerTestBase.java +++ b/server/test-utils/src/main/java/io/deephaven/server/runner/DeephavenApiServerTestBase.java @@ -133,6 +133,10 @@ interface Builder { @Inject RpcServerStateInterceptor serverStateInterceptor; + protected DeephavenApiServerTestBase.TestComponent.Builder testComponentBuilder() { + return DaggerDeephavenApiServerTestBase_TestComponent.builder(); + } + @Before public void setUp() throws Exception { logBuffer = new LogBuffer(128); @@ -149,7 +153,7 @@ public void setUp() throws Exception { .port(-1) .build(); - DaggerDeephavenApiServerTestBase_TestComponent.builder() + testComponentBuilder() .withServerConfig(config) .withAuthorizationProvider(new CommunityAuthorizationProvider()) .withOut(System.out) diff --git a/server/test-utils/src/main/java/io/deephaven/server/test/TestAuthorizationProvider.java b/server/test-utils/src/main/java/io/deephaven/server/test/TestAuthorizationProvider.java index c7877814673..07695a56d06 100644 --- a/server/test-utils/src/main/java/io/deephaven/server/test/TestAuthorizationProvider.java +++ b/server/test-utils/src/main/java/io/deephaven/server/test/TestAuthorizationProvider.java @@ -95,6 +95,14 @@ public HierarchicalTableServiceContextualAuthWiring.TestUseOnly getHierarchicalT @Override public TicketResolver.Authorization getTicketResolverAuthorization() { return new TicketResolver.Authorization() { + @Override + public boolean isDeniedAccess(Object source) { + if (delegateTicketTransformation != null) { + return delegateTicketTransformation.isDeniedAccess(source); + } + return source == null; + } + @Override public T transform(final T source) { if (delegateTicketTransformation != null) { diff --git a/settings.gradle b/settings.gradle index bf2db5d875c..2c4bd554e9d 100644 --- a/settings.gradle +++ b/settings.gradle @@ -424,6 +424,9 @@ include ':clock-impl' include ':sql' +include ':extensions-flight-sql' +project(':extensions-flight-sql').projectDir = file('extensions/flight-sql') + include(':codec-api') project(':codec-api').projectDir = file('codec/api') include(':codec-builtin') diff --git a/sql/src/main/java/io/deephaven/sql/RelNodeVisitorAdapter.java b/sql/src/main/java/io/deephaven/sql/RelNodeVisitorAdapter.java index 100276c55d0..8cd01ba50c5 100644 --- a/sql/src/main/java/io/deephaven/sql/RelNodeVisitorAdapter.java +++ b/sql/src/main/java/io/deephaven/sql/RelNodeVisitorAdapter.java @@ -56,7 +56,7 @@ public RelNode visit(TableFunctionScan scan) { // SELECT * FROM time_table("00:00:01") // // Potentially related to design decisions around SQLTODO(catalog-reader-implementation) - throw new UnsupportedOperationException("SQLTODO(custom-sources)"); + throw new UnsupportedSqlOperation("SQLTODO(custom-sources)", TableFunctionScan.class); } @Override @@ -73,7 +73,7 @@ public RelNode visit(LogicalFilter filter) { @Override public RelNode visit(LogicalCalc calc) { - throw new UnsupportedOperationException(); + throw new UnsupportedSqlOperation(LogicalCalc.class); } @Override @@ -90,7 +90,7 @@ public RelNode visit(LogicalJoin join) { @Override public RelNode visit(LogicalCorrelate correlate) { - throw new UnsupportedOperationException(); + throw new UnsupportedSqlOperation(LogicalCorrelate.class); } @Override @@ -103,14 +103,14 @@ public RelNode visit(LogicalUnion union) { public RelNode visit(LogicalIntersect intersect) { // SQLTODO(logical-intersect) // table.whereIn - throw new UnsupportedOperationException("SQLTODO(logical-intersect)"); + throw new UnsupportedSqlOperation("SQLTODO(logical-intersect)", LogicalIntersect.class); } @Override public RelNode visit(LogicalMinus minus) { // SQLTODO(logical-minus) // table.whereNotIn - throw new UnsupportedOperationException("SQLTODO(logical-minus)"); + throw new UnsupportedSqlOperation("SQLTODO(logical-minus)", LogicalMatch.class); } @Override @@ -121,7 +121,7 @@ public RelNode visit(LogicalAggregate aggregate) { @Override public RelNode visit(LogicalMatch match) { - throw new UnsupportedOperationException(); + throw new UnsupportedSqlOperation(LogicalMatch.class); } @Override @@ -132,16 +132,16 @@ public RelNode visit(LogicalSort sort) { @Override public RelNode visit(LogicalExchange exchange) { - throw new UnsupportedOperationException(); + throw new UnsupportedSqlOperation(LogicalExchange.class); } @Override public RelNode visit(LogicalTableModify modify) { - throw new UnsupportedOperationException(); + throw new UnsupportedSqlOperation(LogicalTableModify.class); } @Override public RelNode visit(RelNode other) { - throw new UnsupportedOperationException(); + throw new UnsupportedSqlOperation(RelNode.class); } } diff --git a/sql/src/main/java/io/deephaven/sql/RexVisitorBase.java b/sql/src/main/java/io/deephaven/sql/RexVisitorBase.java index 574d1b40f65..a15f72ef833 100644 --- a/sql/src/main/java/io/deephaven/sql/RexVisitorBase.java +++ b/sql/src/main/java/io/deephaven/sql/RexVisitorBase.java @@ -23,75 +23,76 @@ class RexVisitorBase implements RexVisitor { @Override public T visitInputRef(RexInputRef inputRef) { - throw unsupported(inputRef); + throw unsupported(inputRef, RexInputRef.class); } @Override public T visitLocalRef(RexLocalRef localRef) { - throw unsupported(localRef); + throw unsupported(localRef, RexLocalRef.class); } @Override public T visitLiteral(RexLiteral literal) { - throw unsupported(literal); + throw unsupported(literal, RexLiteral.class); } @Override public T visitCall(RexCall call) { - throw unsupported(call); + throw unsupported(call, RexCall.class); } @Override public T visitOver(RexOver over) { - throw unsupported(over); + throw unsupported(over, RexOver.class); } @Override public T visitCorrelVariable(RexCorrelVariable correlVariable) { - throw unsupported(correlVariable); + throw unsupported(correlVariable, RexCorrelVariable.class); } @Override public T visitDynamicParam(RexDynamicParam dynamicParam) { - throw unsupported(dynamicParam); + throw unsupported(dynamicParam, RexDynamicParam.class); } @Override public T visitRangeRef(RexRangeRef rangeRef) { - throw unsupported(rangeRef); + throw unsupported(rangeRef, RexRangeRef.class); } @Override public T visitFieldAccess(RexFieldAccess fieldAccess) { - throw unsupported(fieldAccess); + throw unsupported(fieldAccess, RexFieldAccess.class); } @Override public T visitSubQuery(RexSubQuery subQuery) { - throw unsupported(subQuery); + throw unsupported(subQuery, RexSubQuery.class); } @Override public T visitTableInputRef(RexTableInputRef fieldRef) { - throw unsupported(fieldRef); + throw unsupported(fieldRef, RexTableInputRef.class); } @Override public T visitPatternFieldRef(RexPatternFieldRef fieldRef) { - throw unsupported(fieldRef); + throw unsupported(fieldRef, RexPatternFieldRef.class); } @Override public T visitLambda(RexLambda fieldRef) { - throw unsupported(fieldRef); + throw unsupported(fieldRef, RexLambda.class); } @Override public T visitLambdaRef(RexLambdaRef fieldRef) { - throw unsupported(fieldRef); + throw unsupported(fieldRef, RexLambdaRef.class); } - private UnsupportedOperationException unsupported(RexNode node) { - return new UnsupportedOperationException(String.format("%s: %s", getClass().getName(), node.toString())); + private UnsupportedSqlOperation unsupported(T node, Class clazz) { + return new UnsupportedSqlOperation( + String.format("%s: %s %s", getClass().getName(), node.getClass().getName(), node.toString()), clazz); } } diff --git a/sql/src/main/java/io/deephaven/sql/UnsupportedSqlOperation.java b/sql/src/main/java/io/deephaven/sql/UnsupportedSqlOperation.java new file mode 100644 index 00000000000..d09bc8d974e --- /dev/null +++ b/sql/src/main/java/io/deephaven/sql/UnsupportedSqlOperation.java @@ -0,0 +1,23 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.sql; + +import java.util.Objects; + +public class UnsupportedSqlOperation extends UnsupportedOperationException { + private final Class clazz; + + public UnsupportedSqlOperation(Class clazz) { + this.clazz = Objects.requireNonNull(clazz); + } + + public UnsupportedSqlOperation(String message, Class clazz) { + super(message); + this.clazz = Objects.requireNonNull(clazz); + } + + public Class clazz() { + return clazz; + } +}