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

Example 1 with MockConnectorFactory

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

the class TestSetRoleTask method setUp.

@BeforeClass
public void setUp() {
    queryRunner = LocalQueryRunner.create(TEST_SESSION);
    MockConnectorFactory mockConnectorFactory = MockConnectorFactory.builder().withListRoleGrants((connectorSession, roles, grantees, limit) -> ImmutableSet.of(new RoleGrant(new TrinoPrincipal(USER, USER_NAME), ROLE_NAME, false))).build();
    queryRunner.createCatalog(CATALOG_NAME, mockConnectorFactory, ImmutableMap.of());
    MockConnectorFactory systemConnectorFactory = MockConnectorFactory.builder().withName("system_role_connector").build();
    queryRunner.createCatalog(SYSTEM_ROLE_CATALOG_NAME, systemConnectorFactory, ImmutableMap.of());
    transactionManager = queryRunner.getTransactionManager();
    accessControl = queryRunner.getAccessControl();
    metadata = queryRunner.getMetadata();
    parser = queryRunner.getSqlParser();
    executor = newCachedThreadPool(daemonThreadsNamed("test-set-role-task-executor-%s"));
}

Example 2 with MockConnectorFactory

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

the class TestPushProjectionIntoTableScan method testDoesNotFire.

@Test
public void testDoesNotFire() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        String columnName = "input_column";
        Type columnType = ROW_TYPE;
        ColumnHandle inputColumnHandle = column(columnName, columnType);
        MockConnectorFactory factory = createMockFactory(ImmutableMap.of(columnName, inputColumnHandle), Optional.empty());
        ruleTester.getQueryRunner().createCatalog(MOCK_CATALOG, factory, ImmutableMap.of());
        PushProjectionIntoTableScan optimizer = createRule(ruleTester);
        ruleTester.assertThat(optimizer).on(p -> {
            Symbol symbol = p.symbol(columnName, columnType);
            return p.project(Assignments.of(p.symbol("symbol_dereference", BIGINT), new SubscriptExpression(symbol.toSymbolReference(), new LongLiteral("1"))), p.tableScan(TEST_TABLE_HANDLE, ImmutableList.of(symbol), ImmutableMap.of(symbol, inputColumnHandle)));
        }).withSession(MOCK_SESSION).doesNotFire();
    }
}

Example 3 with MockConnectorFactory

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

the class TestPushProjectionIntoTableScan method createMockFactory.

private MockConnectorFactory createMockFactory(Map<String, ColumnHandle> assignments, Optional<MockConnectorFactory.ApplyProjection> applyProjection) {
    List<ColumnMetadata> metadata = assignments.entrySet().stream().map(entry -> new ColumnMetadata(entry.getKey(), ((TpchColumnHandle) entry.getValue()).getType())).collect(toImmutableList());
    MockConnectorFactory.Builder builder = MockConnectorFactory.builder().withListSchemaNames(connectorSession -> ImmutableList.of(TEST_SCHEMA)).withListTables((connectorSession, schema) -> TEST_SCHEMA.equals(schema) ? ImmutableList.of(TEST_SCHEMA_TABLE) : ImmutableList.of()).withGetColumns(schemaTableName -> metadata).withGetTableProperties((session, tableHandle) -> {
        MockConnectorTableHandle mockTableHandle = (MockConnectorTableHandle) tableHandle;
        if (mockTableHandle.getTableName().getTableName().equals(TEST_TABLE)) {
            return new ConnectorTableProperties(TupleDomain.all(), Optional.of(new ConnectorTablePartitioning(PARTITIONING_HANDLE, ImmutableList.of(column("col", VARCHAR)))), Optional.empty(), Optional.empty(), ImmutableList.of());
        }
        return new ConnectorTableProperties();
    });
    if (applyProjection.isPresent()) {
        builder = builder.withApplyProjection(applyProjection.get());
    }
    return builder.build();
}

Example 4 with MockConnectorFactory

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

the class TestPushJoinIntoTableScan method testPushJoinIntoTableRequiresFullColumnHandleMappingInResult.

