1 files changed, 64 insertions, 20 deletions

diff --git a/src/main/java/com/zaxxer/hikari/pool/PoolBase.java b/src/main/java/com/zaxxer/hikari/pool/PoolBase.java
index b463dae..9063d81 100755
--- a/src/main/java/com/zaxxer/hikari/pool/PoolBase.java
+++ b/src/main/java/com/zaxxer/hikari/pool/PoolBase.java
@@ -16,6 +16,7 @@
 
 package com.zaxxer.hikari.pool;
 
+import static com.zaxxer.hikari.pool.ProxyConnection.DIRTY_BIT_SCHEMA;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static java.util.concurrent.TimeUnit.SECONDS;
 
@@ -29,12 +30,16 @@ import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 
 import javax.management.MBeanServer;
 import javax.management.ObjectName;
+import javax.naming.InitialContext;
+import javax.naming.NamingException;
 import javax.sql.DataSource;
 
+import com.zaxxer.hikari.pool.HikariPool.PoolInitializationException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -61,11 +66,11 @@ abstract class PoolBase
 
    protected final HikariConfig config;
    protected final String poolName;
-   protected long connectionTimeout;
-   protected long validationTimeout;
-   protected IMetricsTrackerDelegate metricsTracker;
+   long connectionTimeout;
+   long validationTimeout;
+   IMetricsTrackerDelegate metricsTracker;
 
-   private static final String[] RESET_STATES = {"readOnly", "autoCommit", "isolation", "catalog", "netTimeout"};
+   private static final String[] RESET_STATES = {"readOnly", "autoCommit", "isolation", "catalog", "netTimeout", "schema"};
    private static final int UNINITIALIZED = -1;
    private static final int TRUE = 1;
    private static final int FALSE = 0;
@@ -79,6 +84,7 @@ abstract class PoolBase
    private DataSource dataSource;
 
    private final String catalog;
+   private final String schema;
    private final boolean isReadOnly;
    private final boolean isAutoCommit;
 
@@ -94,6 +100,7 @@ abstract class PoolBase
 
       this.networkTimeout = UNINITIALIZED;
       this.catalog = config.getCatalog();
+      this.schema = config.getSchema();
       this.isReadOnly = config.isReadOnly();
       this.isAutoCommit = config.isAutoCommit();
       this.transactionIsolation = UtilityElf.getTransactionIsolation(config.getTransactionIsolation());
@@ -146,28 +153,30 @@ abstract class PoolBase
    {
       try {
          try {
+            setNetworkTimeout(connection, validationTimeout);
+
+            final long validationSeconds = (int) Math.max(1000L, validationTimeout) / 1000;
+
             if (isUseJdbc4Validation) {
-               return connection.isValid((int) MILLISECONDS.toSeconds(Math.max(1000L, validationTimeout)));
+               return connection.isValid((int) validationSeconds);
             }
 
-            setNetworkTimeout(connection, validationTimeout);
-
             try (Statement statement = connection.createStatement()) {
                if (isNetworkTimeoutSupported != TRUE) {
-                  setQueryTimeout(statement, (int) MILLISECONDS.toSeconds(Math.max(1000L, validationTimeout)));
+                  setQueryTimeout(statement, (int) validationSeconds);
                }
 
                statement.execute(config.getConnectionTestQuery());
             }
          }
          finally {
+            setNetworkTimeout(connection, networkTimeout);
+
             if (isIsolateInternalQueries && !isAutoCommit) {
                connection.rollback();
             }
          }
 
-         setNetworkTimeout(connection, networkTimeout);
-
          return true;
       }
       catch (Exception e) {
@@ -225,6 +234,11 @@ abstract class PoolBase
          resetBits |= DIRTY_BIT_NETTIMEOUT;
       }
 
+      if ((dirtyBits & DIRTY_BIT_SCHEMA) != 0 && schema != null && !schema.equals(proxyConnection.getSchemaState())) {
+         connection.setSchema(schema);
+         resetBits |= DIRTY_BIT_SCHEMA;
+      }
+
       if (resetBits != 0 && LOGGER.isDebugEnabled()) {
          LOGGER.debug("{} - Reset ({}) on connection {}", poolName, stringFromResetBits(resetBits), connection);
       }
@@ -237,6 +251,15 @@ abstract class PoolBase
       }
    }
 
