Fixed incorrect getUpdateCount after getMoreResults

Author: alex268

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

@@ -297,16 +297,13 @@ public int getResultSetConcurrency() {
     protected static class ResultState {
         private final List<YdbResultSet> results;
         private int resultSetIndex;
-        private final int updateCount;
 
         private ResultState() {
             results = null;
             resultSetIndex = -1;
-            updateCount = -1;
         }
 
         private ResultState(List<YdbResultSet> list) {
-            updateCount = (list == null || list.isEmpty()) ? 1 : -1; // TODO: Get update count?
             results = list;
             resultSetIndex = 0;
         }
@@ -315,8 +312,9 @@ public boolean hasResultSets() {
             return results != null && !results.isEmpty();
         }
 
+        // TODO: YDB doesn't return the count of affected rows, so we use little hach to return always 1
         public int getUpdateCount() {
-            return updateCount;
+            return (results != null && results.isEmpty() && resultSetIndex == 0) ? 1 : -1;
         }
 
         public YdbResultSet getCurrentResultSet() throws SQLException {
@@ -336,6 +334,7 @@ public YdbResultSet getResultSet(int index) throws SQLException {
 
         public boolean getMoreResults(int current) throws SQLException {
             if (results == null || results.isEmpty()) {
+                resultSetIndex = -1; // reset updateCount
                 return false;
             }
 

jdbc/src/test/java/tech/ydb/jdbc/impl/YdbStatementImplTest.java

@@ -378,9 +378,13 @@ public void getUpdateCount() throws SQLException {
 
         statement.execute(TEST_UPSERT1_SQL);
         Assertions.assertEquals(1, statement.getUpdateCount());
+        Assertions.assertFalse(statement.getMoreResults());
+        Assertions.assertEquals(-1, statement.getUpdateCount());
 
         statement.execute(TEST_UPSERT1_SQL + ";\n" + TEST_UPSERT2_SQL + ";");
         Assertions.assertEquals(1, statement.getUpdateCount()); // just a single statement
+        Assertions.assertFalse(statement.getMoreResults());
+        Assertions.assertEquals(-1, statement.getUpdateCount());
 
         statement.execute("select 2 + 2");
         Assertions.assertEquals(-1, statement.getUpdateCount());