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

Example 1 with LocalQueryRunner

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

the class TestSpatialJoinPlanning method createLocalQueryRunner.

@Override
protected LocalQueryRunner createLocalQueryRunner() {
    LocalQueryRunner queryRunner = LocalQueryRunner.create(testSessionBuilder().setCatalog("memory").setSchema("default").build());
    queryRunner.installPlugin(new GeoPlugin());
    queryRunner.createCatalog("tpch", new TpchConnectorFactory(1), ImmutableMap.of());
    queryRunner.createCatalog("memory", new MemoryConnectorFactory(), ImmutableMap.of());
    queryRunner.execute(format("CREATE TABLE kdb_tree AS SELECT '%s' AS v", KDB_TREE_JSON));
    queryRunner.execute("CREATE TABLE points (lng, lat, name) AS (VALUES (2.1e0, 2.1e0, 'x'))");
    queryRunner.execute("CREATE TABLE polygons (wkt, name) AS (VALUES ('POLYGON ((30 10, 40 40, 20 40, 10 20, 30 10))', 'a'))");
    return queryRunner;
}

Example 2 with LocalQueryRunner

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

the class TestAccessControlManager method testDenyCatalogAccessControl.

@Test
public void testDenyCatalogAccessControl() {
    try (LocalQueryRunner queryRunner = LocalQueryRunner.create(TEST_SESSION)) {
        TransactionManager transactionManager = queryRunner.getTransactionManager();
        AccessControlManager accessControlManager = createAccessControlManager(transactionManager);
        TestSystemAccessControlFactory accessControlFactory = new TestSystemAccessControlFactory("test");
        accessControlManager.addSystemAccessControlFactory(accessControlFactory);
        accessControlManager.setSystemAccessControl("test", ImmutableMap.of());
        queryRunner.createCatalog("catalog", MockConnectorFactory.create(), ImmutableMap.of());
        accessControlManager.addCatalogAccessControl(new CatalogName("catalog"), new DenyConnectorAccessControl());
        assertThatThrownBy(() -> transaction(transactionManager, accessControlManager).execute(transactionId -> {
            accessControlManager.checkCanSelectFromColumns(context(transactionId), new QualifiedObjectName("catalog", "schema", "table"), ImmutableSet.of("column"));
        })).isInstanceOf(TrinoException.class).hasMessageMatching("Access Denied: Cannot select from columns \\[column\\] in table or view schema.table");
    }
}

Example 3 with LocalQueryRunner

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

the class TestTableScanRedirectionWithPushdown method testRedirectionAfterPredicatePushdownIntoTableScan.

@Test
public void testRedirectionAfterPredicatePushdownIntoTableScan() {
    // the connector can detect a filter
    try (LocalQueryRunner queryRunner = createLocalQueryRunner(getMockApplyRedirectAfterPredicatePushdown(REDIRECTION_MAPPING_A, Optional.empty()), Optional.empty(), Optional.of(getMockApplyFilter(ImmutableSet.of(SOURCE_COLUMN_HANDLE_A, DESTINATION_COLUMN_HANDLE_A))))) {
        assertPlan(queryRunner, "SELECT source_col_a FROM test_table WHERE source_col_a = 1", output(ImmutableList.of("DEST_COL"), tableScan(new MockConnectorTableHandle(DESTINATION_TABLE, TupleDomain.withColumnDomains(ImmutableMap.of(DESTINATION_COLUMN_HANDLE_A, singleValue(INTEGER, 1L))), Optional.empty())::equals, TupleDomain.withColumnDomains(ImmutableMap.of(DESTINATION_COLUMN_HANDLE_A::equals, singleValue(INTEGER, 1L))), ImmutableMap.of("DEST_COL", DESTINATION_COLUMN_HANDLE_A::equals))));
        assertPlan(queryRunner, "SELECT source_col_a FROM test_table", output(ImmutableList.of("SOURCE_COL"), tableScan(TEST_TABLE, ImmutableMap.of("SOURCE_COL", SOURCE_COLUMN_NAME_A))));
    }
}

Example 4 with LocalQueryRunner

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

the class TestTableScanRedirectionWithPushdown method createLocalQueryRunner.

private LocalQueryRunner createLocalQueryRunner(ApplyTableScanRedirect applyTableScanRedirect, Optional<ApplyProjection> applyProjection, Optional<ApplyFilter> applyFilter) {
    LocalQueryRunner queryRunner = LocalQueryRunner.create(MOCK_SESSION);
    MockConnectorFactory.Builder builder = MockConnectorFactory.builder().withGetTableHandle((session, schemaTableName) -> new MockConnectorTableHandle(schemaTableName)).withGetColumns(name -> {
        if (name.equals(SOURCE_TABLE)) {
            return ImmutableList.of(new ColumnMetadata(SOURCE_COLUMN_NAME_A, INTEGER), new ColumnMetadata(SOURCE_COLUMN_NAME_B, INTEGER), new ColumnMetadata(SOURCE_COLUMN_NAME_C, VARCHAR), new ColumnMetadata(SOURCE_COLUMN_NAME_D, ROW_TYPE));
        } else if (name.equals(DESTINATION_TABLE)) {
            return ImmutableList.of(new ColumnMetadata(DESTINATION_COLUMN_NAME_A, INTEGER), new ColumnMetadata(DESTINATION_COLUMN_NAME_B, INTEGER), new ColumnMetadata(DESTINATION_COLUMN_NAME_C, ROW_TYPE), new ColumnMetadata(DESTINATION_COLUMN_NAME_D, BOGUS));
        }
        throw new IllegalArgumentException();
    }).withApplyTableScanRedirect(applyTableScanRedirect);
    applyProjection.ifPresent(builder::withApplyProjection);
    applyFilter.ifPresent(builder::withApplyFilter);
    queryRunner.createCatalog(MOCK_CATALOG, builder.build(), ImmutableMap.of());
    return queryRunner;
}

Example 5 with LocalQueryRunner

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

the class TestTableScanRedirectionWithPushdown method testRedirectionBeforeDeferencePushdown.

@Test
public void testRedirectionBeforeDeferencePushdown() {
    // the connector can detect that source_col_a and source_col_d is projected
    try (LocalQueryRunner queryRunner = createLocalQueryRunner(mockApplyRedirectAfterProjectionPushdown(ROW_TYPE_REDIRECTION_MAPPING_AD, Optional.of(ImmutableSet.of(SOURCE_COLUMN_HANDLE_A, SOURCE_COLUMN_HANDLE_D))), Optional.of(this::mockApplyProjection), Optional.empty())) {
        // Pushdown of dereference for source_col_d.a into table scan results in a new column handle
        // Table scan redirection would not take place if dereference pushdown has already taken place before redirection
        ColumnHandle destinationColumnHandleC0 = new MockConnectorColumnHandle(DESTINATION_COLUMN_NAME_C + "#0", BIGINT);
        assertPlan(queryRunner, "SELECT source_col_a, source_col_d.a FROM test_table", output(ImmutableList.of("DEST_COL_A", "DEST_COL_C#0"), tableScan(new MockConnectorTableHandle(DESTINATION_TABLE, TupleDomain.all(), Optional.of(ImmutableList.of(DESTINATION_COLUMN_HANDLE_A, destinationColumnHandleC0)))::equals, TupleDomain.all(), ImmutableMap.of("DEST_COL_A", DESTINATION_COLUMN_HANDLE_A::equals, "DEST_COL_C#0", destinationColumnHandleC0::equals))));
    }
}