Fix GetInfo test failures

Bit-Quill · Feb 17, 2024 · 2744720 · 2744720
1 parent f540c72
commit 2744720
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 1 deletion.
diff --git a/...ver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/BaseFlightReader.java b/...ver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/BaseFlightReader.java
@@ -30,6 +30,7 @@
 import org.apache.arrow.flight.FlightStream;
 import org.apache.arrow.flight.Location;
 import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.util.AutoCloseables;
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.VectorUnloader;
 import org.apache.arrow.vector.ipc.ArrowReader;
@@ -102,7 +103,7 @@ public long bytesRead() {
   @Override
   protected void closeReadSource() throws IOException {
     try {
-      currentStream.close();
+      AutoCloseables.close(currentStream);
     } catch (Exception e) {
       throw new IOException(e);
     }

diff --git a/...light-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/GetInfoMetadataReader.java b/...light-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/GetInfoMetadataReader.java
@@ -145,6 +145,8 @@ public boolean loadNextBatch() throws IOException {
       try (VectorSchemaRoot root = VectorSchemaRoot.create(readSchema(), allocator)) {
         root.allocateNew();
         this.infoCodes = (UInt4Vector) root.getVector(0);
+        this.infoValues = (DenseUnionVector) root.getVector(1);
+        this.stringValues = this.infoValues.getVarCharVector((byte) 0);
 
         if (requestedCodes.contains(AdbcInfoCode.DRIVER_NAME.getValue())) {
           infoCodes.setSafe(dstIndex, AdbcInfoCode.DRIVER_NAME.getValue());

diff --git a/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/HeaderTest.java b/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/HeaderTest.java
@@ -127,6 +127,7 @@ public void testBearerToken() throws Exception {
       }
     } catch (Exception ex) {
       // Swallow exceptions from the RPC call. Only interested in tracking metadata.
+      // This is expected to fail since GetSqlInfo isn't implemented on this test FlightSqlProducer.
     }
     CallHeaders secondHeaders = headerValidatorFactory.getHeadersReceivedAtRequest(1);
     assertTrue(secondHeaders.get("authorization").contains("Bearer"));