Improve encapsulation of ResultSet tracking

Author: ennuite

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java

@@ -44,8 +44,8 @@ public final class ArrowFlightConnection extends AvaticaConnection {
   private final ArrowFlightSqlClientHandler clientHandler;
   private final ArrowFlightConnectionConfigImpl config;
   private ExecutorService executorService;
-  int metadataResultSetCount;
-  Map<Integer, ArrowFlightJdbcFlightStreamResultSet> metadataResultSetMap = new HashMap<>();
+  private int metadataResultSetCount;
+  private Map<Integer, ArrowFlightJdbcFlightStreamResultSet> metadataResultSetMap = new HashMap<>();
 
   /**
    * Creates a new {@link ArrowFlightConnection}.
@@ -178,6 +178,31 @@ synchronized ExecutorService getExecutorService() {
             : executorService;
   }
 
+  /**
+   * Registers a new metadata ResultSet and assigns it a unique ID. Metadata ResultSets are those
+   * created without an associated Statement.
+   *
+   * @param resultSet the ResultSet to register
+   * @return the assigned ID
+   */
+  int getNewMetadataResultSetId(ArrowFlightJdbcFlightStreamResultSet resultSet) {
+    metadataResultSetMap.put(metadataResultSetCount, resultSet);
+    return metadataResultSetCount++;
+  }
+
+  /**
+   * Unregisters a metadata ResultSet when it is closed. This method is called by metadata
+   * ResultSets during their close operation to remove themselves from the tracking map.
+   *
+   * @param id the ID of the ResultSet to unregister, or null if not a metadata ResultSet
+   */
+  void onResultSetClose(Integer id) {
+    if (id == null) {
+      return;
+    }
+    metadataResultSetMap.remove(id);
+  }
+
   @Override
   public Properties getClientInfo() {
     final Properties copy = new Properties();

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java

@@ -83,8 +83,7 @@ private ArrowFlightJdbcFlightStreamResultSet(
     super(null, state, signature, resultSetMetaData, timeZone, firstFrame);
     this.connection = connection;
     this.flightInfo = flightInfo;
-    this.id = connection.metadataResultSetCount++;
-    connection.metadataResultSetMap.put(id, this);
+    this.id = connection.getNewMetadataResultSetId(this);
   }
 
   /**
@@ -242,9 +241,7 @@ public synchronized void close() {
       if (isClosed()) {
         return;
       }
-      if (id != null) { // id is only set for metadata result sets
-        this.connection.metadataResultSetMap.remove(id);
-      }
+      this.connection.onResultSetClose(id);
       if (flightEndpointDataQueue != null) {
         // flightStreamQueue should close currentFlightStream internally
         flightEndpointDataQueue.close();