Skip to content

Commit f49a537

Browse files
alex268alex268
authored
Added retries for BAD_SESSION on first request (#142)
2 parents 84848a1 + b54f2d8 commit f49a537

File tree

5 files changed

+120
-5
lines changed

5 files changed

+120
-5
lines changed

‎jdbc/src/main/java/tech/ydb/jdbc/context/BaseYdbExecutor.java‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,14 @@
99
import java.util.concurrent.atomic.AtomicReference;
1010

1111
import tech.ydb.core.Result;
12+
import tech.ydb.core.StatusCode;
1213
import tech.ydb.core.grpc.GrpcReadStream;
1314
import tech.ydb.jdbc.YdbConst;
1415
import tech.ydb.jdbc.YdbResultSet;
1516
import tech.ydb.jdbc.YdbStatement;
1617
import tech.ydb.jdbc.YdbTracer;
1718
import tech.ydb.jdbc.common.YdbTypes;
19+
import tech.ydb.jdbc.exception.YdbRetryableException;
1820
import tech.ydb.jdbc.impl.YdbQueryResult;
1921
import tech.ydb.jdbc.impl.YdbStaticResultSet;
2022
import tech.ydb.jdbc.query.QueryType;
@@ -87,6 +89,25 @@ public void ensureOpened() throws SQLException {
8789
}
8890
}
8991

92+
93+
@Override
94+
public YdbQueryResult executeDataQuery(YdbStatement statement, YdbQuery query, String preparedYql, Params params,
95+
long timeout, boolean keepInCache) throws SQLException {
96+
boolean insideTx = isInsideTransaction();
97+
while (true) {
98+
try {
99+
return executeQueryImpl(statement, query, preparedYql, params, timeout, keepInCache);
100+
} catch (YdbRetryableException ex) {
101+
if (insideTx || ex.getStatus().getCode() != StatusCode.BAD_SESSION) {
102+
throw ex;
103+
}
104+
}
105+
}
106+
}
107+
108+
protected abstract YdbQueryResult executeQueryImpl(YdbStatement statement, YdbQuery query, String preparedYql,
109+
Params params, long timeout, boolean keepInCache) throws SQLException;
110+
90111
@Override
91112
public YdbQueryResult executeSchemeQuery(YdbStatement statement, YdbQuery query) throws SQLException {
92113
ensureOpened();

‎jdbc/src/main/java/tech/ydb/jdbc/context/QueryServiceExecutor.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -229,7 +229,7 @@ public void rollback(YdbContext ctx, YdbValidator validator) throws SQLException
229229
}
230230

