Updated implementation of TableTxExecutor

Author: alex268

jdbc/src/main/java/tech/ydb/jdbc/context/TableTxExecutor.java

@@ -87,10 +87,20 @@ public void rollback(YdbContext ctx, YdbValidator validator) throws SQLException
     @Override
     public YdbQueryResult executeDataQuery(YdbStatement statement, YdbQuery query, String preparedYql, Params params,
             long timeout, boolean keepInCache) throws SQLException {
-        YdbQueryResult result = super.executeDataQuery(statement, query, preparedYql, params, timeout, keepInCache);
-        isWriteTx = isInsideTransaction() && (isWriteTx || query.isWriting());
+        try {
+            YdbQueryResult result = super.executeDataQuery(statement, query, preparedYql, params, timeout, keepInCache);
+            isWriteTx = isInsideTransaction() && (isWriteTx || query.isWriting());
+            return result;
+        } catch (YdbConditionallyRetryableException ex) {
+            if (isAutoCommit()) {
+                throw ex;
+            }
 
-        return result;
+            Status status = Status.of(StatusCode.ABORTED).withCause(ex);
+            throw ExceptionFactory.createException("Data query wasn't executed",
+                    new UnexpectedResultException(ex.getMessage(), status)
+            );
+        }
     }
 
     @Override

jdbc/src/test/java/tech/ydb/jdbc/YdbDriverTxValidateTest.java

@@ -18,6 +18,9 @@
 import tech.ydb.core.StatusCode;
 import tech.ydb.core.UnexpectedResultException;
 import tech.ydb.jdbc.context.YdbContext;
+import tech.ydb.jdbc.context.YdbExecutor;
+import tech.ydb.jdbc.exception.YdbConditionallyRetryableException;
+import tech.ydb.jdbc.exception.YdbRetryableException;
 import tech.ydb.jdbc.impl.YdbTracerImpl;
 import tech.ydb.jdbc.impl.helper.ExceptionAssert;
 import tech.ydb.jdbc.impl.helper.JdbcConnectionExtention;
@@ -153,6 +156,52 @@ public void commitedTxTest() throws SQLException {
         }
     }
 
+    @Test
+    public void executeDataQueryTest() throws SQLException {
+        String url = jdbcURL.withArg("withTxValidationTable", "tx1_store").build();
+        try (Connection conn = DriverManager.getConnection(url)) {
+            ErrorTxTracer tracer = YdbTracerImpl.use(new ErrorTxTracer());
+            // table was created automatically
+            assertTxCount("tx1_store", 0);
+
+            conn.setAutoCommit(true);
+            try (Statement st = conn.createStatement()) {
+                st.execute("DELETE FROM tx1_store");
+
+                tracer.throwErrorOn("<-- Status", Status.of(StatusCode.UNDETERMINED));
+                YdbConditionallyRetryableException e = Assertions.assertThrows(YdbConditionallyRetryableException.class,
+                        () -> st.execute("DELETE FROM tx1_store"));
+                Assertions.assertEquals(Status.of(StatusCode.UNDETERMINED), e.getStatus());
+
+                tracer.throwErrorOn("<-- Status", Status.of(StatusCode.ABORTED));
+                YdbRetryableException e2 = Assertions.assertThrows(YdbRetryableException.class,
+                        () -> st.execute("DELETE FROM tx1_store"));
+                Assertions.assertEquals(Status.of(StatusCode.ABORTED), e2.getStatus());
+            }
+
+
+            conn.setAutoCommit(false);
+            try (Statement st = conn.createStatement()) {
+                YdbExecutor executor = st.unwrap(YdbStatement.class).getConnection().getExecutor();
+                st.execute("DELETE FROM tx1_store");
+                conn.commit();
+
+                tracer.throwErrorOn("<-- Status", Status.of(StatusCode.UNDETERMINED));
+                YdbRetryableException e = Assertions.assertThrows(YdbRetryableException.class,
+                        () -> st.execute("DELETE FROM tx1_store"));
+                Assertions.assertEquals(StatusCode.ABORTED, e.getStatus().getCode());
+                Assertions.assertNotNull(e.getStatus().getCause());
+
+                tracer.throwErrorOn("<-- Status", Status.of(StatusCode.ABORTED));
+                YdbRetryableException e2 = Assertions.assertThrows(YdbRetryableException.class,
+                        () -> st.execute("DELETE FROM tx1_store"));
+                Assertions.assertEquals(Status.of(StatusCode.ABORTED), e2.getStatus());
+            }
+        } finally {
+            jdbc.connection().createStatement().execute("DROP TABLE tx1_store");
+        }
+    }
+
     @Test
     public void invalididTxTableTest() throws SQLException {
         String url = jdbcURL.withArg("withTxValidationTable", "tx store").build();