@Test
public void testPushJoinIntoTableRequiresFullColumnHandleMappingInResult() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        MockConnectorFactory connectorFactory = createMockConnectorFactory((session, applyJoinType, left, right, joinConditions, leftAssignments, rightAssignments) -> Optional.of(new JoinApplicationResult<>(JOIN_CONNECTOR_TABLE_HANDLE, ImmutableMap.of(COLUMN_A1_HANDLE, JOIN_COLUMN_A1_HANDLE, COLUMN_A2_HANDLE, JOIN_COLUMN_A2_HANDLE), // mapping for COLUMN_B1_HANDLE is missing
        ImmutableMap.of(), false)));
        ruleTester.getQueryRunner().createCatalog(MOCK_CATALOG, connectorFactory, ImmutableMap.of());
        assertThatThrownBy(() -> {
            ruleTester.assertThat(new PushJoinIntoTableScan(ruleTester.getMetadata())).on(p -> {
                Symbol columnA1Symbol = p.symbol(COLUMN_A1);
                Symbol columnA2Symbol = p.symbol(COLUMN_A2);
                Symbol columnB1Symbol = p.symbol(COLUMN_B1);
                TupleDomain<ColumnHandle> leftContraint = TupleDomain.fromFixedValues(ImmutableMap.of(COLUMN_A2_HANDLE, NullableValue.of(BIGINT, 44L)));
                TupleDomain<ColumnHandle> rightConstraint = TupleDomain.fromFixedValues(ImmutableMap.of(COLUMN_B1_HANDLE, NullableValue.of(BIGINT, 45L)));
                TableScanNode left = p.tableScan(TABLE_A_HANDLE, ImmutableList.of(columnA1Symbol, columnA2Symbol), ImmutableMap.of(columnA1Symbol, COLUMN_A1_HANDLE, columnA2Symbol, COLUMN_A2_HANDLE), leftContraint);
                TableScanNode right = p.tableScan(TABLE_B_HANDLE, ImmutableList.of(columnB1Symbol), ImmutableMap.of(columnB1Symbol, COLUMN_B1_HANDLE), rightConstraint);
                return p.join(INNER, left, right, new JoinNode.EquiJoinClause(columnA1Symbol, columnB1Symbol));
            }).withSession(MOCK_SESSION).matches(anyTree());
        }).isInstanceOf(IllegalStateException.class).hasMessageContaining("Column handle mappings do not match old column handles");
    }
}

Example 5 with MockConnectorFactory

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

the class TestPushJoinIntoTableScan method testPushJoinIntoTableScanDoesNotTriggerWithUnsupportedFilter.

@Test
public void testPushJoinIntoTableScanDoesNotTriggerWithUnsupportedFilter() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        MockConnectorFactory connectorFactory = createMockConnectorFactory((session, applyJoinType, left, right, joinConditions, leftAssignments, rightAssignments) -> {
            throw new IllegalStateException("applyJoin should not be called!");
        });
        ruleTester.getQueryRunner().createCatalog(MOCK_CATALOG, connectorFactory, ImmutableMap.of());
        ruleTester.assertThat(new PushJoinIntoTableScan(ruleTester.getMetadata())).on(p -> {
            Symbol columnA1Symbol = p.symbol(COLUMN_A1);
            Symbol columnA2Symbol = p.symbol(COLUMN_A2);
            Symbol columnB1Symbol = p.symbol(COLUMN_B1);
            TableScanNode left = p.tableScan(TABLE_A_HANDLE, ImmutableList.of(columnA1Symbol, columnA2Symbol), ImmutableMap.of(columnA1Symbol, COLUMN_A1_HANDLE, columnA2Symbol, COLUMN_A2_HANDLE));
            TableScanNode right = p.tableScan(TABLE_B_HANDLE, ImmutableList.of(columnB1Symbol), ImmutableMap.of(columnB1Symbol, COLUMN_B1_HANDLE));
            return p.join(INNER, left, right, new ComparisonExpression(ComparisonExpression.Operator.GREATER_THAN, new ArithmeticBinaryExpression(ArithmeticBinaryExpression.Operator.MULTIPLY, new GenericLiteral("BIGINT", "44"), columnA1Symbol.toSymbolReference()), columnB1Symbol.toSymbolReference()));
        }).withSession(MOCK_SESSION).doesNotFire();
    }
}