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

Example 21 with MockConnectorFactory

use of io.trino.connector.MockConnectorFactory 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;
}

Example 22 with MockConnectorFactory

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

the class TestSystemConnector method createQueryRunner.

@Override
protected QueryRunner createQueryRunner() throws Exception {
    Session defaultSession = testSessionBuilder().setCatalog("mock").setSchema("default").build();
    DistributedQueryRunner queryRunner = DistributedQueryRunner.builder(defaultSession).enableBackupCoordinator().build();
    queryRunner.installPlugin(new Plugin() {

        @Override
        public Iterable<ConnectorFactory> getConnectorFactories() {
            MockConnectorFactory connectorFactory = MockConnectorFactory.builder().withGetViews((session, schemaTablePrefix) -> ImmutableMap.of()).withListTables((session, s) -> ImmutableList.of(SCHEMA_TABLE_NAME)).withGetColumns(tableName -> getColumns.apply(tableName)).build();
            return ImmutableList.of(connectorFactory);
        }
    });
    queryRunner.createCatalog("mock", "mock", ImmutableMap.of());
    return queryRunner;
}

Example 23 with MockConnectorFactory

use of io.trino.connector.MockConnectorFactory 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 24 with MockConnectorFactory

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

the class TestApplyTableScanRedirection method testDoesNotFireForDeleteTableScan.

@Test
public void testDoesNotFireForDeleteTableScan() {
    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), true);
        }).withSession(MOCK_SESSION).doesNotFire();
    }
}

Example 25 with MockConnectorFactory

use of io.trino.connector.MockConnectorFactory 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)));
    }
}