Examples with TEST_SESSION io.trino.SessionTestUtils.TEST_SESSION used on opensource projects

Example 6 with TEST_SESSION

use of io.trino.SessionTestUtils.TEST_SESSION in project trino by trinodb.

the class TestSourcePartitionedScheduler method createStageExecution.

private StageExecution createStageExecution(PlanFragment fragment, NodeTaskMap nodeTaskMap) {
    StageId stageId = new StageId(QUERY_ID, 0);
    SqlStage stage = SqlStage.createSqlStage(stageId, fragment, ImmutableMap.of(TABLE_SCAN_NODE_ID, new TableInfo(new QualifiedObjectName("test", "test", "test"), TupleDomain.all())), new MockRemoteTaskFactory(queryExecutor, scheduledExecutor), TEST_SESSION, true, nodeTaskMap, queryExecutor, new SplitSchedulerStats());
    ImmutableMap.Builder<PlanFragmentId, OutputBufferManager> outputBuffers = ImmutableMap.builder();
    outputBuffers.put(fragment.getId(), new PartitionedOutputBufferManager(FIXED_HASH_DISTRIBUTION, 1));
    fragment.getRemoteSourceNodes().stream().flatMap(node -> node.getSourceFragmentIds().stream()).forEach(fragmentId -> outputBuffers.put(fragmentId, new PartitionedOutputBufferManager(FIXED_HASH_DISTRIBUTION, 10)));
    return createPipelinedStageExecution(stage, outputBuffers.buildOrThrow(), TaskLifecycleListener.NO_OP, new NoOpFailureDetector(), queryExecutor, Optional.of(new int[] { 0 }), 0);
}

Example 7 with TEST_SESSION

use of io.trino.SessionTestUtils.TEST_SESSION in project trino by trinodb.

the class TestAccessControlManager method testDenySystemAccessControl.

@Test
public void testDenySystemAccessControl() {
    try (LocalQueryRunner queryRunner = LocalQueryRunner.create(TEST_SESSION)) {
        TransactionManager transactionManager = queryRunner.getTransactionManager();
        AccessControlManager accessControlManager = createAccessControlManager(transactionManager);
        TestSystemAccessControlFactory accessControlFactory = new TestSystemAccessControlFactory("test");
        accessControlManager.addSystemAccessControlFactory(accessControlFactory);
        accessControlManager.setSystemAccessControl("test", ImmutableMap.of());
        queryRunner.createCatalog("catalog", MockConnectorFactory.create(), ImmutableMap.of());
        accessControlManager.addCatalogAccessControl(new CatalogName("connector"), new DenyConnectorAccessControl());
        assertThatThrownBy(() -> transaction(transactionManager, accessControlManager).execute(transactionId -> {
            accessControlManager.checkCanSelectFromColumns(context(transactionId), new QualifiedObjectName("secured_catalog", "schema", "table"), ImmutableSet.of("column"));
        })).isInstanceOf(TrinoException.class).hasMessageMatching("Access Denied: Cannot select from table secured_catalog.schema.table");
    }
}

Example 8 with TEST_SESSION

use of io.trino.SessionTestUtils.TEST_SESSION in project trino by trinodb.

the class TestWindowClause method testWindowWithFrameCoercions.

@Test
public void testWindowWithFrameCoercions() {
    @Language("SQL") String sql = "SELECT a old_a, 2e0 a FROM (VALUES -100, -99, -98) t(a) WINDOW w AS (ORDER BY a + 1) ORDER BY count(*) OVER (w RANGE BETWEEN CURRENT ROW AND a + 1e0 FOLLOWING)";
    PlanMatchPattern pattern = anyTree(sort(ImmutableList.of(sort("count_result", ASCENDING, LAST)), project(window(windowMatcherBuilder -> windowMatcherBuilder.specification(specification(ImmutableList.of(), ImmutableList.of("sortkey"), ImmutableMap.of("sortkey", SortOrder.ASC_NULLS_LAST))).addFunction("count_result", functionCall("count", ImmutableList.of()), createTestMetadataManager().resolveFunction(TEST_SESSION, QualifiedName.of("count"), fromTypes()), windowFrame(RANGE, CURRENT_ROW, Optional.empty(), Optional.empty(), FOLLOWING, Optional.of("frame_bound"), Optional.of("coerced_sortkey"))), project(ImmutableMap.of("frame_bound", expression(new FunctionCall(QualifiedName.of("$operator$add"), ImmutableList.of(new SymbolReference("coerced_sortkey"), new SymbolReference("frame_offset"))))), project(ImmutableMap.of("coerced_sortkey", expression("CAST(sortkey AS double)")), node(FilterNode.class, project(project(ImmutableMap.of(// sort key based on "a" in source scope
    "sortkey", expression("a + 1"), // frame offset based on "a" in output scope
    "frame_offset", expression("new_a + 1E0")), project(ImmutableMap.of("new_a", expression("2E0")), project(project(values("a")))))))))))));
    assertPlan(sql, CREATED, pattern);
}

