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

Example 16 with MockConnectorFactory

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

the class TestRowFilter method init.

@BeforeAll
public void init() {
    LocalQueryRunner runner = LocalQueryRunner.builder(SESSION).build();
    runner.createCatalog(CATALOG, new TpchConnectorFactory(1), ImmutableMap.of());
    ConnectorViewDefinition view = new ConnectorViewDefinition("SELECT nationkey, name FROM local.tiny.nation", Optional.empty(), Optional.empty(), ImmutableList.of(new ConnectorViewDefinition.ViewColumn("nationkey", BigintType.BIGINT.getTypeId()), new ConnectorViewDefinition.ViewColumn("name", VarcharType.createVarcharType(25).getTypeId())), Optional.empty(), Optional.of(VIEW_OWNER), false);
    MockConnectorFactory mock = MockConnectorFactory.builder().withGetViews((s, prefix) -> ImmutableMap.<SchemaTableName, ConnectorViewDefinition>builder().put(new SchemaTableName("default", "nation_view"), view).buildOrThrow()).withGetColumns(schemaTableName -> {
        if (schemaTableName.equals(new SchemaTableName("tiny", "nation"))) {
            return TPCH_NATION_SCHEMA;
        }
        if (schemaTableName.equals(new SchemaTableName("tiny", "nation_with_hidden_column"))) {
            return TPCH_NATION_WITH_HIDDEN_COLUMN;
        }
        if (schemaTableName.equals(new SchemaTableName("tiny", "nation_with_optional_column"))) {
            return TPCH_NATION_WITH_OPTIONAL_COLUMN;
        }
        throw new UnsupportedOperationException();
    }).withData(schemaTableName -> {
        if (schemaTableName.equals(new SchemaTableName("tiny", "nation"))) {
            return TPCH_NATION_DATA;
        }
        if (schemaTableName.equals(new SchemaTableName("tiny", "nation_with_hidden_column"))) {
            return TPCH_WITH_HIDDEN_COLUMN_DATA;
        }
        if (schemaTableName.equals(new SchemaTableName("tiny", "nation_with_optional_column"))) {
            return TPCH_NATION_DATA;
        }
        throw new UnsupportedOperationException();
    }).build();
    runner.createCatalog(MOCK_CATALOG, mock, ImmutableMap.of());
    MockConnectorFactory mockMissingColumns = MockConnectorFactory.builder().withName("mockmissingcolumns").withGetViews((s, prefix) -> ImmutableMap.<SchemaTableName, ConnectorViewDefinition>builder().put(new SchemaTableName("default", "nation_view"), view).buildOrThrow()).withGetColumns(schemaTableName -> {
        if (schemaTableName.equals(new SchemaTableName("tiny", "nation_with_optional_column"))) {
            return TPCH_NATION_WITH_OPTIONAL_COLUMN;
        }
        throw new UnsupportedOperationException();
    }).withData(schemaTableName -> {
        if (schemaTableName.equals(new SchemaTableName("tiny", "nation_with_optional_column"))) {
            return TPCH_NATION_DATA;
        }
        throw new UnsupportedOperationException();
    }).withAllowMissingColumnsOnInsert(true).build();
    runner.createCatalog(MOCK_CATALOG_MISSING_COLUMNS, mockMissingColumns, ImmutableMap.of());
    assertions = new QueryAssertions(runner);
    accessControl = assertions.getQueryRunner().getAccessControl();
}

Example 17 with MockConnectorFactory

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

the class TestPruneTableScanColumns method testPushColumnPruningProjection.

@Test
public void testPushColumnPruningProjection() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        String mockCatalog = "mock_catalog";
        String testSchema = "test_schema";
        String testTable = "test_table";
        SchemaTableName testSchemaTable = new SchemaTableName(testSchema, testTable);
        ColumnHandle columnHandleA = new MockConnectorColumnHandle("cola", DATE);
        ColumnHandle columnHandleB = new MockConnectorColumnHandle("colb", DOUBLE);
        Map<String, ColumnHandle> assignments = ImmutableMap.of("cola", columnHandleA, "colb", columnHandleB);
        // Create catalog with applyProjection
        MockConnectorFactory factory = MockConnectorFactory.builder().withListSchemaNames(connectorSession -> ImmutableList.of(testSchema)).withListTables((connectorSession, schema) -> testSchema.equals(schema) ? ImmutableList.of(testSchemaTable) : ImmutableList.of()).withGetColumns(schemaTableName -> assignments.entrySet().stream().map(entry -> new ColumnMetadata(entry.getKey(), ((MockConnectorColumnHandle) entry.getValue()).getType())).collect(toImmutableList())).withApplyProjection(this::mockApplyProjection).build();
        ruleTester.getQueryRunner().createCatalog(mockCatalog, factory, ImmutableMap.of());
        ruleTester.assertThat(new PruneTableScanColumns(ruleTester.getMetadata())).on(p -> {
            Symbol symbolA = p.symbol("cola", DATE);
            Symbol symbolB = p.symbol("colb", DOUBLE);
            return p.project(Assignments.of(p.symbol("x"), symbolB.toSymbolReference()), p.tableScan(new TableHandle(new CatalogName(mockCatalog), new MockConnectorTableHandle(testSchemaTable), MockConnectorTransactionHandle.INSTANCE), ImmutableList.of(symbolA, symbolB), ImmutableMap.of(symbolA, columnHandleA, symbolB, columnHandleB)));
        }).withSession(testSessionBuilder().setCatalog(mockCatalog).setSchema(testSchema).build()).matches(strictProject(ImmutableMap.of("expr", PlanMatchPattern.expression("COLB")), tableScan(new MockConnectorTableHandle(testSchemaTable, TupleDomain.all(), Optional.of(ImmutableList.of(columnHandleB)))::equals, TupleDomain.all(), ImmutableMap.of("COLB", columnHandleB::equals))));
    }
}

Example 18 with MockConnectorFactory

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

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

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