Examples with TpchPartitioningHandle io.trino.plugin.tpch.TpchPartitioningHandle used on opensource projects

Example 1 with TpchPartitioningHandle

use of io.trino.plugin.tpch.TpchPartitioningHandle 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;
}