org.apache.calcite.test.ConnectionFactories Maven / Gradle / Ivy
The newest version!
/*
* 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.calcite.test;
import org.apache.calcite.avatica.ConnectionProperty;
import org.apache.calcite.jdbc.CalciteConnection;
import org.apache.calcite.rel.type.RelProtoDataType;
import org.apache.calcite.runtime.FlatLists;
import org.apache.calcite.schema.Schema;
import org.apache.calcite.schema.SchemaPlus;
import org.apache.commons.dbcp2.PoolableConnection;
import org.apache.commons.dbcp2.PoolableConnectionFactory;
import org.apache.commons.dbcp2.PoolingDataSource;
import org.apache.commons.pool2.impl.GenericObjectPool;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.util.Map;
import java.util.Objects;
import java.util.Properties;
import static java.util.Objects.requireNonNull;
/** Utilities for {@link ConnectionFactory} and
* {@link org.apache.calcite.test.CalciteAssert.ConnectionPostProcessor}. */
public abstract class ConnectionFactories {
/** The empty connection factory. */
private static final ConnectionFactory EMPTY =
new MapConnectionFactory(ImmutableMap.of(), ImmutableList.of());
/** Prevent instantiation of utility class. */
private ConnectionFactories() {
}
/** Returns an empty connection factory. */
public static ConnectionFactory empty() {
return EMPTY;
}
/** Creates a connection factory that uses a single pooled connection,
* as opposed to creating a new connection on each invocation. */
public static ConnectionFactory pool(ConnectionFactory connectionFactory) {
return connectionFactory instanceof PoolingConnectionFactory
? connectionFactory
: new PoolingConnectionFactory(connectionFactory);
}
/** Returns a post-processor that adds a {@link CalciteAssert.SchemaSpec}
* (set of schemes) to a connection. */
public static CalciteAssert.ConnectionPostProcessor add(
CalciteAssert.SchemaSpec schemaSpec) {
return new AddSchemaSpecPostProcessor(schemaSpec);
}
/** Returns a post-processor that adds {@link Schema} and sets it as
* default. */
public static CalciteAssert.ConnectionPostProcessor add(String name,
Schema schema) {
return new AddSchemaPostProcessor(name, schema);
}
/** Returns a post-processor that sets a default schema name. */
public static CalciteAssert.ConnectionPostProcessor setDefault(
String schema) {
return new DefaultSchemaPostProcessor(schema);
}
/** Returns a post-processor that adds a type. */
public static CalciteAssert.ConnectionPostProcessor addType(String name,
RelProtoDataType protoDataType) {
return new AddTypePostProcessor(name, protoDataType);
}
/** Connection factory that uses a given map of (name, value) pairs and
* optionally an initial schema. */
private static class MapConnectionFactory implements ConnectionFactory {
private final ImmutableMap map;
private final ImmutableList postProcessors;
MapConnectionFactory(ImmutableMap map,
ImmutableList postProcessors) {
this.map = requireNonNull(map, "map");
this.postProcessors = requireNonNull(postProcessors, "postProcessors");
}
@Override public boolean equals(Object obj) {
return this == obj
|| obj.getClass() == MapConnectionFactory.class
&& ((MapConnectionFactory) obj).map.equals(map)
&& ((MapConnectionFactory) obj).postProcessors.equals(postProcessors);
}
@Override public int hashCode() {
return Objects.hash(map, postProcessors);
}
@Override public Connection createConnection() throws SQLException {
final Properties info = new Properties();
for (Map.Entry entry : map.entrySet()) {
info.setProperty(entry.getKey(), entry.getValue());
}
Connection connection =
DriverManager.getConnection("jdbc:calcite:", info);
for (CalciteAssert.ConnectionPostProcessor postProcessor : postProcessors) {
connection = postProcessor.apply(connection);
}
return connection;
}
@Override public ConnectionFactory with(String property, Object value) {
return new MapConnectionFactory(
FlatLists.append(this.map, property, value.toString()),
postProcessors);
}
@Override public ConnectionFactory with(ConnectionProperty property, Object value) {
if (!property.type().valid(value, property.valueClass())) {
throw new IllegalArgumentException();
}
return with(property.camelName(), value.toString());
}
@Override public ConnectionFactory with(
CalciteAssert.ConnectionPostProcessor postProcessor) {
ImmutableList.Builder builder =
ImmutableList.builder();
builder.addAll(postProcessors);
builder.add(postProcessor);
return new MapConnectionFactory(map, builder.build());
}
}
/** Post-processor that adds a {@link Schema} and sets it as default. */
private static class AddSchemaPostProcessor
implements CalciteAssert.ConnectionPostProcessor {
private final String name;
private final Schema schema;
AddSchemaPostProcessor(String name, Schema schema) {
this.name = requireNonNull(name, "name");
this.schema = requireNonNull(schema, "schema");
}
@Override public Connection apply(Connection connection) throws SQLException {
CalciteConnection con = connection.unwrap(CalciteConnection.class);
SchemaPlus rootSchema = con.getRootSchema();
rootSchema.add(name, schema);
connection.setSchema(name);
return connection;
}
}
/** Post-processor that adds a type. */
private static class AddTypePostProcessor
implements CalciteAssert.ConnectionPostProcessor {
private final String name;
private final RelProtoDataType protoDataType;
AddTypePostProcessor(String name, RelProtoDataType protoDataType) {
this.name = requireNonNull(name, "name");
this.protoDataType = requireNonNull(protoDataType, "protoDataType");
}
@Override public Connection apply(Connection connection) throws SQLException {
CalciteConnection con = connection.unwrap(CalciteConnection.class);
SchemaPlus rootSchema = con.getRootSchema();
rootSchema.add(name, protoDataType);
return connection;
}
}
/** Post-processor that sets a default schema name. */
private static class DefaultSchemaPostProcessor
implements CalciteAssert.ConnectionPostProcessor {
private final String name;
DefaultSchemaPostProcessor(String name) {
this.name = name;
}
@Override public Connection apply(Connection connection) throws SQLException {
connection.setSchema(name);
return connection;
}
}
/** Post-processor that adds a {@link CalciteAssert.SchemaSpec}
* (set of schemes) to a connection. */
private static class AddSchemaSpecPostProcessor
implements CalciteAssert.ConnectionPostProcessor {
private final CalciteAssert.SchemaSpec schemaSpec;
AddSchemaSpecPostProcessor(CalciteAssert.SchemaSpec schemaSpec) {
this.schemaSpec = schemaSpec;
}
@Override public Connection apply(Connection connection) throws SQLException {
CalciteConnection con = connection.unwrap(CalciteConnection.class);
SchemaPlus rootSchema = con.getRootSchema();
switch (schemaSpec) {
case CLONE_FOODMART:
case JDBC_FOODMART_WITH_LATTICE:
CalciteAssert.addSchema(rootSchema, CalciteAssert.SchemaSpec.JDBC_FOODMART);
// fall through
default:
CalciteAssert.addSchema(rootSchema, schemaSpec);
}
con.setSchema(schemaSpec.schemaName);
return connection;
}
}
/** Connection factory that uses the same instance of connections. */
private static class PoolingConnectionFactory implements ConnectionFactory {
private final PoolingDataSource dataSource;
PoolingConnectionFactory(final ConnectionFactory factory) {
final PoolableConnectionFactory connectionFactory =
new PoolableConnectionFactory(factory::createConnection, null);
connectionFactory.setRollbackOnReturn(false);
this.dataSource =
new PoolingDataSource<>(new GenericObjectPool<>(connectionFactory));
}
@Override public Connection createConnection() throws SQLException {
return dataSource.getConnection();
}
}
}