Examples with MockConnectorTableHandle io.trino.connector.MockConnectorTableHandle used on opensource projects

Example 21 with MockConnectorTableHandle

use of io.trino.connector.MockConnectorTableHandle in project trino by trinodb.

the class TestMetadataManager method setUp.

@BeforeClass
public void setUp() throws Exception {
    queryRunner = TpchQueryRunnerBuilder.builder().build();
    queryRunner.installPlugin(new Plugin() {

        @Override
        public Iterable<ConnectorFactory> getConnectorFactories() {
            SchemaTableName viewTableName = new SchemaTableName("UPPER_CASE_SCHEMA", "test_view");
            MockConnectorFactory connectorFactory = MockConnectorFactory.builder().withListSchemaNames(session -> ImmutableList.of("UPPER_CASE_SCHEMA")).withGetTableHandle((session, schemaTableName) -> {
                if (schemaTableName.equals(viewTableName)) {
                    return null;
                }
                return new MockConnectorTableHandle(schemaTableName);
            }).withListTables((session, schemaNameOrNull) -> ImmutableList.of(new SchemaTableName("UPPER_CASE_SCHEMA", "UPPER_CASE_TABLE"))).withGetViews((session, prefix) -> ImmutableMap.of(viewTableName, getConnectorViewDefinition())).build();
            return ImmutableList.of(connectorFactory);
        }
    });
    queryRunner.createCatalog("upper_case_schema_catalog", "mock");
    metadataManager = (MetadataManager) queryRunner.getMetadata();
}

Example 22 with MockConnectorTableHandle

use of io.trino.connector.MockConnectorTableHandle in project trino by trinodb.

the class TestApplyTableScanRedirection method testApplyTableScanRedirection.

@Test
public void testApplyTableScanRedirection() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        // make the mock connector return a table scan on different table
        ApplyTableScanRedirect applyTableScanRedirect = getMockApplyRedirect(ImmutableMap.of(SOURCE_COLUMN_HANDLE_A, DESTINATION_COLUMN_NAME_A));
        MockConnectorFactory mockFactory = createMockFactory(Optional.of(applyTableScanRedirect));
        ruleTester.getQueryRunner().createCatalog(MOCK_CATALOG, mockFactory, ImmutableMap.of());
        ruleTester.assertThat(new ApplyTableScanRedirection(ruleTester.getPlannerContext())).on(p -> {
            Symbol column = p.symbol(SOURCE_COLUMN_NAME_A, VARCHAR);
            return p.tableScan(TEST_TABLE_HANDLE, ImmutableList.of(column), ImmutableMap.of(column, SOURCE_COLUMN_HANDLE_A));
        }).withSession(MOCK_SESSION).matches(tableScan(new MockConnectorTableHandle(DESTINATION_TABLE)::equals, TupleDomain.all(), ImmutableMap.of("DEST_COL", DESTINATION_COLUMN_HANDLE_A::equals)));
    }
}

Example 23 with MockConnectorTableHandle

use of io.trino.connector.MockConnectorTableHandle in project trino by trinodb.

the class TestApplyTableScanRedirection method testMismatchedTypesWithCoercion.

@Test
public void testMismatchedTypesWithCoercion() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        // make the mock connector return a table scan on different table
        ApplyTableScanRedirect applyTableScanRedirect = getMockApplyRedirect(ImmutableMap.of(SOURCE_COLUMN_HANDLE_A, DESTINATION_COLUMN_NAME_C));
        MockConnectorFactory mockFactory = createMockFactory(Optional.of(applyTableScanRedirect));
        LocalQueryRunner runner = ruleTester.getQueryRunner();
        runner.createCatalog(MOCK_CATALOG, mockFactory, ImmutableMap.of());
        ruleTester.assertThat(new ApplyTableScanRedirection(ruleTester.getPlannerContext())).on(p -> {
            Symbol column = p.symbol(SOURCE_COLUMN_NAME_A, VARCHAR);
            return p.tableScan(TEST_TABLE_HANDLE, ImmutableList.of(column), ImmutableMap.of(column, SOURCE_COLUMN_HANDLE_A));
        }).withSession(MOCK_SESSION).matches(project(ImmutableMap.of("COL", expression("CAST(DEST_COL AS VARCHAR)")), tableScan(new MockConnectorTableHandle(DESTINATION_TABLE)::equals, TupleDomain.all(), ImmutableMap.of("DEST_COL", DESTINATION_COLUMN_HANDLE_C::equals))));
    }
}

