org.apache.nifi.provenance.search.QuerySubmission Java Examples

@Override
public QuerySubmission retrieveQuerySubmission(final String queryIdentifier, final NiFiUser user) {
    final QuerySubmission submission = querySubmissionMap.get(queryIdentifier);
    final String userId = submission.getSubmitterIdentity();

    if (user == null && userId == null) {
        return submission;
    }

    if (user == null) {
        throw new AccessDeniedException("Cannot retrieve Provenance Query Submission because no user id was provided in the provenance request.");
    }

    if (userId == null || userId.equals(user.getIdentity())) {
        return submission;
    }

    throw new AccessDeniedException("Cannot retrieve Provenance Query Submission because " + user.getIdentity() + " is not the user who submitted the request.");
}
 

public QueryResult queryEvents(final Query query, final NiFiUser user) throws IOException {
    final QuerySubmission submission = submitQuery(query, user);
    final QueryResult result = submission.getResult();
    while (!result.isFinished()) {
        try {
            Thread.sleep(100L);
        } catch (final InterruptedException ie) {
        }
    }

    if (result.getError() != null) {
        throw new IOException(result.getError());
    }

    return result;
}
 

QueryResult queryEvents(final Query query, final NiFiUser user) throws IOException {
    final QuerySubmission submission = submitQuery(query, user);
    final QueryResult result = submission.getResult();
    while (!result.isFinished()) {
        try {
            Thread.sleep(100L);
        } catch (final InterruptedException ie) {
            Thread.currentThread().interrupt();
        }
    }

    if (result.getError() != null) {
        throw new IOException(result.getError());
    }
    logger.info("{} got {} hits", query, result.getTotalHitCount());
    return result;
}
 

@Override
public QuerySubmission retrieveQuerySubmission(final String queryIdentifier, final NiFiUser user) {
    final QuerySubmission submission = querySubmissionMap.get(queryIdentifier);

    final String userId = submission.getSubmitterIdentity();

    if (user == null && userId == null) {
        return submission;
    }

    if (user == null) {
        throw new AccessDeniedException("Cannot retrieve Provenance Query Submission because no user id was provided in the provenance request.");
    }

    if (userId == null || userId.equals(user.getIdentity())) {
        return submission;
    }

    throw new AccessDeniedException("Cannot retrieve Provenance Query Submission because " + user.getIdentity() + " is not the user who submitted the request.");
}
 

@Override
public QuerySubmission retrieveQuerySubmission(final String queryIdentifier, final NiFiUser user) {
    final QuerySubmission submission = querySubmissionMap.get(queryIdentifier);

    final String userId = submission.getSubmitterIdentity();

    if (user == null && userId == null) {
        return submission;
    }

    if (user == null) {
        throw new AccessDeniedException("Cannot retrieve Provenance Query Submission because no user id was provided");
    }

    if (userId == null || userId.equals(user.getIdentity())) {
        return submission;
    }

    throw new AccessDeniedException("Cannot retrieve Provenance Query Submission because " + user.getIdentity() + " is not the user who submitted the request");
}
 

@Override
public QuerySubmission retrieveQuerySubmission(final String queryIdentifier, final NiFiUser user) {
    final QuerySubmission submission = querySubmissionMap.get(queryIdentifier);

    final String userId = submission.getSubmitterIdentity();

    if (user == null && userId == null) {
        return submission;
    }

    if (user == null) {
        throw new AccessDeniedException("Cannot retrieve Provenance Query Submission because no user id was provided in the provenance request.");
    }

    if (userId == null || userId.equals(user.getIdentity())) {
        return submission;
    }

    throw new AccessDeniedException("Cannot retrieve Provenance Query Submission because " + user.getIdentity() + " is not the user who submitted the request.");
}
 

@Override
public QuerySubmission retrieveQuerySubmission(final String queryIdentifier, final NiFiUser user) {
    final QuerySubmission submission = querySubmissionMap.get(queryIdentifier);

    final String userId = submission.getSubmitterIdentity();

    if (user == null && userId == null) {
        return submission;
    }

    if (user == null) {
        throw new AccessDeniedException("Cannot retrieve Provenance Query Submission because no user id was provided");
    }

    if (userId == null || userId.equals(user.getIdentity())) {
        return submission;
    }

    throw new AccessDeniedException("Cannot retrieve Provenance Query Submission because " + user.getIdentity() + " is not the user who submitted the request");
}
 

QueryResult queryEvents(final Query query, final NiFiUser user) throws IOException {
    final QuerySubmission submission = submitQuery(query, user);
    final QueryResult result = submission.getResult();
    while (!result.isFinished()) {
        try {
            Thread.sleep(100L);
        } catch (final InterruptedException ie) {
            Thread.currentThread().interrupt();
        }
    }

    if (result.getError() != null) {
        throw new IOException(result.getError());
    }
    logger.info("{} got {} hits", query, result.getTotalHitCount());
    return result;
}
 

