Examples with LocalQueryRunner io.trino.testing.LocalQueryRunner used on opensource projects

Example 21 with LocalQueryRunner

use of io.trino.testing.LocalQueryRunner in project trino by trinodb.

the class BaseRuleTest method setUp.

@BeforeClass
public final void setUp() {
    Optional<LocalQueryRunner> localQueryRunner = createLocalQueryRunner();
    if (localQueryRunner.isPresent()) {
        plugins.forEach(plugin -> localQueryRunner.get().installPlugin(plugin));
        tester = new RuleTester(localQueryRunner.get());
    } else {
        tester = defaultRuleTester(plugins, ImmutableMap.of(), Optional.empty());
    }
}

Example 22 with LocalQueryRunner

use of io.trino.testing.LocalQueryRunner in project trino by trinodb.

the class RuleTester method defaultRuleTester.

public static RuleTester defaultRuleTester(List<Plugin> plugins, Map<String, String> sessionProperties, Optional<Integer> nodeCountForStats) {
    Session.SessionBuilder sessionBuilder = testSessionBuilder().setCatalog(CATALOG_ID).setSchema("tiny").setSystemProperty("task_concurrency", // these tests don't handle exchanges from local parallel
    "1");
    for (Map.Entry<String, String> entry : sessionProperties.entrySet()) {
        sessionBuilder.setSystemProperty(entry.getKey(), entry.getValue());
    }
    Session session = sessionBuilder.build();
    LocalQueryRunner queryRunner = nodeCountForStats.map(nodeCount -> LocalQueryRunner.builder(session).withNodeCountForStats(nodeCount).build()).orElseGet(() -> LocalQueryRunner.create(session));
    queryRunner.createCatalog(session.getCatalog().get(), new TpchConnectorFactory(1), ImmutableMap.of());
    plugins.forEach(queryRunner::installPlugin);
    return new RuleTester(queryRunner);
}

Example 23 with LocalQueryRunner

use of io.trino.testing.LocalQueryRunner in project trino by trinodb.

the class TestValidateLimitWithPresortedInput method createLocalQueryRunner.

@Override
protected LocalQueryRunner createLocalQueryRunner() {
    Session session = testSessionBuilder().setCatalog(MOCK_CATALOG).setSchema(TEST_SCHEMA).build();
    LocalQueryRunner queryRunner = LocalQueryRunner.builder(session).build();
    MockConnectorFactory mockFactory = MockConnectorFactory.builder().withGetTableProperties((connectorSession, handle) -> {
        MockConnectorTableHandle tableHandle = (MockConnectorTableHandle) handle;
        if (tableHandle.getTableName().equals(MOCK_TABLE_NAME)) {
            return new ConnectorTableProperties(TupleDomain.all(), Optional.empty(), Optional.empty(), Optional.empty(), ImmutableList.of(new SortingProperty<>(COLUMN_HANDLE_A, ASC_NULLS_FIRST), new SortingProperty<>(COLUMN_HANDLE_C, ASC_NULLS_FIRST)));
        }
        throw new IllegalArgumentException();
    }).withGetColumns(schemaTableName -> {
        if (schemaTableName.equals(MOCK_TABLE_NAME)) {
            return ImmutableList.of(new ColumnMetadata(COLUMN_NAME_A, VARCHAR), new ColumnMetadata(COLUMN_NAME_B, VARCHAR), new ColumnMetadata(COLUMN_NAME_C, VARCHAR));
        }
        throw new IllegalArgumentException();
    }).build();
    queryRunner.createCatalog(MOCK_CATALOG, mockFactory, ImmutableMap.of());
    return queryRunner;
}

Example 24 with LocalQueryRunner

use of io.trino.testing.LocalQueryRunner in project trino by trinodb.

the class TestValidateLimitWithPresortedInput method validatePlan.

private void validatePlan(Function<PlanBuilder, PlanNode> planProvider) {
    LocalQueryRunner queryRunner = getQueryRunner();
    PlanBuilder builder = new PlanBuilder(idAllocator, queryRunner.getMetadata(), queryRunner.getDefaultSession());
    PlanNode planNode = planProvider.apply(builder);
    TypeProvider types = builder.getTypes();
    queryRunner.inTransaction(session -> {
        // metadata.getCatalogHandle() registers the catalog for the transaction
        session.getCatalog().ifPresent(catalog -> queryRunner.getMetadata().getCatalogHandle(session, catalog));
        TypeAnalyzer typeAnalyzer = createTestingTypeAnalyzer(queryRunner.getPlannerContext());
        new ValidateLimitWithPresortedInput().validate(planNode, session, queryRunner.getPlannerContext(), typeAnalyzer, types, WarningCollector.NOOP);
        return null;
    });
}

Example 25 with LocalQueryRunner

use of io.trino.testing.LocalQueryRunner in project trino by trinodb.

the class LocalAtopQueryRunner method createQueryRunner.

public static LocalQueryRunner createQueryRunner(Map<String, String> catalogProperties, Class<? extends AtopFactory> factoryClass) {
    Session session = testSessionBuilder().setCatalog("atop").setSchema("default").setTimeZoneKey(TimeZoneKey.getTimeZoneKey(TimeZone.getDefault().getID())).build();
    LocalQueryRunner queryRunner = LocalQueryRunner.create(session);
    try {
        AtopConnectorFactory connectorFactory = new AtopConnectorFactory(factoryClass, LocalAtopQueryRunner.class.getClassLoader());
        ImmutableMap.Builder<String, String> properties = ImmutableMap.<String, String>builder().putAll(catalogProperties).put("atop.max-history-days", "1");
        queryRunner.createCatalog("atop", connectorFactory, properties.buildOrThrow());
        return queryRunner;
    } catch (Exception e) {
        closeAllSuppress(e, queryRunner);
        throw e;
    }
}