Search in sources :

Example 21 with TransactionId

use of io.trino.transaction.TransactionId in project trino by trinodb.

the class RollbackTask method execute.

@Override
public ListenableFuture<Void> execute(Rollback statement, QueryStateMachine stateMachine, List<Expression> parameters, WarningCollector warningCollector) {
    Session session = stateMachine.getSession();
    if (session.getTransactionId().isEmpty()) {
        throw new TrinoException(NOT_IN_TRANSACTION, "No transaction in progress");
    }
    TransactionId transactionId = session.getTransactionId().get();
    stateMachine.clearTransactionId();
    transactionManager.asyncAbort(transactionId);
    return immediateVoidFuture();
}
Also used : TrinoException(io.trino.spi.TrinoException) Session(io.trino.Session) TransactionId(io.trino.transaction.TransactionId)

Example 22 with TransactionId

use of io.trino.transaction.TransactionId in project trino by trinodb.

the class TestPushProjectionIntoTableScan method testPushProjection.

@Test
public void testPushProjection() {
    try (RuleTester ruleTester = defaultRuleTester()) {
        // Building context for input
        String columnName = "col0";
        Type columnType = ROW_TYPE;
        Symbol baseColumn = new Symbol(columnName);
        ColumnHandle columnHandle = new TpchColumnHandle(columnName, columnType);
        // Create catalog with applyProjection enabled
        MockConnectorFactory factory = createMockFactory(ImmutableMap.of(columnName, columnHandle), Optional.of(this::mockApplyProjection));
        ruleTester.getQueryRunner().createCatalog(MOCK_CATALOG, factory, ImmutableMap.of());
        TypeAnalyzer typeAnalyzer = createTestingTypeAnalyzer(ruleTester.getPlannerContext());
        // Prepare project node symbols and types
        Symbol identity = new Symbol("symbol_identity");
        Symbol dereference = new Symbol("symbol_dereference");
        Symbol constant = new Symbol("symbol_constant");
        Symbol call = new Symbol("symbol_call");
        ImmutableMap<Symbol, Type> types = ImmutableMap.of(baseColumn, ROW_TYPE, identity, ROW_TYPE, dereference, BIGINT, constant, BIGINT, call, VARCHAR);
        // Prepare project node assignments
        ImmutableMap<Symbol, Expression> inputProjections = ImmutableMap.of(identity, baseColumn.toSymbolReference(), dereference, new SubscriptExpression(baseColumn.toSymbolReference(), new LongLiteral("1")), constant, new LongLiteral("5"), call, new FunctionCall(QualifiedName.of("STARTS_WITH"), ImmutableList.of(new StringLiteral("abc"), new StringLiteral("ab"))));
        // Compute expected symbols after applyProjection
        TransactionId transactionId = ruleTester.getQueryRunner().getTransactionManager().beginTransaction(false);
        Session session = MOCK_SESSION.beginTransactionId(transactionId, ruleTester.getQueryRunner().getTransactionManager(), ruleTester.getQueryRunner().getAccessControl());
        ImmutableMap<Symbol, String> connectorNames = inputProjections.entrySet().stream().collect(toImmutableMap(Map.Entry::getKey, e -> translate(session, e.getValue(), typeAnalyzer, viewOf(types), ruleTester.getPlannerContext()).get().toString()));
        ImmutableMap<Symbol, String> newNames = ImmutableMap.of(identity, "projected_variable_" + connectorNames.get(identity), dereference, "projected_dereference_" + connectorNames.get(dereference), constant, "projected_constant_" + connectorNames.get(constant), call, "projected_call_" + connectorNames.get(call));
        Map<String, ColumnHandle> expectedColumns = newNames.entrySet().stream().collect(toImmutableMap(Map.Entry::getValue, e -> column(e.getValue(), types.get(e.getKey()))));
        ruleTester.assertThat(createRule(ruleTester)).on(p -> {
            // Register symbols
            types.forEach((symbol, type) -> p.symbol(symbol.getName(), type));
            return p.project(new Assignments(inputProjections), p.tableScan(tableScan -> tableScan.setTableHandle(TEST_TABLE_HANDLE).setSymbols(ImmutableList.copyOf(types.keySet())).setAssignments(types.keySet().stream().collect(Collectors.toMap(Function.identity(), v -> columnHandle))).setStatistics(Optional.of(PlanNodeStatsEstimate.builder().setOutputRowCount(42).addSymbolStatistics(baseColumn, SymbolStatsEstimate.builder().setNullsFraction(0).setDistinctValuesCount(33).build()).build()))));
        }).withSession(MOCK_SESSION).matches(project(newNames.entrySet().stream().collect(toImmutableMap(e -> e.getKey().getName(), e -> expression(symbolReference(e.getValue())))), tableScan(new MockConnectorTableHandle(new SchemaTableName(TEST_SCHEMA, "projected_" + TEST_TABLE), TupleDomain.all(), Optional.of(ImmutableList.copyOf(expectedColumns.values())))::equals, TupleDomain.all(), expectedColumns.entrySet().stream().collect(toImmutableMap(Map.Entry::getKey, e -> e.getValue()::equals)), Optional.of(PlanNodeStatsEstimate.builder().setOutputRowCount(42).addSymbolStatistics(new Symbol(newNames.get(constant)), SymbolStatsEstimate.builder().setDistinctValuesCount(1).setNullsFraction(0).setLowValue(5).setHighValue(5).build()).addSymbolStatistics(new Symbol(newNames.get(call).toLowerCase(ENGLISH)), SymbolStatsEstimate.builder().setDistinctValuesCount(1).setNullsFraction(0).build()).addSymbolStatistics(new Symbol(newNames.get(identity)), SymbolStatsEstimate.builder().setDistinctValuesCount(33).setNullsFraction(0).build()).addSymbolStatistics(new Symbol(newNames.get(dereference)), SymbolStatsEstimate.unknown()).build())::equals)));
    }
}
Also used : Test(org.testng.annotations.Test) CatalogName(io.trino.connector.CatalogName) MockConnectorFactory(io.trino.connector.MockConnectorFactory) LongLiteral(io.trino.sql.tree.LongLiteral) Arrays.asList(java.util.Arrays.asList) ConnectorTableHandle(io.trino.spi.connector.ConnectorTableHandle) Map(java.util.Map) ProjectionApplicationResult(io.trino.spi.connector.ProjectionApplicationResult) FunctionCall(io.trino.sql.tree.FunctionCall) ENGLISH(java.util.Locale.ENGLISH) PlanNodeStatsEstimate(io.trino.cost.PlanNodeStatsEstimate) PlanMatchPattern.expression(io.trino.sql.planner.assertions.PlanMatchPattern.expression) TpchColumnHandle(io.trino.plugin.tpch.TpchColumnHandle) RowType(io.trino.spi.type.RowType) ImmutableMap(com.google.common.collect.ImmutableMap) Call(io.trino.spi.expression.Call) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) Assignments(io.trino.sql.planner.plan.Assignments) ScalarStatsCalculator(io.trino.cost.ScalarStatsCalculator) Collectors(java.util.stream.Collectors) SchemaTableName(io.trino.spi.connector.SchemaTableName) List(java.util.List) ImmutableMap.toImmutableMap(com.google.common.collect.ImmutableMap.toImmutableMap) TestingSession.testSessionBuilder(io.trino.testing.TestingSession.testSessionBuilder) PlanMatchPattern.anyTree(io.trino.sql.planner.assertions.PlanMatchPattern.anyTree) BIGINT(io.trino.spi.type.BigintType.BIGINT) ConnectorPartitioningHandle(io.trino.spi.connector.ConnectorPartitioningHandle) SymbolReference(io.trino.sql.tree.SymbolReference) Assignment(io.trino.spi.connector.Assignment) Optional(java.util.Optional) Expression(io.trino.sql.tree.Expression) RowType.field(io.trino.spi.type.RowType.field) ConnectorTablePartitioning(io.trino.spi.connector.ConnectorTablePartitioning) ConnectorExpressionTranslator.translate(io.trino.sql.planner.ConnectorExpressionTranslator.translate) Session(io.trino.Session) PlannerContext(io.trino.sql.PlannerContext) TypeAnalyzer.createTestingTypeAnalyzer(io.trino.sql.planner.TypeAnalyzer.createTestingTypeAnalyzer) ColumnMetadata(io.trino.spi.connector.ColumnMetadata) Type(io.trino.spi.type.Type) Variable(io.trino.spi.expression.Variable) Function(java.util.function.Function) SubscriptExpression(io.trino.sql.tree.SubscriptExpression) VARCHAR(io.trino.spi.type.VarcharType.VARCHAR) ImmutableList(com.google.common.collect.ImmutableList) Assertions.assertThatThrownBy(org.assertj.core.api.Assertions.assertThatThrownBy) ColumnHandle(io.trino.spi.connector.ColumnHandle) Constant(io.trino.spi.expression.Constant) TypeProvider.viewOf(io.trino.sql.planner.TypeProvider.viewOf) Symbol(io.trino.sql.planner.Symbol) SymbolStatsEstimate(io.trino.cost.SymbolStatsEstimate) StringLiteral(io.trino.sql.tree.StringLiteral) RuleTester.defaultRuleTester(io.trino.sql.planner.iterative.rule.test.RuleTester.defaultRuleTester) FieldDereference(io.trino.spi.expression.FieldDereference) TestingTransactionHandle(io.trino.testing.TestingTransactionHandle) ConnectorSession(io.trino.spi.connector.ConnectorSession) RuleTester(io.trino.sql.planner.iterative.rule.test.RuleTester) TupleDomain(io.trino.spi.predicate.TupleDomain) QualifiedName(io.trino.sql.tree.QualifiedName) MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle) TableHandle(io.trino.metadata.TableHandle) ConnectorTableProperties(io.trino.spi.connector.ConnectorTableProperties) ConnectorExpression(io.trino.spi.expression.ConnectorExpression) TypeAnalyzer(io.trino.sql.planner.TypeAnalyzer) PlanMatchPattern.project(io.trino.sql.planner.assertions.PlanMatchPattern.project) TransactionId(io.trino.transaction.TransactionId) PlanMatchPattern.tableScan(io.trino.sql.planner.assertions.PlanMatchPattern.tableScan) MockConnectorFactory(io.trino.connector.MockConnectorFactory) TpchColumnHandle(io.trino.plugin.tpch.TpchColumnHandle) RuleTester.defaultRuleTester(io.trino.sql.planner.iterative.rule.test.RuleTester.defaultRuleTester) RuleTester(io.trino.sql.planner.iterative.rule.test.RuleTester) Symbol(io.trino.sql.planner.Symbol) Assignments(io.trino.sql.planner.plan.Assignments) SubscriptExpression(io.trino.sql.tree.SubscriptExpression) FunctionCall(io.trino.sql.tree.FunctionCall) TpchColumnHandle(io.trino.plugin.tpch.TpchColumnHandle) ColumnHandle(io.trino.spi.connector.ColumnHandle) LongLiteral(io.trino.sql.tree.LongLiteral) SchemaTableName(io.trino.spi.connector.SchemaTableName) TypeAnalyzer.createTestingTypeAnalyzer(io.trino.sql.planner.TypeAnalyzer.createTestingTypeAnalyzer) TypeAnalyzer(io.trino.sql.planner.TypeAnalyzer) TransactionId(io.trino.transaction.TransactionId) RowType(io.trino.spi.type.RowType) Type(io.trino.spi.type.Type) StringLiteral(io.trino.sql.tree.StringLiteral) Expression(io.trino.sql.tree.Expression) SubscriptExpression(io.trino.sql.tree.SubscriptExpression) ConnectorExpression(io.trino.spi.expression.ConnectorExpression) MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle) Map(java.util.Map) ImmutableMap(com.google.common.collect.ImmutableMap) ImmutableMap.toImmutableMap(com.google.common.collect.ImmutableMap.toImmutableMap) Session(io.trino.Session) ConnectorSession(io.trino.spi.connector.ConnectorSession) Test(org.testng.annotations.Test)