QueryResult queryEvents(final Query query, final NiFiUser user) throws IOException {
    final QuerySubmission submission = submitQuery(query, user);
    final QueryResult result = submission.getResult();
    while (!result.isFinished()) {
        try {
            Thread.sleep(100L);
        } catch (final InterruptedException ie) {
            Thread.currentThread().interrupt();
        }
    }

    if (result.getError() != null) {
        throw new IOException(result.getError());
    }
    logger.info("{} got {} hits", query, result.getTotalHitCount());
    return result;
}
 

@Override
public QuerySubmission retrieveQuerySubmission(final String queryIdentifier, final NiFiUser user) {
    final QuerySubmission submission = querySubmissionMap.get(queryIdentifier);
    final String userId = submission.getSubmitterIdentity();

    if (user == null && userId == null) {
        return submission;
    }

    if (user == null) {
        throw new AccessDeniedException("Cannot retrieve Provenance Query Submission because no user id was provided in the provenance request.");
    }

    if (userId == null || userId.equals(user.getIdentity())) {
        return submission;
    }

    throw new AccessDeniedException("Cannot retrieve Provenance Query Submission because " + user.getIdentity() + " is not the user who submitted the request.");
}
 

public QueryResult queryEvents(final Query query, final NiFiUser user) throws IOException {
    final QuerySubmission submission = submitQuery(query, user);
    final QueryResult result = submission.getResult();
    while (!result.isFinished()) {
        try {
            Thread.sleep(100L);
        } catch (final InterruptedException ie) {
        }
    }

    if (result.getError() != null) {
        throw new IOException(result.getError());
    }

    return result;
}
 

/**
 * Deletes the query with the specified id.
 *
 * @param provenanceId id
 */
public void deleteProvenanceQuery(final String provenanceId) {
    // get the query to the provenance repository
    final ProvenanceRepository provenanceRepository = flowController.getProvenanceRepository();
    final QuerySubmission querySubmission = provenanceRepository.retrieveQuerySubmission(provenanceId, NiFiUserUtils.getNiFiUser());
    if (querySubmission != null) {
        querySubmission.cancel();
    }
}
 

@Test(timeout = 60000)
public void addThenQueryWithEmptyQuery() throws InterruptedException {
    assumeFalse(isWindowsEnvironment());
    final RepositoryConfiguration repoConfig = createConfig();
    final IndexManager indexManager = new StandardIndexManager(repoConfig);

    final LuceneEventIndex index = new LuceneEventIndex(repoConfig, indexManager, 1, EventReporter.NO_OP);

    final ProvenanceEventRecord event = createEvent();

    index.addEvent(event, new StorageSummary(event.getEventId(), "1.prov", "1", 1, 2L, 2L));

    final Query query = new Query(UUID.randomUUID().toString());

    final ArrayListEventStore eventStore = new ArrayListEventStore();
    eventStore.addEvent(event);
    index.initialize(eventStore);

    // We don't know how long it will take for the event to be indexed, so keep querying until
    // we get a result. The test will timeout after 5 seconds if we've still not succeeded.
    List<ProvenanceEventRecord> matchingEvents = Collections.emptyList();
    while (matchingEvents.isEmpty()) {
        final QuerySubmission submission = index.submitQuery(query, EventAuthorizer.GRANT_ALL, "unit test user");
        assertNotNull(submission);

        final QueryResult result = submission.getResult();
        assertNotNull(result);
        result.awaitCompletion(4000, TimeUnit.MILLISECONDS);

        assertTrue(result.isFinished());
        assertNull(result.getError());

        matchingEvents = result.getMatchingEvents();
        assertNotNull(matchingEvents);
        Thread.sleep(100L); // avoid crushing the CPU
    }

    assertEquals(1, matchingEvents.size());
    assertEquals(event, matchingEvents.get(0));
}
 

/**
 * Deletes the query with the specified id.
 *
 * @param provenanceId id
 */
public void deleteProvenanceQuery(final String provenanceId) {
    // get the query to the provenance repository
    final ProvenanceRepository provenanceRepository = flowController.getProvenanceRepository();
    final QuerySubmission querySubmission = provenanceRepository.retrieveQuerySubmission(provenanceId, NiFiUserUtils.getNiFiUser());
    if (querySubmission != null) {
        querySubmission.cancel();
    }
}
 

