Added retries for BAD_SESSION on first request

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();

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();
 

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();
 

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);
+            }
+        }
+    }
+}

jdbc/src/test/java/tech/ydb/jdbc/YdbDriverTxValidateTest.java → 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;