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) {