/** * To facilitate generating workunits * with the assumption that QueryContext is NOT null * it's not always going to be true, since e.g. QueryContextInfo * may change between ctor and this method * @param options * @param reader * @param rootNode * @param planningSet * @return * @throws ExecutionSetupException */ private List<PlanFragment> generateWorkUnit( OptionList options, PhysicalPlanReader reader, Fragment rootNode, PlanningSet planningSet) throws ExecutionSetupException { Preconditions.checkNotNull(queryContext); return generateWorkUnit(options, queryContext.getCurrentEndpoint(), queryContext.getQueryId(), reader, rootNode, planningSet, queryContext.getSession(), queryContext.getQueryContextInfo(), queryContext.getFunctionRegistry()); }
protected QueryContext mockQueryContext(SabotContext dbContext) throws Exception { final UserSession userSession = UserSession.Builder.newBuilder().withOptionManager(dbContext.getOptionManager()).build(); final SessionOptionManager sessionOptions = (SessionOptionManager) userSession.getOptions(); final QueryOptionManager queryOptions = new QueryOptionManager(sessionOptions); final ExecutionControls executionControls = new ExecutionControls(queryOptions, NodeEndpoint.getDefaultInstance()); final OperatorTable table = new OperatorTable(FUNCTIONS()); final LogicalPlanPersistence lp = dbContext.getLpPersistence(); final CatalogService registry = dbContext.getCatalogService(); final QueryContext context = Mockito.mock(QueryContext.class); when(context.getSession()).thenReturn(userSession); when(context.getLpPersistence()).thenReturn(lp); when(context.getCatalogService()).thenReturn(registry); when(context.getFunctionRegistry()).thenReturn(FUNCTIONS()); when(context.getSession()).thenReturn(UserSession.Builder.newBuilder().setSupportComplexTypes(true).build()); when(context.getCurrentEndpoint()).thenReturn(NodeEndpoint.getDefaultInstance()); when(context.getActiveEndpoints()).thenReturn(ImmutableList.of(NodeEndpoint.getDefaultInstance())); when(context.getPlannerSettings()).thenReturn(new PlannerSettings(dbContext.getConfig(), queryOptions, dbContext.getClusterResourceInformation())); when(context.getOptions()).thenReturn(queryOptions); when(context.getConfig()).thenReturn(DEFAULT_SABOT_CONFIG); when(context.getOperatorTable()).thenReturn(table); when(context.getAllocator()).thenReturn(allocator); when(context.getExecutionControls()).thenReturn(executionControls); when(context.getMaterializationProvider()).thenReturn(Mockito.mock(MaterializationDescriptorProvider.class)); return context; }
uex = UserException.systemError(resultException).addIdentity(queryContext.getCurrentEndpoint()).build(logger); } else { uex = null;
queryContext.getCurrentEndpoint(), queryContext.getQueryId(), endpoints,
@Test public void testResourceSchedulingInProfile() throws Exception { AttemptObservers observers = AttemptObservers.of(observer); final NodeEndpoint endpoint = NodeEndpoint.newBuilder().setAddress("host1").setFabricPort(12345).build(); when(context.getCurrentEndpoint()).thenReturn(endpoint); QueryManager queryManager = new QueryManager(queryId, context, null, completionListener, new Pointer<>(), observers, true, true, catalog); ResourceSchedulingDecisionInfo result = new ResourceSchedulingDecisionInfo(); result.setQueueId("abcd"); result.setQueueName("queue.abcd"); observers.resourcesScheduled(result); UserBitShared.QueryProfile queryProfile = queryManager.getQueryProfile("my description", UserBitShared.QueryResult.QueryState.RUNNING, null, "some reason"); assertNotNull(queryProfile.getResourceSchedulingProfile()); assertEquals("abcd", queryProfile.getResourceSchedulingProfile().getQueueId()); assertEquals("queue.abcd", queryProfile.getResourceSchedulingProfile().getQueueName()); } }
.setId(queryId) .setState(state) .setForeman(context.getCurrentEndpoint()) .setStart(startTime) .setEnd(endTime)