From 09f8d35730ba7765404325291a2d4b3723785f21 Mon Sep 17 00:00:00 2001 From: Devin Smith Date: Mon, 25 Nov 2024 13:16:21 -0800 Subject: [PATCH] test: Add Flight SQL ADBC tests This is adding Flight SQL Java ADBC tests, mainly in support of a documentation effort on how to connect different clients to Deephaven Flight SQL, see https://github.com/deephaven/deephaven-docs-community/issues/365 --- extensions/flight-sql/build.gradle | 37 +++++- .../server/DeephavenServerTestBase.java | 62 +++++++++ .../flightsql/FlightSqlAdbcTestBase.java | 115 ++++++++++++++++ .../server/flightsql/FlightSqlTestModule.java | 123 ++++++++++++++++++ .../server/flightsql/JettyTestComponent.java | 38 ++++++ .../jetty/FlightSqlAdbcTestJetty.java | 15 +++ gradle/libs.versions.toml | 3 + 7 files changed, 392 insertions(+), 1 deletion(-) create mode 100644 extensions/flight-sql/src/adbcTest/java/io/deephaven/server/DeephavenServerTestBase.java create mode 100644 extensions/flight-sql/src/adbcTest/java/io/deephaven/server/flightsql/FlightSqlAdbcTestBase.java create mode 100644 extensions/flight-sql/src/adbcTest/java/io/deephaven/server/flightsql/FlightSqlTestModule.java create mode 100644 extensions/flight-sql/src/adbcTest/java/io/deephaven/server/flightsql/JettyTestComponent.java create mode 100644 extensions/flight-sql/src/adbcTest/java/io/deephaven/server/flightsql/jetty/FlightSqlAdbcTestJetty.java diff --git a/extensions/flight-sql/build.gradle b/extensions/flight-sql/build.gradle index 3c4e155e411..3c1c2a0b37c 100644 --- a/extensions/flight-sql/build.gradle +++ b/extensions/flight-sql/build.gradle @@ -6,6 +6,11 @@ plugins { description = 'The Deephaven Flight SQL library' sourceSets { + adbcTest { + compileClasspath += sourceSets.main.output + runtimeClasspath += sourceSets.main.output + } + jdbcTest { compileClasspath += sourceSets.main.output runtimeClasspath += sourceSets.main.output @@ -13,6 +18,9 @@ sourceSets { } configurations { + adbcTestImplementation.extendsFrom implementation + adbcTestRuntimeOnly.extendsFrom runtimeOnly + jdbcTestImplementation.extendsFrom implementation jdbcTestRuntimeOnly.extendsFrom runtimeOnly } @@ -42,8 +50,23 @@ dependencies { testRuntimeOnly project(':log-to-slf4j') testRuntimeOnly libs.slf4j.simple + // ADBC testing needs an actually server instance bound to a port because it can only connect over ADBC URIs like + // grpc://localhost:10000 + adbcTestImplementation project(':server-jetty') + adbcTestImplementation libs.adbc.flight.sql + + adbcTestImplementation project(':server-test-utils') + adbcTestAnnotationProcessor libs.dagger.compiler + adbcTestImplementation libs.assertj + adbcTestImplementation platform(libs.junit.bom) + adbcTestImplementation libs.junit.jupiter + adbcTestRuntimeOnly libs.junit.platform.launcher + adbcTestRuntimeOnly libs.junit.vintage.engine + adbcTestRuntimeOnly project(':log-to-slf4j') + adbcTestRuntimeOnly 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. + // jdbc:arrow-flight-sql://localhost:10000. jdbcTestImplementation project(':server-jetty') jdbcTestRuntimeOnly libs.arrow.flight.sql.jdbc @@ -62,6 +85,17 @@ test { useJUnitPlatform() } +def adbcTest = tasks.register('adbcTest', Test) { + description = 'Runs ADBC tests.' + group = 'verification' + + testClassesDirs = sourceSets.adbcTest.output.classesDirs + classpath = sourceSets.adbcTest.runtimeClasspath + shouldRunAfter test + + useJUnitPlatform() +} + def jdbcTest = tasks.register('jdbcTest', Test) { description = 'Runs JDBC tests.' group = 'verification' @@ -73,6 +107,7 @@ def jdbcTest = tasks.register('jdbcTest', Test) { useJUnitPlatform() } +check.dependsOn adbcTest check.dependsOn jdbcTest apply plugin: 'io.deephaven.java-open-nio' diff --git a/extensions/flight-sql/src/adbcTest/java/io/deephaven/server/DeephavenServerTestBase.java b/extensions/flight-sql/src/adbcTest/java/io/deephaven/server/DeephavenServerTestBase.java new file mode 100644 index 00000000000..76458d91a57 --- /dev/null +++ b/extensions/flight-sql/src/adbcTest/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/adbcTest/java/io/deephaven/server/flightsql/FlightSqlAdbcTestBase.java b/extensions/flight-sql/src/adbcTest/java/io/deephaven/server/flightsql/FlightSqlAdbcTestBase.java new file mode 100644 index 00000000000..d035e790781 --- /dev/null +++ b/extensions/flight-sql/src/adbcTest/java/io/deephaven/server/flightsql/FlightSqlAdbcTestBase.java @@ -0,0 +1,115 @@ +// +// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending +// +package io.deephaven.server.flightsql; + +import io.deephaven.server.DeephavenServerTestBase; +import org.apache.arrow.adbc.core.AdbcConnection; +import org.apache.arrow.adbc.core.AdbcDatabase; +import org.apache.arrow.adbc.core.AdbcDriver; +import org.apache.arrow.adbc.core.AdbcException; +import org.apache.arrow.adbc.core.AdbcStatement; +import org.apache.arrow.adbc.driver.flightsql.FlightSqlConnectionProperties; +import org.apache.arrow.adbc.driver.flightsql.FlightSqlDriverFactory; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.apache.arrow.vector.types.Types; +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.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +public abstract class FlightSqlAdbcTestBase extends DeephavenServerTestBase { + + 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"); + + BufferAllocator allocator; + AdbcDatabase database; + AdbcConnection connection; + + @BeforeEach + void setUp() throws AdbcException { + final Map options = new HashMap<>(); + AdbcDriver.PARAM_URI.set(options, String.format("grpc://localhost:%d", localPort)); + FlightSqlConnectionProperties.WITH_COOKIE_MIDDLEWARE.set(options, true); + options.put(FlightSqlConnectionProperties.RPC_CALL_HEADER_PREFIX + "Authorization", "Anonymous"); + options.put(FlightSqlConnectionProperties.RPC_CALL_HEADER_PREFIX + "x-deephaven-auth-cookie-request", "true"); + allocator = new RootAllocator(); + database = new FlightSqlDriverFactory().getDriver(allocator).open(options); + connection = database.connect(); + } + + @AfterEach + void tearDown() throws Exception { + connection.close(); + database.close(); + allocator.close(); + } + + @Test + void executeSchema() throws Exception { + final Schema expectedSchema = new Schema(List + .of(new Field("Foo", new FieldType(true, Types.MinorType.INT.getType(), null, DEEPHAVEN_INT), null))); + try (final AdbcStatement statement = connection.createStatement()) { + statement.setSqlQuery("SELECT 42 as Foo"); + assertThat(statement.executeSchema()).isEqualTo(expectedSchema); + } + } + + @Test + void executeQuery() throws Exception { + final Schema expectedSchema = new Schema(List + .of(new Field("Foo", new FieldType(true, Types.MinorType.INT.getType(), null, DEEPHAVEN_INT), null))); + try (final AdbcStatement statement = connection.createStatement()) { + statement.setSqlQuery("SELECT 42 as Foo"); + try (final AdbcStatement.QueryResult result = statement.executeQuery()) { + final ArrowReader reader = result.getReader(); + assertThat(reader.loadNextBatch()).isTrue(); + final VectorSchemaRoot root = reader.getVectorSchemaRoot(); + assertThat(root.getSchema()).isEqualTo(expectedSchema); + final IntVector vector = (IntVector) root.getVector(0); + assertThat(vector.isNull(0)).isFalse(); + assertThat(vector.get(0)).isEqualTo(42); + assertThat(reader.loadNextBatch()).isFalse(); + } + } + } + + @Test + void preparedExecuteQuery() throws Exception { + final Schema expectedSchema = new Schema(List + .of(new Field("Foo", new FieldType(true, Types.MinorType.INT.getType(), null, DEEPHAVEN_INT), null))); + try (final AdbcStatement statement = connection.createStatement()) { + statement.setSqlQuery("SELECT 42 as Foo"); + statement.prepare(); + try (final AdbcStatement.QueryResult result = statement.executeQuery()) { + final ArrowReader reader = result.getReader(); + assertThat(reader.loadNextBatch()).isTrue(); + final VectorSchemaRoot root = reader.getVectorSchemaRoot(); + assertThat(root.getSchema()).isEqualTo(expectedSchema); + final IntVector vector = (IntVector) root.getVector(0); + assertThat(vector.isNull(0)).isFalse(); + assertThat(vector.get(0)).isEqualTo(42); + assertThat(reader.loadNextBatch()).isFalse(); + } + } + } +} diff --git a/extensions/flight-sql/src/adbcTest/java/io/deephaven/server/flightsql/FlightSqlTestModule.java b/extensions/flight-sql/src/adbcTest/java/io/deephaven/server/flightsql/FlightSqlTestModule.java new file mode 100644 index 00000000000..df7b7e3c9c2 --- /dev/null +++ b/extensions/flight-sql/src/adbcTest/java/io/deephaven/server/flightsql/FlightSqlTestModule.java @@ -0,0 +1,123 @@ +// +// 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 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/adbcTest/java/io/deephaven/server/flightsql/JettyTestComponent.java b/extensions/flight-sql/src/adbcTest/java/io/deephaven/server/flightsql/JettyTestComponent.java new file mode 100644 index 00000000000..add6799ad28 --- /dev/null +++ b/extensions/flight-sql/src/adbcTest/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/adbcTest/java/io/deephaven/server/flightsql/jetty/FlightSqlAdbcTestJetty.java b/extensions/flight-sql/src/adbcTest/java/io/deephaven/server/flightsql/jetty/FlightSqlAdbcTestJetty.java new file mode 100644 index 00000000000..87ac886c782 --- /dev/null +++ b/extensions/flight-sql/src/adbcTest/java/io/deephaven/server/flightsql/jetty/FlightSqlAdbcTestJetty.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.FlightSqlAdbcTestBase; + +public class FlightSqlAdbcTestJetty extends FlightSqlAdbcTestBase { + + @Override + protected TestComponent component() { + return DaggerJettyTestComponent.create(); + } +} diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index 1b33f7f1fe2..bb9212544ce 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -1,4 +1,5 @@ [versions] +adbc = "0.15.0" airlift = "2.0.2" arrow = "18.0.0" autoservice = "1.1.1" @@ -92,6 +93,8 @@ jmh = "1.37" spockframework = "2.3-groovy-3.0" [libraries] +adbc-flight-sql = { module = "org.apache.arrow.adbc:adbc-driver-flight-sql", version.ref = "adbc" } + airlift-aircompressor = { module = "io.airlift:aircompressor", version.ref = "airlift"} arrow-compression = { module = "org.apache.arrow:arrow-compression", version.ref = "arrow" }