Examples with QueryBody io.trino.sql.tree.QueryBody used on opensource projects

Example 1 with QueryBody

use of io.trino.sql.tree.QueryBody in project trino by trinodb.

the class AstBuilder method visitSetOperation.

@Override
public Node visitSetOperation(SqlBaseParser.SetOperationContext context) {
    QueryBody left = (QueryBody) visit(context.left);
    QueryBody right = (QueryBody) visit(context.right);
    boolean distinct = context.setQuantifier() == null || context.setQuantifier().DISTINCT() != null;
    switch(context.operator.getType()) {
        case SqlBaseLexer.UNION:
            return new Union(getLocation(context.UNION()), ImmutableList.of(left, right), distinct);
        case SqlBaseLexer.INTERSECT:
            return new Intersect(getLocation(context.INTERSECT()), ImmutableList.of(left, right), distinct);
        case SqlBaseLexer.EXCEPT:
            return new Except(getLocation(context.EXCEPT()), left, right, distinct);
    }
    throw new IllegalArgumentException("Unsupported set operation: " + context.operator.getText());
}

Example 2 with QueryBody

use of io.trino.sql.tree.QueryBody in project trino by trinodb.

the class AstBuilder method visitQueryNoWith.

@Override
public Node visitQueryNoWith(SqlBaseParser.QueryNoWithContext context) {
    QueryBody term = (QueryBody) visit(context.queryTerm());
    Optional<OrderBy> orderBy = Optional.empty();
    if (context.ORDER() != null) {
        orderBy = Optional.of(new OrderBy(getLocation(context.ORDER()), visit(context.sortItem(), SortItem.class)));
    }
    Optional<Offset> offset = Optional.empty();
    if (context.OFFSET() != null) {
        Expression rowCount;
        if (context.offset.INTEGER_VALUE() != null) {
            rowCount = new LongLiteral(getLocation(context.offset.INTEGER_VALUE()), context.offset.getText());
        } else {
            rowCount = new Parameter(getLocation(context.offset.QUESTION_MARK()), parameterPosition);
            parameterPosition++;
        }
        offset = Optional.of(new Offset(Optional.of(getLocation(context.OFFSET())), rowCount));
    }
    Optional<Node> limit = Optional.empty();
    if (context.FETCH() != null) {
        Optional<Expression> rowCount = Optional.empty();
        if (context.fetchFirst != null) {
            if (context.fetchFirst.INTEGER_VALUE() != null) {
                rowCount = Optional.of(new LongLiteral(getLocation(context.fetchFirst.INTEGER_VALUE()), context.fetchFirst.getText()));
            } else {
                rowCount = Optional.of(new Parameter(getLocation(context.fetchFirst.QUESTION_MARK()), parameterPosition));
                parameterPosition++;
            }
        }
        limit = Optional.of(new FetchFirst(Optional.of(getLocation(context.FETCH())), rowCount, context.TIES() != null));
    } else if (context.LIMIT() != null) {
        if (context.limit == null) {
            throw new IllegalStateException("Missing LIMIT value");
        }
        Expression rowCount;
        if (context.limit.ALL() != null) {
            rowCount = new AllRows(getLocation(context.limit.ALL()));
        } else if (context.limit.rowCount().INTEGER_VALUE() != null) {
            rowCount = new LongLiteral(getLocation(context.limit.rowCount().INTEGER_VALUE()), context.limit.getText());
        } else {
            rowCount = new Parameter(getLocation(context.limit.rowCount().QUESTION_MARK()), parameterPosition);
            parameterPosition++;
        }
        limit = Optional.of(new Limit(Optional.of(getLocation(context.LIMIT())), rowCount));
    }
    if (term instanceof QuerySpecification) {
        // When we have a simple query specification
        // followed by order by, offset, limit or fetch,
        // fold the order by, limit, offset or fetch clauses
        // into the query specification (analyzer/planner
        // expects this structure to resolve references with respect
        // to columns defined in the query specification)
        QuerySpecification query = (QuerySpecification) term;
        return new Query(getLocation(context), Optional.empty(), new QuerySpecification(getLocation(context), query.getSelect(), query.getFrom(), query.getWhere(), query.getGroupBy(), query.getHaving(), query.getWindows(), orderBy, offset, limit), Optional.empty(), Optional.empty(), Optional.empty());
    }
    return new Query(getLocation(context), Optional.empty(), term, orderBy, offset, limit);
}

Example 3 with QueryBody

use of io.trino.sql.tree.QueryBody in project trino by trinodb.

the class QueryRewriter method getColumns.

private List<Column> getColumns(Connection connection, CreateTableAsSelect createTableAsSelect) throws SQLException {
    io.trino.sql.tree.Query createSelectClause = createTableAsSelect.getQuery();
    // Rewrite the query to select zero rows, so that we can get the column names and types
    QueryBody innerQuery = createSelectClause.getQueryBody();
    io.trino.sql.tree.Query zeroRowsQuery;
    if (innerQuery instanceof QuerySpecification) {
        QuerySpecification querySpecification = (QuerySpecification) innerQuery;
        innerQuery = new QuerySpecification(querySpecification.getSelect(), querySpecification.getFrom(), querySpecification.getWhere(), querySpecification.getGroupBy(), querySpecification.getHaving(), querySpecification.getWindows(), querySpecification.getOrderBy(), querySpecification.getOffset(), Optional.of(new Limit(new LongLiteral("0"))));
        zeroRowsQuery = new io.trino.sql.tree.Query(createSelectClause.getWith(), innerQuery, Optional.empty(), Optional.empty(), Optional.empty());
    } else {
        zeroRowsQuery = new io.trino.sql.tree.Query(createSelectClause.getWith(), innerQuery, Optional.empty(), Optional.empty(), Optional.of(new Limit(new LongLiteral("0"))));
    }
    ImmutableList.Builder<Column> columns = ImmutableList.builder();
    try (java.sql.Statement jdbcStatement = connection.createStatement()) {
        ExecutorService executor = newSingleThreadExecutor();
        TimeLimiter limiter = SimpleTimeLimiter.create(executor);
        java.sql.Statement limitedStatement = limiter.newProxy(jdbcStatement, java.sql.Statement.class, timeout.toMillis(), TimeUnit.MILLISECONDS);
        try (ResultSet resultSet = limitedStatement.executeQuery(formatSql(zeroRowsQuery))) {
            ResultSetMetaData metaData = resultSet.getMetaData();
            for (int i = 1; i <= metaData.getColumnCount(); i++) {
                String name = metaData.getColumnName(i);
                int type = metaData.getColumnType(i);
                columns.add(new Column(name, APPROXIMATE_TYPES.contains(type)));
            }
        } catch (UncheckedTimeoutException e) {
            throw new SQLException("SQL statement execution timed out", e);
        } finally {
            executor.shutdownNow();
        }
    }
    return columns.build();
}