Merge pull request apache#88 from rafael-telles/make-flight-sql-close…

…-resources Make FlightSqlClient AutoCloseable
rafael-telles · Aug 23, 2021 · 65204f1 · 65204f1
2 parents 1d2699e + 0f19a60
commit 65204f1
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/java/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java b/java/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java
@@ -56,6 +56,7 @@
 import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementResult;
 import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementQuery;
 import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.util.AutoCloseables;
 import org.apache.arrow.util.Preconditions;
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.types.pojo.Schema;
@@ -68,7 +69,7 @@
 /**
  * Flight client with Flight SQL semantics.
  */
-public class FlightSqlClient {
+public class FlightSqlClient implements AutoCloseable {
   private final FlightClient client;
 
   public FlightSqlClient(final FlightClient client) {
@@ -353,6 +354,11 @@ public PreparedStatement prepare(final String query, final CallOption... options
     return new PreparedStatement(client, query, options);
   }
 
+  @Override
+  public void close() throws Exception {
+    AutoCloseables.close(client);
+  }
+
   /**
    * Helper class to encapsulate Flight SQL prepared statement logic.
    */

diff --git a/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/TestFlightSql.java b/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/TestFlightSql.java
@@ -99,7 +99,6 @@ public class TestFlightSql {
   private static int port;
   private static BufferAllocator allocator;
   private static FlightServer server;
-  private static FlightClient client;
   private static FlightSqlClient sqlClient;
   @Rule
   public final ErrorCollector collector = new ErrorCollector();
@@ -121,8 +120,7 @@ public static void setUp() throws Exception {
         .start();
 
     final Location clientLocation = Location.forGrpcInsecure(LOCALHOST, server.getPort());
-    client = FlightClient.builder(allocator, clientLocation).build();
-    sqlClient = new FlightSqlClient(client);
+    sqlClient = new FlightSqlClient(FlightClient.builder(allocator, clientLocation).build());
 
     GET_SQL_INFO_EXPECTED_RESULTS_MAP
         .put(Integer.toString(ALL_SQL_INFO_ARGS[0]), "Apache Derby");
@@ -148,7 +146,7 @@ public static void setUp() throws Exception {
 
   @AfterClass
   public static void tearDown() throws Exception {
-    close(client, server, allocator);
+    close(sqlClient, server, allocator);
   }
 
   private static List<List<String>> getNonConformingResultsForGetSqlInfo(final List<? extends List<String>> results) {