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

Example 11 with MockConnectorFactory

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

the class TestPushTopNIntoTableScan method testDoesNotFire.

@Test
public void testDoesNotFire() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        MockConnectorFactory mockFactory = createMockFactory(assignments, Optional.empty());
        ruleTester.getQueryRunner().createCatalog(MOCK_CATALOG, mockFactory, ImmutableMap.of());
        ruleTester.assertThat(new PushTopNIntoTableScan(ruleTester.getMetadata())).on(p -> {
            Symbol dimension = p.symbol(dimensionName, VARCHAR);
            Symbol metric = p.symbol(metricName, BIGINT);
            return p.topN(1, ImmutableList.of(dimension), p.tableScan(TEST_TABLE_HANDLE, ImmutableList.of(dimension, metric), ImmutableMap.of(dimension, dimensionColumn, metric, metricColumn)));
        }).withSession(MOCK_SESSION).doesNotFire();
    }
}

Example 12 with MockConnectorFactory

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

the class TestApplyTableScanRedirection method doesNotFireIfNoTableScan.

@Test
public void doesNotFireIfNoTableScan() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        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 -> p.values(p.symbol("a", BIGINT))).withSession(MOCK_SESSION).doesNotFire();
    }
}

Example 13 with MockConnectorFactory

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

the class TestApplyTableScanRedirection method testApplyTableScanRedirectionWithFilter.

@Test
public void testApplyTableScanRedirectionWithFilter() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        // make the mock connector return a table scan on different table
        // source table handle has a pushed down predicate
        ApplyTableScanRedirect applyTableScanRedirect = getMockApplyRedirect(ImmutableMap.of(SOURCE_COLUMN_HANDLE_A, DESTINATION_COLUMN_NAME_A, SOURCE_COLUMN_HANDLE_B, DESTINATION_COLUMN_NAME_B));
        MockConnectorFactory mockFactory = createMockFactory(Optional.of(applyTableScanRedirect));
        ruleTester.getQueryRunner().createCatalog(MOCK_CATALOG, mockFactory, ImmutableMap.of());
        ApplyTableScanRedirection applyTableScanRedirection = new ApplyTableScanRedirection(ruleTester.getPlannerContext());
        TupleDomain<ColumnHandle> constraint = TupleDomain.withColumnDomains(ImmutableMap.of(SOURCE_COLUMN_HANDLE_A, singleValue(VARCHAR, utf8Slice("foo"))));
        ruleTester.assertThat(applyTableScanRedirection).on(p -> {
            Symbol column = p.symbol(SOURCE_COLUMN_NAME_A, VARCHAR);
            return p.tableScan(createTableHandle(new MockConnectorTableHandle(SOURCE_TABLE, constraint, Optional.empty())), ImmutableList.of(column), ImmutableMap.of(column, SOURCE_COLUMN_HANDLE_A), constraint);
        }).withSession(MOCK_SESSION).matches(filter("DEST_COL = VARCHAR 'foo'", tableScan(new MockConnectorTableHandle(DESTINATION_TABLE)::equals, TupleDomain.all(), ImmutableMap.of("DEST_COL", DESTINATION_COLUMN_HANDLE_A::equals))));
        ruleTester.assertThat(applyTableScanRedirection).on(p -> {
            Symbol column = p.symbol(SOURCE_COLUMN_NAME_B, VARCHAR);
            return p.tableScan(createTableHandle(new MockConnectorTableHandle(SOURCE_TABLE, constraint, Optional.empty())), ImmutableList.of(column), // predicate on non-projected column
            ImmutableMap.of(column, SOURCE_COLUMN_HANDLE_B), TupleDomain.all());
        }).withSession(MOCK_SESSION).matches(project(ImmutableMap.of("expr", expression("DEST_COL_B")), filter("DEST_COL_A = VARCHAR 'foo'", tableScan(new MockConnectorTableHandle(DESTINATION_TABLE)::equals, TupleDomain.all(), ImmutableMap.of("DEST_COL_A", DESTINATION_COLUMN_HANDLE_A::equals, "DEST_COL_B", DESTINATION_COLUMN_HANDLE_B::equals)))));
    }
}

Example 14 with MockConnectorFactory

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

the class TestApplyTableScanRedirection method testMismatchedTypesWithMissingCoercion.

