Use getSchema rpc to get schema (#39)

sabbasani · Dec 18, 2024 · 2d46518 · 2d46518
1 parent 9b67550
commit 2d46518
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/...rrow-flight/src/main/java/com/facebook/plugin/arrow/AbstractArrowFlightClientHandler.java b/...rrow-flight/src/main/java/com/facebook/plugin/arrow/AbstractArrowFlightClientHandler.java
@@ -23,12 +23,14 @@
 import org.apache.arrow.vector.types.pojo.Schema;
 
 import java.io.FileInputStream;
+import java.io.IOException;
 import java.io.InputStream;
 import java.util.Objects;
 import java.util.Optional;
 
 import static com.facebook.plugin.arrow.ArrowErrorCode.ARROW_FLIGHT_CLIENT_ERROR;
 import static com.facebook.plugin.arrow.ArrowErrorCode.ARROW_FLIGHT_INFO_ERROR;
+import static com.facebook.plugin.arrow.ArrowErrorCode.ARROW_FLIGHT_METADATA_ERROR;
 import static java.util.Objects.requireNonNull;
 
 public abstract class AbstractArrowFlightClientHandler
@@ -107,14 +109,20 @@ public FlightInfo getFlightInfo(FlightDescriptor flightDescriptor, ConnectorSess
             FlightInfo flightInfo = client.getFlightClient().getInfo(flightDescriptor, auth);
             return flightInfo;
         }
-        catch (Exception e) {
+        catch (InterruptedException | IOException e) {
             throw new ArrowException(ARROW_FLIGHT_INFO_ERROR, "The flight information could not be obtained from the flight server." + e.getMessage(), e);
         }
     }
 
-    public Optional<Schema> getSchema(FlightDescriptor flightDescriptor, ConnectorSession connectorSession)
+    public Schema getSchema(FlightDescriptor flightDescriptor, ConnectorSession connectorSession)
     {
-        return getFlightInfo(flightDescriptor, connectorSession).getSchemaOptional();
+        try (ArrowFlightClient client = createArrowFlightClient()) {
+            CredentialCallOption[] auth = this.getCallOptions(connectorSession);
+            return client.getFlightClient().getSchema(flightDescriptor, auth).getSchema();
+        }
+        catch (InterruptedException | IOException e) {
+            throw new ArrowException(ARROW_FLIGHT_METADATA_ERROR, "The flight information could not be obtained from the flight server." + e.getMessage(), e);
+        }
     }
 
     public void closeRootAllocator()

diff --git a/presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/AbstractArrowMetadata.java b/presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/AbstractArrowMetadata.java
@@ -88,9 +88,8 @@ public List<Field> getColumnsList(String schema, String table, ConnectorSession
             FlightDescriptor flightDescriptor = getFlightDescriptor(
                     dataSourceSpecificSchemaName, dataSourceSpecificTableName);
 
-            Optional<Schema> flightSchema = clientHandler.getSchema(flightDescriptor, connectorSession);
-            List<Field> fields = flightSchema.map(Schema::getFields).orElse(Collections.emptyList());
-            return fields;
+            Schema flightSchema = clientHandler.getSchema(flightDescriptor, connectorSession);
+            return flightSchema.getFields();
         }
         catch (Exception e) {
             throw new ArrowException(ARROW_FLIGHT_METADATA_ERROR, "The table columns could not be listed for the table " + table, e);