From 9a2f3fc01748224fc8084fbf2b4a0223426b1603 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Mon, 14 Aug 2023 15:45:14 +0200
Subject: [PATCH] fix: getColumnCount would fail for empty partititioned result
 sets (#2588)

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.
---
 .../spanner/connection/MergedResultSet.java   | 20 +++++++++++++++++++
 .../connection/MergedResultSetTest.java       | 15 ++++++++++++++
 2 files changed, 35 insertions(+)

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
index 1f7d15365f..36451bc8f4 100644
--- 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
index b9d9578d6b..c8295fd5ec 100644
--- 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.