From b54f2d8fbf7b8b80cb3a98a5c4bd4c74c82b845a Mon Sep 17 00:00:00 2001 From: Alexandr Gorshenin Date: Tue, 21 Oct 2025 11:02:54 +0100 Subject: [PATCH] Added retries for BAD_SESSION on first request --- .../ydb/jdbc/context/BaseYdbExecutor.java | 21 ++++ .../jdbc/context/QueryServiceExecutor.java | 2 +- .../jdbc/context/TableServiceExecutor.java | 2 +- .../ydb/jdbc/context/BadSessionRetryTest.java | 95 +++++++++++++++++++ .../YdbDriverTxValidateTest.java | 5 +- 5 files changed, 120 insertions(+), 5 deletions(-) create mode 100644 jdbc/src/test/java/tech/ydb/jdbc/context/BadSessionRetryTest.java rename jdbc/src/test/java/tech/ydb/jdbc/{ => context}/YdbDriverTxValidateTest.java (98%) diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/BaseYdbExecutor.java b/jdbc/src/main/java/tech/ydb/jdbc/context/BaseYdbExecutor.java index f1f335b..16c90eb 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/context/BaseYdbExecutor.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/context/BaseYdbExecutor.java @@ -9,12 +9,14 @@ import java.util.concurrent.atomic.AtomicReference; import tech.ydb.core.Result; +import tech.ydb.core.StatusCode; import tech.ydb.core.grpc.GrpcReadStream; import tech.ydb.jdbc.YdbConst; import tech.ydb.jdbc.YdbResultSet; import tech.ydb.jdbc.YdbStatement; import tech.ydb.jdbc.YdbTracer; import tech.ydb.jdbc.common.YdbTypes; +import tech.ydb.jdbc.exception.YdbRetryableException; import tech.ydb.jdbc.impl.YdbQueryResult; import tech.ydb.jdbc.impl.YdbStaticResultSet; import tech.ydb.jdbc.query.QueryType; @@ -87,6 +89,25 @@ public void ensureOpened() throws SQLException { } } + + @Override + public YdbQueryResult executeDataQuery(YdbStatement statement, YdbQuery query, String preparedYql, Params params, + long timeout, boolean keepInCache) throws SQLException { + boolean insideTx = isInsideTransaction(); + while (true) { + try { + return executeQueryImpl(statement, query, preparedYql, params, timeout, keepInCache); + } catch (YdbRetryableException ex) { + if (insideTx || ex.getStatus().getCode() != StatusCode.BAD_SESSION) { + throw ex; + } + } + } + } + + protected abstract YdbQueryResult executeQueryImpl(YdbStatement statement, YdbQuery query, String preparedYql, + Params params, long timeout, boolean keepInCache) throws SQLException; + @Override public YdbQueryResult executeSchemeQuery(YdbStatement statement, YdbQuery query) throws SQLException { ensureOpened(); diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/QueryServiceExecutor.java b/jdbc/src/main/java/tech/ydb/jdbc/context/QueryServiceExecutor.java index ab32014..b020a72 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/context/QueryServiceExecutor.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/context/QueryServiceExecutor.java @@ -229,7 +229,7 @@ public void rollback(YdbContext ctx, YdbValidator validator) throws SQLException } @Override - public YdbQueryResult executeDataQuery(YdbStatement statement, YdbQuery query, String preparedYql, Params params, + protected YdbQueryResult executeQueryImpl(YdbStatement statement, YdbQuery query, String preparedYql, Params params, long timeout, boolean keepInCache) throws SQLException { ensureOpened(); diff --git a/jdbc/src/main/java/tech/ydb/jdbc/context/TableServiceExecutor.java b/jdbc/src/main/java/tech/ydb/jdbc/context/TableServiceExecutor.java index 02240fb..df718e9 100644 --- a/jdbc/src/main/java/tech/ydb/jdbc/context/TableServiceExecutor.java +++ b/jdbc/src/main/java/tech/ydb/jdbc/context/TableServiceExecutor.java @@ -200,7 +200,7 @@ public YdbQueryResult executeExplainQuery(YdbStatement statement, YdbQuery query } @Override - public YdbQueryResult executeDataQuery(YdbStatement statement, YdbQuery query, String preparedYql, Params params, + protected YdbQueryResult executeQueryImpl(YdbStatement statement, YdbQuery query, String preparedYql, Params params, long timeout, boolean keepInCache) throws SQLException { ensureOpened(); diff --git a/jdbc/src/test/java/tech/ydb/jdbc/context/BadSessionRetryTest.java b/jdbc/src/test/java/tech/ydb/jdbc/context/BadSessionRetryTest.java new file mode 100644 index 0000000..47f8f98 --- /dev/null +++ b/jdbc/src/test/java/tech/ydb/jdbc/context/BadSessionRetryTest.java @@ -0,0 +1,95 @@ +package tech.ydb.jdbc.context; + +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.SQLException; +import java.sql.Statement; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import tech.ydb.core.Status; +import tech.ydb.core.StatusCode; +import tech.ydb.core.UnexpectedResultException; +import tech.ydb.jdbc.impl.YdbTracerImpl; +import tech.ydb.jdbc.impl.helper.ExceptionAssert; +import tech.ydb.jdbc.impl.helper.JdbcUrlHelper; +import tech.ydb.test.junit5.YdbHelperExtension; + +/** + * + * @author Aleksandr Gorshenin + */ +public class BadSessionRetryTest { + + @RegisterExtension + private static final YdbHelperExtension ydb = new YdbHelperExtension(); + + private static final JdbcUrlHelper jdbcURL = new JdbcUrlHelper(ydb) + .withArg("enableTxTracer", "true"); + + @ParameterizedTest + @ValueSource(strings = { "true", "false" }) + public void badSessionRetryTest(String useQueryService) throws SQLException { + String url = jdbcURL.withArg("useQueryService", useQueryService).build(); + try (Connection conn = DriverManager.getConnection(url)) { + ErrorTxTracer tracer = YdbTracerImpl.use(new ErrorTxTracer()); + + // BAD_SESSION will be retried + tracer.throwErrorOn(3, "<-- Status{code = SUCCESS}", Status.of(StatusCode.BAD_SESSION)); + try (Statement st = conn.createStatement()) { + Assertions.assertTrue(st.execute("SELECT 1 + 2")); + } + + // BAD_REQUEST will not be retried + tracer.throwErrorOn(1, "<-- Status{code = SUCCESS}", Status.of(StatusCode.BAD_REQUEST)); + try (Statement st = conn.createStatement()) { + ExceptionAssert.ydbException("" + + "Cannot call 'DATA_QUERY >>\n" + + "SELECT 1 + 2' with Status{code = BAD_REQUEST(code=400010)", + () -> st.execute("SELECT 1 + 2")); + } + + conn.setAutoCommit(false); + + // BAD_SESSION will be retried + tracer.throwErrorOn(3, "<-- Status{code = SUCCESS}", Status.of(StatusCode.BAD_SESSION)); + try (Statement st = conn.createStatement()) { + Assertions.assertTrue(st.execute("SELECT 1 + 2")); + } + + // BAD_SESSION will not be retried, transaction is already started + tracer.throwErrorOn(1, "<-- Status{code = SUCCESS}", Status.of(StatusCode.BAD_SESSION)); + try (Statement st = conn.createStatement()) { + ExceptionAssert.sqlRecoverable("" + + "Cannot call 'DATA_QUERY >>\n" + + "SELECT 1 + 2' with Status{code = BAD_SESSION(code=400100)", + () -> st.execute("SELECT 1 + 2")); + } + } + } + + private class ErrorTxTracer extends YdbTracerImpl { + private String traceMsg = null; + private Status error = null; + private int count = 0; + + public void throwErrorOn(int count, String traceMsg, Status error) { + this.count = count; + this.traceMsg = traceMsg; + this.error = error; + } + + @Override + public void trace(String message) { + super.trace(message); + if (count > 0 && message.startsWith(traceMsg)) { + Status status = error; + count -= 1; + throw new UnexpectedResultException("Test error", status); + } + } + } +} diff --git a/jdbc/src/test/java/tech/ydb/jdbc/YdbDriverTxValidateTest.java b/jdbc/src/test/java/tech/ydb/jdbc/context/YdbDriverTxValidateTest.java similarity index 98% rename from jdbc/src/test/java/tech/ydb/jdbc/YdbDriverTxValidateTest.java rename to jdbc/src/test/java/tech/ydb/jdbc/context/YdbDriverTxValidateTest.java index 8bf70a1..82cbcb5 100644 --- a/jdbc/src/test/java/tech/ydb/jdbc/YdbDriverTxValidateTest.java +++ b/jdbc/src/test/java/tech/ydb/jdbc/context/YdbDriverTxValidateTest.java @@ -1,4 +1,4 @@ -package tech.ydb.jdbc; +package tech.ydb.jdbc.context; import java.sql.Connection; import java.sql.DriverManager; @@ -17,7 +17,7 @@ import tech.ydb.core.Status; import tech.ydb.core.StatusCode; import tech.ydb.core.UnexpectedResultException; -import tech.ydb.jdbc.context.YdbContext; +import tech.ydb.jdbc.YdbConnection; import tech.ydb.jdbc.exception.YdbConditionallyRetryableException; import tech.ydb.jdbc.exception.YdbRetryableException; import tech.ydb.jdbc.impl.YdbTracerImpl; @@ -221,7 +221,6 @@ public void throwErrorOn(String traceMsg, Status error) { @Override public void trace(String message) { super.trace(message); - System.out.println("TRACE " + message); if (traceMsg != null && error != null && message.startsWith(traceMsg)) { Status status = error; error = null;