Examples with Hook org.apache.calcite.runtime.Hook used on opensource projects

Example 1 with Hook

use of org.apache.calcite.runtime.Hook in project calcite by apache.

the class CalciteAssert method assertQuery.

static void assertQuery(Connection connection, String sql, int limit, boolean materializationsEnabled, List<Pair<Hook, Function>> hooks, Function<ResultSet, Void> resultChecker, Function<Integer, Void> updateChecker, Function<Throwable, Void> exceptionChecker) throws Exception {
    final String message = "With materializationsEnabled=" + materializationsEnabled + ", limit=" + limit;
    try (final Closer closer = new Closer()) {
        if (connection instanceof CalciteConnection) {
            CalciteConnection calciteConnection = (CalciteConnection) connection;
            calciteConnection.getProperties().setProperty(CalciteConnectionProperty.MATERIALIZATIONS_ENABLED.camelName(), Boolean.toString(materializationsEnabled));
            calciteConnection.getProperties().setProperty(CalciteConnectionProperty.CREATE_MATERIALIZATIONS.camelName(), Boolean.toString(materializationsEnabled));
            if (!calciteConnection.getProperties().containsKey(CalciteConnectionProperty.TIME_ZONE.camelName())) {
                // Do not override id some test has already set this property.
                calciteConnection.getProperties().setProperty(CalciteConnectionProperty.TIME_ZONE.camelName(), DateTimeUtils.UTC_ZONE.getID());
            }
        }
        for (Pair<Hook, Function> hook : hooks) {
            closer.add(hook.left.addThread(hook.right));
        }
        Statement statement = connection.createStatement();
        statement.setMaxRows(limit <= 0 ? limit : Math.max(limit, 1));
        ResultSet resultSet = null;
        Integer updateCount = null;
        try {
            if (updateChecker == null) {
                resultSet = statement.executeQuery(sql);
            } else {
                updateCount = statement.executeUpdate(sql);
            }
            if (exceptionChecker != null) {
                exceptionChecker.apply(null);
                return;
            }
        } catch (Exception | Error e) {
            if (exceptionChecker != null) {
                exceptionChecker.apply(e);
                return;
            }
            throw e;
        }
        if (resultChecker != null) {
            resultChecker.apply(resultSet);
        }
        if (updateChecker != null) {
            updateChecker.apply(updateCount);
        }
        if (resultSet != null) {
            resultSet.close();
        }
        statement.close();
        connection.close();
    } catch (Error | RuntimeException e) {
        // at the very top level of the exception stack.
        throw e;
    } catch (Throwable e) {
        throw new RuntimeException(message, e);
    }
}