Initial changes

pgjdbc/src/main/java/com/yugabyte/ysql/YBDriver.java

@@ -0,0 +1,33 @@
+package com.yugabyte.ysql;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.util.Properties;
+
+import org.postgresql.CustomDriver;
+import org.postgresql.jdbc.PgConnection;
+import static org.postgresql.Driver.*;
+
+public class YBDriver implements CustomDriver{
+
+    /**
+   * Create a connection from URL and properties. Always does the connection work in the current
+   * thread without enforcing a timeout, regardless of any timeout specified in the properties.
+   *
+   * @param url the original URL
+   * @param properties the parsed/defaulted connection properties
+   * @return a new connection
+   * @throws SQLException if the connection could not be made
+   */
+    @Override
+    public Connection makeConnection(String url, Properties properties) throws SQLException {
+        Connection connection = LoadBalanceManager.getConnection(url, properties, user(properties),
+        database(properties));
+    if (connection != null) {
+      return connection;
+    }
+    return new PgConnection(hostSpecs(properties), user(properties), database(properties),
+        properties, url);
+}
+
+}

pgjdbc/src/main/java/org/postgresql/CustomDriver.java

@@ -0,0 +1,24 @@
+package org.postgresql;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.util.Properties;
+
+/**
+ * An interface for any custom driver to implement. connect() invokes the
+ * implemented method while processing a new connection request.
+ */
+public interface CustomDriver {
+
+ /**
+   * Create a connection from URL and properties. Always does the connection work in the current
+   * thread without enforcing a timeout, regardless of any timeout specified in the properties.
+   *
+   * @param url the original URL
+   * @param properties the parsed/defaulted connection properties
+   * @return a new connection
+   * @throws SQLException if the connection could not be made
+   */
+    Connection makeConnection(String url, Properties properties) throws SQLException;
+
+}

pgjdbc/src/main/java/org/postgresql/DefaultCustomDriver.java

@@ -0,0 +1,26 @@
+package org.postgresql;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.util.Properties;
+import org.postgresql.jdbc.PgConnection;
+import static org.postgresql.Driver.*;
+
+public class DefaultCustomDriver implements CustomDriver{
+
+    /**
+   * Create a connection from URL and properties. Always does the connection work in the current
+   * thread without enforcing a timeout, regardless of any timeout specified in the properties.
+   *
+   * @param url the original URL
+   * @param properties the parsed/defaulted connection properties
+   * @return a new connection
+   * @throws SQLException if the connection could not be made
+   */
+    @Override
+    public Connection makeConnection(String url, Properties properties) throws SQLException {
+    return new PgConnection(hostSpecs(properties), user(properties), database(properties),
+        properties, url);
+    }
+
+}

pgjdbc/src/main/java/org/postgresql/Driver.java

@@ -38,6 +38,8 @@
 import org.postgresql.util.URLCoder;
 
 import com.yugabyte.ysql.LoadBalanceManager;
