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"))))))));
}
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");
}
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))");
});
}
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();
}
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()));
}
Aggregations