use of datawave.microservice.querymetric.QueryMetricFactoryImpl in project datawave by NationalSecurityAgency.
the class ExtendedQueryExecutorBeanTest method testEnableTracing_NullChecksAndHappyPath.
@Test
public void testEnableTracing_NullChecksAndHappyPath() throws Exception {
// Set local test input
String user = "user";
String queryRegex = "queryRegex";
PowerMock.resetAll();
// Set expectations
expect(traceInfos.containsEntry(user, PatternWrapper.wrap(queryRegex))).andReturn(false);
expect(traceInfos.put(user, PatternWrapper.wrap(queryRegex))).andReturn(true);
traceCache.put("traceInfos", traceInfos);
// Run the test
PowerMock.replayAll();
QueryExecutorBean subject = new QueryExecutorBean();
setInternalState(subject, Multimap.class, traceInfos);
setInternalState(subject, QueryTraceCache.class, traceCache);
setInternalState(subject, QueryMetricFactory.class, new QueryMetricFactoryImpl());
Exception result1 = null;
try {
subject.enableTracing(null, null);
} catch (BadRequestException e) {
result1 = e;
}
VoidResponse result2 = subject.enableTracing(queryRegex, user);
PowerMock.verifyAll();
// Verify results
assertNotNull("Expected a BadRequestException due to null regex and user values", result1);
assertNotNull("Expected a non-null response", result2);
}
use of datawave.microservice.querymetric.QueryMetricFactoryImpl in project datawave by NationalSecurityAgency.
the class ExtendedQueryExecutorBeanTest method testList_HappyPath.
@SuppressWarnings({ "unchecked", "rawtypes" })
@Test
public void testList_HappyPath() throws Exception {
// Set local test input
String queryName = "queryName";
String queryLogicName = "queryLogicName";
String userName = "userName";
String sid = "sid";
UUID queryId = UUID.randomUUID();
List<String> dnList = Collections.singletonList("userDn");
HashMap<String, Collection<String>> authsMap = new HashMap<>();
authsMap.put("USERDN", Arrays.asList("AUTH_1"));
// Set expectations
expect(this.context.getCallerPrincipal()).andReturn(this.principal);
expect(this.principal.getName()).andReturn(userName);
expect(this.principal.getShortName()).andReturn(sid);
expect(this.principal.getPrimaryUser()).andReturn(dwUser);
expect(this.dwUser.getAuths()).andReturn(Collections.singletonList("AUTH_1"));
expect(this.principal.getProxiedUsers()).andReturn(Collections.singletonList(dwUser));
expect(this.principal.getAuthorizations()).andReturn((Collection) Arrays.asList(Arrays.asList("AUTH_1")));
expect(this.persister.findByName(queryName)).andReturn(Arrays.asList((Query) this.query));
expect(this.query.getOwner()).andReturn(sid).anyTimes();
expect(this.query.getUserDN()).andReturn(sid).anyTimes();
expect(this.query.getQueryLogicName()).andReturn(queryLogicName).anyTimes();
expect(this.queryLogicFactory.getQueryLogic(queryLogicName, principal)).andReturn((QueryLogic) this.queryLogic1);
expect(this.queryLogic1.getConnectionPriority()).andReturn(Priority.HIGH);
expect(this.query.getQueryAuthorizations()).andReturn(null).anyTimes();
expect(this.queryLogic1.getCollectQueryMetrics()).andReturn(false);
expect(this.query.getId()).andReturn(queryId).anyTimes();
expect(this.query.getQuery()).andReturn(queryName).anyTimes();
expect(this.query.getBeginDate()).andReturn(null).anyTimes();
expect(this.query.getEndDate()).andReturn(null).anyTimes();
expect(this.query.getColumnVisibility()).andReturn(null).anyTimes();
expect(this.cache.containsKey(queryId.toString())).andReturn(false);
expect(this.query.getQueryName()).andReturn(null).anyTimes();
expect(this.query.getPagesize()).andReturn(0).anyTimes();
expect(this.query.getPageTimeout()).andReturn(-1).anyTimes();
expect(this.query.getExpirationDate()).andReturn(null).anyTimes();
expect(this.query.getParameters()).andReturn((Set) Collections.emptySet()).anyTimes();
expect(this.query.getDnList()).andReturn(dnList).anyTimes();
expect(this.queryLogic1.getResultLimit(dnList)).andReturn(-1L);
expect(this.queryLogic1.getMaxResults()).andReturn(-1L);
this.cache.put(eq(queryId.toString()), isA(RunningQuery.class));
// Run the test
PowerMock.replayAll();
QueryExecutorBean subject = new QueryExecutorBean();
setInternalState(subject, EJBContext.class, context);
setInternalState(subject, QueryCache.class, cache);
setInternalState(subject, ClosedQueryCache.class, closedCache);
setInternalState(subject, Persister.class, persister);
setInternalState(subject, QueryLogicFactory.class, queryLogicFactory);
setInternalState(subject, QueryExpirationConfiguration.class, queryExpirationConf);
setInternalState(subject, QueryParameters.class, new QueryParametersImpl());
setInternalState(subject, QueryMetricFactory.class, new QueryMetricFactoryImpl());
QueryImplListResponse result1 = subject.list(queryName);
PowerMock.verifyAll();
// Verify results
assertNotNull("QueryLogicResponse should not be returned null", result1);
}
use of datawave.microservice.querymetric.QueryMetricFactoryImpl in project datawave by NationalSecurityAgency.
the class QueryExecutorBeanTest method testPredict.
@SuppressWarnings("unchecked")
@Test
public void testPredict() throws Exception {
QueryImpl q = createNewQuery();
MultivaluedMap p = createNewQueryParameterMap();
p.putSingle(QueryParameters.QUERY_LOGIC_NAME, queryLogicName);
MultivaluedMap<String, String> optionalParameters = createNewQueryParameters(q, p);
@SuppressWarnings("rawtypes") QueryLogic logic = createMock(BaseQueryLogic.class);
DatawaveUser user = new DatawaveUser(SubjectIssuerDNPair.of(userDN, "<CN=MY_CA, OU=MY_SUBDIVISION, OU=MY_DIVISION, O=ORG, C=US>"), UserType.USER, Arrays.asList(auths), null, null, 0L);
DatawavePrincipal principal = new DatawavePrincipal(Collections.singletonList(user));
String[] dns = principal.getDNs();
Arrays.sort(dns);
List<String> dnList = Arrays.asList(dns);
PowerMock.resetAll();
EasyMock.expect(ctx.getCallerPrincipal()).andReturn(principal).anyTimes();
suppress(constructor(QueryParametersImpl.class));
EasyMock.expect(persister.create(principal.getUserDN().subjectDN(), dnList, (SecurityMarking) Whitebox.getField(bean.getClass(), "marking").get(bean), queryLogicName, (QueryParameters) Whitebox.getField(bean.getClass(), "qp").get(bean), optionalParameters)).andReturn(q);
EasyMock.expect(queryLogicFactory.getQueryLogic(queryLogicName, principal)).andReturn(logic);
EasyMock.expect(logic.getRequiredQueryParameters()).andReturn(Collections.EMPTY_SET);
EasyMock.expect(logic.containsDNWithAccess(dnList)).andReturn(true);
EasyMock.expect(logic.getMaxPageSize()).andReturn(0);
BaseQueryMetric metric = new QueryMetricFactoryImpl().createMetric();
metric.populate(q);
metric.setQueryType(RunningQuery.class.getSimpleName());
QueryMetric testMetric = new QueryMetric((QueryMetric) metric) {
@Override
public boolean equals(Object o) {
// test for equality except for the create date
if (null == o) {
return false;
}
if (this == o) {
return true;
}
if (o instanceof QueryMetric) {
QueryMetric other = (QueryMetric) o;
return new EqualsBuilder().append(this.getQueryId(), other.getQueryId()).append(this.getQueryType(), other.getQueryType()).append(this.getQueryAuthorizations(), other.getQueryAuthorizations()).append(this.getColumnVisibility(), other.getColumnVisibility()).append(this.getBeginDate(), other.getBeginDate()).append(this.getEndDate(), other.getEndDate()).append(this.getCreateDate(), other.getCreateDate()).append(this.getSetupTime(), other.getSetupTime()).append(this.getUser(), other.getUser()).append(this.getUserDN(), other.getUserDN()).append(this.getQuery(), other.getQuery()).append(this.getQueryLogic(), other.getQueryLogic()).append(this.getQueryName(), other.getQueryName()).append(this.getParameters(), other.getParameters()).append(this.getHost(), other.getHost()).append(this.getPageTimes(), other.getPageTimes()).append(this.getProxyServers(), other.getProxyServers()).append(this.getLifecycle(), other.getLifecycle()).append(this.getErrorMessage(), other.getErrorMessage()).append(this.getErrorCode(), other.getErrorCode()).append(this.getSourceCount(), other.getSourceCount()).append(this.getNextCount(), other.getNextCount()).append(this.getSeekCount(), other.getSeekCount()).append(this.getYieldCount(), other.getYieldCount()).append(this.getDocRanges(), other.getDocRanges()).append(this.getFiRanges(), other.getFiRanges()).append(this.getPlan(), other.getPlan()).append(this.getVersion(), other.getVersion()).append(this.getLoginTime(), other.getLoginTime()).append(this.getPredictions(), other.getPredictions()).isEquals();
} else {
return false;
}
}
};
Set<Prediction> predictions = new HashSet<>();
predictions.add(new Prediction("source", 1));
EasyMock.expect(predictor.predict(EasyMock.eq(testMetric))).andReturn(predictions);
PowerMock.replayAll();
GenericResponse<String> response = bean.predictQuery(queryLogicName, p);
PowerMock.verifyAll();
Object cachedRunningQuery = cache.get(q.getId().toString());
Assert.assertNull(cachedRunningQuery);
Assert.assertEquals(predictions.toString(), response.getResult());
}
use of datawave.microservice.querymetric.QueryMetricFactoryImpl in project datawave by NationalSecurityAgency.
the class QueryExecutorBeanTest method setup.
@Before
public void setup() throws Exception {
System.setProperty(NpeUtils.NPE_OU_PROPERTY, "iamnotaperson");
System.setProperty("dw.metadatahelper.all.auths", "A,B,C,D");
QueryTraceCache traceCache = new QueryTraceCache();
Whitebox.invokeMethod(traceCache, "init");
cache = new QueryCache();
cache.init();
closedCache = new ClosedQueryCache();
bean = new QueryExecutorBean();
connectionFactory = createStrictMock(AccumuloConnectionFactory.class);
auditor = new AuditBean();
auditService = createStrictMock(AuditService.class);
metrics = createStrictMock(QueryMetricsBean.class);
queryLogicFactory = createStrictMock(QueryLogicFactoryImpl.class);
persister = createStrictMock(Persister.class);
predictor = createStrictMock(QueryPredictor.class);
ctx = createStrictMock(EJBContext.class);
qlCache = new CreatedQueryLogicCacheBean();
queryExpirationConf = new QueryExpirationConfiguration();
queryExpirationConf.setPageSizeShortCircuitCheckTime(45);
queryExpirationConf.setPageShortCircuitTimeout(58);
queryExpirationConf.setCallTime(60);
connectionRequestBean = createStrictMock(AccumuloConnectionRequestBean.class);
setInternalState(auditor, AuditService.class, auditService);
setInternalState(auditor, AuditParameterBuilder.class, new DefaultAuditParameterBuilder());
setInternalState(connectionRequestBean, EJBContext.class, ctx);
setInternalState(bean, QueryCache.class, cache);
setInternalState(bean, ClosedQueryCache.class, closedCache);
setInternalState(bean, AccumuloConnectionFactory.class, connectionFactory);
setInternalState(bean, AuditBean.class, auditor);
setInternalState(bean, QueryMetricsBean.class, metrics);
setInternalState(bean, QueryLogicFactory.class, queryLogicFactory);
setInternalState(bean, QueryExpirationConfiguration.class, queryExpirationConf);
setInternalState(bean, Persister.class, persister);
setInternalState(bean, QueryPredictor.class, predictor);
setInternalState(bean, EJBContext.class, ctx);
setInternalState(bean, CreatedQueryLogicCacheBean.class, qlCache);
setInternalState(bean, QueryTraceCache.class, traceCache);
setInternalState(bean, Multimap.class, HashMultimap.create());
setInternalState(bean, LookupUUIDConfiguration.class, new LookupUUIDConfiguration());
setInternalState(bean, SecurityMarking.class, new ColumnVisibilitySecurityMarking());
setInternalState(bean, QueryParameters.class, new QueryParametersImpl());
setInternalState(bean, QueryMetricFactory.class, new QueryMetricFactoryImpl());
setInternalState(bean, AccumuloConnectionRequestBean.class, connectionRequestBean);
// RESTEasy mock stuff
dispatcher = MockDispatcherFactory.createDispatcher();
dispatcher.getRegistry().addSingletonResource(bean, "/DataWave/Query");
response = new MockHttpResponse();
}
use of datawave.microservice.querymetric.QueryMetricFactoryImpl in project datawave by NationalSecurityAgency.
the class QueryExecutorBeanTest method testCloseActuallyCloses.
@SuppressWarnings("unchecked")
@Test(timeout = 5000)
public void testCloseActuallyCloses() throws Exception {
QueryImpl q = createNewQuery();
final MultivaluedMap<String, String> queryParameters = createNewQueryParameterMap();
queryParameters.putSingle(QueryParameters.QUERY_LOGIC_NAME, "EventQueryLogic");
final Thread createQuery = new Thread(() -> {
try {
bean.createQuery("EventQueryLogic", queryParameters);
} catch (Exception e) {
// ok if we fail the call
log.debug("createQuery terminated with " + e);
}
});
final Throwable[] createQueryException = { null };
createQuery.setUncaughtExceptionHandler((t, e) -> createQueryException[0] = e);
@SuppressWarnings("rawtypes") QueryLogic logic = createMock(BaseQueryLogic.class);
DatawaveUser user = new DatawaveUser(SubjectIssuerDNPair.of(userDN, "<CN=MY_CA, OU=MY_SUBDIVISION, OU=MY_DIVISION, O=ORG, C=US>"), UserType.USER, Arrays.asList(auths), null, null, 0L);
DatawavePrincipal principal = new DatawavePrincipal(Collections.singletonList(user));
principal.getShortName();
String[] dns = principal.getDNs();
Arrays.sort(dns);
List<String> dnList = Arrays.asList(dns);
InMemoryInstance instance = new InMemoryInstance();
Connector c = instance.getConnector("root", new PasswordToken(""));
MultivaluedMap<String, String> optionalParameters = createNewQueryParameters(q, queryParameters);
PowerMock.resetAll();
EasyMock.expect(ctx.getCallerPrincipal()).andReturn(principal).anyTimes();
EasyMock.expect(logic.getAuditType(null)).andReturn(AuditType.NONE);
EasyMock.expect(persister.create(principal.getUserDN().subjectDN(), dnList, Whitebox.getInternalState(bean, SecurityMarking.class), queryLogicName, Whitebox.getInternalState(bean, QueryParameters.class), optionalParameters)).andReturn(q);
EasyMock.expect(persister.findById(EasyMock.anyString())).andReturn(null).anyTimes();
EasyMock.expect(connectionFactory.getTrackingMap(anyObject())).andReturn(Maps.newHashMap()).anyTimes();
BaseQueryMetric metric = new QueryMetricFactoryImpl().createMetric();
metric.populate(q);
EasyMock.expectLastCall();
metric.setQueryType(RunningQuery.class.getSimpleName());
metric.setLifecycle(Lifecycle.DEFINED);
System.out.println(metric);
Set<Prediction> predictions = new HashSet<>();
predictions.add(new Prediction("source", 1));
EasyMock.expect(predictor.predict(metric)).andReturn(predictions);
connectionRequestBean.requestBegin(q.getId().toString());
EasyMock.expectLastCall();
EasyMock.expect(connectionFactory.getConnection(eq("connPool1"), anyObject(), anyObject())).andReturn(c).anyTimes();
connectionRequestBean.requestEnd(q.getId().toString());
EasyMock.expectLastCall();
connectionFactory.returnConnection(c);
EasyMock.expectLastCall();
EasyMock.expect(queryLogicFactory.getQueryLogic(queryLogicName, principal)).andReturn(logic);
EasyMock.expect(logic.getRequiredQueryParameters()).andReturn(Collections.emptySet());
EasyMock.expect(logic.getConnectionPriority()).andReturn(AccumuloConnectionFactory.Priority.NORMAL).atLeastOnce();
EasyMock.expect(logic.containsDNWithAccess(dnList)).andReturn(true);
EasyMock.expect(logic.getMaxPageSize()).andReturn(0);
EasyMock.expect(logic.getAuditType(q)).andReturn(AuditType.NONE);
EasyMock.expect(logic.getConnPoolName()).andReturn("connPool1");
EasyMock.expect(logic.getResultLimit(eq(q.getDnList()))).andReturn(-1L).anyTimes();
EasyMock.expect(logic.getMaxResults()).andReturn(-1L).anyTimes();
EasyMock.expect(connectionRequestBean.cancelConnectionRequest(q.getId().toString(), principal)).andReturn(false).anyTimes();
connectionFactory.returnConnection(EasyMock.isA(Connector.class));
final AtomicBoolean initializeLooping = new AtomicBoolean(false);
// During initialize, mark that we get here, and then sleep
final IAnswer<GenericQueryConfiguration> initializeAnswer = () -> {
initializeLooping.set(true);
try {
while (true) {
Thread.sleep(1000);
log.debug("Initialize: woke up");
}
} catch (InterruptedException e) {
throw new QueryException("EXPECTED EXCEPTION: initialize interrupted");
}
};
EasyMock.expect(logic.initialize(anyObject(Connector.class), anyObject(Query.class), anyObject(Set.class))).andAnswer(initializeAnswer);
EasyMock.expect(logic.getCollectQueryMetrics()).andReturn(Boolean.FALSE);
// On close, interrupt the thread to simulate the ScannerFactory cleaning up
final IAnswer<Object> closeAnswer = () -> {
if (null != createQuery) {
log.debug("createQuery thread is not null. interrupting");
createQuery.interrupt();
} else {
log.debug("createQuery thread is null. not interrupting");
}
return null;
};
logic.close();
EasyMock.expectLastCall().andAnswer(closeAnswer).anyTimes();
// Make the QueryLogic mock not threadsafe, otherwise it will be blocked infinitely
// trying to get the lock on the infinite loop
EasyMock.makeThreadSafe(logic, false);
metrics.updateMetric(EasyMock.isA(QueryMetric.class));
PowerMock.replayAll();
try {
createQuery.start();
// Wait for the create call to get to initialize
while (!initializeLooping.get()) {
if (!createQuery.isAlive() && !initializeLooping.get()) {
Assert.fail("createQuery thread died before reaching initialize: " + createQueryException[0]);
}
Thread.sleep(50);
}
// initialize has not completed yet so it will not appear in the cache
Object cachedRunningQuery = cache.get(q.getId().toString());
Assert.assertNull(cachedRunningQuery);
Pair<QueryLogic<?>, Connector> pair = qlCache.poll(q.getId().toString());
Assert.assertNotNull(pair);
Assert.assertEquals(logic, pair.getFirst());
Assert.assertEquals(c, pair.getSecond());
// Have to add these back because poll was destructive
qlCache.add(q.getId().toString(), principal.getShortName(), pair.getFirst(), pair.getSecond());
// Call close
bean.close(q.getId().toString());
// Make sure that it's gone from the qlCache
pair = qlCache.poll(q.getId().toString());
Assert.assertNull("Still found an entry in the qlCache: " + pair, pair);
// Should have already joined by now, but just to be sure
createQuery.join();
} finally {
if (null != createQuery && createQuery.isAlive()) {
createQuery.interrupt();
}
}
}
Aggregations