Added wait of statement closing

Author: alex268

jdbc/src/main/java/tech/ydb/jdbc/YdbConnection.java

@@ -9,13 +9,8 @@
 import tech.ydb.jdbc.context.YdbExecutor;
 
 public interface YdbConnection extends Connection {
-    /**
-     * Return current YDB transaction, if exists
-     *
-     * @return YDB transaction ID or null, if no transaction started
-     */
     @Nullable
-    String getYdbTxId();
+    String getYdbTxId() throws SQLException;
 
     YdbContext getCtx();
 

jdbc/src/main/java/tech/ydb/jdbc/YdbStatement.java

@@ -43,4 +43,6 @@ public interface YdbStatement extends Statement {
 
     @Override
     YdbConnection getConnection() throws SQLException;
+
+    void waitReady() throws SQLException;
 }

jdbc/src/main/java/tech/ydb/jdbc/impl/BaseYdbStatement.java

@@ -148,11 +148,18 @@ private void ensureOpened() throws SQLException {
         }
     }
 
+    @Override
+    public void waitReady() throws SQLException {
+        state.close();
+    }
+
     protected void cleanState() throws SQLException {
         ensureOpened();
-        clearWarnings();
+
         state.close();
         state = YdbQueryResult.EMPTY;
+
+        clearWarnings();
     }
 
     protected boolean updateState(YdbQueryResult result) throws SQLException {

jdbc/src/main/java/tech/ydb/jdbc/impl/YdbConnectionImpl.java

@@ -17,6 +17,7 @@
 import java.util.Map;
 import java.util.Properties;
 import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -38,6 +39,7 @@ public class YdbConnectionImpl implements YdbConnection {
     private final YdbContext ctx;
     private final YdbValidator validator;
     private final YdbExecutor executor;
+    private final AtomicReference<YdbStatement> currState = new AtomicReference<>();
 
     public YdbConnectionImpl(YdbContext context) throws SQLException {
         this.ctx = context;
@@ -83,18 +85,35 @@ public void setAutoCommit(boolean autoCommit) throws SQLException {
         executor.setAutoCommit(autoCommit);
     }
 
+    private <T extends YdbStatement> T updateStatement(T statement) throws SQLException {
+        Statement prev = currState.getAndSet(statement);
+        if (prev != null) {
+            prev.close();
+        }
+        return statement;
+    }
+
+    private void waitStatementReady() throws SQLException {
+        YdbStatement curr = currState.get();
+        if (curr != null) {
+            curr.waitReady();
+        }
+    }
+
     @Override
     public boolean getAutoCommit() throws SQLException {
         return executor.isAutoCommit();
     }
 
     @Override
     public void commit() throws SQLException {
+        waitStatementReady();
         executor.commit(ctx, validator);
     }
 
     @Override
     public void rollback() throws SQLException {
+        waitStatementReady();
         executor.rollback(ctx, validator);
     }
 
@@ -167,6 +186,8 @@ public SQLWarning getWarnings() throws SQLException {
     @Override
     public void clearWarnings() throws SQLException {
         executor.ensureOpened();
+
+        waitStatementReady();
         validator.clearWarnings();
     }
 
@@ -210,7 +231,7 @@ public YdbStatement createStatement(int resultSetType, int resultSetConcurrency,
             int resultSetHoldability) throws SQLException {
         executor.ensureOpened();
         checkStatementParams(resultSetType, resultSetConcurrency, resultSetHoldability);
-        return new YdbStatementImpl(this, resultSetType);
+        return updateStatement(new YdbStatementImpl(this, resultSetType));
     }
 
     @Override
@@ -242,7 +263,7 @@ private YdbPreparedStatement prepareStatement(String sql, int resultSetType, Ydb
         YdbQuery query = ctx.findOrParseYdbQuery(sql);
 
         YdbPreparedQuery params = ctx.findOrPrepareParams(query, mode);
-        return new YdbPreparedStatementImpl(this, query, params, resultSetType);
+        return updateStatement(new YdbPreparedStatementImpl(this, query, params, resultSetType));
     }
 
     @Override
@@ -287,7 +308,8 @@ public int getNetworkTimeout() throws SQLException {
     }
 
     @Override
-    public String getYdbTxId() {
+    public String getYdbTxId() throws SQLException {
+        waitStatementReady();
         return executor.txID();
     }