Search in sources :

Example 1 with Session

use of io.prestosql.Session in project hetu-core by openlookeng.

the class TestDynamicFilter method testJoinExceedMaxSize.

@Test
public void testJoinExceedMaxSize() {
    Session session = Session.builder(this.getQueryRunner().getDefaultSession()).setSystemProperty(DYNAMIC_FILTERING_MAX_SIZE, "10").build();
    assertPlanWithSession("SELECT o.orderkey FROM orders o, lineitem l WHERE l.orderkey = o.orderkey AND l.partkey < 100", session, true, anyTree(join(INNER, ImmutableList.of(equiJoinClause("ORDERS_OK", "LINEITEM_OK")), Optional.empty(), project(tableScan("orders", ImmutableMap.of("ORDERS_OK", "orderkey"))), exchange(project(node(FilterNode.class, tableScan("lineitem", ImmutableMap.of("LINEITEM_OK", "orderkey"))))))));
}
Also used : FilterNode(io.prestosql.spi.plan.FilterNode) Session(io.prestosql.Session) Test(org.testng.annotations.Test) BasePlanTest(io.prestosql.sql.planner.assertions.BasePlanTest)

Example 2 with Session

use of io.prestosql.Session in project hetu-core by openlookeng.

the class TestDynamicFiltersCollector method TestCollectingGlobalDynamicFilters.

@Test
public void TestCollectingGlobalDynamicFilters() throws InterruptedException {
    final QueryId queryId = new QueryId("test_query");
    final String filterId = "1";
    final String columnName = "column";
    final TestingColumnHandle columnHandle = new TestingColumnHandle(columnName);
    final Set<String> valueSet = ImmutableSet.of("1", "2", "3");
    TaskContext taskContext = mock(TaskContext.class);
    Session session = testSessionBuilder().setQueryId(queryId).setSystemProperty(ENABLE_DYNAMIC_FILTERING, "true").setSystemProperty(DYNAMIC_FILTERING_DATA_TYPE, "HASHSET").build();
    when(taskContext.getSession()).thenReturn(session);
    // set up state store and merged dynamic filters map
    Map mockMap = new HashMap<>();
    StateStoreProvider stateStoreProvider = mock(StateStoreProvider.class);
    StateStore stateStore = mock(StateStore.class);
    StateMap stateMap = new MockStateMap<>("test-map", mockMap);
    when(stateStoreProvider.getStateStore()).thenReturn(stateStore);
    when(stateStore.getStateCollection(any())).thenReturn(stateMap);
    when(stateStore.createStateMap(any())).thenReturn(stateMap);
    when(stateStore.getOrCreateStateCollection(any(), any())).thenReturn(stateMap);
    // set up state store listener and dynamic filter cache
    StateStoreListenerManager stateStoreListenerManager = new StateStoreListenerManager(stateStoreProvider);
    DynamicFilterCacheManager dynamicFilterCacheManager = new DynamicFilterCacheManager();
    stateStoreListenerManager.addStateStoreListener(new DynamicFilterListener(dynamicFilterCacheManager), MERGED_DYNAMIC_FILTERS);
    LocalDynamicFiltersCollector collector = new LocalDynamicFiltersCollector(taskContext, Optional.empty(), dynamicFilterCacheManager);
    TableScanNode tableScan = mock(TableScanNode.class);
    when(tableScan.getAssignments()).thenReturn(ImmutableMap.of(new Symbol(columnName), columnHandle));
    List<DynamicFilters.Descriptor> dynamicFilterDescriptors = ImmutableList.of(new DynamicFilters.Descriptor(filterId, new VariableReferenceExpression(columnName, BIGINT)));
    collector.initContext(ImmutableList.of(dynamicFilterDescriptors), SymbolUtils.toLayOut(tableScan.getOutputSymbols()));
    assertTrue(collector.getDynamicFilters(tableScan).isEmpty(), "there should be no dynamic filter available");
    // put some values in state store as a new dynamic filter
    // and wait for the listener to process the event
    stateMap.put(createKey(DynamicFilterUtils.FILTERPREFIX, filterId, queryId.getId()), valueSet);
    TimeUnit.MILLISECONDS.sleep(100);
    // get available dynamic filter and verify it
    List<Map<ColumnHandle, DynamicFilter>> dynamicFilters = collector.getDynamicFilters(tableScan);
    assertEquals(dynamicFilters.size(), 1, "there should be a new dynamic filter");
    assertEquals(dynamicFilters.size(), 1);
    DynamicFilter dynamicFilter = dynamicFilters.get(0).get(columnHandle);
    assertTrue(dynamicFilter instanceof HashSetDynamicFilter, "new dynamic filter should be hashset");
    assertEquals(dynamicFilter.getSize(), valueSet.size(), "new dynamic filter should have correct size");
    for (String value : valueSet) {
        assertTrue(dynamicFilter.contains(value), "new dynamic filter should contain correct values");
    }
    // clean up when task finishes
    collector.removeDynamicFilter(true);
    DynamicFilter cachedFilter = dynamicFilterCacheManager.getDynamicFilter(DynamicFilterCacheManager.createCacheKey(filterId, queryId.getId()));
    assertNull(cachedFilter, "cached dynamic filter should have been removed");
}
Also used : HashMap(java.util.HashMap) Symbol(io.prestosql.spi.plan.Symbol) MockStateMap(io.prestosql.statestore.MockStateMap) StateMap(io.prestosql.spi.statestore.StateMap) StateStoreListenerManager(io.prestosql.statestore.listener.StateStoreListenerManager) StateStoreProvider(io.prestosql.statestore.StateStoreProvider) MockStateMap(io.prestosql.statestore.MockStateMap) TestingColumnHandle(io.prestosql.spi.connector.TestingColumnHandle) DynamicFilters(io.prestosql.sql.DynamicFilters) DynamicFilterListener(io.prestosql.dynamicfilter.DynamicFilterListener) DynamicFilterCacheManager(io.prestosql.dynamicfilter.DynamicFilterCacheManager) TaskContext(io.prestosql.operator.TaskContext) DynamicFilter(io.prestosql.spi.dynamicfilter.DynamicFilter) HashSetDynamicFilter(io.prestosql.spi.dynamicfilter.HashSetDynamicFilter) QueryId(io.prestosql.spi.QueryId) StateStore(io.prestosql.spi.statestore.StateStore) HashSetDynamicFilter(io.prestosql.spi.dynamicfilter.HashSetDynamicFilter) TableScanNode(io.prestosql.spi.plan.TableScanNode) VariableReferenceExpression(io.prestosql.spi.relation.VariableReferenceExpression) MockStateMap(io.prestosql.statestore.MockStateMap) HashMap(java.util.HashMap) Map(java.util.Map) ImmutableMap(com.google.common.collect.ImmutableMap) StateMap(io.prestosql.spi.statestore.StateMap) Session(io.prestosql.Session) Test(org.testng.annotations.Test)

