fix(bq jdbc): ensure getMoreResults() always moves the cursor (#12353)

Author: logachev

java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryStatement.java

@@ -643,6 +643,10 @@ private QueryJobConfiguration setDestinationDatasetAndTableInJobConfig(
   }
 
   Job getNextJob() {
+    if (this.parentJobId == null) {
+      return null;
+    }
+
     while (this.currentJobIdIndex + 1 < this.parentJobId.getJobs().size()) {
       this.currentJobIdIndex += 1;
       Job currentJob = this.parentJobId.getJobs().get(this.currentJobIdIndex);
@@ -1448,10 +1452,6 @@ public boolean getMoreResults(int current) throws SQLException {
           "The JDBC driver only supports Statement.CLOSE_CURRENT_RESULT.");
     }
 
-    if (this.parentJobId == null) {
-      return false;
-    }
-
     try {
       if (this.currentResultSet != null) {
         this.currentResultSet.close();

java-bigquery/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java

@@ -87,10 +87,7 @@ public class ITBigQueryJDBCTest extends ITBase {
       "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;PROJECTID="
           + PROJECT_ID
           + ";OAUTHTYPE=3";
-  static final String session_enabled_connection_uri =
-      "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;PROJECTID="
-          + PROJECT_ID
-          + ";OAUTHTYPE=3;EnableSession=1";
+  static final String session_enabled_connection_uri = connection_uri + ";EnableSession=1";
   private static final String BASE_QUERY =
       "SELECT * FROM bigquery-public-data.new_york_taxi_trips.tlc_yellow_trips_2017 order by"
           + " trip_distance asc LIMIT %s";
@@ -3717,6 +3714,63 @@ public void testExecuteScriptWithExpession() throws SQLException {
     assertEquals(-1, bigQueryStatement.getUpdateCount());
   }
 
+  @Test
+  public void testExecuteSingleSelectMoreResultsBehavior() throws SQLException {
+    String sql = "SELECT 1;";
+    assertTrue(bigQueryStatement.execute(sql));
+    ResultSet rs = bigQueryStatement.getResultSet();
+    assertNotNull(rs);
+    assertTrue(rs.next());
+    assertEquals(1, rs.getInt(1));
+    assertFalse(rs.next());
+
+    // Validate no more results reset update count to -1
+    assertFalse(bigQueryStatement.getMoreResults());
+    assertEquals(-1, bigQueryStatement.getUpdateCount());
+  }
+
+  @Test
+  public void testExecuteSingleInsertMoreResultsBehavior() throws SQLException {
+    String tableName = "test_insert_more_results_" + System.currentTimeMillis();
+    String createSql =
+        String.format("CREATE OR REPLACE TABLE %s.%s (id INT64);", DATASET, tableName);
+    String insertSql = String.format("INSERT INTO %s.%s (id) VALUES (1);", DATASET, tableName);
+    String dropSql = String.format("DROP TABLE IF EXISTS %s.%s;", DATASET, tableName);
+
+    try {
+      bigQueryStatement.execute(createSql);
+
+      assertFalse(bigQueryStatement.execute(insertSql));
+      assertNull(bigQueryStatement.getResultSet());
+      assertTrue(bigQueryStatement.getUpdateCount() > 0);
+
+      // Validate no more results reset update count to -1
+      assertFalse(bigQueryStatement.getMoreResults());
+      assertEquals(-1, bigQueryStatement.getUpdateCount());
+    } finally {
+      bigQueryStatement.execute(dropSql);
+    }
+  }
+
+  @Test
+  public void testExecuteSingleTclMoreResultsBehavior() throws SQLException {
+    Connection sessionConnection = DriverManager.getConnection(session_enabled_connection_uri);
+    Statement sessionStatement = sessionConnection.createStatement();
+
+    try {
+      assertFalse(sessionStatement.execute("START TRANSACTION;"));
+      assertNull(sessionStatement.getResultSet());
+      assertEquals(0, sessionStatement.getUpdateCount());
+
+      // Validate no more results reset update count to -1
+      assertFalse(sessionStatement.getMoreResults());
+      assertEquals(-1, sessionStatement.getUpdateCount());
+    } finally {
+      sessionStatement.close();
+      sessionConnection.close();
+    }
+  }
+
   @Test
   public void testInformationSchemaTables() throws SQLException {
     String query = String.format("SELECT * FROM %s.INFORMATION_SCHEMA.TABLES", DATASET);