Examples with LocalQueryRunner io.trino.testing.LocalQueryRunner used on opensource projects

Example 11 with LocalQueryRunner

use of io.trino.testing.LocalQueryRunner 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 12 with LocalQueryRunner

use of io.trino.testing.LocalQueryRunner in project trino by trinodb.

the class TestTransactionManager method testAbortedTransactionWorkflow.

@Test
public void testAbortedTransactionWorkflow() {
    try (LocalQueryRunner queryRunner = LocalQueryRunner.create(TEST_SESSION)) {
        TransactionManager transactionManager = queryRunner.getTransactionManager();
        queryRunner.createCatalog(CATALOG, new TpchConnectorFactory(), ImmutableMap.of());
        TransactionId transactionId = transactionManager.beginTransaction(false);
        assertEquals(transactionManager.getAllTransactionInfos().size(), 1);
        TransactionInfo transactionInfo = transactionManager.getTransactionInfo(transactionId);
        assertFalse(transactionInfo.isAutoCommitContext());
        assertTrue(transactionInfo.getCatalogNames().isEmpty());
        assertFalse(transactionInfo.getWrittenConnectorId().isPresent());
        ConnectorMetadata metadata = transactionManager.getOptionalCatalogMetadata(transactionId, CATALOG).get().getMetadata(TEST_SESSION);
        metadata.listSchemaNames(TEST_SESSION.toConnectorSession(CATALOG_NAME));
        transactionInfo = transactionManager.getTransactionInfo(transactionId);
        assertEquals(transactionInfo.getCatalogNames(), ImmutableList.of(CATALOG_NAME, INFORMATION_SCHEMA_ID, SYSTEM_TABLES_ID));
        assertFalse(transactionInfo.getWrittenConnectorId().isPresent());
        getFutureValue(transactionManager.asyncAbort(transactionId));
        assertTrue(transactionManager.getAllTransactionInfos().isEmpty());
    }
}

Example 13 with LocalQueryRunner

use of io.trino.testing.LocalQueryRunner in project trino by trinodb.

the class TestTransactionManager method testFailedTransactionWorkflow.

@Test
public void testFailedTransactionWorkflow() {
    try (LocalQueryRunner queryRunner = LocalQueryRunner.create(TEST_SESSION)) {
        TransactionManager transactionManager = queryRunner.getTransactionManager();
        queryRunner.createCatalog(CATALOG, new TpchConnectorFactory(), ImmutableMap.of());
        TransactionId transactionId = transactionManager.beginTransaction(false);
        assertEquals(transactionManager.getAllTransactionInfos().size(), 1);
        TransactionInfo transactionInfo = transactionManager.getTransactionInfo(transactionId);
        assertFalse(transactionInfo.isAutoCommitContext());
        assertTrue(transactionInfo.getCatalogNames().isEmpty());
        assertFalse(transactionInfo.getWrittenConnectorId().isPresent());
        ConnectorMetadata metadata = transactionManager.getOptionalCatalogMetadata(transactionId, CATALOG).get().getMetadata(TEST_SESSION);
        metadata.listSchemaNames(TEST_SESSION.toConnectorSession(CATALOG_NAME));
        transactionInfo = transactionManager.getTransactionInfo(transactionId);
        assertEquals(transactionInfo.getCatalogNames(), ImmutableList.of(CATALOG_NAME, INFORMATION_SCHEMA_ID, SYSTEM_TABLES_ID));
        assertFalse(transactionInfo.getWrittenConnectorId().isPresent());
        transactionManager.fail(transactionId);
        assertEquals(transactionManager.getAllTransactionInfos().size(), 1);
        assertTrinoExceptionThrownBy(() -> transactionManager.getCatalogMetadata(transactionId, CATALOG_NAME)).hasErrorCode(TRANSACTION_ALREADY_ABORTED);
        assertEquals(transactionManager.getAllTransactionInfos().size(), 1);
        getFutureValue(transactionManager.asyncAbort(transactionId));
        assertTrue(transactionManager.getAllTransactionInfos().isEmpty());
    }
}

Example 14 with LocalQueryRunner

use of io.trino.testing.LocalQueryRunner 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();
}

Example 15 with LocalQueryRunner

use of io.trino.testing.LocalQueryRunner 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();
}