fix!(java): require supplying BufferAllocator to create drivers

The old way was prone to memory leaks. Fixes apache#563. BREAKING CHANGE: removes static driver instances.
lidavidm · Apr 28, 2023 · b0dc8f8 · b0dc8f8
1 parent a175420
commit b0dc8f8
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 21 deletions.
diff --git a/java/driver-manager/src/main/java/org/apache/arrow/adbc/drivermanager/AdbcDriverManager.java b/java/driver-manager/src/main/java/org/apache/arrow/adbc/drivermanager/AdbcDriverManager.java
@@ -20,15 +20,23 @@
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.function.Function;
 import org.apache.arrow.adbc.core.AdbcDatabase;
 import org.apache.arrow.adbc.core.AdbcDriver;
 import org.apache.arrow.adbc.core.AdbcException;
 import org.apache.arrow.adbc.core.AdbcStatusCode;
+import org.apache.arrow.memory.BufferAllocator;
 
+/**
+ * Load an ADBC driver based on name.
+ *
+ * <p>This class is EXPERIMENTAL. It will be rewritten before 1.0.0 (see <a
+ * href="https://github.com/apache/arrow-adbc/issues/48">issue #48</a>).
+ */
 public final class AdbcDriverManager {
   private static final AdbcDriverManager INSTANCE = new AdbcDriverManager();
 
-  private final ConcurrentMap<String, AdbcDriver> drivers;
+  private final ConcurrentMap<String, Function<BufferAllocator, AdbcDriver>> drivers;
 
   public AdbcDriverManager() {
     drivers = new ConcurrentHashMap<>();
@@ -38,18 +46,20 @@ public AdbcDriverManager() {
    * Connect to a database.
    *
    * @param driverName The driver to use.
+   * @param allocator The allocator to use.
    * @param parameters Parameters for the driver.
    * @return The AdbcDatabase instance.
    * @throws AdbcException if the driver was not found or if connection fails.
    */
-  public AdbcDatabase connect(String driverName, Map<String, Object> parameters)
+  public AdbcDatabase connect(
+      String driverName, BufferAllocator allocator, Map<String, Object> parameters)
       throws AdbcException {
-    final AdbcDriver driver = lookupDriver(driverName);
+    final Function<BufferAllocator, AdbcDriver> driver = lookupDriver(driverName);
     if (driver == null) {
       throw new AdbcException(
           "Driver not found for '" + driverName + "'", null, AdbcStatusCode.NOT_FOUND, null, 0);
     }
-    return driver.open(parameters);
+    return driver.apply(allocator).open(parameters);
   }
 
   /**
@@ -58,11 +68,11 @@ public AdbcDatabase connect(String driverName, Map<String, Object> parameters)
    * @param driverName The driver to lookup.
    * @return The driver instance, or null if not found.
    */
-  public AdbcDriver lookupDriver(String driverName) {
+  public Function<BufferAllocator, AdbcDriver> lookupDriver(String driverName) {
     return drivers.get(driverName);
   }
 
-  public void registerDriver(String driverName, AdbcDriver driver) {
+  public void registerDriver(String driverName, Function<BufferAllocator, AdbcDriver> driver) {
     if (drivers.putIfAbsent(driverName, driver) != null) {
       throw new IllegalStateException(
           "[DriverManager] Driver already registered for '" + driverName + "'");

diff --git a/...iver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlDriver.java b/...iver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlDriver.java
@@ -26,24 +26,17 @@
 import org.apache.arrow.adbc.sql.SqlQuirks;
 import org.apache.arrow.flight.Location;
 import org.apache.arrow.memory.BufferAllocator;
-import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.util.Preconditions;
 
 /** An ADBC driver wrapping Arrow Flight SQL. */
 public class FlightSqlDriver implements AdbcDriver {
-  public static final FlightSqlDriver INSTANCE = new FlightSqlDriver();
-
   static {
     AdbcDriverManager.getInstance()
-        .registerDriver("org.apache.arrow.adbc.driver.flightsql", INSTANCE);
+        .registerDriver("org.apache.arrow.adbc.driver.flightsql", FlightSqlDriver::new);
   }
 
   private final BufferAllocator allocator;
 
-  FlightSqlDriver() {
-    this(new RootAllocator());
-  }
-
   FlightSqlDriver(BufferAllocator allocator) {
     this.allocator = Objects.requireNonNull(allocator);
   }

diff --git a/java/driver/jdbc/src/main/java/org/apache/arrow/adbc/driver/jdbc/JdbcDriver.java b/java/driver/jdbc/src/main/java/org/apache/arrow/adbc/driver/jdbc/JdbcDriver.java
@@ -26,11 +26,9 @@
 import org.apache.arrow.adbc.drivermanager.AdbcDriverManager;
 import org.apache.arrow.adbc.sql.SqlQuirks;
 import org.apache.arrow.memory.BufferAllocator;
-import org.apache.arrow.memory.RootAllocator;
 
 /** An ADBC driver wrapping the JDBC API. */
 public class JdbcDriver implements AdbcDriver {
-  public static final JdbcDriver INSTANCE = new JdbcDriver();
   /** A parameter for creating an {@link AdbcDatabase} from a {@link DataSource}. */
   public static final String PARAM_DATASOURCE = "adbc.jdbc.datasource";
   /**
@@ -41,15 +39,12 @@ public class JdbcDriver implements AdbcDriver {
   public static final String PARAM_URI = "uri";
 
   static {
-    AdbcDriverManager.getInstance().registerDriver("org.apache.arrow.adbc.driver.jdbc", INSTANCE);
+    AdbcDriverManager.getInstance()
+        .registerDriver("org.apache.arrow.adbc.driver.jdbc", JdbcDriver::new);
   }
 
   private final BufferAllocator allocator;
 
-  public JdbcDriver() {
-    this(new RootAllocator());
-  }
-
   public JdbcDriver(BufferAllocator allocator) {
     this.allocator = Objects.requireNonNull(allocator);
   }