(pinot) branch master updated: Fix file handle leaks in Pinot Driver (apache#12263) (#12356)

[xi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

xiangfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git


The following commit(s) were added to refs/heads/master by this push:
     new e1e22bd4ed Fix file handle leaks in Pinot Driver (apache#12263) (#12356)
e1e22bd4ed is described below

commit e1e22bd4edac5794f8bea104bf43afa43f551641
Author: Brendan Stansfield <63...@users.noreply.github.com>
AuthorDate: Tue Feb 13 02:45:47 2024 -0700

    Fix file handle leaks in Pinot Driver (apache#12263) (#12356)
---
 .../java/org/apache/pinot/client/PinotDriver.java  | 35 ++++++++++++++++++++--
 .../org/apache/pinot/client/PinotDriverTest.java   | 20 +++++++++++++
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotDriver.java b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotDriver.java
index 2e44fadeb5..d24e880b38 100644
--- a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotDriver.java
+++ b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/PinotDriver.java
@@ -80,7 +80,13 @@ public class PinotDriver implements Driver {
   @Override
   public Connection connect(String url, Properties info)
       throws SQLException {
+
+      PinotClientTransport pinotClientTransport = null;
+      PinotControllerTransport pinotControllerTransport = null;
     try {
+      if (!this.acceptsURL(url)) {
+        return null;
+      }
       LOGGER.info("Initiating connection to database for url: " + url);
 
       Map<String, String> urlParams = DriverUtils.getURLParams(url);
@@ -112,18 +118,43 @@ public class PinotDriver implements Driver {
         pinotControllerTransportFactory.setHeaders(headers);
       }
 
-      PinotClientTransport pinotClientTransport = factory.withConnectionProperties(info).buildTransport();
-      PinotControllerTransport pinotControllerTransport = pinotControllerTransportFactory
+      pinotClientTransport = factory.withConnectionProperties(info).buildTransport();
+      pinotControllerTransport = pinotControllerTransportFactory
               .withConnectionProperties(info)
               .buildTransport();
       String controllerUrl = DriverUtils.getControllerFromURL(url);
       String tenant = info.getProperty(INFO_TENANT, DEFAULT_TENANT);
       return new PinotConnection(info, controllerUrl, pinotClientTransport, tenant, pinotControllerTransport);
     } catch (Exception e) {
+      closeResourcesSafely(pinotClientTransport, pinotControllerTransport);
       throw new SQLException(String.format("Failed to connect to url : %s", url), e);
     }
   }
 
+  /**
+   * If something goes wrong generating the connection, the transports need to be safely closed to prevent leaks.
+   * @param pinotClientTransport connection client transport
+   * @param pinotControllerTransport connection controller transport
+   */
+  private static void closeResourcesSafely(PinotClientTransport pinotClientTransport,
+      PinotControllerTransport pinotControllerTransport) {
+    try {
+      if (pinotClientTransport != null) {
+        pinotClientTransport.close();
+      }
+    } catch (PinotClientException ignored) {
+      // ignored
+    }
+
+    try {
+      if (pinotControllerTransport != null) {
+        pinotControllerTransport.close();
+      }
+    } catch (PinotClientException ignored) {
+      // ignored
+    }
+  }
+
   private Map<String, String> getHeadersFromProperties(Properties info) {
     return info.entrySet().stream().filter(entry -> entry.getKey().toString().startsWith(INFO_HEADERS + ".")).map(
             entry -> Pair.of(entry.getKey().toString().substring(INFO_HEADERS.length() + 1),
diff --git a/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotDriverTest.java b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotDriverTest.java
index 36dccaa5e9..05691534c0 100644
--- a/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotDriverTest.java
+++ b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/PinotDriverTest.java
@@ -30,6 +30,8 @@ import org.testng.annotations.Test;
 
 public class PinotDriverTest {
   static final String DB_URL = "jdbc:pinot://localhost:8000?controller=localhost:9000";
+  static final String BAD_URL = "jdbc:someOtherDB://localhost:8000?controller=localhost:9000";
+  static final String GOOD_URL_NO_CONNECTION = "jdbc:pinot://localhost:1111?controller=localhost:2222";
 
   @Test(enabled = false)
   public void testDriver()
@@ -59,4 +61,22 @@ public class PinotDriverTest {
     ;
     conn.close();
   }
+
+  @Test
+  public void testDriverBadURL() {
+    try {
+      DriverManager.getConnection(BAD_URL);
+    } catch (Exception e) {
+      Assert.assertTrue(e.getMessage().contains("No suitable driver found"));
+    }
+  }
+
+  @Test
+  public void testDriverGoodURLNoConnection() {
+    try {
+      DriverManager.getConnection(GOOD_URL_NO_CONNECTION);
+    } catch (Exception e) {
+      Assert.assertTrue(e.getMessage().contains("Failed to connect to url"));
+    }
+  }
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org