@Test(timeout = 5000)
public void addThenQueryWithEmptyQuery() throws InterruptedException {
    final RepositoryConfiguration repoConfig = createConfig();
    final IndexManager indexManager = new SimpleIndexManager(repoConfig);

    final LuceneEventIndex index = new LuceneEventIndex(repoConfig, indexManager, 1, EventReporter.NO_OP);

    final ProvenanceEventRecord event = createEvent();

    index.addEvent(event, new StorageSummary(event.getEventId(), "1.prov", "1", 1, 2L, 2L));

    final Query query = new Query(UUID.randomUUID().toString());

    final ArrayListEventStore eventStore = new ArrayListEventStore();
    eventStore.addEvent(event);
    index.initialize(eventStore);

    // We don't know how long it will take for the event to be indexed, so keep querying until
    // we get a result. The test will timeout after 5 seconds if we've still not succeeded.
    List<ProvenanceEventRecord> matchingEvents = Collections.emptyList();
    while (matchingEvents.isEmpty()) {
        final QuerySubmission submission = index.submitQuery(query, EventAuthorizer.GRANT_ALL, "unit test user");
        assertNotNull(submission);

        final QueryResult result = submission.getResult();
        assertNotNull(result);
        result.awaitCompletion(100, TimeUnit.MILLISECONDS);

        assertTrue(result.isFinished());
        assertNull(result.getError());

        matchingEvents = result.getMatchingEvents();
        assertNotNull(matchingEvents);
        Thread.sleep(100L); // avoid crushing the CPU
    }

    assertEquals(1, matchingEvents.size());
    assertEquals(event, matchingEvents.get(0));
}
 

@Override
public QuerySubmission submitQuery(final Query query, final NiFiUser user) {
    return eventIndex.submitQuery(query, createEventAuthorizer(user), user.getIdentity());
}
 

@Override
public QuerySubmission submitQuery(Query query, NiFiUser user) {
    throw new UnsupportedOperationException("MockProvenanceRepository does not support querying");
}
 

/**
 * Retrieves the results of a provenance query.
 *
 * @param provenanceId id
 * @return the results of a provenance query
 */
public ProvenanceDTO getProvenanceQuery(String provenanceId, Boolean summarize, Boolean incrementalResults) {
    try {
        // get the query to the provenance repository
        final ProvenanceRepository provenanceRepository = flowController.getProvenanceRepository();
        final QuerySubmission querySubmission = provenanceRepository.retrieveQuerySubmission(provenanceId, NiFiUserUtils.getNiFiUser());

        // ensure the query results could be found
        if (querySubmission == null) {
            throw new ResourceNotFoundException("Cannot find the results for the specified provenance requests. Results may have been purged.");
        }

        // get the original query and the results
        final Query query = querySubmission.getQuery();
        final QueryResult queryResult = querySubmission.getResult();

        // build the response
        final ProvenanceDTO provenanceDto = new ProvenanceDTO();
        final ProvenanceRequestDTO requestDto = new ProvenanceRequestDTO();
        final ProvenanceResultsDTO resultsDto = new ProvenanceResultsDTO();

        // include the original request and results
        provenanceDto.setRequest(requestDto);
        provenanceDto.setResults(resultsDto);

        // convert the original request
        requestDto.setStartDate(query.getStartDate());
        requestDto.setEndDate(query.getEndDate());
        requestDto.setMinimumFileSize(query.getMinFileSize());
        requestDto.setMaximumFileSize(query.getMaxFileSize());
        requestDto.setMaxResults(query.getMaxResults());
        if (query.getSearchTerms() != null) {
            final Map<String, String> searchTerms = new HashMap<>();
            for (final SearchTerm searchTerm : query.getSearchTerms()) {
                searchTerms.put(searchTerm.getSearchableField().getFriendlyName(), searchTerm.getValue());
            }
            requestDto.setSearchTerms(searchTerms);
        }

        // convert the provenance
        provenanceDto.setId(query.getIdentifier());
        provenanceDto.setSubmissionTime(querySubmission.getSubmissionTime());
        provenanceDto.setExpiration(queryResult.getExpiration());
        provenanceDto.setFinished(queryResult.isFinished());
        provenanceDto.setPercentCompleted(queryResult.getPercentComplete());

        // convert each event
        final boolean includeResults = incrementalResults == null || Boolean.TRUE.equals(incrementalResults);
        if (includeResults || queryResult.isFinished()) {
            final List<ProvenanceEventDTO> events = new ArrayList<>();
            for (final ProvenanceEventRecord record : queryResult.getMatchingEvents()) {
                events.add(createProvenanceEventDto(record, Boolean.TRUE.equals(summarize)));
            }
            resultsDto.setProvenanceEvents(events);
        }

        if (requestDto.getMaxResults() != null && queryResult.getTotalHitCount() >= requestDto.getMaxResults()) {
            resultsDto.setTotalCount(requestDto.getMaxResults().longValue());
            resultsDto.setTotal(FormatUtils.formatCount(requestDto.getMaxResults().longValue()) + "+");
        } else {
            resultsDto.setTotalCount(queryResult.getTotalHitCount());
            resultsDto.setTotal(FormatUtils.formatCount(queryResult.getTotalHitCount()));
        }

        // include any errors
        if (queryResult.getError() != null) {
            final Set<String> errors = new HashSet<>();
            errors.add(queryResult.getError());
            resultsDto.setErrors(errors);
        }

        // set the generated timestamp
        final Date now = new Date();
        resultsDto.setGenerated(now);
        resultsDto.setTimeOffset(TimeZone.getDefault().getOffset(now.getTime()));

        // get the oldest available event time
        final List<ProvenanceEventRecord> firstEvent = provenanceRepository.getEvents(0, 1);
        if (!firstEvent.isEmpty()) {
            resultsDto.setOldestEvent(new Date(firstEvent.get(0).getEventTime()));
        }

        provenanceDto.setResults(resultsDto);
        return provenanceDto;
    } catch (final IOException ioe) {
        throw new NiFiCoreException("An error occurred while searching the provenance events.", ioe);
    }
}
 