Example 3 with Session

use of io.prestosql.Session in project hetu-core by openlookeng.

the class TestColumnMask method testView.

@Test
public void testView() {
    // mask on the underlying table for view owner when running query as different user
    assertions.executeExclusively(() -> {
        accessControl.reset();
        accessControl.columnMask(new QualifiedObjectName(CATALOG, "tiny", "nation"), "name", VIEW_OWNER, new ViewExpression(VIEW_OWNER, Optional.empty(), Optional.empty(), "reverse(name)"));
        Session session = Session.builder(SESSION).setIdentity(new Identity(RUN_AS_USER, Optional.empty())).build();
        assertions.assertQuery(session, "SELECT name FROM mock.default.nation_view WHERE nationkey = 1", "VALUES CAST('ANITNEGRA' AS VARCHAR(25))");
    });
    // mask on the underlying table for view owner when running as themselves
    assertions.executeExclusively(() -> {
        accessControl.reset();
        accessControl.columnMask(new QualifiedObjectName(CATALOG, "tiny", "nation"), "name", VIEW_OWNER, new ViewExpression(VIEW_OWNER, Optional.of(CATALOG), Optional.of("tiny"), "reverse(name)"));
        Session session = Session.builder(SESSION).setIdentity(new Identity(VIEW_OWNER, Optional.empty())).build();
        assertions.assertQuery(session, "SELECT name FROM mock.default.nation_view WHERE nationkey = 1", "VALUES CAST('ANITNEGRA' AS VARCHAR(25))");
    });
    // mask on the underlying table for user running the query (different from view owner) should not be applied
    assertions.executeExclusively(() -> {
        accessControl.reset();
        accessControl.columnMask(new QualifiedObjectName(CATALOG, "tiny", "nation"), "name", RUN_AS_USER, new ViewExpression(RUN_AS_USER, Optional.of(CATALOG), Optional.of("tiny"), "reverse(name)"));
        Session session = Session.builder(SESSION).setIdentity(new Identity(RUN_AS_USER, Optional.empty())).build();
        assertions.assertQuery(session, "SELECT name FROM mock.default.nation_view WHERE nationkey = 1", "VALUES CAST('ARGENTINA' AS VARCHAR(25))");
    });
}
Also used : Identity(io.prestosql.spi.security.Identity) QualifiedObjectName(io.prestosql.spi.connector.QualifiedObjectName) ViewExpression(io.prestosql.spi.security.ViewExpression) Session(io.prestosql.Session) Test(org.testng.annotations.Test)

Example 4 with Session

use of io.prestosql.Session in project hetu-core by openlookeng.

the class TestColumnMask method init.

