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..b4890a0b245 --- /dev/null +++ b/extensions/flight-sql/build.gradle @@ -0,0 +1,85 @@ +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') + + // Another reason why we have a separate source set is because flight-sql-jdbc-driver breaks logging until we can + // upgrade to a newer version. + // See https://github.com/apache/arrow/pull/40908 + // See https://github.com/deephaven/deephaven-core/issues/5947 + // Although, even when we upgrade, it may make sense to keep JDBC separate so that the + // FlightSqlClient testing does not need to depend on 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..71247f2a01e --- /dev/null +++ b/extensions/flight-sql/src/jdbcTest/java/io/deephaven/server/flightsql/FlightSqlJdbcTestBase.java @@ -0,0 +1,172 @@ +// +// 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 (SQLException e) { + assertThat((Throwable) e).getRootCause() + .hasMessageContaining( + "Flight SQL: Must use the original session; is the client echoing the authentication token properly?"); + } + // 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) { + // Note: this is arguably a JDBC implementation bug; it should be throwing SQLException, but it's + // exposing shadowed internal error from Flight. + assertThat(e.getClass().getName()).isEqualTo("cfjd.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..c38e7821a74 --- /dev/null +++ b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlCommandHelper.java @@ -0,0 +1,227 @@ +// +// 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.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); + } + + public static boolean handlesCommand(FlightDescriptor descriptor) { + if (descriptor.getType() != FlightDescriptor.DescriptorType.CMD) { + // noinspection DataFlowIssue + throw Assert.statementNeverExecuted(); + } + // 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 (descriptor.getType() != FlightDescriptor.DescriptorType.CMD) { + // If we get here, there is an error with io.deephaven.server.session.TicketRouter.getPathResolver / + // handlesPath + // noinspection DataFlowIssue + throw Assert.statementNeverExecuted(); + } + final Any command = parseOrNull(descriptor.getCmd()); + if (command == null) { + // If we get here, there is an error with io.deephaven.server.session.TicketRouter.getCommandResolver / + // handlesCommand + // noinspection DataFlowIssue + throw Assert.statementNeverExecuted(); + } + final String typeUrl = command.getTypeUrl(); + if (!typeUrl.startsWith(FlightSqlSharedConstants.FLIGHT_SQL_COMMAND_TYPE_PREFIX)) { + // If we get here, there is an error with io.deephaven.server.session.TicketRouter.getCommandResolver / + // handlesCommand + // noinspection DataFlowIssue + throw Assert.statementNeverExecuted(); + } + 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)); + } + 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); + } + } +} 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..13fdc5e08c4 --- /dev/null +++ b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlResolver.java @@ -0,0 +1,1747 @@ +// +// 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.Set; +import java.util.function.Consumer; +import java.util.function.Function; +import java.util.function.Predicate; +import java.util.regex.Pattern; + +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 (!handlesActionType(action.getType())) { + // If we get here, there is an error with io.deephaven.server.session.ActionRouter.doAction / + // handlesActionType + // noinspection DataFlowIssue + throw Assert.statementNeverExecuted(); + } + 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 (table.isRefreshing()) { + throw new IllegalStateException( + "TicketHandler implementation error; should only override totalRecords for non-refreshing tables"); + } + if (table.size() != totalRecords) { + throw new IllegalStateException( + "Ticket handler implementation error; totalRecords does not match the table size"); + } + } + 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(); + if (table.isRefreshing()) { + throw new IllegalArgumentException("Expected static table"); + } + 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) { + // 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 int L = flightSqlPattern.length(); + 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); + } + }; + for (int i = 0; i < L; ++i) { + final char c = flightSqlPattern.charAt(i); + if (c == '%') { + appendQuoted.run(); + pattern.append(".*"); + } else if (c == '_') { + appendQuoted.run(); + pattern.append('.'); + } else { + quoted.append(c); + } + } + 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..7fe3a332848 --- /dev/null +++ b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlSharedConstants.java @@ -0,0 +1,46 @@ +// +// 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"; + + // Need to update to newer FlightSql version for this + // 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..a3aa98a87c1 --- /dev/null +++ b/extensions/flight-sql/src/main/java/io/deephaven/server/flightsql/FlightSqlTicketHelper.java @@ -0,0 +1,182 @@ +// +// 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 (ticket.get() != TICKET_PREFIX) { + // If we get here, it means there is an error with FlightSqlResolver.ticketRoute / + // io.deephaven.server.session.TicketRouter.getResolver + // noinspection DataFlowIssue + throw Assert.statementNeverExecuted(); + } + 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..3db141d4e11 --- /dev/null +++ b/extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlFilterPredicateTest.java @@ -0,0 +1,160 @@ +// +// 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); + } + + 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..a2b1a93d627 --- /dev/null +++ b/extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlTest.java @@ -0,0 +1,1092 @@ +// +// 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.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() { + // 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); + getSchemaUnknown(() -> flightClient.getSchema(descriptor), typeUrl); + commandUnknown(() -> flightClient.getInfo(descriptor), typeUrl); + } + + @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 0dea94ab38d..5fbe93b6854 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -99,6 +99,8 @@ 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-grpc = { module = "org.apache.arrow:flight-grpc", 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 ac3bedf066d..f2656ca7b3e 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; @Singleton public class FlightServiceGrpcImpl extends FlightServiceGrpc.FlightServiceImplBase { @@ -53,6 +60,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; @@ -64,6 +72,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; @@ -71,6 +80,7 @@ public FlightServiceGrpcImpl( this.sessionService = sessionService; this.errorTransformer = errorTransformer; this.ticketRouter = ticketRouter; + this.actionRouter = actionRouter; this.doExchangeFactory = doExchangeFactory; this.authRequestHandlers = authRequestHandlers; } @@ -199,14 +209,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, @@ -332,4 +363,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 f9ccce190d2..a7195f53a01 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); @@ -266,12 +269,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 @@ -307,6 +315,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); + } } } } @@ -332,8 +343,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 @@ -352,20 +363,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 74b5aafc304..5a1ba49487d 100644 --- a/server/src/main/java/io/deephaven/server/session/SessionState.java +++ b/server/src/main/java/io/deephaven/server/session/SessionState.java @@ -718,7 +718,7 @@ private synchronized void setWork( return; } - this.exportMain = exportMain; + this.exportMain = Objects.requireNonNull(exportMain); this.errorHandler = errorHandler; this.successHandler = successHandler; @@ -800,6 +800,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. * @@ -1368,7 +1375,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 54133d912b0..2fbef0ee2c4 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 d2bad06572b..dfc7e40f6eb 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; + } +}