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

Example 16 with MockConnectorTableHandle

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

the class TestPartialTopNWithPresortedInput 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(tableA)) {
            return new ConnectorTableProperties(TupleDomain.all(), Optional.empty(), Optional.empty(), Optional.empty(), ImmutableList.of(new SortingProperty<>(columnHandleA, ASC_NULLS_FIRST)));
        }
        throw new IllegalArgumentException();
    }).withGetColumns(schemaTableName -> {
        if (schemaTableName.equals(tableA)) {
            return ImmutableList.of(new ColumnMetadata(columnNameA, VARCHAR), new ColumnMetadata(columnNameB, VARCHAR));
        }
        throw new IllegalArgumentException();
    }).build();
    queryRunner.createCatalog(MOCK_CATALOG, mockFactory, ImmutableMap.of());
    return queryRunner;
}

Example 17 with MockConnectorTableHandle

use of io.trino.connector.MockConnectorTableHandle 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 18 with MockConnectorTableHandle

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

the class TestRefreshMaterializedView method createQueryRunner.

@Override
protected QueryRunner createQueryRunner() throws Exception {
    Session session = testSessionBuilder().setCatalog("mock").setSchema("default").build();
    DistributedQueryRunner queryRunner = DistributedQueryRunner.builder(session).build();
    queryRunner.installPlugin(new MockConnectorPlugin(MockConnectorFactory.builder().withListSchemaNames(connectionSession -> ImmutableList.of("default")).withGetColumns(schemaTableName -> ImmutableList.of(new ColumnMetadata("nationkey", BIGINT))).withGetTableHandle((connectorSession, tableName) -> new MockConnectorTableHandle(tableName)).withGetMaterializedViews((connectorSession, schemaTablePrefix) -> ImmutableMap.of(new SchemaTableName("default", "delegate_refresh_to_connector"), new ConnectorMaterializedViewDefinition("SELECT nationkey FROM mock.default.test_table", Optional.of(new CatalogSchemaTableName("mock", "default", "test_storage")), Optional.of("mock"), Optional.of("default"), ImmutableList.of(new ConnectorMaterializedViewDefinition.Column("nationkey", BIGINT.getTypeId())), Optional.empty(), Optional.of("alice"), ImmutableMap.of()))).withDelegateMaterializedViewRefreshToConnector((connectorSession, schemaTableName) -> true).withRefreshMaterializedView(((connectorSession, schemaTableName) -> {
        startRefreshMaterializedView.set(null);
        SettableFuture<Void> refreshMaterializedView = SettableFuture.create();
        finishRefreshMaterializedView.addListener(() -> refreshMaterializedView.set(null), directExecutor());
        addExceptionCallback(refreshMaterializedView, () -> refreshInterrupted.set(null));
        return toCompletableFuture(refreshMaterializedView);
    })).build()));
    queryRunner.createCatalog("mock", "mock");
    return queryRunner;
}

Example 19 with MockConnectorTableHandle

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

the class TestRevokeOnTable method initClass.

@BeforeClass
public void initClass() throws Exception {
    SchemaTableName table = new SchemaTableName("default", "table_one");
    queryRunner = DistributedQueryRunner.builder(userWithAllPrivileges).build();
    Grants<SchemaTableName> tableGrants = new MutableGrants<>();
    tableGrants.grant(new TrinoPrincipal(USER, admin.getUser()), table, EnumSet.allOf(Privilege.class), true);
    tableGrants.grant(new TrinoPrincipal(USER, userWithAllPrivileges.getUser()), table, EnumSet.allOf(Privilege.class), true);
    tableGrants.grant(new TrinoPrincipal(USER, userWithCreate.getUser()), table, ImmutableSet.of(Privilege.CREATE), true);
    tableGrants.grant(new TrinoPrincipal(USER, userWithSelect.getUser()), table, ImmutableSet.of(Privilege.SELECT), true);
    tableGrants.grant(new TrinoPrincipal(USER, userWithInsert.getUser()), table, ImmutableSet.of(Privilege.INSERT), true);
    tableGrants.grant(new TrinoPrincipal(USER, userWithUpdate.getUser()), table, ImmutableSet.of(Privilege.UPDATE), true);
    tableGrants.grant(new TrinoPrincipal(USER, userWithDelete.getUser()), table, ImmutableSet.of(Privilege.DELETE), true);
    MockConnectorFactory connectorFactory = MockConnectorFactory.builder().withListSchemaNames(session -> ImmutableList.of("default")).withListTables((session, schemaName) -> "default".equalsIgnoreCase(schemaName) ? ImmutableList.of(table) : ImmutableList.of()).withGetTableHandle((session, tableName) -> tableName.equals(table) ? new MockConnectorTableHandle(tableName) : null).withSchemaGrants(new MutableGrants<>()).withTableGrants(tableGrants).build();
    queryRunner.installPlugin(new MockConnectorPlugin(connectorFactory));
    queryRunner.createCatalog("local", "mock");
    assertions = new QueryAssertions(queryRunner);
}

Example 20 with MockConnectorTableHandle

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

the class TestErrorThrowableInQuery method createQueryRunner.

@Override
protected DistributedQueryRunner createQueryRunner() throws Exception {
    Session session = testSessionBuilder().setSystemProperty("task_concurrency", "1").setCatalog("mock").setSchema("default").setClientInfo("{\"clientVersion\":\"testVersion\"}").build();
    DistributedQueryRunner queryRunner = DistributedQueryRunner.builder(session).setNodeCount(1).build();
    try {
        queryRunner.installPlugin(new TpchPlugin());
        queryRunner.installPlugin(new ResourceGroupManagerPlugin());
        queryRunner.installPlugin(new Plugin() {

            @Override
            public Iterable<ConnectorFactory> getConnectorFactories() {
                SchemaTableName stackOverflowErrorTableName = new SchemaTableName("default", "stack_overflow_during_planning");
                SchemaTableName classFormatErrorTableName = new SchemaTableName("default", "class_format_error_during_planning");
                MockConnectorFactory connectorFactory = MockConnectorFactory.builder().withListTables((session, s) -> ImmutableList.of(stackOverflowErrorTableName)).withGetColumns(schemaTableName -> ImmutableList.of(new ColumnMetadata("test_varchar", createUnboundedVarcharType()), new ColumnMetadata("test_bigint", BIGINT))).withGetTableHandle((session, schemaTableName) -> new MockConnectorTableHandle(schemaTableName)).withApplyProjection((session, handle, projections, assignments) -> {
                    MockConnectorTableHandle mockTableHandle = (MockConnectorTableHandle) handle;
                    if (stackOverflowErrorTableName.equals(mockTableHandle.getTableName())) {
                        throw new StackOverflowError("We run out of stack!!!!!!!!!!!");
                    }
                    if (classFormatErrorTableName.equals(mockTableHandle.getTableName())) {
                        throw new ClassFormatError("Bad class format!!!!!!!!!!");
                    }
                    throw new TrinoException(NOT_FOUND, "Unknown table: " + mockTableHandle.getTableName());
                }).build();
                return ImmutableList.of(connectorFactory);
            }
        });
        queryRunner.createCatalog("mock", "mock", ImmutableMap.of());
    } catch (Exception e) {
        queryRunner.close();
        throw e;
    }
    return queryRunner;
}