Skip to content

Commit b61abcf

Browse files
devinrsmithdevinrsmith
authored
test: Add Flight SQL ADBC Java tests (deephaven#6432)
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 deephaven/deephaven-docs-community#365
1 parent 7a9526c commit b61abcf

File tree

7 files changed

+392
-1
lines changed

7 files changed

+392
-1
lines changed

‎extensions/flight-sql/build.gradle

+36-1
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,21 @@ plugins {
66
description = 'The Deephaven Flight SQL library'
77

88
sourceSets {
9+
adbcTest {
10+
compileClasspath += sourceSets.main.output
11+
runtimeClasspath += sourceSets.main.output
12+
}
13+
914
jdbcTest {
1015
compileClasspath += sourceSets.main.output
1116
runtimeClasspath += sourceSets.main.output
1217
}
1318
}
1419

1520
configurations {
21+
adbcTestImplementation.extendsFrom implementation
22+
adbcTestRuntimeOnly.extendsFrom runtimeOnly
23+
1624
jdbcTestImplementation.extendsFrom implementation
1725
jdbcTestRuntimeOnly.extendsFrom runtimeOnly
1826
}
@@ -42,8 +50,23 @@ dependencies {
4250
testRuntimeOnly project(':log-to-slf4j')
4351
testRuntimeOnly libs.slf4j.simple
4452

53+
// ADBC testing needs an actually server instance bound to a port because it can only connect over ADBC URIs like
54+
// grpc://localhost:10000
55+
adbcTestImplementation project(':server-jetty')
56+
adbcTestImplementation libs.adbc.flight.sql
57+
58+
adbcTestImplementation project(':server-test-utils')
59+
adbcTestAnnotationProcessor libs.dagger.compiler
60+
adbcTestImplementation libs.assertj
61+
adbcTestImplementation platform(libs.junit.bom)
62+
adbcTestImplementation libs.junit.jupiter
63+
adbcTestRuntimeOnly libs.junit.platform.launcher
64+
adbcTestRuntimeOnly libs.junit.vintage.engine
65+
adbcTestRuntimeOnly project(':log-to-slf4j')
66+
adbcTestRuntimeOnly libs.slf4j.simple
67+
4568
// JDBC testing needs an actually server instance bound to a port because it can only connect over JDBC URIs like
46-
// jdbc:arrow-flight-sql://localhost:1000.
69+
// jdbc:arrow-flight-sql://localhost:10000.
4770
jdbcTestImplementation project(':server-jetty')
4871
jdbcTestRuntimeOnly libs.arrow.flight.sql.jdbc
4972

@@ -62,6 +85,17 @@ test {
6285
useJUnitPlatform()
6386
}
6487

88+
def adbcTest = tasks.register('adbcTest', Test) {
89+
description = 'Runs ADBC tests.'
90+
group = 'verification'
91+
92+
testClassesDirs = sourceSets.adbcTest.output.classesDirs
93+
classpath = sourceSets.adbcTest.runtimeClasspath
94+
shouldRunAfter test
95+
96+
useJUnitPlatform()
97+
}
98+
6599
def jdbcTest = tasks.register('jdbcTest', Test) {
66100
description = 'Runs JDBC tests.'
67101
group = 'verification'
@@ -73,6 +107,7 @@ def jdbcTest = tasks.register('jdbcTest', Test) {
73107
useJUnitPlatform()
74108
}
75109

110+
check.dependsOn adbcTest
76111
check.dependsOn jdbcTest
77112

78113
apply plugin: 'io.deephaven.java-open-nio'

‎extensions/flight-sql/src/adbcTest/java/io/deephaven/server/DeephavenServerTestBase.java

+62
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
//
2+
// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending
3+
//
4+
package io.deephaven.server;
5+
6+
import io.deephaven.engine.context.ExecutionContext;
7+
import io.deephaven.io.logger.LogBuffer;
8+
import io.deephaven.io.logger.LogBufferGlobal;
9+
import io.deephaven.server.runner.GrpcServer;
10+
import io.deephaven.server.runner.MainHelper;
11+
import io.deephaven.util.SafeCloseable;
12+
import org.junit.jupiter.api.AfterEach;
13+
import org.junit.jupiter.api.BeforeAll;
14+
import org.junit.jupiter.api.BeforeEach;
15+
import org.junit.jupiter.api.Timeout;
16+
17+
import java.io.IOException;
18+
import java.util.concurrent.TimeUnit;
19+
20+
@Timeout(30)
21+
public abstract class DeephavenServerTestBase {
22+
23+
public interface TestComponent {
24+
25+
GrpcServer server();
26+
27+
ExecutionContext executionContext();
28+
}
29+
30+
protected TestComponent component;
31+
32+
private LogBuffer logBuffer;
33+
private SafeCloseable executionContext;
34+
private GrpcServer server;
35+
protected int localPort;
36+
37+
protected abstract TestComponent component();
38+
39+
@BeforeAll
40+
static void setupOnce() throws IOException {
41+
MainHelper.bootstrapProjectDirectories();
42+
}
43+
44+
@BeforeEach
45+
void setup() throws IOException {
46+
logBuffer = new LogBuffer(128);
47+
LogBufferGlobal.setInstance(logBuffer);
48+
component = component();
49+
executionContext = component.executionContext().open();
50+
server = component.server();
51+
server.start();
52+
localPort = server.getPort();
53+
}
54+
55+
@AfterEach
56+
void tearDown() throws InterruptedException {
57+
server.stopWithTimeout(10, TimeUnit.SECONDS);
58+
server.join();
59+
executionContext.close();
60+
LogBufferGlobal.clear(logBuffer);
61+
}
62+
}

‎extensions/flight-sql/src/adbcTest/java/io/deephaven/server/flightsql/FlightSqlAdbcTestBase.java

+115
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
//
2+
// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending
3+
//
4+
package io.deephaven.server.flightsql;
5+
6+
import io.deephaven.server.DeephavenServerTestBase;
7+
import org.apache.arrow.adbc.core.AdbcConnection;
8+
import org.apache.arrow.adbc.core.AdbcDatabase;
9+
import org.apache.arrow.adbc.core.AdbcDriver;
10+
import org.apache.arrow.adbc.core.AdbcException;
11+
import org.apache.arrow.adbc.core.AdbcStatement;
12+
import org.apache.arrow.adbc.driver.flightsql.FlightSqlConnectionProperties;
13+
import org.apache.arrow.adbc.driver.flightsql.FlightSqlDriverFactory;
14+
import org.apache.arrow.memory.BufferAllocator;
15+
import org.apache.arrow.memory.RootAllocator;
16+
import org.apache.arrow.vector.IntVector;
17+
import org.apache.arrow.vector.VectorSchemaRoot;
18+
import org.apache.arrow.vector.ipc.ArrowReader;
19+
import org.apache.arrow.vector.types.Types;
20+
import org.apache.arrow.vector.types.pojo.Field;
21+
import org.apache.arrow.vector.types.pojo.FieldType;
22+
import org.apache.arrow.vector.types.pojo.Schema;
23+
import org.junit.jupiter.api.AfterEach;
24+
import org.junit.jupiter.api.BeforeEach;
25+
import org.junit.jupiter.api.Test;
26+
27+
import java.util.HashMap;
28+
import java.util.List;
29+
import java.util.Map;
30+
31+
import static org.assertj.core.api.Assertions.assertThat;
32+
33+
public abstract class FlightSqlAdbcTestBase extends DeephavenServerTestBase {
34+
35+
private static final Map<String, String> DEEPHAVEN_INT = Map.of(
36+
"deephaven:isSortable", "true",
37+
"deephaven:isRowStyle", "false",
38+
"deephaven:isPartitioning", "false",
39+
"deephaven:type", "int",
40+
"deephaven:isNumberFormat", "false",
41+
"deephaven:isStyle", "false",
42+
"deephaven:isDateFormat", "false");
43+
44+
BufferAllocator allocator;
45+
AdbcDatabase database;
46+
AdbcConnection connection;
47+
48+
@BeforeEach
49+
void setUp() throws AdbcException {
50+
final Map<String, Object> options = new HashMap<>();
51+
AdbcDriver.PARAM_URI.set(options, String.format("grpc://localhost:%d", localPort));
52+
FlightSqlConnectionProperties.WITH_COOKIE_MIDDLEWARE.set(options, true);
53+
options.put(FlightSqlConnectionProperties.RPC_CALL_HEADER_PREFIX + "Authorization", "Anonymous");
54+
options.put(FlightSqlConnectionProperties.RPC_CALL_HEADER_PREFIX + "x-deephaven-auth-cookie-request", "true");
55+
allocator = new RootAllocator();
56+
database = new FlightSqlDriverFactory().getDriver(allocator).open(options);
57+
connection = database.connect();
58+
}
59+
60+
@AfterEach
61+
void tearDown() throws Exception {
62+
connection.close();
63+
database.close();
64+
allocator.close();
65+
}
66+
67+
@Test
68+
void executeSchema() throws Exception {
69+
final Schema expectedSchema = new Schema(List
70+
.of(new Field("Foo", new FieldType(true, Types.MinorType.INT.getType(), null, DEEPHAVEN_INT), null)));
71+
try (final AdbcStatement statement = connection.createStatement()) {
72+
statement.setSqlQuery("SELECT 42 as Foo");
73+
assertThat(statement.executeSchema()).isEqualTo(expectedSchema);
74+
}
75+
}
76+
77+
@Test
78+
void executeQuery() throws Exception {
79+
final Schema expectedSchema = new Schema(List
80+
.of(new Field("Foo", new FieldType(true, Types.MinorType.INT.getType(), null, DEEPHAVEN_INT), null)));
81+
try (final AdbcStatement statement = connection.createStatement()) {
82+
statement.setSqlQuery("SELECT 42 as Foo");
83+
try (final AdbcStatement.QueryResult result = statement.executeQuery()) {
84+
final ArrowReader reader = result.getReader();
85+
assertThat(reader.loadNextBatch()).isTrue();
86+
final VectorSchemaRoot root = reader.getVectorSchemaRoot();
87+
assertThat(root.getSchema()).isEqualTo(expectedSchema);
88+
final IntVector vector = (IntVector) root.getVector(0);
89+
assertThat(vector.isNull(0)).isFalse();
90+
assertThat(vector.get(0)).isEqualTo(42);
91+
assertThat(reader.loadNextBatch()).isFalse();
92+
}
93+
}
94+
}
95+
96+
@Test
97+
void preparedExecuteQuery() throws Exception {
98+
final Schema expectedSchema = new Schema(List
99+
.of(new Field("Foo", new FieldType(true, Types.MinorType.INT.getType(), null, DEEPHAVEN_INT), null)));
100+
try (final AdbcStatement statement = connection.createStatement()) {
101+
statement.setSqlQuery("SELECT 42 as Foo");
102+
statement.prepare();
103+
try (final AdbcStatement.QueryResult result = statement.executeQuery()) {
104+
final ArrowReader reader = result.getReader();
105+
assertThat(reader.loadNextBatch()).isTrue();
106+
final VectorSchemaRoot root = reader.getVectorSchemaRoot();
107+
assertThat(root.getSchema()).isEqualTo(expectedSchema);
108+
final IntVector vector = (IntVector) root.getVector(0);
109+
assertThat(vector.isNull(0)).isFalse();
110+
assertThat(vector.get(0)).isEqualTo(42);
111+
assertThat(reader.loadNextBatch()).isFalse();
112+
}
113+
}
114+
}
115+
}

‎extensions/flight-sql/src/adbcTest/java/io/deephaven/server/flightsql/FlightSqlTestModule.java

+123
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
//
2+
// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending
3+
//
4+
package io.deephaven.server.flightsql;
5+
6+
import dagger.Module;
7+
import dagger.Provides;
8+
import dagger.multibindings.IntoSet;
9+
import io.deephaven.base.clock.Clock;
10+
import io.deephaven.engine.context.ExecutionContext;
11+
import io.deephaven.engine.updategraph.OperationInitializer;
12+
import io.deephaven.engine.updategraph.UpdateGraph;
13+
import io.deephaven.engine.util.AbstractScriptSession;
14+
import io.deephaven.engine.util.NoLanguageDeephavenSession;
15+
import io.deephaven.engine.util.ScriptSession;
16+
import io.deephaven.server.arrow.ArrowModule;
17+
import io.deephaven.server.auth.AuthorizationProvider;
18+
import io.deephaven.server.config.ConfigServiceModule;
19+
import io.deephaven.server.console.ConsoleModule;
20+
import io.deephaven.server.log.LogModule;
21+
import io.deephaven.server.plugin.PluginsModule;
22+
import io.deephaven.server.session.ExportTicketResolver;
23+
import io.deephaven.server.session.ObfuscatingErrorTransformerModule;
24+
import io.deephaven.server.session.SessionModule;
25+
import io.deephaven.server.session.TicketResolver;
26+
import io.deephaven.server.table.TableModule;
27+
import io.deephaven.server.test.TestAuthModule;
28+
import io.deephaven.server.test.TestAuthorizationProvider;
29+
import io.deephaven.server.util.Scheduler;
30+
31+
import javax.inject.Named;
32+
import javax.inject.Singleton;
33+
import java.util.concurrent.Executors;
34+
import java.util.concurrent.ScheduledExecutorService;
35+
36+
@Module(includes = {
37+
ArrowModule.class,
38+
ConfigServiceModule.class,
39+
ConsoleModule.class,
40+
LogModule.class,
41+
SessionModule.class,
42+
TableModule.class,
43+
TestAuthModule.class,
44+
ObfuscatingErrorTransformerModule.class,
45+
PluginsModule.class,
46+
FlightSqlModule.class
47+
})
48+
public class FlightSqlTestModule {
49+
@IntoSet
50+
@Provides
51+
TicketResolver ticketResolver(ExportTicketResolver resolver) {
52+
return resolver;
53+
}
54+
55+
@Singleton
56+
@Provides
57+
AbstractScriptSession<?> provideAbstractScriptSession(
58+
final UpdateGraph updateGraph,
59+
final OperationInitializer operationInitializer) {
60+
return new NoLanguageDeephavenSession(
61+
updateGraph, operationInitializer, "non-script-session");
62+
}
63+
64+
@Provides
65+
ScriptSession provideScriptSession(AbstractScriptSession<?> scriptSession) {
66+
return scriptSession;
67+
}
68+
69+
@Provides
70+
@Singleton
71+
ScheduledExecutorService provideExecutorService() {
72+
return Executors.newScheduledThreadPool(1);
73+
}
74+
75+
@Provides
76+
Scheduler provideScheduler(ScheduledExecutorService concurrentExecutor) {
77+
return new Scheduler.DelegatingImpl(
78+
Executors.newSingleThreadExecutor(),
79+
concurrentExecutor,
80+
Clock.system());
81+
}
82+
83+
@Provides
84+
@Named("session.tokenExpireMs")
85+
long provideTokenExpireMs() {
86+
return 60_000_000;
87+
}
88+
89+
@Provides
90+
@Named("http.port")
91+
int provideHttpPort() {
92+
return 0;// 'select first available'
93+
}
94+
95+
@Provides
96+
@Named("grpc.maxInboundMessageSize")
97+
int provideMaxInboundMessageSize() {
98+
return 1024 * 1024;
99+
}
100+
101+
@Provides
102+
AuthorizationProvider provideAuthorizationProvider(TestAuthorizationProvider provider) {
103+
return provider;
104+
}
105+
106+
@Provides
107+
@Singleton
108+
TestAuthorizationProvider provideTestAuthorizationProvider() {
109+
return new TestAuthorizationProvider();
110+
}
111+
112+
@Provides
113+
@Singleton
114+
static UpdateGraph provideUpdateGraph() {
115+
return ExecutionContext.getContext().getUpdateGraph();
116+
}
117+
118+
@Provides
119+
@Singleton
120+
static OperationInitializer provideOperationInitializer() {
121+
return ExecutionContext.getContext().getOperationInitializer();
122+
}
123+
}

0 commit comments

Comments
 (0)