/**
 * Submits a provenance query.
 *
 * @param provenanceDto dto
 * @return provenance info
 */
public ProvenanceDTO submitProvenance(ProvenanceDTO provenanceDto) {
    final ProvenanceRequestDTO requestDto = provenanceDto.getRequest();

    // create the query
    final Query query = new Query(provenanceDto.getId());

    // if the request was specified
    if (requestDto != null) {
        // add each search term specified
        final Map<String, String> searchTerms = requestDto.getSearchTerms();
        if (searchTerms != null) {
            for (final Map.Entry<String, String> searchTerm : searchTerms.entrySet()) {
                SearchableField field;

                field = SearchableFields.getSearchableField(searchTerm.getKey());
                if (field == null) {
                    field = SearchableFields.newSearchableAttribute(searchTerm.getKey());
                }
                query.addSearchTerm(SearchTerms.newSearchTerm(field, searchTerm.getValue()));
            }
        }

        // specify the start date if specified
        if (requestDto.getStartDate() != null) {
            query.setStartDate(requestDto.getStartDate());
        }

        // ensure an end date is populated
        if (requestDto.getEndDate() != null) {
            query.setEndDate(requestDto.getEndDate());
        }

        // set the min/max file size
        query.setMinFileSize(requestDto.getMinimumFileSize());
        query.setMaxFileSize(requestDto.getMaximumFileSize());

        // set the max results desired
        query.setMaxResults(requestDto.getMaxResults());
    }

    // submit the query to the provenance repository
    final ProvenanceRepository provenanceRepository = flowController.getProvenanceRepository();
    final QuerySubmission querySubmission = provenanceRepository.submitQuery(query, NiFiUserUtils.getNiFiUser());

    // return the query with the results populated at this point
    return getProvenanceQuery(querySubmission.getQueryIdentifier(), requestDto.getSummarize(), requestDto.getIncrementalResults());
}
 

@Test
public void testIndexAndCompressOnRolloverAndSubsequentSearchAsync() throws InterruptedException {
    repo = new VolatileProvenanceRepository(NiFiProperties.createBasicNiFiProperties(null));

    final String uuid = "00000000-0000-0000-0000-000000000000";
    final Map<String, String> attributes = new HashMap<>();
    attributes.put("abc", "xyz");
    attributes.put("xyz", "abc");
    attributes.put("filename", "file-" + uuid);

    final ProvenanceEventBuilder builder = new StandardProvenanceEventRecord.Builder();
    builder.setEventTime(System.currentTimeMillis());
    builder.setEventType(ProvenanceEventType.RECEIVE);
    builder.setTransitUri("nifi://unit-test");
    builder.fromFlowFile(createFlowFile(3L, 3000L, attributes));
    builder.setComponentId("1234");
    builder.setComponentType("dummy processor");

    for (int i = 0; i < 10; i++) {
        attributes.put("uuid", "00000000-0000-0000-0000-00000000000" + i);
        builder.fromFlowFile(createFlowFile(i, 3000L, attributes));
        repo.registerEvent(builder.build());
    }

    final Query query = new Query(UUID.randomUUID().toString());
    query.addSearchTerm(SearchTerms.newSearchTerm(SearchableFields.FlowFileUUID, "00000*"));
    query.addSearchTerm(SearchTerms.newSearchTerm(SearchableFields.Filename, "file-*"));
    query.addSearchTerm(SearchTerms.newSearchTerm(SearchableFields.ComponentID, "12?4"));
    query.addSearchTerm(SearchTerms.newSearchTerm(SearchableFields.TransitURI, "nifi://*"));
    query.setMaxResults(100);

    final QuerySubmission submission = repo.submitQuery(query, createUser());
    while (!submission.getResult().isFinished()) {
        Thread.sleep(100L);
    }

    assertEquals(10, submission.getResult().getMatchingEvents().size());
    for (final ProvenanceEventRecord match : submission.getResult().getMatchingEvents()) {
        System.out.println(match);
    }
}
 

