com.facebook.presto.tests.H2QueryRunner Maven / Gradle / Ivy
/*
* Licensed 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 com.facebook.presto.tests;
import com.facebook.presto.Session;
import com.facebook.presto.spi.ColumnMetadata;
import com.facebook.presto.spi.ConnectorTableMetadata;
import com.facebook.presto.spi.RecordCursor;
import com.facebook.presto.spi.RecordSet;
import com.facebook.presto.spi.SchemaTableName;
import com.facebook.presto.spi.type.DecimalType;
import com.facebook.presto.spi.type.Type;
import com.facebook.presto.spi.type.VarcharType;
import com.facebook.presto.testing.MaterializedResult;
import com.facebook.presto.testing.MaterializedRow;
import com.facebook.presto.tpch.TpchMetadata;
import com.facebook.presto.tpch.TpchTableHandle;
import com.google.common.base.Joiner;
import org.intellij.lang.annotations.Language;
import org.joda.time.DateTimeZone;
import org.skife.jdbi.v2.DBI;
import org.skife.jdbi.v2.Handle;
import org.skife.jdbi.v2.PreparedBatch;
import org.skife.jdbi.v2.PreparedBatchPart;
import org.skife.jdbi.v2.StatementContext;
import org.skife.jdbi.v2.tweak.ResultSetMapper;
import java.math.BigDecimal;
import java.sql.Date;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Time;
import java.sql.Timestamp;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;
import static com.facebook.presto.spi.type.BigintType.BIGINT;
import static com.facebook.presto.spi.type.BooleanType.BOOLEAN;
import static com.facebook.presto.spi.type.DateType.DATE;
import static com.facebook.presto.spi.type.DoubleType.DOUBLE;
import static com.facebook.presto.spi.type.IntegerType.INTEGER;
import static com.facebook.presto.spi.type.SmallintType.SMALLINT;
import static com.facebook.presto.spi.type.TimeType.TIME;
import static com.facebook.presto.spi.type.TimeWithTimeZoneType.TIME_WITH_TIME_ZONE;
import static com.facebook.presto.spi.type.TimestampType.TIMESTAMP;
import static com.facebook.presto.spi.type.TimestampWithTimeZoneType.TIMESTAMP_WITH_TIME_ZONE;
import static com.facebook.presto.spi.type.TinyintType.TINYINT;
import static com.facebook.presto.spi.type.Varchars.isVarcharType;
import static com.facebook.presto.tpch.TpchMetadata.TINY_SCHEMA_NAME;
import static com.facebook.presto.tpch.TpchRecordSet.createTpchRecordSet;
import static com.facebook.presto.type.UnknownType.UNKNOWN;
import static com.facebook.presto.util.DateTimeZoneIndex.getDateTimeZone;
import static com.facebook.presto.util.ImmutableCollectors.toImmutableList;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import static io.airlift.tpch.TpchTable.LINE_ITEM;
import static io.airlift.tpch.TpchTable.ORDERS;
import static java.lang.String.format;
import static java.util.Collections.nCopies;
public class H2QueryRunner
{
private final Handle handle;
public H2QueryRunner()
{
handle = DBI.open("jdbc:h2:mem:test" + System.nanoTime());
TpchMetadata tpchMetadata = new TpchMetadata("");
handle.execute("CREATE TABLE orders (\n" +
" orderkey BIGINT PRIMARY KEY,\n" +
" custkey BIGINT NOT NULL,\n" +
" orderstatus CHAR(1) NOT NULL,\n" +
" totalprice DOUBLE NOT NULL,\n" +
" orderdate DATE NOT NULL,\n" +
" orderpriority CHAR(15) NOT NULL,\n" +
" clerk CHAR(15) NOT NULL,\n" +
" shippriority INTEGER NOT NULL,\n" +
" comment VARCHAR(79) NOT NULL\n" +
")");
handle.execute("CREATE INDEX custkey_index ON orders (custkey)");
TpchTableHandle ordersHandle = tpchMetadata.getTableHandle(null, new SchemaTableName(TINY_SCHEMA_NAME, ORDERS.getTableName()));
insertRows(tpchMetadata.getTableMetadata(null, ordersHandle), handle, createTpchRecordSet(ORDERS, ordersHandle.getScaleFactor()));
handle.execute("CREATE TABLE lineitem (\n" +
" orderkey BIGINT,\n" +
" partkey BIGINT NOT NULL,\n" +
" suppkey BIGINT NOT NULL,\n" +
" linenumber INTEGER,\n" +
" quantity DOUBLE NOT NULL,\n" +
" extendedprice DOUBLE NOT NULL,\n" +
" discount DOUBLE NOT NULL,\n" +
" tax DOUBLE NOT NULL,\n" +
" returnflag CHAR(1) NOT NULL,\n" +
" linestatus CHAR(1) NOT NULL,\n" +
" shipdate DATE NOT NULL,\n" +
" commitdate DATE NOT NULL,\n" +
" receiptdate DATE NOT NULL,\n" +
" shipinstruct VARCHAR(25) NOT NULL,\n" +
" shipmode VARCHAR(10) NOT NULL,\n" +
" comment VARCHAR(44) NOT NULL,\n" +
" PRIMARY KEY (orderkey, linenumber)" +
")");
TpchTableHandle lineItemHandle = tpchMetadata.getTableHandle(null, new SchemaTableName(TINY_SCHEMA_NAME, LINE_ITEM.getTableName()));
insertRows(tpchMetadata.getTableMetadata(null, lineItemHandle), handle, createTpchRecordSet(LINE_ITEM, lineItemHandle.getScaleFactor()));
}
public void close()
{
handle.close();
}
public MaterializedResult execute(Session session, @Language("SQL") String sql, List extends Type> resultTypes)
{
MaterializedResult materializedRows = new MaterializedResult(
handle.createQuery(sql)
.map(rowMapper(resultTypes))
.list(),
resultTypes
);
// H2 produces dates in the JVM time zone instead of the session timezone
materializedRows = materializedRows.toTimeZone(DateTimeZone.getDefault(), getDateTimeZone(session.getTimeZoneKey()));
return materializedRows;
}
private static ResultSetMapper rowMapper(final List extends Type> types)
{
return new ResultSetMapper()
{
@Override
public MaterializedRow map(int index, ResultSet resultSet, StatementContext ctx)
throws SQLException
{
int count = resultSet.getMetaData().getColumnCount();
checkArgument(types.size() == count, "type does not match result");
List
© 2015 - 2025 Weber Informatics LLC | Privacy Policy