Start line:  
End line:  

Snippet Preview

Snippet HTML Code

Stack Overflow Questions
  /*
   * 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.server;
 
 
 import java.net.URI;
 import java.util.List;
 import java.util.Map;
 
 import static com.facebook.presto.sql.analyzer.Session.DEFAULT_CATALOG;
 import static com.facebook.presto.sql.analyzer.Session.DEFAULT_SCHEMA;
 import static com.facebook.presto.util.MaterializedResult.DEFAULT_PRECISION;
 import static com.facebook.presto.util.Types.checkType;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.Iterables.transform;
 import static io.airlift.json.JsonCodec.jsonCodec;
 import static io.airlift.testing.Assertions.assertLessThan;
 import static io.airlift.units.Duration.nanosSince;
 import static java.lang.String.format;
 import static java.util.Collections.nCopies;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertTrue;
 
 public class TestDistributedQueries
         extends AbstractTestSampledQueries
 {
     private static final Session SESSION = new Session("user""test""test"nullnull);
 
     private static final String ENVIRONMENT = "testing";
     private static final Logger log = Logger.get(TestDistributedQueries.class.getSimpleName());
     private final JsonCodec<QueryResultsqueryResultsCodec = jsonCodec(QueryResults.class);
 
     private TestingPrestoServer coordinator;
     private List<TestingPrestoServerservers;
     private AsyncHttpClient httpClient;
 
     @Test(expectedExceptions = RuntimeException.class, expectedExceptionsMessageRegExp = "statement is too large \\(stack overflow during analysis\\)")
     public void testLargeQueryFailure()
             throws Exception
     {
         assertQuery("SELECT " + Joiner.on(" AND ").join(nCopies(1000, "1 = 1")), "SELECT true");
     }
 
     @Test
     public void testLargeQuerySuccess()
             throws Exception
     {
         assertQuery("SELECT " + Joiner.on(" AND ").join(nCopies(500, "1 = 1")), "SELECT true");
    }
    @Test
    public void testTableSampleSystem()
            throws Exception
    {
        int total = computeActual("SELECT orderkey FROM orders").getMaterializedTuples().size();
        boolean sampleSizeFound = false;
        for (int i = 0; i < 100; i++) {
            int sampleSize = computeActual("SELECT orderkey FROM ORDERS TABLESAMPLE SYSTEM (50)").getMaterializedTuples().size();
            if (sampleSize > 0 && sampleSize < total) {
                sampleSizeFound = true;
                break;
            }
        }
        assertTrue(sampleSizeFound"Table sample returned unexpected number of rows");
    }
    @Test
            throws Exception
    {
        MaterializedResult fullSample = computeActual("SELECT orderkey FROM orders TABLESAMPLE SYSTEM (100)");
        MaterializedResult emptySample = computeActual("SELECT orderkey FROM orders TABLESAMPLE SYSTEM (0)");
        MaterializedResult all = computeActual("SELECT orderkey FROM orders");
        assertEquals(emptySample.getMaterializedTuples().size(), 0);
    }
    @Test
    public void testCreateTableAsSelect()
            throws Exception
    {
        assertCreateTable(
                "test_simple",
                "SELECT orderkey, totalprice, orderdate FROM orders",
                "SELECT count(*) FROM orders");
    }
    @Test
    public void testCreateTableAsSelectGroupBy()
            throws Exception
    {
        assertCreateTable(
                "test_group",
                "SELECT orderstatus, sum(totalprice) x FROM orders GROUP BY orderstatus",
                "SELECT count(DISTINCT orderstatus) FROM orders");
    }
    @Test
    public void testCreateTableAsSelectLimit()
            throws Exception
    {
        assertCreateTable(
                "test_limit",
                "SELECT orderkey FROM orders ORDER BY orderkey LIMIT 10",
                "SELECT 10");
    }
    @Test
    public void testCreateTableAsSelectJoin()
            throws Exception
    {
        assertCreateTable(
                "test_join",
                "SELECT count(*) x FROM lineitem JOIN orders ON lineitem.orderkey = orders.orderkey",
                "SELECT 1");
    }
    @Test
            throws Exception
    {
        assertCreateTable(
                "test_limit_sampled",
                "SELECT orderkey FROM tpch_sampled.tiny.orders ORDER BY orderkey LIMIT 10",
                "SELECT orderkey FROM (SELECT orderkey FROM orders) UNION ALL (SELECT orderkey FROM orders) ORDER BY orderkey LIMIT 10",
                "SELECT 10");
    }
    private void assertCreateTable(String table, @Language("SQL"String query, @Language("SQL"String rowCountQuery)
            throws Exception
    {
        assertCreateTable(tablequeryqueryrowCountQuery);
    }
    private void assertCreateTable(String table, @Language("SQL"String query, @Language("SQL"String expectedQuery, @Language("SQL"String rowCountQuery)
           throws Exception
    {
        try {
            assertQuery("CREATE TABLE " +  table + " AS " + queryrowCountQuery);
            assertQuery("SELECT * FROM " + tableexpectedQuery);
        }
        finally {
            QualifiedTableName name = new QualifiedTableName(table);
            Optional<TableHandlehandle = .getMetadata().getTableHandle(name);
            if (handle.isPresent()) {
                .getMetadata().dropTable(handle.get());
            }
        }
    }
    @Override
    protected int getNodeCount()
    {
        return 3;
    }
    @Override
    protected Session setUpQueryFramework()
            throws Exception
    {
        try {
             = new TestingDiscoveryServer();
             = ImmutableList.<TestingPrestoServer>builder()
                    .add()
                    .add(createTestingPrestoServer(.getBaseUrl(), false))
                    .add(createTestingPrestoServer(.getBaseUrl(), false))
                    .build();
        }
        catch (Exception e) {
            tearDownQueryFramework();
            throw e;
        }
        this. = new JettyHttpClient(
                new HttpClientConfig()
                        .setConnectTimeout(new Duration(1, .))
                        .setReadTimeout(new Duration(10, .)));
        long start = System.nanoTime();
        while (!allNodesGloballyVisible()) {
            assertLessThan(nanosSince(start), new Duration(10, ));
            .sleep(10);
        }
        for (TestingPrestoServer server : ) {
            server.getMetadata().addFunctions();
        }
        .info("Loading data...");
        long startTime = System.nanoTime();
        .info("Loading complete in %s"nanosSince(startTime).toString());
        return ;
    }
    private boolean allNodesGloballyVisible()
    {
        for (TestingPrestoServer server : ) {
            AllNodes allNodes = server.refreshNodes();
            if (!allNodes.getInactiveNodes().isEmpty() ||
                    (allNodes.getActiveNodes().size() != .size())) {
                return false;
            }
        }
        return true;
    }
    @SuppressWarnings("deprecation")
    @Override
    protected void tearDownQueryFramework()
            throws Exception
    {
        if ( != null) {
            for (TestingPrestoServer server : ) {
                Closeables.closeQuietly(server);
            }
        }
        Closeables.closeQuietly();
    }
    private void distributeData(String catalogString schemaClientSession session)
            throws Exception
    {
        for (QualifiedTableName table : .getMetadata().listTables(new QualifiedTablePrefix(catalogschema))) {
            if (table.getTableName().equalsIgnoreCase("dual")) {
                continue;
            }
            .info("Running import for %s"table.getTableName());
            @Language("SQL"String sql = format("CREATE TABLE %s AS SELECT * FROM %s"table.getTableName(), table);
            long rows = checkType(compute(sqlsession).getMaterializedTuples().get(0).getField(0), Long.class"rows");
            .info("Imported %s rows for %s"rowstable.getTableName());
        }
    }
    protected ClientSession getClientSession()
    {
    }
    {
        return new ClientSession(.getBaseUrl(), .getUser(), .getSource(), .getCatalog(), "sampled"true);
    }
    @Override
    protected MaterializedResult computeActualSampled(@Language("SQL"String sql)
    {
        return compute(sqlgetSampledClientSession());
    }
    @Override
    protected MaterializedResult computeActual(@Language("SQL"String sql)
    {
        return compute(sqlgetClientSession());
    }
    private MaterializedResult compute(@Language("SQL"String sqlClientSession session)
    {
        try (StatementClient client = new StatementClient(sessionsql)) {
            AtomicBoolean loggedUri = new AtomicBoolean(false);
            ImmutableList.Builder<MaterializedTuplerows = ImmutableList.builder();
            List<TupleInfotypes = null;
            while (client.isValid()) {
                QueryResults results = client.current();
                if (!loggedUri.getAndSet(true)) {
                    .info("Query %s: %s?pretty"results.getId(), results.getInfoUri());
                }
                if ((types == null) && (results.getColumns() != null)) {
                    types = getTupleInfos(results.getColumns());
                }
                if (results.getData() != null) {
                    rows.addAll(transform(results.getData(), dataToTuple(types)));
                }
                client.advance();
            }
            if (!client.isFailed()) {
                return new MaterializedResult(rows.build(), types);
            }
            QueryError error = client.finalResults().getError();
            assert error != null;
            if (error.getFailureInfo() != null) {
                throw error.getFailureInfo().toException();
            }
            throw new RuntimeException("Query failed: " + error.getMessage());
            // dump query info to console for debugging (NOTE: not pretty printed)
            // JsonCodec<QueryInfo> queryInfoJsonCodec = createCodecFactory().prettyPrint().jsonCodec(QueryInfo.class);
            // log.info("\n" + queryInfoJsonCodec.toJson(queryInfo));
        }
    }
    private static List<TupleInfogetTupleInfos(List<Columncolumns)
    {
        return ImmutableList.copyOf(transform(columnscolumnTupleInfoGetter()));
    }
    private static Function<ColumnTupleInfocolumnTupleInfoGetter()
    {
        return new Function<ColumnTupleInfo>()
        {
            @Override
            public TupleInfo apply(Column column)
            {
                String type = column.getType();
                switch (type) {
                    case "boolean":
                        return .;
                    case "bigint":
                        return .;
                    case "double":
                        return .;
                    case "varchar":
                        return .;
                }
                throw new AssertionError("Unhandled type: " + type);
            }
        };
    }
    private static Function<List<Object>, MaterializedTupledataToTuple(final List<TupleInfotupleInfos)
    {
        return new Function<List<Object>, MaterializedTuple>()
        {
            @Override
            public MaterializedTuple apply(List<Objectdata)
            {
                checkArgument(data.size() == tupleInfos.size(), "columns size does not match tuple infos");
                List<Objectrow = new ArrayList<>();
                for (int i = 0; i < data.size(); i++) {
                    Object value = data.get(i);
                    if (value == null) {
                        row.add(null);
                        continue;
                    }
                    Type type = tupleInfos.get(i).getType();
                    switch (type) {
                        case :
                            row.add(value);
                            break;
                        case :
                            row.add(((Numbervalue).longValue());
                            break;
                        case :
                            row.add(((Numbervalue).doubleValue());
                            break;
                        case :
                            row.add(value);
                            break;
                        default:
                            throw new AssertionError("unhandled type: " + type);
                    }
                }
                return new MaterializedTuple(row);
            }
        };
    }
    private static TestingPrestoServer createTestingPrestoServer(URI discoveryUriboolean coordinator)
            throws Exception
    {
        Map<StringStringproperties = ImmutableMap.<StringString>builder()
                .put("query.client.timeout""10m")
                .put("exchange.http-client.read-timeout""1h")
                .put("datasources""native,tpch,tpch_sampled")
                .build();
        TestingPrestoServer server = new TestingPrestoServer(coordinatorpropertiesdiscoveryUri, ImmutableList.<Module>of());
        server.installPlugin(new TpchPlugin(), "tpch""tpch");
        server.installPlugin(new SampledTpchPlugin(), "tpch_sampled""tpch_sampled");
        return server;
    }
New to GrepCode? Check out our FAQ X