Example 9 with TEST_SESSION

use of io.trino.SessionTestUtils.TEST_SESSION in project trino by trinodb.

the class TestWindowFrameRange method testFrameFollowingWithOffsetCoercion.

@Test
public void testFrameFollowingWithOffsetCoercion() {
    @Language("SQL") String sql = "SELECT array_agg(key) OVER(ORDER BY key RANGE BETWEEN CURRENT ROW AND x FOLLOWING) " + "FROM (VALUES (1.1, 1), (2.2, 2)) t(key, x)";
    PlanMatchPattern pattern = anyTree(window(windowMatcherBuilder -> windowMatcherBuilder.specification(specification(ImmutableList.of(), ImmutableList.of("key"), ImmutableMap.of("key", SortOrder.ASC_NULLS_LAST))).addFunction("array_agg_result", functionCall("array_agg", ImmutableList.of("key")), createTestMetadataManager().resolveFunction(TEST_SESSION, QualifiedName.of("array_agg"), fromTypes(createDecimalType(2, 1))), windowFrame(RANGE, CURRENT_ROW, Optional.empty(), Optional.empty(), FOLLOWING, Optional.of("frame_end_value"), Optional.of("key_for_frame_end_comparison"))), project(ImmutableMap.of("key_for_frame_end_comparison", expression("CAST(key AS decimal(12, 1))")), project(ImmutableMap.of("frame_end_value", expression(new FunctionCall(QualifiedName.of("$operator$add"), ImmutableList.of(new SymbolReference("key"), new SymbolReference("offset"))))), filter("IF((offset >= CAST(0 AS DECIMAL(10, 0))), " + "true, " + "CAST(fail(CAST('Window frame offset value must not be negative or null' AS varchar)) AS boolean))", project(ImmutableMap.of("offset", expression("CAST(x AS decimal(10, 0))")), anyTree(values(ImmutableList.of("key", "x"), ImmutableList.of(ImmutableList.of(new DecimalLiteral("1.1"), new LongLiteral("1")), ImmutableList.of(new DecimalLiteral("2.2"), new LongLiteral("2")))))))))));
    assertPlan(sql, CREATED, pattern);
}

Example 10 with TEST_SESSION

use of io.trino.SessionTestUtils.TEST_SESSION in project trino by trinodb.

the class TestPushProjectionThroughJoin method testPushesProjectionThroughJoin.

@Test
public void testPushesProjectionThroughJoin() {
    PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator();
    PlanBuilder p = new PlanBuilder(idAllocator, dummyMetadata(), TEST_SESSION);
    Symbol a0 = p.symbol("a0");
    Symbol a1 = p.symbol("a1");
    Symbol a2 = p.symbol("a2");
    Symbol a3 = p.symbol("a3");
    Symbol b0 = p.symbol("b0");
    Symbol b1 = p.symbol("b1");
    Symbol b2 = p.symbol("b2");
    ProjectNode planNode = p.project(Assignments.of(a3, new ArithmeticUnaryExpression(MINUS, a2.toSymbolReference()), b2, new ArithmeticUnaryExpression(PLUS, b1.toSymbolReference())), p.join(INNER, // intermediate non-identity projections should be fully inlined
    p.project(Assignments.of(a2, new ArithmeticUnaryExpression(PLUS, a0.toSymbolReference()), a1, a1.toSymbolReference()), p.project(Assignments.builder().putIdentity(a0).putIdentity(a1).build(), p.values(a0, a1))), p.values(b0, b1), new JoinNode.EquiJoinClause(a1, b1)));
    Session session = testSessionBuilder().build();
    Optional<PlanNode> rewritten = pushProjectionThroughJoin(PLANNER_CONTEXT, planNode, noLookup(), idAllocator, session, createTestingTypeAnalyzer(PLANNER_CONTEXT), p.getTypes());
    assertTrue(rewritten.isPresent());
    assertPlan(session, dummyMetadata(), createTestingFunctionManager(), node -> unknown(), new Plan(rewritten.get(), p.getTypes(), empty()), noLookup(), join(INNER, ImmutableList.of(aliases -> new JoinNode.EquiJoinClause(new Symbol("a1"), new Symbol("b1"))), strictProject(ImmutableMap.of("a3", expression("-(+a0)"), "a1", expression("a1")), strictProject(ImmutableMap.of("a0", expression("a0"), "a1", expression("a1")), PlanMatchPattern.values("a0", "a1"))), strictProject(ImmutableMap.of("b2", expression("+b1"), "b1", expression("b1")), PlanMatchPattern.values("b0", "b1"))).withExactOutputs("a3", "b2"));
}