+import com.yugabyte.ysql.YBDriver;
+
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 import java.io.IOException;
@@ -87,6 +89,7 @@ public class Driver implements java.sql.Driver {
   private static final Logger LOGGER = Logger.getLogger("org.postgresql.Driver");
   private static final SharedTimer SHARED_TIMER = new SharedTimer();
   private static final String DEFAULT_PORT = "5433";
+  public static CustomDriver custDriver;
 
   static {
     try {
@@ -266,6 +269,7 @@ private Properties loadDefaultProperties() throws IOException {
     if ((props = parseURL(url, props)) == null) {
       return null;
     }
+    custDriver = loadCustomDriver(props);
     try {
       // Setup java.util.logging.Logger using connection properties.
       setupLoggerFromProperties(props);
@@ -282,7 +286,7 @@ private Properties loadDefaultProperties() throws IOException {
       // more details.
       long timeout = timeout(props);
       if (timeout <= 0) {
-        return makeConnection(url, props);
+        return custDriver.makeConnection(url, props);
       }
 
       ConnectThread ct = new ConnectThread(url, props);
@@ -393,7 +397,7 @@ public void run() {
       Throwable error;
 
       try {
-        conn = makeConnection(url, props);
+        conn = custDriver.makeConnection(url, props);
         error = null;
       } catch (Throwable t) {
         conn = null;
@@ -473,25 +477,6 @@ public Connection getResult(long timeout) throws SQLException {
     private boolean abandoned;
   }
 
-  /**
-   * Create a connection from URL and properties. Always does the connection work in the current
-   * thread without enforcing a timeout, regardless of any timeout specified in the properties.
-   *
-   * @param url the original URL
-   * @param properties the parsed/defaulted connection properties
-   * @return a new connection
-   * @throws SQLException if the connection could not be made
-   */
-  private static Connection makeConnection(String url, Properties properties) throws SQLException {
-    Connection connection = LoadBalanceManager.getConnection(url, properties, user(properties),
-        database(properties));
-    if (connection != null) {
-      return connection;
-    }
-    return new PgConnection(hostSpecs(properties), user(properties), database(properties),
-        properties, url);
-  }
-
   /**
    * Returns true if the driver thinks it can open a connection to the given URL. Typically, drivers
    * will return true if they understand the subprotocol specified in the URL and false if they
@@ -760,14 +745,14 @@ public static HostSpec[] hostSpecs(Properties props) {
   /**
    * @return the username of the URL
    */
-  private static String user(Properties props) {
+  public static String user(Properties props) {
     return props.getProperty("user", "");
   }
 
   /**
    * @return the database name of the URL
    */
-  private static String database(Properties props) {
+  public static String database(Properties props) {
     return props.getProperty("PGDBNAME", "");
   }
 
@@ -831,6 +816,22 @@ public static void register() throws SQLException {
     Driver.registeredDriver = registeredDriver;
   }
 
+  public static CustomDriver loadCustomDriver(Properties props) throws SQLException {
+    try {
+      String custom_driver = PGProperty.CUSTOM_DRIVER_CLASS.get(props);
+      if (custom_driver != null){
+        if (custom_driver.equalsIgnoreCase("com.yugabyte.ysql.YBDriver"))
+        return new YBDriver();
+      } else {
+        return new DefaultCustomDriver();
+      }
+    } catch (Exception e) {
+      //stripurl();
+      return new DefaultCustomDriver();
+    }
+    return new DefaultCustomDriver();
+  }
+
   /**
    * According to JDBC specification, this driver is registered against {@link DriverManager} when
    * the class is loaded. To avoid leaks, this method allow unregistering the driver so that the

pgjdbc/src/main/java/org/postgresql/PGProperty.java

@@ -539,6 +539,11 @@ public enum PGProperty {
       null,
       "Load Balance connections only in the given topologies"),
 
+  CUSTOM_DRIVER_CLASS(
+        "customdriverclass",
+        null,
+        "Class name of the custom driver, class should implement CustomDriver interface"),
+
   YB_SERVERS_REFRESH_INTERVAL(
       "yb-servers-refresh-interval",
       "300",

pgjdbc/src/main/java/org/postgresql/jdbc/PgConnection.java

@@ -59,6 +59,8 @@
 
 import com.yugabyte.ysql.LoadBalanceManager;
 import com.yugabyte.ysql.LoadBalancer;
+import com.yugabyte.ysql.YBDriver;
+
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.checkerframework.checker.nullness.qual.PolyNull;
 import org.checkerframework.dataflow.qual.Pure;
@@ -771,9 +773,11 @@ public void close() throws SQLException {
     releaseTimer();
     queryExecutor.close();
     openStackTrace = null;
-    String host = queryExecutor.getHostSpec().getHost();
-    if (loadBalancer != null && host != null) {
-      LoadBalanceManager.decrementConnectionCount(host);
+    if (Driver.custDriver instanceof YBDriver ) {
+      String host = queryExecutor.getHostSpec().getHost();
+      if (loadBalancer != null && host != null) {
+        LoadBalanceManager.decrementConnectionCount(host);
+      }
     }
   }