Example 24 with MockConnectorTableHandle

use of io.trino.connector.MockConnectorTableHandle in project trino by trinodb.

the class TestInsert method createLocalQueryRunner.

@Override
protected LocalQueryRunner createLocalQueryRunner() {
    Session.SessionBuilder sessionBuilder = testSessionBuilder().setCatalog("mock").setSchema("schema");
    LocalQueryRunner queryRunner = LocalQueryRunner.create(sessionBuilder.build());
    queryRunner.createCatalog("mock", MockConnectorFactory.builder().withGetTableHandle((session, schemaTableName) -> {
        if (schemaTableName.getTableName().equals("test_table_preferred_partitioning")) {
            return new MockConnectorTableHandle(schemaTableName);
        }
        if (schemaTableName.getTableName().equals("test_table_required_partitioning")) {
            return new MockConnectorTableHandle(schemaTableName);
        }
        return null;
    }).withGetColumns(name -> ImmutableList.of(new ColumnMetadata("column1", INTEGER), new ColumnMetadata("column2", INTEGER))).withGetInsertLayout((session, tableName) -> {
        if (tableName.getTableName().equals("test_table_preferred_partitioning")) {
            return Optional.of(new ConnectorTableLayout(ImmutableList.of("column1")));
        }
        if (tableName.getTableName().equals("test_table_required_partitioning")) {
            return Optional.of(new ConnectorTableLayout(new TpchPartitioningHandle("orders", 10), ImmutableList.of("column1")));
        }
        return Optional.empty();
    }).withGetNewTableLayout((session, tableMetadata) -> {
        if (tableMetadata.getTable().getTableName().equals("new_test_table_preferred_partitioning")) {
            return Optional.of(new ConnectorTableLayout(ImmutableList.of("column1")));
        }
        if (tableMetadata.getTable().getTableName().equals("new_test_table_required_partitioning")) {
            return Optional.of(new ConnectorTableLayout(new TpchPartitioningHandle("orders", 10), ImmutableList.of("column1")));
        }
        if (tableMetadata.getTable().getTableName().equals("new_test_table_unpartitioned")) {
            return Optional.empty();
        }
        return Optional.empty();
    }).build(), ImmutableMap.of());
    return queryRunner;
}

Example 25 with MockConnectorTableHandle

use of io.trino.connector.MockConnectorTableHandle in project trino by trinodb.

the class TestTableScanRedirectionWithPushdown method testRedirectionAfterProjectionPushdown.

@Test
public void testRedirectionAfterProjectionPushdown() {
    // the connector can detect that source_col_a is projected
    try (LocalQueryRunner queryRunner = createLocalQueryRunner(mockApplyRedirectAfterProjectionPushdown(REDIRECTION_MAPPING_A, Optional.of(ImmutableSet.of(SOURCE_COLUMN_HANDLE_A))), Optional.of(this::mockApplyProjection), Optional.empty())) {
        assertPlan(queryRunner, "SELECT source_col_a FROM test_table", output(ImmutableList.of("DEST_COL"), tableScan("target_table", ImmutableMap.of("DEST_COL", DESTINATION_COLUMN_NAME_A))));
        assertPlan(queryRunner, "SELECT source_col_a, source_col_b FROM test_table", output(ImmutableList.of("SOURCE_COLA", "SOURCE_COLB"), tableScan(TEST_TABLE, ImmutableMap.of("SOURCE_COLA", SOURCE_COLUMN_NAME_A, "SOURCE_COLB", SOURCE_COLUMN_NAME_B))));
        assertPlan(queryRunner, "SELECT source_col_a FROM test_table WHERE source_col_a > 0", output(ImmutableList.of("DEST_COL"), filter("DEST_COL > 0", tableScan(new MockConnectorTableHandle(DESTINATION_TABLE)::equals, TupleDomain.all(), ImmutableMap.of("DEST_COL", DESTINATION_COLUMN_HANDLE_A::equals)))));
    }
}