This is an automated email from the ASF dual-hosted git repository.
JNSimba pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris-flink-connector.git
The following commit(s) were added to refs/heads/master by this push:
new 4a3c81fc [fix] Close Flight SQL sessions when readers finish (#719)
4a3c81fc is described below
commit 4a3c81fcaff2b21e8edcf3573e2263eec5295075
Author: wudi <[email protected]>
AuthorDate: Fri Nov 6 10:26:47 2026 +0800
[fix] Close Flight SQL sessions when readers finish (#719)
---
.../flink/source/reader/DorisFlightSession.java | 135 +++++++++++++
.../source/reader/DorisFlightValueReader.java | 7 +-
.../source/reader/DorisFlightSessionTest.java | 218 +++++++++++++++++++++
.../source/reader/DorisFlightValueReaderTest.java | 11 +-
.../flink-doris-connector-flink1/pom.xml | 3 +
.../flink-doris-connector-flink2/pom.xml | 3 +
flink-doris-connector/pom.xml | 18 +-
7 files changed, 385 insertions(+), 10 deletions(-)
diff --git
a/flink-doris-connector/flink-doris-connector-base/src/main/java/org/apache/doris/flink/source/reader/DorisFlightSession.java
b/flink-doris-connector/flink-doris-connector-base/src/main/java/org/apache/doris/flink/source/reader/DorisFlightSession.java
new file mode 100644
index 00000000..1a60cc4e
--- /dev/null
+++
b/flink-doris-connector/flink-doris-connector-base/src/main/java/org/apache/doris/flink/source/reader/DorisFlightSession.java
@@ -0,0 +1,135 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.flink.source.reader;
+
+import org.apache.arrow.adbc.core.AdbcConnection;
+import org.apache.arrow.adbc.core.AdbcDatabase;
+import org.apache.arrow.adbc.core.AdbcException;
+import org.apache.arrow.adbc.driver.flightsql.FlightSqlClientWithCallOptions;
+import org.apache.arrow.adbc.driver.flightsql.FlightSqlConnection;
+import org.apache.arrow.flight.CallOption;
+import org.apache.arrow.flight.CallOptions;
+import org.apache.arrow.flight.CloseSessionRequest;
+import org.apache.arrow.flight.CloseSessionResult;
+import org.apache.arrow.flight.sql.FlightSqlClient;
+import org.apache.doris.flink.exception.DorisRuntimeException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.lang.reflect.Field;
+import java.util.Arrays;
+import java.util.Objects;
+import java.util.concurrent.TimeUnit;
+
+/** Closes the server session without taking ownership of the ADBC connection.
*/
+final class DorisFlightSession implements AutoCloseable {
+ private static final Logger LOG =
LoggerFactory.getLogger(DorisFlightSession.class);
+
+ private final AdbcConnection connection;
+ private final FlightSqlClient client;
+ private final CallOption[] callOptions;
+ private boolean closed;
+
+ private DorisFlightSession(
+ AdbcConnection connection, FlightSqlClient client, CallOption[]
callOptions) {
+ this.connection = connection;
+ this.client = client;
+ this.callOptions = Arrays.copyOf(callOptions, callOptions.length + 1);
+ this.callOptions[callOptions.length] = CallOptions.timeout(5,
TimeUnit.SECONDS);
+ }
+
+ static DorisFlightSession open(AdbcDatabase database) throws AdbcException
{
+ // shortcut: ADBC 0.14.0 hides the client; remove once its close()
sends CloseSession.
+ Field connectionClient =
+ field(FlightSqlConnection.class, "client",
FlightSqlClientWithCallOptions.class);
+ Field flightClient =
+ field(FlightSqlClientWithCallOptions.class, "client",
FlightSqlClient.class);
+ Field connectionOptions =
+ field(
+ FlightSqlClientWithCallOptions.class,
+ "connectionOptions",
+ CallOption[].class);
+
+ // Validate the driver layout before authentication creates a server
session.
+ AdbcConnection connection = database.connect();
+ try {
+ Object wrapper = connectionClient.get(connection);
+ FlightSqlClient client =
+ Objects.requireNonNull((FlightSqlClient)
flightClient.get(wrapper));
+ CallOption[] options =
+ Objects.requireNonNull((CallOption[])
connectionOptions.get(wrapper));
+ return new DorisFlightSession(connection, client, options);
+ } catch (ReflectiveOperationException | RuntimeException failure) {
+ try {
+ connection.close();
+ } catch (Exception closeFailure) {
+ failure.addSuppressed(closeFailure);
+ }
+ throw new DorisRuntimeException("Cannot bind ADBC Flight SQL
session", failure);
+ }
+ }
+
+ private static Field field(Class<?> owner, String name, Class<?> type) {
+ try {
+ Field field = owner.getDeclaredField(name);
+ if (field.getType() != type) {
+ throw new NoSuchFieldException(
+ owner.getSimpleName() + "." + name + " has changed
type");
+ }
+ field.setAccessible(true);
+ return field;
+ } catch (ReflectiveOperationException | RuntimeException failure) {
+ throw new DorisRuntimeException(
+ "Unsupported ADBC Flight SQL driver: cannot access session
for CloseSession",
+ failure);
+ }
+ }
+
+ AdbcConnection getConnection() {
+ return connection;
+ }
+
+ public synchronized void closeSession() {
+ if (closed) {
+ return;
+ }
+ closed = true;
+ // Remote session cleanup must not turn a completed read into a
failure.
+ try {
+ CloseSessionResult.Status status =
+ client.closeSession(new CloseSessionRequest(),
callOptions).getStatus();
+ if (status != CloseSessionResult.Status.CLOSED
+ && status != CloseSessionResult.Status.CLOSING) {
+ LOG.warn(
+ "Doris Flight SQL session was not closed, status: {};
it may remain open"
+ + " until the server idle timeout",
+ status);
+ }
+ } catch (Exception e) {
+ LOG.warn(
+ "Failed to close Doris Flight SQL session; it may remain
open until the server"
+ + " idle timeout",
+ e);
+ }
+ }
+
+ @Override
+ public void close() {
+ closeSession();
+ }
+}
diff --git
a/flink-doris-connector/flink-doris-connector-base/src/main/java/org/apache/doris/flink/source/reader/DorisFlightValueReader.java
b/flink-doris-connector/flink-doris-connector-base/src/main/java/org/apache/doris/flink/source/reader/DorisFlightValueReader.java
index 9856db53..f553480c 100644
---
a/flink-doris-connector/flink-doris-connector-base/src/main/java/org/apache/doris/flink/source/reader/DorisFlightValueReader.java
+++
b/flink-doris-connector/flink-doris-connector-base/src/main/java/org/apache/doris/flink/source/reader/DorisFlightValueReader.java
@@ -71,6 +71,7 @@ public class DorisFlightValueReader extends ValueReader
implements AutoCloseable
protected AdbcConnection client;
private RootAllocator allocator;
private AdbcDatabase database;
+ private DorisFlightSession session;
protected Lock clientLock = new ReentrantLock();
private final DorisSourceSplit split;
@@ -274,7 +275,8 @@ public class DorisFlightValueReader extends ValueReader
implements AutoCloseable
createConnectionParameters(
String.valueOf(split[0]), flightSqlPort, options,
rootCertificates);
database = driver.open(parameters);
- return database.connect();
+ session = DorisFlightSession.open(database);
+ return session.getConnection();
} catch (AdbcException | IOException e) {
LOG.debug("Open Flight Connection error: {}", e.getMessage());
throw new RuntimeException(e);
@@ -415,10 +417,11 @@ public class DorisFlightValueReader extends ValueReader
implements AutoCloseable
private void closeFlightResources() throws Exception {
try {
- closeAll(arrowReader, statement, client, database, allocator);
+ closeAll(arrowReader, statement, session, client, database,
allocator);
} finally {
arrowReader = null;
statement = null;
+ session = null;
client = null;
database = null;
allocator = null;
diff --git
a/flink-doris-connector/flink-doris-connector-base/src/test/java/org/apache/doris/flink/source/reader/DorisFlightSessionTest.java
b/flink-doris-connector/flink-doris-connector-base/src/test/java/org/apache/doris/flink/source/reader/DorisFlightSessionTest.java
new file mode 100644
index 00000000..81be277b
--- /dev/null
+++
b/flink-doris-connector/flink-doris-connector-base/src/test/java/org/apache/doris/flink/source/reader/DorisFlightSessionTest.java
@@ -0,0 +1,218 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.flink.source.reader;
+
+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.driver.flightsql.FlightSqlDriver;
+import org.apache.arrow.flight.CallHeaders;
+import org.apache.arrow.flight.CallStatus;
+import org.apache.arrow.flight.CloseSessionRequest;
+import org.apache.arrow.flight.CloseSessionResult;
+import org.apache.arrow.flight.FlightServer;
+import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.auth2.CallHeaderAuthenticator;
+import org.apache.arrow.flight.sql.NoOpFlightSqlProducer;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.doris.flink.exception.DorisRuntimeException;
+import org.junit.jupiter.api.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+class DorisFlightSessionTest {
+ @Test
+ void closesTheSessionSharedWithAdbcOnlyOnce() throws Exception {
+ SessionProducer producer = new SessionProducer();
+ try (RootAllocator allocator = new RootAllocator();
+ FlightServer server = server(allocator, producer)) {
+ Map<String, Object> parameters = parameters(server);
+ try (AdbcDatabase database = new
FlightSqlDriver(allocator).open(parameters)) {
+ DorisFlightSession session = DorisFlightSession.open(database);
+ try (AdbcConnection connection = session.getConnection();
+ DorisFlightSession closingSession = session) {
+ assertThat(producer.opens.get()).isEqualTo(1);
+ assertThat(parameters)
+ .containsEntry(AdbcDriver.PARAM_USERNAME.getKey(),
"root")
+ .containsEntry(AdbcDriver.PARAM_PASSWORD.getKey(),
"password");
+ closingSession.closeSession();
+ closingSession.closeSession();
+ assertThat(producer.closes.get()).isEqualTo(1);
+ assertThat(producer.bearerCalls.get()).isEqualTo(1);
+ assertThat(producer.active.get()).isZero();
+ }
+ }
+ assertThat(producer.closes.get()).isEqualTo(1);
+ }
+ }
+
+ @Test
+ void releasesAdbcResourcesEvenWhenCloseSessionFails() throws Exception {
+ SessionProducer producer = new SessionProducer();
+ producer.rejectClose = true;
+ try (RootAllocator allocator = new RootAllocator();
+ FlightServer server = server(allocator, producer);
+ AdbcDatabase database = new
FlightSqlDriver(allocator).open(parameters(server))) {
+ DorisFlightSession session = DorisFlightSession.open(database);
+ assertThatCode(() -> DorisFlightValueReader.closeAll(session,
session.getConnection()))
+ .doesNotThrowAnyException();
+ session.close();
+ assertThat(producer.closes.get()).isEqualTo(1);
+ }
+ }
+
+ @Test
+ void boundsUnresponsiveSessionCloseAndContinuesCleanup() throws Exception {
+ SessionProducer producer = new SessionProducer();
+ producer.hangClose = true;
+ try (RootAllocator allocator = new RootAllocator();
+ FlightServer server = server(allocator, producer);
+ AdbcDatabase database = new
FlightSqlDriver(allocator).open(parameters(server))) {
+ DorisFlightSession session = DorisFlightSession.open(database);
+ long start = System.nanoTime();
+ assertThatCode(() -> DorisFlightValueReader.closeAll(session,
session.getConnection()))
+ .doesNotThrowAnyException();
+ assertThat(TimeUnit.NANOSECONDS.toMillis(System.nanoTime() -
start))
+ .isBetween(4000L, 10000L);
+ session.close();
+ assertThat(producer.closes.get()).isEqualTo(1);
+ }
+ }
+
+ @Test
+ void closesConnectionWhenSessionBindingFails() throws Exception {
+ AdbcDatabase database = mock(AdbcDatabase.class);
+ AdbcConnection connection = mock(AdbcConnection.class);
+ when(database.connect()).thenReturn(connection);
+
+ assertThatThrownBy(() -> DorisFlightSession.open(database))
+ .isInstanceOf(DorisRuntimeException.class)
+ .hasMessageContaining("Cannot bind ADBC Flight SQL session");
+ verify(connection).close();
+ }
+
+ @Test
+ void acceptsSessionCloseInProgress() throws Exception {
+ SessionProducer producer = new SessionProducer();
+ producer.closeStatus = CloseSessionResult.Status.CLOSING;
+ try (RootAllocator allocator = new RootAllocator();
+ FlightServer server = server(allocator, producer);
+ AdbcDatabase database = new
FlightSqlDriver(allocator).open(parameters(server))) {
+ DorisFlightSession session = DorisFlightSession.open(database);
+ DorisFlightValueReader.closeAll(session, session.getConnection());
+ assertThat(producer.closes.get()).isEqualTo(1);
+ }
+ }
+
+ @Test
+ void continuesCleanupWhenTheServerCannotCloseSession() throws Exception {
+ SessionProducer producer = new SessionProducer();
+ producer.closeStatus = CloseSessionResult.Status.NOT_CLOSABLE;
+ try (RootAllocator allocator = new RootAllocator();
+ FlightServer server = server(allocator, producer);
+ AdbcDatabase database = new
FlightSqlDriver(allocator).open(parameters(server))) {
+ DorisFlightSession session = DorisFlightSession.open(database);
+ assertThatCode(() -> DorisFlightValueReader.closeAll(session,
session.getConnection()))
+ .doesNotThrowAnyException();
+ assertThat(producer.closes.get()).isEqualTo(1);
+ assertThat(producer.active.get()).isEqualTo(1);
+ }
+ }
+
+ private FlightServer server(RootAllocator allocator, SessionProducer
producer)
+ throws Exception {
+ return FlightServer.builder(allocator,
Location.forGrpcInsecure("127.0.0.1", 0), producer)
+ .headerAuthenticator(producer::authenticate)
+ .build()
+ .start();
+ }
+
+ private Map<String, Object> parameters(FlightServer server) {
+ Map<String, Object> parameters = new HashMap<>();
+ AdbcDriver.PARAM_URI.set(
+ parameters,
+ Location.forGrpcInsecure("127.0.0.1",
server.getPort()).getUri().toString());
+ AdbcDriver.PARAM_USERNAME.set(parameters, "root");
+ AdbcDriver.PARAM_PASSWORD.set(parameters, "password");
+ return parameters;
+ }
+
+ private static class SessionProducer extends NoOpFlightSqlProducer {
+ private final AtomicInteger opens = new AtomicInteger();
+ private final AtomicInteger active = new AtomicInteger();
+ private final AtomicInteger closes = new AtomicInteger();
+ private final AtomicInteger bearerCalls = new AtomicInteger();
+ private boolean rejectClose;
+ private boolean hangClose;
+ private CloseSessionResult.Status closeStatus =
CloseSessionResult.Status.CLOSED;
+
+ private CallHeaderAuthenticator.AuthResult authenticate(CallHeaders
headers) {
+ String authorization = headers.get("authorization");
+ if ("Basic cm9vdDpwYXNzd29yZA==".equals(authorization)) {
+ opens.incrementAndGet();
+ active.incrementAndGet();
+ } else if ("Bearer original-session".equals(authorization) &&
active.get() == 1) {
+ bearerCalls.incrementAndGet();
+ } else {
+ throw CallStatus.UNAUTHENTICATED.toRuntimeException();
+ }
+ return new CallHeaderAuthenticator.AuthResult() {
+ @Override
+ public String getPeerIdentity() {
+ return "original-session";
+ }
+
+ @Override
+ public void appendToOutgoingHeaders(CallHeaders outgoing) {
+ outgoing.insert("authorization", "Bearer
original-session");
+ }
+ };
+ }
+
+ @Override
+ public void closeSession(
+ CloseSessionRequest request,
+ CallContext context,
+ StreamListener<CloseSessionResult> listener) {
+ closes.incrementAndGet();
+ if (hangClose) {
+ return;
+ }
+ if (rejectClose) {
+ listener.onError(
+ CallStatus.INTERNAL.withDescription("close
refused").toRuntimeException());
+ return;
+ }
+ if (closeStatus == CloseSessionResult.Status.CLOSED) {
+ active.decrementAndGet();
+ }
+ listener.onNext(new CloseSessionResult(closeStatus));
+ listener.onCompleted();
+ }
+ }
+}
diff --git
a/flink-doris-connector/flink-doris-connector-base/src/test/java/org/apache/doris/flink/source/reader/DorisFlightValueReaderTest.java
b/flink-doris-connector/flink-doris-connector-base/src/test/java/org/apache/doris/flink/source/reader/DorisFlightValueReaderTest.java
index 059714d6..23663cad 100644
---
a/flink-doris-connector/flink-doris-connector-base/src/test/java/org/apache/doris/flink/source/reader/DorisFlightValueReaderTest.java
+++
b/flink-doris-connector/flink-doris-connector-base/src/test/java/org/apache/doris/flink/source/reader/DorisFlightValueReaderTest.java
@@ -108,10 +108,14 @@ class DorisFlightValueReaderTest {
DorisReadOptions readOptions =
DorisReadOptions.builder().setFlightSqlPort(8815).build();
AdbcDatabase database = mock(AdbcDatabase.class);
AdbcConnection connection = mock(AdbcConnection.class);
- when(database.connect()).thenReturn(connection);
+ DorisFlightSession session = mock(DorisFlightSession.class);
+ when(session.getConnection()).thenReturn(connection);
when(connection.createStatement()).thenThrow(new
RuntimeException("create failed"));
+ RuntimeException closeFailure = new RuntimeException("close connection
failed");
+ doThrow(closeFailure).when(connection).close();
try (MockedStatic<RestService> restService =
mockStatic(RestService.class);
+ MockedStatic<DorisFlightSession> sessions =
mockStatic(DorisFlightSession.class);
MockedConstruction<RootAllocator> allocators =
mockConstruction(RootAllocator.class);
MockedConstruction<FlightSqlDriver> drivers =
@@ -119,6 +123,7 @@ class DorisFlightValueReaderTest {
FlightSqlDriver.class,
(driver, context) ->
when(driver.open(anyMap())).thenReturn(database))) {
+ sessions.when(() ->
DorisFlightSession.open(database)).thenReturn(session);
restService
.when(() -> RestService.getSchema(eq(options),
eq(readOptions), any()))
.thenReturn(new Schema());
@@ -128,8 +133,10 @@ class DorisFlightValueReaderTest {
assertThatThrownBy(() -> new DorisFlightValueReader(split,
options, readOptions))
.isInstanceOf(RuntimeException.class)
- .hasMessage("create failed");
+ .hasMessage("create failed")
+ .hasSuppressedException(closeFailure);
+ verify(session).close();
verify(connection).close();
verify(database).close();
verify(allocators.constructed().get(0)).close();
diff --git a/flink-doris-connector/flink-doris-connector-flink1/pom.xml
b/flink-doris-connector/flink-doris-connector-flink1/pom.xml
index 95e3835c..6d92937e 100644
--- a/flink-doris-connector/flink-doris-connector-flink1/pom.xml
+++ b/flink-doris-connector/flink-doris-connector-flink1/pom.xml
@@ -229,6 +229,9 @@ under the License.
<artifactId>maven-shade-plugin</artifactId>
<version>3.4.1</version>
<configuration>
+ <transformers>
+ <transformer
implementation="org.apache.maven.plugins.shade.resource.ServicesResourceTransformer"/>
+ </transformers>
<relocations>
<relocation>
<pattern>org.apache.arrow</pattern>
diff --git a/flink-doris-connector/flink-doris-connector-flink2/pom.xml
b/flink-doris-connector/flink-doris-connector-flink2/pom.xml
index 0343790a..1dc27563 100644
--- a/flink-doris-connector/flink-doris-connector-flink2/pom.xml
+++ b/flink-doris-connector/flink-doris-connector-flink2/pom.xml
@@ -144,6 +144,9 @@ under the License.
<artifactId>maven-shade-plugin</artifactId>
<version>3.4.1</version>
<configuration>
+ <transformers>
+ <transformer
implementation="org.apache.maven.plugins.shade.resource.ServicesResourceTransformer"/>
+ </transformers>
<relocations>
<relocation>
<pattern>org.apache.arrow</pattern>
diff --git a/flink-doris-connector/pom.xml b/flink-doris-connector/pom.xml
index d672c1fa..5c2196ff 100644
--- a/flink-doris-connector/pom.xml
+++ b/flink-doris-connector/pom.xml
@@ -97,8 +97,9 @@ under the License.
<mysql.driver.version>8.0.26</mysql.driver.version>
<ojdbc.version>19.3.0.0</ojdbc.version>
<log4j.version>2.17.1</log4j.version>
- <arrow.version>15.0.2</arrow.version>
- <adbc.version>0.12.0</adbc.version>
+ <arrow.version>17.0.0</arrow.version>
+ <adbc.version>0.14.0</adbc.version>
+ <protobuf.version>3.25.1</protobuf.version>
<awssdk.version>2.29.52</awssdk.version>
</properties>
@@ -176,13 +177,18 @@ under the License.
</dependency>
<dependency>
<groupId>org.apache.arrow</groupId>
- <artifactId>arrow-vector</artifactId>
+ <artifactId>arrow-bom</artifactId>
<version>${arrow.version}</version>
+ <type>pom</type>
+ <scope>import</scope>
</dependency>
+ <!-- Align protobuf runtime and utilities used by ADBC and Flight.
-->
<dependency>
- <groupId>org.apache.arrow</groupId>
- <artifactId>arrow-memory-netty</artifactId>
- <version>${arrow.version}</version>
+ <groupId>com.google.protobuf</groupId>
+ <artifactId>protobuf-bom</artifactId>
+ <version>${protobuf.version}</version>
+ <type>pom</type>
+ <scope>import</scope>
</dependency>
<!-- jackson -->
<dependency>
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
