| /* |
| * 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.druid.sql.avatica; |
| |
| import com.google.common.base.Function; |
| import com.google.common.collect.ImmutableSet; |
| import com.google.common.collect.Lists; |
| import org.apache.calcite.avatica.ColumnMetaData; |
| import org.apache.calcite.avatica.Meta; |
| import org.apache.calcite.avatica.remote.TypedValue; |
| import org.apache.druid.common.config.NullHandling; |
| import org.apache.druid.java.util.common.DateTimes; |
| import org.apache.druid.java.util.common.io.Closer; |
| import org.apache.druid.math.expr.ExprMacroTable; |
| import org.apache.druid.query.QueryRunnerFactoryConglomerate; |
| import org.apache.druid.segment.join.JoinableFactoryWrapper; |
| import org.apache.druid.server.QueryStackTests; |
| import org.apache.druid.server.SpecificSegmentsQuerySegmentWalker; |
| import org.apache.druid.server.security.AllowAllAuthenticator; |
| import org.apache.druid.server.security.AuthConfig; |
| import org.apache.druid.server.security.AuthTestUtils; |
| import org.apache.druid.sql.SqlQueryPlus; |
| import org.apache.druid.sql.SqlStatementFactory; |
| import org.apache.druid.sql.avatica.DruidJdbcResultSet.ResultFetcherFactory; |
| import org.apache.druid.sql.calcite.planner.CalciteRulesManager; |
| import org.apache.druid.sql.calcite.planner.CatalogResolver; |
| import org.apache.druid.sql.calcite.planner.DruidOperatorTable; |
| import org.apache.druid.sql.calcite.planner.PlannerConfig; |
| import org.apache.druid.sql.calcite.planner.PlannerFactory; |
| import org.apache.druid.sql.calcite.schema.DruidSchemaCatalog; |
| import org.apache.druid.sql.calcite.util.CalciteTestBase; |
| import org.apache.druid.sql.calcite.util.CalciteTests; |
| import org.junit.Assert; |
| import org.junit.jupiter.api.AfterAll; |
| import org.junit.jupiter.api.AfterEach; |
| import org.junit.jupiter.api.BeforeAll; |
| import org.junit.jupiter.api.BeforeEach; |
| import org.junit.jupiter.api.Test; |
| import org.junit.jupiter.api.io.TempDir; |
| |
| import java.io.File; |
| import java.io.IOException; |
| import java.util.Collections; |
| import java.util.List; |
| |
| public class DruidStatementTest extends CalciteTestBase |
| { |
| private static String SUB_QUERY_WITH_ORDER_BY = |
| "select T20.F13 as F22\n" |
| + "from (SELECT DISTINCT dim1 as F13 FROM druid.foo T10) T20\n" |
| + "order by T20.F13 ASC"; |
| private static String SELECT_FROM_FOO = |
| "SELECT __time, cnt, dim1, dim2, m1 FROM druid.foo"; |
| private static String SELECT_STAR_FROM_FOO = |
| "SELECT * FROM druid.foo"; |
| |
| private static SpecificSegmentsQuerySegmentWalker walker; |
| private static QueryRunnerFactoryConglomerate conglomerate; |
| private static Closer resourceCloser; |
| |
| @BeforeAll |
| public static void setUpClass(@TempDir File tempDir) |
| { |
| resourceCloser = Closer.create(); |
| conglomerate = QueryStackTests.createQueryRunnerFactoryConglomerate(resourceCloser); |
| walker = CalciteTests.createMockWalker(conglomerate, tempDir); |
| resourceCloser.register(walker); |
| } |
| |
| @AfterAll |
| public static void tearDownClass() throws IOException |
| { |
| resourceCloser.close(); |
| } |
| |
| private SqlStatementFactory sqlStatementFactory; |
| |
| @BeforeEach |
| public void setUp() |
| { |
| final PlannerConfig plannerConfig = new PlannerConfig(); |
| final DruidOperatorTable operatorTable = CalciteTests.createOperatorTable(); |
| final ExprMacroTable macroTable = CalciteTests.createExprMacroTable(); |
| DruidSchemaCatalog rootSchema = |
| CalciteTests.createMockRootSchema(conglomerate, walker, plannerConfig, AuthTestUtils.TEST_AUTHORIZER_MAPPER); |
| final JoinableFactoryWrapper joinableFactoryWrapper = CalciteTests.createJoinableFactoryWrapper(); |
| final PlannerFactory plannerFactory = new PlannerFactory( |
| rootSchema, |
| operatorTable, |
| macroTable, |
| plannerConfig, |
| AuthTestUtils.TEST_AUTHORIZER_MAPPER, |
| CalciteTests.getJsonMapper(), |
| CalciteTests.DRUID_SCHEMA_NAME, |
| new CalciteRulesManager(ImmutableSet.of()), |
| joinableFactoryWrapper, |
| CatalogResolver.NULL_RESOLVER, |
| new AuthConfig() |
| ); |
| this.sqlStatementFactory = CalciteTests.createSqlStatementFactory( |
| CalciteTests.createMockSqlEngine(walker, conglomerate), |
| plannerFactory |
| ); |
| } |
| |
| @AfterEach |
| public void tearDown() |
| { |
| |
| } |
| |
| //----------------------------------------------------------------- |
| // Druid JDBC Statement |
| // |
| // The JDBC Statement class starts "empty", then allows executing |
| // one statement at a time. Executing a second automatically closes |
| // the result set from the first. Each statement takes a new query. |
| // Parameters are not generally used in this pattern. |
| |
| private DruidJdbcStatement jdbcStatement() |
| { |
| return new DruidJdbcStatement( |
| "", |
| 0, |
| Collections.emptyMap(), |
| sqlStatementFactory, |
| new ResultFetcherFactory(AvaticaServerConfig.DEFAULT_FETCH_TIMEOUT_MS) |
| ); |
| } |
| |
| @Test |
| public void testSubQueryWithOrderByDirect() |
| { |
| SqlQueryPlus queryPlus = new SqlQueryPlus( |
| SUB_QUERY_WITH_ORDER_BY, |
| null, |
| null, |
| AllowAllAuthenticator.ALLOW_ALL_RESULT |
| ); |
| try (final DruidJdbcStatement statement = jdbcStatement()) { |
| // First frame, ask for all rows. |
| statement.execute(queryPlus, -1); |
| Meta.Frame frame = statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 6); |
| Assert.assertEquals( |
| subQueryWithOrderByResults(), |
| frame |
| ); |
| Assert.assertTrue(statement.isDone()); |
| } |
| } |
| |
| @Test |
| public void testFetchPastEOFDirect() |
| { |
| SqlQueryPlus queryPlus = new SqlQueryPlus( |
| SUB_QUERY_WITH_ORDER_BY, |
| null, |
| null, |
| AllowAllAuthenticator.ALLOW_ALL_RESULT |
| ); |
| try (final DruidJdbcStatement statement = jdbcStatement()) { |
| // First frame, ask for all rows. |
| statement.execute(queryPlus, -1); |
| Meta.Frame frame = statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 6); |
| Assert.assertEquals( |
| subQueryWithOrderByResults(), |
| frame |
| ); |
| Assert.assertTrue(statement.isDone()); |
| try { |
| statement.nextFrame(6, 6); |
| Assert.fail(); |
| } |
| catch (Exception e) { |
| // Expected: can't work with an auto-closed result set. |
| } |
| } |
| } |
| |
| /** |
| * Ensure an error is thrown if the execution step is skipped. |
| */ |
| @Test |
| public void testSkipExecuteDirect() |
| { |
| try (final DruidJdbcStatement statement = jdbcStatement()) { |
| // Error: no call to execute; |
| statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 6); |
| Assert.fail(); |
| } |
| catch (Exception e) { |
| // Expected |
| } |
| } |
| |
| /** |
| * Ensure an error is thrown if the client attempts to fetch from a |
| * statement after its result set is closed. |
| */ |
| @Test |
| public void testFetchAfterResultCloseDirect() |
| { |
| SqlQueryPlus queryPlus = new SqlQueryPlus( |
| SUB_QUERY_WITH_ORDER_BY, |
| null, |
| null, |
| AllowAllAuthenticator.ALLOW_ALL_RESULT |
| ); |
| try (final DruidJdbcStatement statement = jdbcStatement()) { |
| // First frame, ask for all rows. |
| statement.execute(queryPlus, -1); |
| statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 6); |
| statement.closeResultSet(); |
| statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 6); |
| Assert.fail(); |
| } |
| catch (Exception e) { |
| // Expected |
| } |
| } |
| |
| @Test |
| public void testSubQueryWithOrderByDirectTwice() |
| { |
| SqlQueryPlus queryPlus = new SqlQueryPlus( |
| SUB_QUERY_WITH_ORDER_BY, |
| null, |
| null, |
| AllowAllAuthenticator.ALLOW_ALL_RESULT |
| ); |
| try (final DruidJdbcStatement statement = jdbcStatement()) { |
| statement.execute(queryPlus, -1); |
| Meta.Frame frame = statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 6); |
| Assert.assertEquals( |
| subQueryWithOrderByResults(), |
| frame |
| ); |
| |
| // Do it again. JDBC says we can reuse statements sequentially. |
| Assert.assertTrue(statement.isDone()); |
| statement.execute(queryPlus, -1); |
| frame = statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 6); |
| Assert.assertEquals( |
| subQueryWithOrderByResults(), |
| frame |
| ); |
| Assert.assertTrue(statement.isDone()); |
| } |
| } |
| |
| private Meta.Frame subQueryWithOrderByResults() |
| { |
| return Meta.Frame.create( |
| 0, |
| true, |
| Lists.newArrayList( |
| new Object[]{""}, |
| new Object[]{"1"}, |
| new Object[]{"10.1"}, |
| new Object[]{"2"}, |
| new Object[]{"abc"}, |
| new Object[]{"def"} |
| ) |
| ); |
| } |
| |
| @Test |
| public void testSelectAllInFirstFrameDirect() |
| { |
| SqlQueryPlus queryPlus = new SqlQueryPlus( |
| SELECT_FROM_FOO, |
| null, |
| null, |
| AllowAllAuthenticator.ALLOW_ALL_RESULT |
| ); |
| try (final DruidJdbcStatement statement = jdbcStatement()) { |
| // First frame, ask for all rows. |
| statement.execute(queryPlus, -1); |
| Meta.Frame frame = statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 6); |
| Assert.assertEquals( |
| Meta.Frame.create( |
| 0, |
| true, |
| Lists.newArrayList( |
| new Object[]{DateTimes.of("2000-01-01").getMillis(), 1L, "", "a", 1.0f}, |
| new Object[]{ |
| DateTimes.of("2000-01-02").getMillis(), |
| 1L, |
| "10.1", |
| NullHandling.defaultStringValue(), |
| 2.0f |
| }, |
| new Object[]{DateTimes.of("2000-01-03").getMillis(), 1L, "2", "", 3.0f}, |
| new Object[]{DateTimes.of("2001-01-01").getMillis(), 1L, "1", "a", 4.0f}, |
| new Object[]{DateTimes.of("2001-01-02").getMillis(), 1L, "def", "abc", 5.0f}, |
| new Object[]{DateTimes.of("2001-01-03").getMillis(), 1L, "abc", NullHandling.defaultStringValue(), 6.0f} |
| ) |
| ), |
| frame |
| ); |
| Assert.assertTrue(statement.isDone()); |
| } |
| } |
| |
| /** |
| * Test results spread over two frames. Also checks various state-related |
| * methods. |
| */ |
| @Test |
| public void testSelectSplitOverTwoFramesDirect() |
| { |
| SqlQueryPlus queryPlus = new SqlQueryPlus( |
| SELECT_FROM_FOO, |
| null, |
| null, |
| AllowAllAuthenticator.ALLOW_ALL_RESULT |
| ); |
| try (final DruidJdbcStatement statement = jdbcStatement()) { |
| |
| // First frame, ask for 2 rows. |
| statement.execute(queryPlus, -1); |
| Assert.assertEquals(0, statement.getCurrentOffset()); |
| Assert.assertFalse(statement.isDone()); |
| Meta.Frame frame = statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 2); |
| Assert.assertEquals( |
| firstFrameResults(), |
| frame |
| ); |
| Assert.assertFalse(statement.isDone()); |
| Assert.assertEquals(2, statement.getCurrentOffset()); |
| |
| // Last frame, ask for all remaining rows. |
| frame = statement.nextFrame(2, 10); |
| Assert.assertEquals( |
| secondFrameResults(), |
| frame |
| ); |
| Assert.assertTrue(statement.isDone()); |
| } |
| } |
| |
| /** |
| * Verify that JDBC automatically closes the first result set when we |
| * open a second for the same statement. |
| */ |
| @Test |
| public void testTwoFramesAutoCloseDirect() |
| { |
| SqlQueryPlus queryPlus = new SqlQueryPlus( |
| SELECT_FROM_FOO, |
| null, |
| null, |
| AllowAllAuthenticator.ALLOW_ALL_RESULT |
| ); |
| try (final DruidJdbcStatement statement = jdbcStatement()) { |
| // First frame, ask for 2 rows. |
| statement.execute(queryPlus, -1); |
| Meta.Frame frame = statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 2); |
| Assert.assertEquals( |
| firstFrameResults(), |
| frame |
| ); |
| Assert.assertFalse(statement.isDone()); |
| |
| // Do it again. Closes the prior result set. |
| statement.execute(queryPlus, -1); |
| frame = statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 2); |
| Assert.assertEquals( |
| firstFrameResults(), |
| frame |
| ); |
| Assert.assertFalse(statement.isDone()); |
| |
| // Last frame, ask for all remaining rows. |
| frame = statement.nextFrame(2, 10); |
| Assert.assertEquals( |
| secondFrameResults(), |
| frame |
| ); |
| Assert.assertTrue(statement.isDone()); |
| } |
| } |
| |
| /** |
| * Test that closing a statement with pending results automatically |
| * closes the underlying result set. |
| */ |
| @Test |
| public void testTwoFramesCloseWithResultSetDirect() |
| { |
| SqlQueryPlus queryPlus = new SqlQueryPlus( |
| SELECT_FROM_FOO, |
| null, |
| null, |
| AllowAllAuthenticator.ALLOW_ALL_RESULT |
| ); |
| try (final DruidJdbcStatement statement = jdbcStatement()) { |
| // First frame, ask for 2 rows. |
| statement.execute(queryPlus, -1); |
| Meta.Frame frame = statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 2); |
| Assert.assertEquals( |
| firstFrameResults(), |
| frame |
| ); |
| Assert.assertFalse(statement.isDone()); |
| |
| // Leave result set open; close statement. |
| } |
| } |
| |
| private Meta.Frame firstFrameResults() |
| { |
| return Meta.Frame.create( |
| 0, |
| false, |
| Lists.newArrayList( |
| new Object[]{DateTimes.of("2000-01-01").getMillis(), 1L, "", "a", 1.0f}, |
| new Object[]{ |
| DateTimes.of("2000-01-02").getMillis(), |
| 1L, |
| "10.1", |
| NullHandling.defaultStringValue(), |
| 2.0f |
| } |
| ) |
| ); |
| } |
| |
| private Meta.Frame secondFrameResults() |
| { |
| return Meta.Frame.create( |
| 2, |
| true, |
| Lists.newArrayList( |
| new Object[]{DateTimes.of("2000-01-03").getMillis(), 1L, "2", "", 3.0f}, |
| new Object[]{DateTimes.of("2001-01-01").getMillis(), 1L, "1", "a", 4.0f}, |
| new Object[]{DateTimes.of("2001-01-02").getMillis(), 1L, "def", "abc", 5.0f}, |
| new Object[]{DateTimes.of("2001-01-03").getMillis(), 1L, "abc", NullHandling.defaultStringValue(), 6.0f} |
| ) |
| ); |
| } |
| |
| @Test |
| public void testSignatureDirect() |
| { |
| SqlQueryPlus queryPlus = new SqlQueryPlus( |
| SELECT_STAR_FROM_FOO, |
| null, |
| null, |
| AllowAllAuthenticator.ALLOW_ALL_RESULT |
| ); |
| try (final DruidJdbcStatement statement = jdbcStatement()) { |
| // Check signature. |
| statement.execute(queryPlus, -1); |
| verifySignature(statement.getSignature()); |
| } |
| } |
| |
| @SuppressWarnings("unchecked") |
| private void verifySignature(Meta.Signature signature) |
| { |
| Assert.assertEquals(Meta.CursorFactory.ARRAY, signature.cursorFactory); |
| Assert.assertEquals(Meta.StatementType.SELECT, signature.statementType); |
| Assert.assertEquals(SELECT_STAR_FROM_FOO, signature.sql); |
| Assert.assertEquals( |
| Lists.newArrayList( |
| Lists.newArrayList("__time", "TIMESTAMP", "java.math.BigDecimal"), |
| Lists.newArrayList("dim1", "VARCHAR", "java.lang.String"), |
| Lists.newArrayList("dim2", "VARCHAR", "java.lang.String"), |
| Lists.newArrayList("dim3", "VARCHAR", "java.lang.String"), |
| Lists.newArrayList("cnt", "BIGINT", "java.math.BigDecimal"), |
| Lists.newArrayList("m1", "FLOAT", "java.math.BigDecimal"), |
| Lists.newArrayList("m2", "DOUBLE", "java.math.BigDecimal"), |
| Lists.newArrayList("unique_dim1", "OTHER", "java.lang.Object") |
| ), |
| Lists.transform( |
| signature.columns, |
| new Function<ColumnMetaData, List<String>>() |
| { |
| @Override |
| public List<String> apply(final ColumnMetaData columnMetaData) |
| { |
| return Lists.newArrayList( |
| columnMetaData.label, |
| columnMetaData.type.name, |
| columnMetaData.type.rep.clazz.getName() |
| ); |
| } |
| } |
| ) |
| ); |
| } |
| |
| //----------------------------------------------------------------- |
| // Druid JDBC Prepared Statement |
| // |
| // The JDBC PreparedStatement class starts with, then allows executing |
| // the statement sequentially, typically with a set of parameters. |
| |
| private DruidJdbcPreparedStatement jdbcPreparedStatement(SqlQueryPlus queryPlus) |
| { |
| return new DruidJdbcPreparedStatement( |
| "", |
| 0, |
| sqlStatementFactory.preparedStatement(queryPlus), |
| Long.MAX_VALUE, |
| new ResultFetcherFactory(AvaticaServerConfig.DEFAULT_FETCH_TIMEOUT_MS) |
| ); |
| } |
| |
| @Test |
| public void testSubQueryWithOrderByPrepared() |
| { |
| final String sql = "select T20.F13 as F22 from (SELECT DISTINCT dim1 as F13 FROM druid.foo T10) T20 order by T20.F13 ASC"; |
| SqlQueryPlus queryPlus = new SqlQueryPlus( |
| sql, |
| null, |
| null, |
| AllowAllAuthenticator.ALLOW_ALL_RESULT |
| ); |
| try (final DruidJdbcPreparedStatement statement = jdbcPreparedStatement(queryPlus)) { |
| statement.prepare(); |
| // First frame, ask for all rows. |
| statement.execute(Collections.emptyList()); |
| Meta.Frame frame = statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 6); |
| Assert.assertEquals( |
| subQueryWithOrderByResults(), |
| frame |
| ); |
| Assert.assertTrue(statement.isDone()); |
| } |
| } |
| |
| @Test |
| public void testSubQueryWithOrderByPreparedTwice() |
| { |
| final String sql = "select T20.F13 as F22 from (SELECT DISTINCT dim1 as F13 FROM druid.foo T10) T20 order by T20.F13 ASC"; |
| SqlQueryPlus queryPlus = new SqlQueryPlus( |
| sql, |
| null, |
| null, |
| AllowAllAuthenticator.ALLOW_ALL_RESULT |
| ); |
| try (final DruidJdbcPreparedStatement statement = jdbcPreparedStatement(queryPlus)) { |
| statement.prepare(); |
| statement.execute(Collections.emptyList()); |
| Meta.Frame frame = statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 6); |
| Assert.assertEquals( |
| subQueryWithOrderByResults(), |
| frame |
| ); |
| |
| // Do it again. JDBC says we can reuse prepared statements sequentially. |
| Assert.assertTrue(statement.isDone()); |
| statement.execute(Collections.emptyList()); |
| frame = statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 6); |
| Assert.assertEquals( |
| subQueryWithOrderByResults(), |
| frame |
| ); |
| Assert.assertTrue(statement.isDone()); |
| } |
| } |
| |
| @Test |
| public void testSignaturePrepared() |
| { |
| SqlQueryPlus queryPlus = new SqlQueryPlus( |
| SELECT_STAR_FROM_FOO, |
| null, |
| null, |
| AllowAllAuthenticator.ALLOW_ALL_RESULT |
| ); |
| try (final DruidJdbcPreparedStatement statement = jdbcPreparedStatement(queryPlus)) { |
| statement.prepare(); |
| verifySignature(statement.getSignature()); |
| } |
| } |
| |
| @Test |
| public void testParameters() |
| { |
| SqlQueryPlus queryPlus = new SqlQueryPlus( |
| "SELECT COUNT(*) AS cnt FROM sys.servers WHERE servers.host = ?", |
| null, |
| null, |
| AllowAllAuthenticator.ALLOW_ALL_RESULT |
| ); |
| Meta.Frame expected = Meta.Frame.create(0, true, Collections.singletonList(new Object[] {1L})); |
| List<TypedValue> matchingParams = Collections.singletonList(TypedValue.ofLocal(ColumnMetaData.Rep.STRING, "dummy")); |
| try (final DruidJdbcPreparedStatement statement = jdbcPreparedStatement(queryPlus)) { |
| |
| // PreparedStatement protocol: prepare once... |
| statement.prepare(); |
| |
| // Execute many times. First time. |
| statement.execute(matchingParams); |
| Meta.Frame frame = statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 6); |
| Assert.assertEquals( |
| expected, |
| frame |
| ); |
| |
| // Again, same value. |
| statement.execute(matchingParams); |
| frame = statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 6); |
| Assert.assertEquals( |
| expected, |
| frame |
| ); |
| |
| // Again, no matches. |
| statement.execute( |
| Collections.singletonList( |
| TypedValue.ofLocal(ColumnMetaData.Rep.STRING, "foo"))); |
| frame = statement.nextFrame(AbstractDruidJdbcStatement.START_OFFSET, 6); |
| Assert.assertEquals( |
| Meta.Frame.create(0, true, Collections.emptyList()), |
| frame |
| ); |
| } |
| } |
| } |