@Override
public QuerySubmission retrieveQuerySubmission(String queryIdentifier, NiFiUser user) {
    throw new UnsupportedOperationException("MockProvenanceRepository does not support querying");
}
 

@Test(timeout = 50000)
public void testQuerySpecificField() throws InterruptedException {
    final RepositoryConfiguration repoConfig = createConfig();
    final IndexManager indexManager = new StandardIndexManager(repoConfig);

    final LuceneEventIndex index = new LuceneEventIndex(repoConfig, indexManager, 2, EventReporter.NO_OP);

    // add 2 events, one of which we will query for.
    final ProvenanceEventRecord event = createEvent();
    index.addEvent(event, new StorageSummary(event.getEventId(), "1.prov", "1", 1, 2L, 2L));
    index.addEvent(createEvent(), new StorageSummary(2L, "1.prov", "1", 1, 2L, 2L));

    // Create a query that searches for the event with the FlowFile UUID equal to the first event's.
    final Query query = new Query(UUID.randomUUID().toString());
    query.addSearchTerm(SearchTerms.newSearchTerm(SearchableFields.FlowFileUUID, event.getFlowFileUuid()));

    final ArrayListEventStore eventStore = new ArrayListEventStore();
    eventStore.addEvent(event);
    index.initialize(eventStore);

    // We don't know how long it will take for the event to be indexed, so keep querying until
    // we get a result. The test will timeout after 5 seconds if we've still not succeeded.
    List<ProvenanceEventRecord> matchingEvents = Collections.emptyList();
    while (matchingEvents.isEmpty()) {
        final QuerySubmission submission = index.submitQuery(query, EventAuthorizer.GRANT_ALL, "unit test user");
        assertNotNull(submission);

        final QueryResult result = submission.getResult();
        assertNotNull(result);
        result.awaitCompletion(4000, TimeUnit.MILLISECONDS);

        assertTrue(result.isFinished());
        assertNull(result.getError());

        matchingEvents = result.getMatchingEvents();
        assertNotNull(matchingEvents);
        Thread.sleep(100L); // avoid crushing the CPU
    }

    assertEquals(1, matchingEvents.size());
    assertEquals(event, matchingEvents.get(0));
}
 

@Test(timeout = 10000)
public void testNotAuthorizedQuery() throws IOException, InterruptedException {
    assumeFalse(isWindowsEnvironment());
    final RepositoryConfiguration config = createConfiguration();
    config.setMaxRecordLife(5, TimeUnit.MINUTES);
    config.setMaxStorageCapacity(1024L * 1024L);
    config.setMaxEventFileLife(500, TimeUnit.MILLISECONDS);
    config.setMaxEventFileCapacity(1024L * 1024L);
    config.setSearchableFields(new ArrayList<>(SearchableFields.getStandardFields()));
    config.setDesiredIndexSize(10); // force new index to be created for each rollover

    repo = new PersistentProvenanceRepository(config, DEFAULT_ROLLOVER_MILLIS) {
        @Override
        public boolean isAuthorized(ProvenanceEventRecord event, NiFiUser user) {
            return event.getEventId() > 2;
        }
    };

    repo.initialize(getEventReporter(), null, null, IdentifierLookup.EMPTY);

    final String uuid = "00000000-0000-0000-0000-000000000000";
    final Map<String, String> attributes = new HashMap<>();
    attributes.put("abc", "xyz");
    attributes.put("xyz", "abc");
    attributes.put("filename", "file-" + uuid);

    final ProvenanceEventBuilder builder = new StandardProvenanceEventRecord.Builder();
    builder.setEventTime(System.currentTimeMillis());
    builder.setEventType(ProvenanceEventType.RECEIVE);
    builder.setTransitUri("nifi://unit-test");
    builder.fromFlowFile(createFlowFile(3L, 3000L, attributes));
    builder.setComponentId("1234");
    builder.setComponentType("dummy processor");

    for (int i = 0; i < 10; i++) {
        attributes.put("uuid", "00000000-0000-0000-0000-00000000000" + i);
        builder.fromFlowFile(createFlowFile(i, 3000L, attributes));
        builder.setEventTime(10L); // make sure the events are destroyed when we call purge
        repo.registerEvent(builder.build());
    }

    repo.waitForRollover();

    final Query query = new Query("1234");
    query.addSearchTerm(SearchTerms.newSearchTerm(SearchableFields.ComponentID, "1234"));
    final QuerySubmission submission = repo.submitQuery(query, createUser());

    final QueryResult result = submission.getResult();
    while (!result.isFinished()) {
        Thread.sleep(100L);
    }

    // Ensure that we gets events with ID's 3 through 10.
    final List<ProvenanceEventRecord> events = result.getMatchingEvents();
    assertEquals(7, events.size());
    final List<Long> eventIds = events.stream().map(event -> event.getEventId()).sorted().collect(Collectors.toList());
    for (int i = 0; i < 7; i++) {
        Assert.assertEquals(i + 3, eventIds.get(i).intValue());
    }
}
 

