DRILL-2127: make JDBC tests re-use bit instances; reset cached bits o…

…n failure

common/src/main/java/org/apache/drill/common/config/DrillConfig.java

@@ -96,9 +96,9 @@ public static DrillConfig create() {
   }
 
   /**
-   * Creates a {{@link DrillConfig configuration}} disabling server specific configuration options.
+   * Creates a {@link DrillConfig configuration} disabling server specific configuration options.
    *
-   * @return {{@link DrillConfig}} instance
+   * @return {@link DrillConfig} instance
    */
   public static DrillConfig forClient() {
     return create(null, false);

exec/jdbc/src/test/java/org/apache/drill/jdbc/CachingConnectionFactory.java

@@ -0,0 +1,31 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.drill.jdbc;
+
+import java.sql.SQLException;
+
+/**
+ * A connection factory that caches connections.
+ */
+public interface CachingConnectionFactory extends ConnectionFactory {
+
+  /**
+   * Closes all active connections in the cache.
+   */
+  void close() throws SQLException;
+}

exec/jdbc/src/test/java/org/apache/drill/jdbc/ConnectionFactory.java

@@ -0,0 +1,33 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.drill.jdbc;
+
+import java.sql.Connection;
+
+/**
+ * A factory used to create {@link java.sql.Connection} instances.
+ */
+public interface ConnectionFactory {
+  /**
+   * Creates a new {@link java.sql.Connection} based on the given {@link org.apache.drill.jdbc.ConnectionInfo info}
+   *
+   * @param info connection parameters
+   * @throws Exception if factory fails to create a connection.
+   */
+  Connection createConnection(ConnectionInfo info) throws Exception;
+}

exec/jdbc/src/test/java/org/apache/drill/jdbc/ConnectionInfo.java

@@ -0,0 +1,81 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.drill.jdbc;
+
+import java.util.Map;
+import java.util.Properties;
+
+import com.google.common.base.Objects;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableMap;
+
+/**
+ * An immutable bag of parameters that describes a {@link java.sql.Connection}.
+ */
+public class ConnectionInfo {
+  private final String url;
+  private final Map<Object, Object> params;
+
+  public ConnectionInfo(String url, Properties params) {
+    this(url, ImmutableMap.copyOf(params));
+  }
+
+  public ConnectionInfo(String url, ImmutableMap<Object, Object> params) {
+    this.url = Preconditions.checkNotNull(url, "URL cannot be null");
+    this.params = params;
+  }
+
+  /**
+   * Returns connection url.
+   */
+  public String getUrl() {
+    return url;
+  }
+
+  /**
+   * Returns connection parameters.
+   */
+  public Map<Object, Object> getParameters() {
+    return params;
+  }
+
+  /**
+   * Creates a new {@link java.util.Properties} instance from underlying parameters.
+   */
+  public Properties getParamsAsProperties() {
+    final Properties props = new Properties();
+    for (Object key:params.keySet()) {
+      props.put(key, params.get(key));
+    }
+    return props;
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hashCode(url, params);
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (obj instanceof ConnectionInfo) {
+      final ConnectionInfo info = ConnectionInfo.class.cast(obj);
+      return Objects.equal(url, info.getUrl()) && Objects.equal(params, info.getParameters());
+    }
+    return false;
+  }
+}

exec/jdbc/src/test/java/org/apache/drill/jdbc/JdbcTest.java

@@ -17,8 +17,104 @@
  */
 package org.apache.drill.jdbc;
 
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Properties;
+
+import com.google.common.base.Strings;
 import org.apache.drill.exec.ExecTest;
+import org.apache.drill.jdbc.test.JdbcAssert;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.rules.TestRule;
+import org.junit.rules.TestWatcher;
+import org.junit.runner.Description;
 
-public class JdbcTest extends ExecTest{
+public class JdbcTest extends ExecTest {
   static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(JdbcTest.class);
+
+  @Rule
+  public final TestRule watcher = new TestWatcher() {
+    @Override
+    protected void failed(Throwable e, Description description) {
+      reset();
+    }
+  };
+
+  private static CachingConnectionFactory factory;
+
+  @BeforeClass
+  public static void setUp() {
+    factory = new SingleConnectionCachingFactory(new ConnectionFactory() {
+      @Override
+      public Connection createConnection(ConnectionInfo info) throws Exception {
+        Class.forName("org.apache.drill.jdbc.Driver");
+        return DriverManager.getConnection(info.getUrl(), info.getParamsAsProperties());
+      }
+    });
+    JdbcAssert.setFactory(factory);
+  }
+
+  /**
+   * Creates a {@link java.sql.Connection connection} using default parameters.
+   * @param url connection URL
+   * @throws Exception if connection fails
+   */
+  protected static Connection connect(String url) throws Exception {
+    return connect(url, JdbcAssert.getDefaultProperties());
+  }
+
+
+  /**
+   * Creates a {@link java.sql.Connection connection} using the given parameters.
+   * @param url connection URL
+   * @param info connection info
+   * @throws Exception if connection fails
+   */
+  protected static Connection connect(String url, Properties info) throws Exception {
+    final Connection conn = factory.createConnection(new ConnectionInfo(url, info));
+    changeSchemaIfSupplied(conn, info);
+    return conn;
+  }
+
+  /**
+   * Changes schema of the given connection if the field "schema" is present in {@link java.util.Properties info}.
+   * Does nothing otherwise.
+   */
+  protected static void changeSchemaIfSupplied(Connection conn, Properties info) {
+    final String schema = info.getProperty("schema", null);
+    if (!Strings.isNullOrEmpty(schema)) {
+      changeSchema(conn, schema);
+    }
+  }
+
+  protected static void changeSchema(Connection conn, String schema) {
+    final String query = String.format("use %s", schema);
+    try {
+      Statement s = conn.createStatement();
+      ResultSet r = s.executeQuery(query);
+    } catch (SQLException e) {
+      throw new RuntimeException("unable to change schema", e);
+    }
+  }
+
+  /**
+   * Resets the factory closing all of the active connections.
+   */
+  protected static void reset() {
+    try {
+      factory.close();
+    } catch (SQLException e) {
+      throw new RuntimeException("error while closing connection factory", e);
+    }
+  }
+
+  @AfterClass
+  public static void clearUp() throws Exception {
+    factory.close();
+  }
 }

exec/jdbc/src/test/java/org/apache/drill/jdbc/MultiConnectionCachingFactory.java

@@ -0,0 +1,65 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.drill.jdbc;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.util.Map;
+
+import com.google.common.collect.Maps;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A connection factory that caches connections based on given {@link org.apache.drill.jdbc.ConnectionInfo}.
+ */
+public class MultiConnectionCachingFactory implements CachingConnectionFactory {
+  private static final Logger logger = LoggerFactory.getLogger(MultiConnectionCachingFactory.class);
+
+  private final ConnectionFactory delegate;
+  private final Map<ConnectionInfo, Connection> cache = Maps.newHashMap();
+
+  public MultiConnectionCachingFactory(ConnectionFactory delegate) {
+    this.delegate = delegate;
+  }
+
+  /**
+   * Creates a {@link org.apache.drill.jdbc.NonClosableConnection connection} and caches it.
+   *
+   * The returned {@link org.apache.drill.jdbc.NonClosableConnection connection} does not support
+   * {@link java.sql.Connection#close()}. Consumer must call {#close} to close the cached connections.
+   */
+  @Override
+  public Connection createConnection(ConnectionInfo info) throws Exception {
+    Connection conn = cache.get(info);
+    if (conn == null) {
+      conn = delegate.createConnection(info);
+      cache.put(info, conn);
+    }
+    return new NonClosableConnection(conn);
+  }
+
+  /**
+   * Closes all active connections in the cache.
+   */
+  public void close() throws SQLException {
+    for (Connection conn:cache.values()) {
+      conn.close();
+    }
+  }
+}