diff --git a/pgjdbc/src/main/java/com/yugabyte/ysql/YBDriver.java b/pgjdbc/src/main/java/com/yugabyte/ysql/YBDriver.java
index d2b1bd9e87..eebeaff57c 100644
--- 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);
-}
+ }
+
+ /**
+ * Note: 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
index c101f86213..4c54247675 100644
--- 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;
+ /**
+ * Note: 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
index 1bf50ea9a9..a63d0af8b3 100644
--- 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);
}
+ /**
+ * Note: 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
index 6db2e879f5..e3f9376588 100644
--- 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) {