@BeforeClass
public void init() {
    LocalQueryRunner runner = new LocalQueryRunner(SESSION);
    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.getTypeSignature()), new ConnectorViewDefinition.ViewColumn("name", VarcharType.createVarcharType(25).getTypeSignature())), Optional.of(VIEW_OWNER), false);
    MockConnectorFactory mock = MockConnectorFactory.builder().withGetViews((s, prefix) -> ImmutableMap.<SchemaTableName, ConnectorViewDefinition>builder().put(new SchemaTableName("default", "nation_view"), view).build()).build();
    runner.createCatalog(MOCK_CATALOG, mock, ImmutableMap.of());
    assertions = new QueryAssertions(runner);
    accessControl = assertions.getQueryRunner().getAccessControl();
}
Also used : BigintType(io.prestosql.spi.type.BigintType) AfterClass(org.testng.annotations.AfterClass) ImmutableMap(com.google.common.collect.ImmutableMap) TINY_SCHEMA_NAME(io.prestosql.plugin.tpch.TpchMetadata.TINY_SCHEMA_NAME) ConnectorViewDefinition(io.prestosql.spi.connector.ConnectorViewDefinition) BeforeClass(org.testng.annotations.BeforeClass) TpchConnectorFactory(io.prestosql.plugin.tpch.TpchConnectorFactory) ViewExpression(io.prestosql.spi.security.ViewExpression) Identity(io.prestosql.spi.security.Identity) Test(org.testng.annotations.Test) MockConnectorFactory(io.prestosql.connector.MockConnectorFactory) QualifiedObjectName(io.prestosql.spi.connector.QualifiedObjectName) TestingSession.testSessionBuilder(io.prestosql.testing.TestingSession.testSessionBuilder) TestingAccessControlManager(io.prestosql.testing.TestingAccessControlManager) SchemaTableName(io.prestosql.spi.connector.SchemaTableName) LocalQueryRunner(io.prestosql.testing.LocalQueryRunner) ImmutableList(com.google.common.collect.ImmutableList) Session(io.prestosql.Session) Optional(java.util.Optional) VarcharType(io.prestosql.spi.type.VarcharType) TpchConnectorFactory(io.prestosql.plugin.tpch.TpchConnectorFactory) MockConnectorFactory(io.prestosql.connector.MockConnectorFactory) SchemaTableName(io.prestosql.spi.connector.SchemaTableName) LocalQueryRunner(io.prestosql.testing.LocalQueryRunner) ConnectorViewDefinition(io.prestosql.spi.connector.ConnectorViewDefinition) BeforeClass(org.testng.annotations.BeforeClass)

Example 5 with Session

use of io.prestosql.Session in project hetu-core by openlookeng.

the class TestQueuesDb method assertResourceGroupWithClientTags.

private void assertResourceGroupWithClientTags(Set<String> clientTags, ResourceGroupId expectedResourceGroup) throws InterruptedException {
    Session session = testSessionBuilder().setCatalog("tpch").setSchema("sf100000").setSource("client_tags").setClientTags(clientTags).build();
    QueryId queryId = createQuery(queryRunner, session, LONG_LASTING_QUERY);
    waitForQueryState(queryRunner, queryId, ImmutableSet.of(RUNNING, FINISHED));
    Optional<ResourceGroupId> resourceGroupId = queryRunner.getCoordinator().getQueryManager().getFullQueryInfo(queryId).getResourceGroupId();
    assertTrue(resourceGroupId.isPresent(), "Query should have a resource group");
    assertEquals(resourceGroupId.get(), expectedResourceGroup, format("Expected: '%s' resource group, found: %s", expectedResourceGroup, resourceGroupId.get()));
}
Also used : ResourceGroupId(io.prestosql.spi.resourcegroups.ResourceGroupId) TestQueues.createResourceGroupId(io.prestosql.execution.TestQueues.createResourceGroupId) QueryId(io.prestosql.spi.QueryId) H2TestUtil.adhocSession(io.prestosql.execution.resourcegroups.db.H2TestUtil.adhocSession) H2TestUtil.dashboardSession(io.prestosql.execution.resourcegroups.db.H2TestUtil.dashboardSession) H2TestUtil.rejectingSession(io.prestosql.execution.resourcegroups.db.H2TestUtil.rejectingSession) Session(io.prestosql.Session)

Aggregations

Session (io.prestosql.Session)254 Test (org.testng.annotations.Test)131 List (java.util.List)52 Optional (java.util.Optional)50 Map (java.util.Map)43 Metadata (io.prestosql.metadata.Metadata)41 QualifiedObjectName (io.prestosql.spi.connector.QualifiedObjectName)41 MaterializedResult (io.prestosql.testing.MaterializedResult)40 Objects.requireNonNull (java.util.Objects.requireNonNull)39 ConnectorSession (io.prestosql.spi.connector.ConnectorSession)35 PrestoException (io.prestosql.spi.PrestoException)31 TableHandle (io.prestosql.spi.metadata.TableHandle)31 Set (java.util.Set)31 ImmutableList (com.google.common.collect.ImmutableList)29 Type (io.prestosql.spi.type.Type)29 HashMap (java.util.HashMap)29 HeuristicIndexerManager (io.prestosql.heuristicindex.HeuristicIndexerManager)27 Symbol (io.prestosql.spi.plan.Symbol)27 ImmutableList.toImmutableList (com.google.common.collect.ImmutableList.toImmutableList)26 ImmutableSet (com.google.common.collect.ImmutableSet)26