Examples with Prediction datawave.microservice.querymetric.BaseQueryMetric.Prediction used on opensource projects

Example 1 with Prediction

use of datawave.microservice.querymetric.BaseQueryMetric.Prediction in project datawave by NationalSecurityAgency.

the class RunningQuery method applyPrediction.

protected void applyPrediction(String context) {
    if (getPredictor() != null) {
        try {
            Set<Prediction> predictions = getPredictor().predict(this.getMetric());
            if (predictions != null) {
                // now we have a predictions, lets broadcast
                log.info("Model predictions: " + predictions);
                context = (context == null ? "" : context + ' ');
                for (Prediction prediction : predictions) {
                    // append the context to the predictions
                    this.getMetric().addPrediction(new Prediction(context + prediction.getName(), prediction.getPrediction()));
                }
            }
        } catch (Exception e) {
            log.warn("Unable to apply query prediction", e);
        }
    }
}

Example 2 with Prediction

use of datawave.microservice.querymetric.BaseQueryMetric.Prediction in project datawave by NationalSecurityAgency.

the class QueryExecutorBean method predictions.

/**
 * Pulls back the current predictions for a query.
 *
 * @param id
 *            - (@Required)
 *
 * @return GenericResponse containing predictions
 * @RequestHeader X-ProxiedEntitiesChain use when proxying request for user, by specifying a chain of DNs of the identities to proxy
 * @RequestHeader X-ProxiedIssuersChain required when using X-ProxiedEntitiesChain, specify one issuer DN per subject DN listed in X-ProxiedEntitiesChain
 * @RequestHeader query-session-id session id value used for load balancing purposes. query-session-id can be placed in the request in a Cookie header or as
 *                a query parameter
 * @ResponseHeader X-OperationTimeInMS time spent on the server performing the operation, does not account for network or result serialization
 *
 * @HTTP 200 success
 * @HTTP 204 success and no results
 * @HTTP 404 if id not found
 * @HTTP 412 if the query is no longer alive, client should call {@link #reset(String)} and try again
 * @HTTP 500 internal server error
 */
@GET
@Path("/{id}/predictions")
@Produces({ "application/xml", "text/xml", "application/json", "text/yaml", "text/x-yaml", "application/x-yaml", "application/x-protobuf", "application/x-protostuff" })
@GZIP
@Interceptors({ ResponseInterceptor.class, RequiredInterceptor.class })
@Override
@Timed(name = "dw.query.predictions", absolute = true)
public GenericResponse<String> predictions(@Required("id") @PathParam("id") String id) {
    // in case we don't make it to creating the response from the QueryLogic
    GenericResponse<String> response = new GenericResponse<>();
    Principal p = ctx.getCallerPrincipal();
    String userid = p.getName();
    if (p instanceof DatawavePrincipal) {
        DatawavePrincipal dp = (DatawavePrincipal) p;
        userid = dp.getShortName();
    }
    try {
        // Not calling getQueryById() here. We don't want to pull the persisted definition.
        RunningQuery query = queryCache.get(id);
        // an error.
        if (null == query || null == query.getConnection()) {
            // status code.
            if (null == query) {
                List<Query> queries = persister.findById(id);
                if (queries == null || queries.size() != 1) {
                    throw new NotFoundQueryException(DatawaveErrorCode.NO_QUERY_OBJECT_MATCH, MessageFormat.format("{0}", id));
                }
            }
            throw new PreConditionFailedQueryException(DatawaveErrorCode.QUERY_TIMEOUT_OR_SERVER_ERROR, MessageFormat.format("id = {0}", id));
        } else {
            // Validate the query belongs to the caller
            if (!query.getSettings().getOwner().equals(userid)) {
                throw new UnauthorizedQueryException(DatawaveErrorCode.QUERY_OWNER_MISMATCH, MessageFormat.format("{0} != {1}", userid, query.getSettings().getOwner()));
            }
            // pull the predictions out of the query metric
            Set<Prediction> predictions = query.getMetric().getPredictions();
            if (predictions != null && !predictions.isEmpty()) {
                response.setResult(predictions.toString());
                response.setHasResults(true);
            }
        }
    } catch (Exception e) {
        log.error("Failed to get query predictions", e);
        QueryException qe = new QueryException(DatawaveErrorCode.QUERY_PREDICTIONS_ERROR, e, MessageFormat.format("query id: {0}", id));
        log.error(qe, e);
        response.addException(qe.getBottomQueryException());
        int statusCode = qe.getBottomQueryException().getStatusCode();
        throw new DatawaveWebApplicationException(qe, response, statusCode);
    }
    return response;
}

Example 3 with Prediction

use of datawave.microservice.querymetric.BaseQueryMetric.Prediction 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());
}

Example 4 with Prediction

use of datawave.microservice.querymetric.BaseQueryMetric.Prediction 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();
        }
    }
}

Example 5 with Prediction

use of datawave.microservice.querymetric.BaseQueryMetric.Prediction in project datawave by NationalSecurityAgency.

the class QueryExecutorBean method predictQuery.

/**
 * @param queryLogicName
 * @param queryParameters
 * @return query predictions
 */
@POST
@Path("/{logicName}/predict")
@Produces({ "application/xml", "text/xml", "application/json", "text/yaml", "text/x-yaml", "application/x-yaml", "application/x-protobuf", "application/x-protostuff" })
@Interceptors({ RequiredInterceptor.class, ResponseInterceptor.class })
@Timed(name = "dw.query.predictQuery", absolute = true)
public GenericResponse<String> predictQuery(@Required("logicName") @PathParam("logicName") String queryLogicName, MultivaluedMap<String, String> queryParameters) {
    CreateQuerySessionIDFilter.QUERY_ID.set(null);
    QueryData qd = validateQuery(queryLogicName, queryParameters, null);
    GenericResponse<String> response = new GenericResponse<>();
    if (predictor != null) {
        try {
            qp.setPersistenceMode(QueryPersistence.TRANSIENT);
            MultivaluedMap<String, String> optionalQueryParameters = new MultivaluedMapImpl<>();
            optionalQueryParameters.putAll(qp.getUnknownParameters(queryParameters));
            Query q = persister.create(qd.userDn, qd.dnList, marking, queryLogicName, qp, optionalQueryParameters);
            BaseQueryMetric metric = metricFactory.createMetric();
            metric.populate(q);
            metric.setQueryType(RunningQuery.class.getSimpleName());
            Set<Prediction> predictions = predictor.predict(metric);
            if (predictions != null && !predictions.isEmpty()) {
                String predictionString = predictions.toString();
                // now we have a predictions, lets broadcast
                log.info("Model predictions: " + predictionString);
                response.setHasResults(true);
                response.setResult(predictionString);
            } else {
                response.setHasResults(false);
            }
        } catch (Throwable t) {
            response.setHasResults(false);
            if (t instanceof Error && !(t instanceof TokenMgrError)) {
                log.error(t.getMessage(), t);
                throw (Error) t;
            } else if (t instanceof WebApplicationException) {
                log.error(t.getMessage(), t);
                throw ((WebApplicationException) t);
            } else {
                log.error(t.getMessage(), t);
                QueryException qe = new QueryException(DatawaveErrorCode.QUERY_PREDICTIONS_ERROR, t);
                response.addException(qe.getBottomQueryException());
                int statusCode = qe.getBottomQueryException().getStatusCode();
                throw new DatawaveWebApplicationException(qe, response, statusCode);
            }
        }
    } else {
        response.setHasResults(false);
    }
    return response;
}