@Test
public void testIndexAndCompressOnRolloverAndSubsequentSearchMultipleStorageDirs() throws IOException, InterruptedException, ParseException {
    assumeFalse(isWindowsEnvironment());
    final RepositoryConfiguration config = createConfiguration();
    config.addStorageDirectory("2", new File("target/storage/" + UUID.randomUUID().toString()));
    config.setMaxRecordLife(30, TimeUnit.SECONDS);
    config.setMaxStorageCapacity(1024L * 1024L);
    config.setMaxEventFileLife(1, TimeUnit.SECONDS);
    config.setMaxEventFileCapacity(1024L * 1024L);
    config.setSearchableFields(new ArrayList<>(SearchableFields.getStandardFields()));

    repo = new PersistentProvenanceRepository(config, DEFAULT_ROLLOVER_MILLIS);
    repo.initialize(getEventReporter(), null, null, IdentifierLookup.EMPTY);

    final String uuid = "00000000-0000-0000-0000-000000000000";
    final Map<String, String> attributes = new HashMap<>();
    attributes.put("abc", "xyz");
    attributes.put("xyz", "abc");
    attributes.put("filename", "file-" + uuid);

    final ProvenanceEventBuilder builder = new StandardProvenanceEventRecord.Builder();
    for (int j = 0; j < 3; j++) {
        attributes.put("iteration", String.valueOf(j));

        builder.setEventTime(System.currentTimeMillis());
        builder.setEventType(ProvenanceEventType.RECEIVE);
        builder.setTransitUri("nifi://unit-test");
        builder.setComponentId("1234");
        builder.setComponentType("dummy processor");
        builder.fromFlowFile(createFlowFile(3L, 3000L, attributes));

        for (int i = 0; i < 10; i++) {
            String uuidSuffix = String.valueOf(i + j * 10);
            if (uuidSuffix.length() < 2) {
                uuidSuffix = "0" + uuidSuffix;
            }

            attributes.put("uuid", "00000000-0000-0000-0000-0000000000" + uuidSuffix);
            builder.fromFlowFile(createFlowFile(i + j * 10, 3000L, attributes));
            repo.registerEvent(builder.build());
        }

        repo.waitForRollover();
    }

    final Query query = new Query(UUID.randomUUID().toString());
    query.addSearchTerm(SearchTerms.newSearchTerm(SearchableFields.Filename, "file-*"));
    query.addSearchTerm(SearchTerms.newSearchTerm(SearchableFields.ComponentID, "12?4"));
    query.addSearchTerm(SearchTerms.newSearchTerm(SearchableFields.TransitURI, "nifi://*"));
    query.setMaxResults(100);

    final QuerySubmission submission = repo.submitQuery(query, createUser());
    while (!submission.getResult().isFinished()) {
        Thread.sleep(100L);
    }

    assertEquals(30, submission.getResult().getMatchingEvents().size());
    final Map<String, Integer> counts = new HashMap<>();
    for (final ProvenanceEventRecord match : submission.getResult().getMatchingEvents()) {
        System.out.println(match);

        final String index = match.getAttributes().get("iteration");
        Integer count = counts.get(index);
        if (count == null) {
            count = 0;
        }
        counts.put(index, count + 1);
    }

    assertEquals(3, counts.size());
    assertEquals(10, counts.get("0").intValue());
    assertEquals(10, counts.get("1").intValue());
    assertEquals(10, counts.get("2").intValue());

    config.setMaxRecordLife(1, TimeUnit.MILLISECONDS);

    repo.purgeOldEvents();

    Thread.sleep(2000L); // purge is async. Give it time to do its job.

    query.setMaxResults(100);
    final QuerySubmission noResultSubmission = repo.submitQuery(query, createUser());
    while (!noResultSubmission.getResult().isFinished()) {
        Thread.sleep(10L);
    }

    assertEquals(0, noResultSubmission.getResult().getTotalHitCount());
}
 