Aggregations

TransactionId (io.trino.transaction.TransactionId)22 Test (org.testng.annotations.Test)12 SchemaTableName (io.trino.spi.connector.SchemaTableName)11 ImmutableMap (com.google.common.collect.ImmutableMap)8 ConnectorSession (io.trino.spi.connector.ConnectorSession)8 Session (io.trino.Session)7 InformationSchemaMetadata (io.trino.connector.informationschema.InformationSchemaMetadata)7 InformationSchemaTableHandle (io.trino.connector.informationschema.InformationSchemaTableHandle)7 ConnectorMetadata (io.trino.spi.connector.ConnectorMetadata)7 Constraint (io.trino.spi.connector.Constraint)7 QualifiedObjectName (io.trino.metadata.QualifiedObjectName)6 TrinoException (io.trino.spi.TrinoException)6 CatalogName (io.trino.connector.CatalogName)5 ImmutableList (com.google.common.collect.ImmutableList)4 MockConnectorFactory (io.trino.connector.MockConnectorFactory)4 TransactionManager (io.trino.transaction.TransactionManager)4 Map (java.util.Map)4 ImmutableSet (com.google.common.collect.ImmutableSet)3 TEST_SESSION (io.trino.SessionTestUtils.TEST_SESSION)3 InformationSchemaColumnHandle (io.trino.connector.informationschema.InformationSchemaColumnHandle)3