fix: getColumnCount would fail for empty partititioned result sets (#2588)

olavloite · web-flow · commit 9a2f3fc01748 · 2023-08-14T15:45:14.000+02:00
Calling getColumnCount() would fail for ResultSets that were returned for partitioned
queries. The reason for this was that the number of columns would be calculated based
on the last seen row. This would not work for ResultSets without any rows at all.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/MergedResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/MergedResultSet.java
@@ -394,6 +394,26 @@ public Type getType() {
     return rowProducer.getType();
   }
 
+  @Override
+  public int getColumnCount() {
+    return getType().getStructFields().size();
+  }
+
+  @Override
+  public int getColumnIndex(String columnName) {
+    return getType().getFieldIndex(columnName);
+  }
+
+  @Override
+  public Type getColumnType(int columnIndex) {
+    return getType().getStructFields().get(columnIndex).getType();
+  }
+
+  @Override
+  public Type getColumnType(String columnName) {
+    return getType().getStructFields().get(getColumnIndex(columnName)).getType();
+  }
+
   @Override
   public int getNumPartitions() {
     return rowProducer.getNumPartitions();
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/MergedResultSetTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/MergedResultSetTest.java
@@ -17,6 +17,7 @@
 package com.google.cloud.spanner.connection;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
@@ -30,6 +31,7 @@
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.cloud.spanner.Struct;
+import com.google.cloud.spanner.Type;
 import java.util.ArrayList;
 import java.util.BitSet;
 import java.util.Collection;
@@ -140,6 +142,19 @@ public void testAllResultsAreReturned() {
       while (resultSet.next()) {
         assertRowExists(results.allRows, resultSet.getCurrentRowAsStruct(), rowsFound);
       }
+      // Verify that we can get the metadata after having gotten all rows.
+      // This failed in the initial release of this feature for result sets that were empty. The
+      // reason for that was that the initial implementation would do a call to currentRowAsStruct,
+      // which would always be null for result sets that never returned any data.
+      assertNotNull(resultSet.getMetadata());
+      if (numPartitions == 0) {
+        assertEquals(0, resultSet.getColumnCount());
+      } else {
+        assertEquals(18, resultSet.getColumnCount());
+        assertEquals(Type.bool(), resultSet.getColumnType(0));
+        assertEquals(Type.bool(), resultSet.getColumnType("COL0"));
+        assertEquals(10, resultSet.getColumnIndex("COL10"));
+      }
       // Check that all rows were found.
       assertEquals(results.allRows.size(), rowsFound.nextClearBit(0));
       // Check extended metadata.