@Override
public QuerySubmission retrieveQuerySubmission(final String queryIdentifier, final NiFiUser user) {
    return eventIndex.retrieveQuerySubmission(queryIdentifier, user);
}
 

@Override
public QuerySubmission submitQuery(final Query query, final NiFiUser user) {
    return eventIndex.submitQuery(query, createEventAuthorizer(user), user == null ? null : user.getIdentity());
}
 

@Test(timeout = 10000)
public void testNotAuthorizedQuery() throws IOException, InterruptedException {
    final RepositoryConfiguration config = createConfiguration();
    config.setMaxRecordLife(5, TimeUnit.MINUTES);
    config.setMaxStorageCapacity(1024L * 1024L);
    config.setMaxEventFileLife(500, TimeUnit.MILLISECONDS);
    config.setMaxEventFileCapacity(1024L * 1024L);
    config.setSearchableFields(new ArrayList<>(SearchableFields.getStandardFields()));
    config.setDesiredIndexSize(10); // force new index to be created for each rollover

    repo = new PersistentProvenanceRepository(config, DEFAULT_ROLLOVER_MILLIS) {
        @Override
        public boolean isAuthorized(ProvenanceEventRecord event, NiFiUser user) {
            return event.getEventId() > 2;
        }
    };

    repo.initialize(getEventReporter(), null, null, IdentifierLookup.EMPTY);

    final String uuid = "00000000-0000-0000-0000-000000000000";
    final Map<String, String> attributes = new HashMap<>();
    attributes.put("abc", "xyz");
    attributes.put("xyz", "abc");
    attributes.put("filename", "file-" + uuid);

    final ProvenanceEventBuilder builder = new StandardProvenanceEventRecord.Builder();
    builder.setEventTime(System.currentTimeMillis());
    builder.setEventType(ProvenanceEventType.RECEIVE);
    builder.setTransitUri("nifi://unit-test");
    builder.fromFlowFile(createFlowFile(3L, 3000L, attributes));
    builder.setComponentId("1234");
    builder.setComponentType("dummy processor");

    for (int i = 0; i < 10; i++) {
        attributes.put("uuid", "00000000-0000-0000-0000-00000000000" + i);
        builder.fromFlowFile(createFlowFile(i, 3000L, attributes));
        builder.setEventTime(10L); // make sure the events are destroyed when we call purge
        repo.registerEvent(builder.build());
    }

    repo.waitForRollover();

    final Query query = new Query("1234");
    query.addSearchTerm(SearchTerms.newSearchTerm(SearchableFields.ComponentID, "1234"));
    final QuerySubmission submission = repo.submitQuery(query, createUser());

    final QueryResult result = submission.getResult();
    while (!result.isFinished()) {
        Thread.sleep(100L);
    }

    // Ensure that we gets events with ID's 3 through 10.
    final List<ProvenanceEventRecord> events = result.getMatchingEvents();
    assertEquals(7, events.size());
    final List<Long> eventIds = events.stream().map(event -> event.getEventId()).sorted().collect(Collectors.toList());
    for (int i = 0; i < 7; i++) {
        Assert.assertEquals(i + 3, eventIds.get(i).intValue());
    }
}
 

@Override
public QuerySubmission retrieveQuerySubmission(final String queryIdentifier, final NiFiUser user) {
    return eventIndex.retrieveQuerySubmission(queryIdentifier, user);
}
 