231231
@Override
232-
public YdbQueryResult executeDataQuery(YdbStatement statement, YdbQuery query, String preparedYql, Params params,
232+
protected YdbQueryResult executeQueryImpl(YdbStatement statement, YdbQuery query, String preparedYql, Params params,
233233
long timeout, boolean keepInCache) throws SQLException {
234234
ensureOpened();
235235

‎jdbc/src/main/java/tech/ydb/jdbc/context/TableServiceExecutor.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -200,7 +200,7 @@ public YdbQueryResult executeExplainQuery(YdbStatement statement, YdbQuery query
200200
}
201201

202202
@Override
203-
public YdbQueryResult executeDataQuery(YdbStatement statement, YdbQuery query, String preparedYql, Params params,
203+
protected YdbQueryResult executeQueryImpl(YdbStatement statement, YdbQuery query, String preparedYql, Params params,
204204
long timeout, boolean keepInCache) throws SQLException {
205205
ensureOpened();
206206

‎jdbc/src/test/java/tech/ydb/jdbc/context/BadSessionRetryTest.java‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
package tech.ydb.jdbc.context;
2+
3+
import java.sql.Connection;
4+
import java.sql.DriverManager;
5+
import java.sql.SQLException;
6+
import java.sql.Statement;
7+
8+
import org.junit.jupiter.api.Assertions;
9+
import org.junit.jupiter.api.extension.RegisterExtension;
10+
import org.junit.jupiter.params.ParameterizedTest;
11+
import org.junit.jupiter.params.provider.ValueSource;
12+
13+
import tech.ydb.core.Status;
14+
import tech.ydb.core.StatusCode;
15+
import tech.ydb.core.UnexpectedResultException;
16+
import tech.ydb.jdbc.impl.YdbTracerImpl;
17+
import tech.ydb.jdbc.impl.helper.ExceptionAssert;
18+
import tech.ydb.jdbc.impl.helper.JdbcUrlHelper;
19+
import tech.ydb.test.junit5.YdbHelperExtension;
20+
21+
/**
22+
*
23+
* @author Aleksandr Gorshenin
24+
*/
25+
public class BadSessionRetryTest {
26+
27+
@RegisterExtension
28+
private static final YdbHelperExtension ydb = new YdbHelperExtension();
29+
30+
private static final JdbcUrlHelper jdbcURL = new JdbcUrlHelper(ydb)
31+
.withArg("enableTxTracer", "true");
32+
33+
@ParameterizedTest
34+
@ValueSource(strings = { "true", "false" })
35+
public void badSessionRetryTest(String useQueryService) throws SQLException {
36+
String url = jdbcURL.withArg("useQueryService", useQueryService).build();
37+
try (Connection conn = DriverManager.getConnection(url)) {
38+
ErrorTxTracer tracer = YdbTracerImpl.use(new ErrorTxTracer());
39+
40+
// BAD_SESSION will be retried
41+
tracer.throwErrorOn(3, "<-- Status{code = SUCCESS}", Status.of(StatusCode.BAD_SESSION));
42+
try (Statement st = conn.createStatement()) {
43+
Assertions.assertTrue(st.execute("SELECT 1 + 2"));
44+
}
45+
46+
// BAD_REQUEST will not be retried
47+
tracer.throwErrorOn(1, "<-- Status{code = SUCCESS}", Status.of(StatusCode.BAD_REQUEST));
48+
try (Statement st = conn.createStatement()) {
49+
ExceptionAssert.ydbException(""
50+
+ "Cannot call 'DATA_QUERY >>\n"
51+
+ "SELECT 1 + 2' with Status{code = BAD_REQUEST(code=400010)",
52+
() -> st.execute("SELECT 1 + 2"));
53+
}
54+
55+
conn.setAutoCommit(false);
56+
57+
// BAD_SESSION will be retried
58+
tracer.throwErrorOn(3, "<-- Status{code = SUCCESS}", Status.of(StatusCode.BAD_SESSION));
59+
try (Statement st = conn.createStatement()) {
60+
Assertions.assertTrue(st.execute("SELECT 1 + 2"));
61+
}
62+
63+
// BAD_SESSION will not be retried, transaction is already started
64+
tracer.throwErrorOn(1, "<-- Status{code = SUCCESS}", Status.of(StatusCode.BAD_SESSION));
65+
try (Statement st = conn.createStatement()) {
66+
ExceptionAssert.sqlRecoverable(""
67+
+ "Cannot call 'DATA_QUERY >>\n"
68+
+ "SELECT 1 + 2' with Status{code = BAD_SESSION(code=400100)",
69+
() -> st.execute("SELECT 1 + 2"));
70+
}
71+
}
72+
}
73+
74+
private class ErrorTxTracer extends YdbTracerImpl {
75+
private String traceMsg = null;
76+
private Status error = null;
77+
private int count = 0;
78+
79+
public void throwErrorOn(int count, String traceMsg, Status error) {
80+
this.count = count;
81+
this.traceMsg = traceMsg;
82+
this.error = error;
83+
}
84+
85+
@Override
86+
public void trace(String message) {
87+
super.trace(message);
88+
if (count > 0 && message.startsWith(traceMsg)) {
89+
Status status = error;
90+
count -= 1;
91+
throw new UnexpectedResultException("Test error", status);
92+
}
93+
}
94+
}
95+
}

‎jdbc/src/test/java/tech/ydb/jdbc/YdbDriverTxValidateTest.java‎ renamed to ‎jdbc/src/test/java/tech/ydb/jdbc/context/YdbDriverTxValidateTest.java‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package tech.ydb.jdbc;
1+
package tech.ydb.jdbc.context;
22

33
import java.sql.Connection;
44
import java.sql.DriverManager;
@@ -17,7 +17,7 @@
1717
import tech.ydb.core.Status;
1818
import tech.ydb.core.StatusCode;
1919
import tech.ydb.core.UnexpectedResultException;
20-
import tech.ydb.jdbc.context.YdbContext;
20+
import tech.ydb.jdbc.YdbConnection;
2121
import tech.ydb.jdbc.exception.YdbConditionallyRetryableException;
2222
import tech.ydb.jdbc.exception.YdbRetryableException;
2323
import tech.ydb.jdbc.impl.YdbTracerImpl;
@@ -221,7 +221,6 @@ public void throwErrorOn(String traceMsg, Status error) {
221221
@Override
222222
public void trace(String message) {
223223
super.trace(message);
224-
System.out.println("TRACE " + message);
225224
if (traceMsg != null && error != null && message.startsWith(traceMsg)) {
226225
Status status = error;
227226
error = null;

0 commit comments

Comments
 (0)