+   long getLoginTimeout()
+   {
+      try {
+         return (dataSource != null) ? dataSource.getLoginTimeout() : SECONDS.toSeconds(5);
+      } catch (SQLException e) {
+         return SECONDS.toSeconds(5);
+      }
+   }
+
    // ***********************************************************************
    //                       JMX methods
    // ***********************************************************************
@@ -244,7 +267,7 @@ abstract class PoolBase
    /**
     * Register MBeans for HikariConfig and HikariPool.
     *
-    * @param pool a HikariPool instance
+    * @param hikariPool a HikariPool instance
     */
    void registerMBeans(final HikariPool hikariPool)
    {
@@ -300,8 +323,6 @@ abstract class PoolBase
 
    /**
     * Create/initialize the underlying DataSource.
-    *
-    * @return a DataSource instance
     */
    private void initializeDataSource()
    {
@@ -310,6 +331,7 @@ abstract class PoolBase
       final String password = config.getPassword();
       final String dsClassName = config.getDataSourceClassName();
       final String driverClassName = config.getDriverClassName();
+      final String dataSourceJNDI = config.getDataSourceJNDI();
       final Properties dataSourceProperties = config.getDataSourceProperties();
 
       DataSource dataSource = config.getDataSource();
@@ -320,6 +342,14 @@ abstract class PoolBase
       else if (jdbcUrl != null && dataSource == null) {
          dataSource = new DriverDataSource(jdbcUrl, driverClassName, dataSourceProperties, username, password);
       }
+      else if (dataSourceJNDI != null && dataSource == null) {
+         try {
+            InitialContext ic = new InitialContext();
+            dataSource = (DataSource) ic.lookup(dataSourceJNDI);
+         } catch (NamingException e) {
+            throw new PoolInitializationException(e);
+         }
+      }
 
       if (dataSource != null) {
          setLoginTimeout(dataSource);
@@ -334,7 +364,7 @@ abstract class PoolBase
     *
     * @return a Connection connection
     */
-   Connection newConnection() throws Exception
+   private Connection newConnection() throws Exception
    {
       final long start = currentTime();
 
@@ -375,7 +405,7 @@ abstract class PoolBase
     * Setup a connection initial state.
     *
     * @param connection a Connection
-    * @throws SQLException thrown from driver
+    * @throws ConnectionSetupException thrown if any exception is encountered
     */
    private void setupConnection(final Connection connection) throws ConnectionSetupException
    {
@@ -400,6 +430,10 @@ abstract class PoolBase
             connection.setCatalog(catalog);
          }
 
+         if (schema != null) {
+            connection.setSchema(schema);
+         }
+
          executeSql(connection, config.getConnectionInitSql(), true);
 
          setNetworkTimeout(connection, networkTimeout);
@@ -438,10 +472,12 @@ abstract class PoolBase
          }
          catch (SQLException e) {
             LOGGER.warn("{} - Default transaction isolation level detection failed ({}).", poolName, e.getMessage());
+            if (e.getSQLState() != null && !e.getSQLState().startsWith("08")) {
+               throw e;
+            }
          }
-         finally {
-            isValidChecked = true;
-         }
+
+         isValidChecked = true;
       }
    }
 
@@ -610,7 +646,7 @@ abstract class PoolBase
    {
       private static final long serialVersionUID = 929872118275916521L;
 
-      public ConnectionSetupException(Throwable t)
+      ConnectionSetupException(Throwable t)
       {
          super(t);
       }
@@ -635,12 +671,14 @@ abstract class PoolBase
       }
    }
 
-   static interface IMetricsTrackerDelegate extends AutoCloseable
+   interface IMetricsTrackerDelegate extends AutoCloseable
    {
       default void recordConnectionUsage(PoolEntry poolEntry) {}
 
       default void recordConnectionCreated(long connectionCreatedMillis) {}
 
+      default void recordBorrowTimeoutStats(long startTime) {}
+
       default void recordBorrowStats(final PoolEntry poolEntry, final long startTime) {}
 
       default void recordConnectionTimeout() {}
@@ -676,6 +714,12 @@ abstract class PoolBase
       }
 
       @Override
+      public void recordBorrowTimeoutStats(long startTime)
+      {
+         tracker.recordConnectionAcquiredNanos(elapsedNanos(startTime));
+      }
+
+      @Override
       public void recordBorrowStats(final PoolEntry poolEntry, final long startTime)
       {
          final long now = currentTime();