@Test
public void testIndexAndCompressOnRolloverAndSubsequentSearchMultipleStorageDirs() throws IOException, InterruptedException, ParseException {
    final RepositoryConfiguration config = createConfiguration();
    config.addStorageDirectory("2", new File("target/storage/" + UUID.randomUUID().toString()));
    config.setMaxRecordLife(30, TimeUnit.SECONDS);
    config.setMaxStorageCapacity(1024L * 1024L);
    config.setMaxEventFileLife(1, TimeUnit.SECONDS);
    config.setMaxEventFileCapacity(1024L * 1024L);
    config.setSearchableFields(new ArrayList<>(SearchableFields.getStandardFields()));

    repo = new PersistentProvenanceRepository(config, DEFAULT_ROLLOVER_MILLIS);
    repo.initialize(getEventReporter(), null, null, IdentifierLookup.EMPTY);

    final String uuid = "00000000-0000-0000-0000-000000000000";
    final Map<String, String> attributes = new HashMap<>();
    attributes.put("abc", "xyz");
    attributes.put("xyz", "abc");
    attributes.put("filename", "file-" + uuid);

    final ProvenanceEventBuilder builder = new StandardProvenanceEventRecord.Builder();
    for (int j = 0; j < 3; j++) {
        attributes.put("iteration", String.valueOf(j));

        builder.setEventTime(System.currentTimeMillis());
        builder.setEventType(ProvenanceEventType.RECEIVE);
        builder.setTransitUri("nifi://unit-test");
        builder.setComponentId("1234");
        builder.setComponentType("dummy processor");
        builder.fromFlowFile(createFlowFile(3L, 3000L, attributes));

        for (int i = 0; i < 10; i++) {
            String uuidSuffix = String.valueOf(i + j * 10);
            if (uuidSuffix.length() < 2) {
                uuidSuffix = "0" + uuidSuffix;
            }

            attributes.put("uuid", "00000000-0000-0000-0000-0000000000" + uuidSuffix);
            builder.fromFlowFile(createFlowFile(i + j * 10, 3000L, attributes));
            repo.registerEvent(builder.build());
        }

        repo.waitForRollover();
    }

    final Query query = new Query(UUID.randomUUID().toString());
    query.addSearchTerm(SearchTerms.newSearchTerm(SearchableFields.Filename, "file-*"));
    query.addSearchTerm(SearchTerms.newSearchTerm(SearchableFields.ComponentID, "12?4"));
    query.addSearchTerm(SearchTerms.newSearchTerm(SearchableFields.TransitURI, "nifi://*"));
    query.setMaxResults(100);

    final QuerySubmission submission = repo.submitQuery(query, createUser());
    while (!submission.getResult().isFinished()) {
        Thread.sleep(100L);
    }

    assertEquals(30, submission.getResult().getMatchingEvents().size());
    final Map<String, Integer> counts = new HashMap<>();
    for (final ProvenanceEventRecord match : submission.getResult().getMatchingEvents()) {
        System.out.println(match);

        final String index = match.getAttributes().get("iteration");
        Integer count = counts.get(index);
        if (count == null) {
            count = 0;
        }
        counts.put(index, count + 1);
    }

    assertEquals(3, counts.size());
    assertEquals(10, counts.get("0").intValue());
    assertEquals(10, counts.get("1").intValue());
    assertEquals(10, counts.get("2").intValue());

    config.setMaxRecordLife(1, TimeUnit.MILLISECONDS);

    repo.purgeOldEvents();

    Thread.sleep(2000L); // purge is async. Give it time to do its job.

    query.setMaxResults(100);
    final QuerySubmission noResultSubmission = repo.submitQuery(query, createUser());
    while (!noResultSubmission.getResult().isFinished()) {
        Thread.sleep(10L);
    }

    assertEquals(0, noResultSubmission.getResult().getTotalHitCount());
}
 

@Test(timeout = 50000)
public void testQuerySpecificField() throws InterruptedException {
    final RepositoryConfiguration repoConfig = createConfig();
    final IndexManager indexManager = new SimpleIndexManager(repoConfig);

    final LuceneEventIndex index = new LuceneEventIndex(repoConfig, indexManager, 2, EventReporter.NO_OP);

    // add 2 events, one of which we will query for.
    final ProvenanceEventRecord event = createEvent();
    index.addEvent(event, new StorageSummary(event.getEventId(), "1.prov", "1", 1, 2L, 2L));
    index.addEvent(createEvent(), new StorageSummary(2L, "1.prov", "1", 1, 2L, 2L));

    // Create a query that searches for the event with the FlowFile UUID equal to the first event's.
    final Query query = new Query(UUID.randomUUID().toString());
    query.addSearchTerm(SearchTerms.newSearchTerm(SearchableFields.FlowFileUUID, event.getFlowFileUuid()));

    final ArrayListEventStore eventStore = new ArrayListEventStore();
    eventStore.addEvent(event);
    index.initialize(eventStore);

    // We don't know how long it will take for the event to be indexed, so keep querying until
    // we get a result. The test will timeout after 5 seconds if we've still not succeeded.
    List<ProvenanceEventRecord> matchingEvents = Collections.emptyList();
    while (matchingEvents.isEmpty()) {
        final QuerySubmission submission = index.submitQuery(query, EventAuthorizer.GRANT_ALL, "unit test user");
        assertNotNull(submission);

        final QueryResult result = submission.getResult();
        assertNotNull(result);
        result.awaitCompletion(100, TimeUnit.MILLISECONDS);

        assertTrue(result.isFinished());
        assertNull(result.getError());

        matchingEvents = result.getMatchingEvents();
        assertNotNull(matchingEvents);
        Thread.sleep(100L); // avoid crushing the CPU
    }

    assertEquals(1, matchingEvents.size());
    assertEquals(event, matchingEvents.get(0));
}