@Test
public void testMismatchedTypesWithMissingCoercion() {
    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_D));
        MockConnectorFactory mockFactory = createMockFactory(Optional.of(applyTableScanRedirect));
        LocalQueryRunner runner = ruleTester.getQueryRunner();
        runner.createCatalog(MOCK_CATALOG, mockFactory, ImmutableMap.of());
        transaction(runner.getTransactionManager(), runner.getAccessControl()).execute(MOCK_SESSION, session -> {
            assertThatThrownBy(() -> runner.createPlan(session, "SELECT source_col_a FROM test_table", WarningCollector.NOOP)).isInstanceOf(TrinoException.class).hasMessageMatching("Cast not possible from redirected column mock_catalog.target_schema.target_table.destination_col_d with type Bogus to source column .*mock_catalog.test_schema.test_table.*source_col_a.* with type: varchar");
        });
    }
}

Example 15 with MockConnectorFactory

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

the class TestColumnMask 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);
    ConnectorMaterializedViewDefinition materializedView = new ConnectorMaterializedViewDefinition("SELECT * FROM local.tiny.nation", Optional.empty(), Optional.empty(), Optional.empty(), ImmutableList.of(new ConnectorMaterializedViewDefinition.Column("nationkey", BigintType.BIGINT.getTypeId()), new ConnectorMaterializedViewDefinition.Column("name", VarcharType.createVarcharType(25).getTypeId()), new ConnectorMaterializedViewDefinition.Column("regionkey", BigintType.BIGINT.getTypeId()), new ConnectorMaterializedViewDefinition.Column("comment", VarcharType.createVarcharType(152).getTypeId())), Optional.empty(), Optional.of(VIEW_OWNER), ImmutableMap.of());
    ConnectorMaterializedViewDefinition freshMaterializedView = new ConnectorMaterializedViewDefinition("SELECT * FROM local.tiny.nation", Optional.of(new CatalogSchemaTableName("local", "tiny", "nation")), Optional.empty(), Optional.empty(), ImmutableList.of(new ConnectorMaterializedViewDefinition.Column("nationkey", BigintType.BIGINT.getTypeId()), new ConnectorMaterializedViewDefinition.Column("name", VarcharType.createVarcharType(25).getTypeId()), new ConnectorMaterializedViewDefinition.Column("regionkey", BigintType.BIGINT.getTypeId()), new ConnectorMaterializedViewDefinition.Column("comment", VarcharType.createVarcharType(152).getTypeId())), Optional.empty(), Optional.of(VIEW_OWNER), ImmutableMap.of());
    ConnectorMaterializedViewDefinition materializedViewWithCasts = new ConnectorMaterializedViewDefinition("SELECT nationkey, cast(name as varchar(1)) as name, regionkey, comment FROM local.tiny.nation", Optional.of(new CatalogSchemaTableName("local", "tiny", "nation")), Optional.empty(), Optional.empty(), ImmutableList.of(new ConnectorMaterializedViewDefinition.Column("nationkey", BigintType.BIGINT.getTypeId()), new ConnectorMaterializedViewDefinition.Column("name", VarcharType.createVarcharType(2).getTypeId()), new ConnectorMaterializedViewDefinition.Column("regionkey", BigintType.BIGINT.getTypeId()), new ConnectorMaterializedViewDefinition.Column("comment", VarcharType.createVarcharType(152).getTypeId())), Optional.empty(), Optional.of(VIEW_OWNER), ImmutableMap.of());
    MockConnectorFactory mock = MockConnectorFactory.builder().withGetColumns(schemaTableName -> {
        if (schemaTableName.equals(new SchemaTableName("tiny", "nation_with_hidden_column"))) {
            return TPCH_NATION_WITH_HIDDEN_COLUMN;
        }
        throw new UnsupportedOperationException();
    }).withData(schemaTableName -> {
        if (schemaTableName.equals(new SchemaTableName("tiny", "nation_with_hidden_column"))) {
            return TPCH_WITH_HIDDEN_COLUMN_DATA;
        }
        throw new UnsupportedOperationException();
    }).withGetViews((s, prefix) -> ImmutableMap.of(new SchemaTableName("default", "nation_view"), view)).withGetMaterializedViews((s, prefix) -> ImmutableMap.of(new SchemaTableName("default", "nation_materialized_view"), materializedView, new SchemaTableName("default", "nation_fresh_materialized_view"), freshMaterializedView, new SchemaTableName("default", "materialized_view_with_casts"), materializedViewWithCasts)).build();
    runner.createCatalog(MOCK_CATALOG, mock, ImmutableMap.of());
    assertions = new QueryAssertions(runner);
    accessControl = assertions.getQueryRunner().getAccessControl();
}