added close() for connection

yugabyte · May 7, 2024 · a891cb7 · a891cb7
1 parent 5fbfc91
commit a891cb7
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 7 deletions.
diff --git a/pgjdbc/src/main/java/com/yugabyte/ysql/YBDriver.java b/pgjdbc/src/main/java/com/yugabyte/ysql/YBDriver.java
@@ -5,6 +5,7 @@
 import java.util.Properties;
 
 import org.postgresql.CustomDriver;
+import org.postgresql.core.QueryExecutor;
 import org.postgresql.jdbc.PgConnection;
 import static org.postgresql.Driver.*;
 
@@ -28,6 +29,22 @@ public Connection makeConnection(String url, Properties properties) throws SQLEx
     }
     return new PgConnection(hostSpecs(properties), user(properties), database(properties),
         properties, url);
-}
+    }
+
+    /**
+   * <B>Note:</B> even though {@code Statement} is automatically closed when it is garbage
+   * collected, it is better to close it explicitly to lower resource consumption.
+   * The spec says that calling close on a closed connection is a no-op.
+   *
+   * {@inheritDoc}
+   */
+  @Override
+  public void close(QueryExecutor queryExecutor) throws SQLException {
+
+    String host = queryExecutor.getHostSpec().getHost();
+    if (host != null) {
+      LoadBalanceManager.decrementConnectionCount(host);
+    }
+  }
 
 }
diff --git a/pgjdbc/src/main/java/org/postgresql/CustomDriver.java b/pgjdbc/src/main/java/org/postgresql/CustomDriver.java
@@ -4,6 +4,8 @@
 import java.sql.SQLException;
 import java.util.Properties;
 
+import org.postgresql.core.QueryExecutor;
+
 /**
  * An interface for any custom driver to implement. connect() invokes the
  * implemented method while processing a new connection request.
@@ -21,4 +23,13 @@ public interface CustomDriver {
    */
     Connection makeConnection(String url, Properties properties) throws SQLException;
 
+    /**
+   * <B>Note:</B> even though {@code Statement} is automatically closed when it is garbage
+   * collected, it is better to close it explicitly to lower resource consumption.
+   * The spec says that calling close on a closed connection is a no-op.
+   *
+   * {@inheritDoc}
+   */
+  public void close(QueryExecutor queryExecutor) throws SQLException;
+
 }
diff --git a/pgjdbc/src/main/java/org/postgresql/DefaultCustomDriver.java b/pgjdbc/src/main/java/org/postgresql/DefaultCustomDriver.java
@@ -3,6 +3,8 @@
 import java.sql.Connection;
 import java.sql.SQLException;
 import java.util.Properties;
+
+import org.postgresql.core.QueryExecutor;
 import org.postgresql.jdbc.PgConnection;
 import static org.postgresql.Driver.*;
 
@@ -23,4 +25,15 @@ public Connection makeConnection(String url, Properties properties) throws SQLEx
         properties, url);
     }
 
+    /**
+   * <B>Note:</B> even though {@code Statement} is automatically closed when it is garbage
+   * collected, it is better to close it explicitly to lower resource consumption.
+   * The spec says that calling close on a closed connection is a no-op.
+   *
+   * {@inheritDoc}
+   */
+  @Override
+  public void close(QueryExecutor queryExecutor) throws SQLException {
+    //No op
+  }
 }
diff --git a/pgjdbc/src/main/java/org/postgresql/jdbc/PgConnection.java b/pgjdbc/src/main/java/org/postgresql/jdbc/PgConnection.java
@@ -22,6 +22,7 @@
 
 import static org.postgresql.util.internal.Nullness.castNonNull;
 
+import org.postgresql.CustomDriver;
 import org.postgresql.Driver;
 import org.postgresql.PGNotification;
 import org.postgresql.PGProperty;
@@ -773,12 +774,7 @@ public void close() throws SQLException {
     releaseTimer();
     queryExecutor.close();
     openStackTrace = null;
-    if (Driver.custDriver instanceof YBDriver ) {
-      String host = queryExecutor.getHostSpec().getHost();
-      if (loadBalancer != null && host != null) {
-        LoadBalanceManager.decrementConnectionCount(host);
-      }
-    }
+    Driver.custDriver.close(queryExecutor);
   }
 
   public void setLoadBalancer(LoadBalancer lb) {