org.apache.nifi.provenance.ProvenanceRepository Java Examples

public StandardExtensionDiscoveringManager() {
    definitionMap.put(Processor.class, new HashSet<>());
    definitionMap.put(FlowFilePrioritizer.class, new HashSet<>());
    definitionMap.put(ReportingTask.class, new HashSet<>());
    definitionMap.put(ControllerService.class, new HashSet<>());
    definitionMap.put(Authorizer.class, new HashSet<>());
    definitionMap.put(UserGroupProvider.class, new HashSet<>());
    definitionMap.put(AccessPolicyProvider.class, new HashSet<>());
    definitionMap.put(LoginIdentityProvider.class, new HashSet<>());
    definitionMap.put(ProvenanceRepository.class, new HashSet<>());
    definitionMap.put(ComponentStatusRepository.class, new HashSet<>());
    definitionMap.put(FlowFileRepository.class, new HashSet<>());
    definitionMap.put(FlowFileSwapManager.class, new HashSet<>());
    definitionMap.put(ContentRepository.class, new HashSet<>());
    definitionMap.put(StateProvider.class, new HashSet<>());
    definitionMap.put(StatusAnalyticsModel.class, new HashSet<>());
}
 

private void consumeNiFiProvenanceEvents(ReportingContext context, NiFiFlow nifiFlow) {
    final EventAccess eventAccess = context.getEventAccess();
    final AnalysisContext analysisContext = new StandardAnalysisContext(nifiFlow, namespaceResolvers,
            // FIXME: This class cast shouldn't be necessary to query lineage. Possible refactor target in next major update.
            (ProvenanceRepository)eventAccess.getProvenanceRepository());
    consumer.consumeEvents(context, (componentMapHolder, events) -> {
        for (ProvenanceEventRecord event : events) {
            try {
                lineageStrategy.processEvent(analysisContext, nifiFlow, event);
            } catch (Exception e) {
                // If something went wrong, log it and continue with other records.
                getLogger().error("Skipping failed analyzing event {} due to {}.", new Object[]{event, e, e});
            }
        }
        nifiAtlasHook.commitMessages();
    });
}
 

public RepositoryContextFactory(final ContentRepository contentRepository, final FlowFileRepository flowFileRepository,
        final FlowFileEventRepository flowFileEventRepository, final CounterRepository counterRepository,
        final ProvenanceRepository provenanceRepository) {

    this.contentRepo = contentRepository;
    this.flowFileRepo = flowFileRepository;
    this.flowFileEventRepo = flowFileEventRepository;
    this.counterRepo = counterRepository;
    this.provenanceRepo = provenanceRepository;
}
 

/**
 * Deletes the lineage with the specified id.
 *
 * @param lineageId id
 */
public void deleteLineage(final String lineageId) {
    // get the query to the provenance repository
    final ProvenanceRepository provenanceRepository = flowController.getProvenanceRepository();
    final ComputeLineageSubmission computeLineageSubmission = provenanceRepository.retrieveLineageSubmission(lineageId, NiFiUserUtils.getNiFiUser());
    if (computeLineageSubmission != null) {
        computeLineageSubmission.cancel();
    }
}
 

/**
 * 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();
    }
}
 

/**
 * Gets the lineage with the specified id.
 *
 * @param lineageId id
 * @return the lineage with the specified id
 */
public LineageDTO getLineage(final String lineageId) {
    // get the query to the provenance repository
    final ProvenanceRepository provenanceRepository = flowController.getProvenanceRepository();
    final ComputeLineageSubmission computeLineageSubmission = provenanceRepository.retrieveLineageSubmission(lineageId, NiFiUserUtils.getNiFiUser());

    // ensure the submission was found
    if (computeLineageSubmission == null) {
        throw new ResourceNotFoundException("Cannot find the results for the specified lineage request. Results may have been purged.");
    }

    return dtoFactory.createLineageDto(computeLineageSubmission);
}
 

/**
 * Submits the specified lineage request.
 *
 * @param lineageDto dto
 * @return updated lineage
 */
public LineageDTO submitLineage(LineageDTO lineageDto) {
    final LineageRequestDTO requestDto = lineageDto.getRequest();

    // get the provenance repo
    final ProvenanceRepository provenanceRepository = flowController.getProvenanceRepository();
    final ComputeLineageSubmission result;

    if (LineageRequestType.FLOWFILE.equals(requestDto.getLineageRequestType())) {
        if (requestDto.getUuid() != null) {
            // submit uuid if it is specified
            result = provenanceRepository.submitLineageComputation(requestDto.getUuid(), NiFiUserUtils.getNiFiUser());
        } else {
            // submit the event if the flowfile uuid needs to be looked up
            result = provenanceRepository.submitLineageComputation(requestDto.getEventId(), NiFiUserUtils.getNiFiUser());
        }
    } else {
        // submit event... (parents or children)
        if (LineageRequestType.PARENTS.equals(requestDto.getLineageRequestType())) {
            result = provenanceRepository.submitExpandParents(requestDto.getEventId(), NiFiUserUtils.getNiFiUser());
        } else {
            result = provenanceRepository.submitExpandChildren(requestDto.getEventId(), NiFiUserUtils.getNiFiUser());
        }
    }

    return getLineage(result.getLineageIdentifier());
}
 

/**
 * Gets the available options for searching provenance.
 *
 * @return the available options for searching provenance
 */
public ProvenanceOptionsDTO getProvenanceSearchOptions() {
    final ProvenanceRepository provenanceRepository = flowController.getProvenanceRepository();

    // create the search options dto
    final ProvenanceOptionsDTO searchOptions = new ProvenanceOptionsDTO();
    final List<ProvenanceSearchableFieldDTO> searchableFieldNames = new ArrayList<>();
    final List<SearchableField> fields = provenanceRepository.getSearchableFields();
    for (final SearchableField field : fields) {
        final ProvenanceSearchableFieldDTO searchableField = new ProvenanceSearchableFieldDTO();
        searchableField.setId(field.getIdentifier());
        searchableField.setField(field.getSearchableFieldName());
        searchableField.setLabel(field.getFriendlyName());
        searchableField.setType(field.getFieldType().name());
        searchableFieldNames.add(searchableField);
    }
    final List<SearchableField> searchableAttributes = provenanceRepository.getSearchableAttributes();
    for (final SearchableField searchableAttr : searchableAttributes) {
        final ProvenanceSearchableFieldDTO searchableAttribute = new ProvenanceSearchableFieldDTO();
        searchableAttribute.setId(searchableAttr.getIdentifier());
        searchableAttribute.setField(searchableAttr.getSearchableFieldName());
        searchableAttribute.setLabel(searchableAttr.getFriendlyName());
        searchableAttribute.setType(searchableAttr.getFieldType().name());
        searchableFieldNames.add(searchableAttribute);
    }
    searchOptions.setSearchableFields(searchableFieldNames);
    return searchOptions;
}
 

private ProvenanceRepository createProvenanceRepository(final NiFiProperties properties) throws InstantiationException, IllegalAccessException, ClassNotFoundException {
    final String implementationClassName = properties.getProperty(NiFiProperties.PROVENANCE_REPO_IMPLEMENTATION_CLASS, DEFAULT_PROVENANCE_REPO_IMPLEMENTATION);
    if (implementationClassName == null) {
        throw new RuntimeException("Cannot create Provenance Repository because the NiFi Properties is missing the following property: "
                + NiFiProperties.PROVENANCE_REPO_IMPLEMENTATION_CLASS);
    }

    try {
        return NarThreadContextClassLoader.createInstance(extensionManager, implementationClassName, ProvenanceRepository.class, properties);
    } catch (final Exception e) {
        throw new RuntimeException(e);
    }
}
 

/**
 * Gets the available options for searching provenance.
 *
 * @return the available options for searching provenance
 */
public ProvenanceOptionsDTO getProvenanceSearchOptions() {
    final ProvenanceRepository provenanceRepository = flowController.getProvenanceRepository();

    // create the search options dto
    final ProvenanceOptionsDTO searchOptions = new ProvenanceOptionsDTO();
    final List<ProvenanceSearchableFieldDTO> searchableFieldNames = new ArrayList<>();
    final List<SearchableField> fields = provenanceRepository.getSearchableFields();
    for (final SearchableField field : fields) {
        // we exclude the Event Time because it is always searchable but don't want support querying it this way...
        // we prefer the user queries using startDate and endDate
        if (SearchableFields.EventTime.equals(field)) {
            continue;
        }

        final ProvenanceSearchableFieldDTO searchableField = new ProvenanceSearchableFieldDTO();
        searchableField.setId(field.getIdentifier());
        searchableField.setField(field.getSearchableFieldName());
        searchableField.setLabel(field.getFriendlyName());
        searchableField.setType(field.getFieldType().name());
        searchableFieldNames.add(searchableField);
    }
    final List<SearchableField> searchableAttributes = provenanceRepository.getSearchableAttributes();
    for (final SearchableField searchableAttr : searchableAttributes) {
        final ProvenanceSearchableFieldDTO searchableAttribute = new ProvenanceSearchableFieldDTO();
        searchableAttribute.setId(searchableAttr.getIdentifier());
        searchableAttribute.setField(searchableAttr.getSearchableFieldName());
        searchableAttribute.setLabel(searchableAttr.getFriendlyName());
        searchableAttribute.setType(searchableAttr.getFieldType().name());
        searchableFieldNames.add(searchableAttribute);
    }
    searchOptions.setSearchableFields(searchableFieldNames);
    return searchOptions;
}
 

public StandardLoadBalanceProtocol(final FlowFileRepository flowFileRepository, final ContentRepository contentRepository, final ProvenanceRepository provenanceRepository,
                                   final FlowController flowController, final LoadBalanceAuthorizer authorizer) {
    this.flowFileRepository = flowFileRepository;
    this.contentRepository = contentRepository;
    this.provenanceRepository = provenanceRepository;
    this.flowController = flowController;
    this.authorizer = authorizer;
}
 

private void captureDiagnostics(final ProvenanceRepository repository, final List<String> details) throws IOException {
    details.add("Provenance Repository Implementation: " + repository.getClass().getName());
    for (final String containerName : repository.getContainerNames()) {
        details.add("Provenance Repository <" + containerName + "> File Store: " + repository.getContainerFileStoreName(containerName));
        details.add("Provenance Repository <" + containerName + "> Storage Capacity: " + FormatUtils.formatDataSize(repository.getContainerCapacity(containerName)));
        details.add("Provenance Repository <" + containerName + "> Usable Space: " + FormatUtils.formatDataSize(repository.getContainerUsableSpace(containerName)));
    }
}
 

private ProvenanceRepository createProvenanceRepository(final NiFiProperties properties) throws InstantiationException, IllegalAccessException, ClassNotFoundException {
    final String implementationClassName = properties.getProperty(NiFiProperties.PROVENANCE_REPO_IMPLEMENTATION_CLASS, DEFAULT_PROVENANCE_REPO_IMPLEMENTATION);
    if (implementationClassName == null) {
        throw new RuntimeException("Cannot create Provenance Repository because the NiFi Properties is missing the following property: "
                + NiFiProperties.PROVENANCE_REPO_IMPLEMENTATION_CLASS);
    }

    try {
        return NarThreadContextClassLoader.createInstance(implementationClassName, ProvenanceRepository.class, properties);
    } catch (final Exception e) {
        throw new RuntimeException(e);
    }
}
 

/**
 * Deletes the lineage with the specified id.
 *
 * @param lineageId id
 */
public void deleteLineage(final String lineageId) {
    // get the query to the provenance repository
    final ProvenanceRepository provenanceRepository = flowController.getProvenanceRepository();
    final ComputeLineageSubmission computeLineageSubmission = provenanceRepository.retrieveLineageSubmission(lineageId, NiFiUserUtils.getNiFiUser());
    if (computeLineageSubmission != null) {
        computeLineageSubmission.cancel();
    }
}
 

/**
 * 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();
    }
}
 

/**
 * Gets the lineage with the specified id.
 *
 * @param lineageId id
 * @return the lineage with the specified id
 */
public LineageDTO getLineage(final String lineageId) {
    // get the query to the provenance repository
    final ProvenanceRepository provenanceRepository = flowController.getProvenanceRepository();
    final ComputeLineageSubmission computeLineageSubmission = provenanceRepository.retrieveLineageSubmission(lineageId, NiFiUserUtils.getNiFiUser());

    // ensure the submission was found
    if (computeLineageSubmission == null) {
        throw new ResourceNotFoundException("Cannot find the results for the specified lineage request. Results may have been purged.");
    }

    return dtoFactory.createLineageDto(computeLineageSubmission);
}
 

/**
 * Submits the specified lineage request.
 *
 * @param lineageDto dto
 * @return updated lineage
 */
public LineageDTO submitLineage(LineageDTO lineageDto) {
    final LineageRequestDTO requestDto = lineageDto.getRequest();

    // get the provenance repo
    final ProvenanceRepository provenanceRepository = flowController.getProvenanceRepository();
    final ComputeLineageSubmission result;

    // submit the event
    if (LineageRequestType.FLOWFILE.equals(requestDto.getLineageRequestType())) {
        // submit uuid
        if (requestDto.getEventId() == null) {
            result = provenanceRepository.submitLineageComputation(requestDto.getUuid(), NiFiUserUtils.getNiFiUser());
        } else {
            result = provenanceRepository.submitLineageComputation(requestDto.getEventId(), NiFiUserUtils.getNiFiUser());
        }
    } else {
        // submit event... (parents or children)
        if (LineageRequestType.PARENTS.equals(requestDto.getLineageRequestType())) {
            result = provenanceRepository.submitExpandParents(requestDto.getEventId(), NiFiUserUtils.getNiFiUser());
        } else {
            result = provenanceRepository.submitExpandChildren(requestDto.getEventId(), NiFiUserUtils.getNiFiUser());
        }
    }

    return getLineage(result.getLineageIdentifier());
}
 

public ProvenanceRepository getProvenanceRepository() {
    return provenanceRepository;
}
 

protected final void initialize(final NiFiProperties nifiProperties) throws IOException {
    this.nifiProperties = nifiProperties;

    final FlowFileEventRepository flowFileEventRepository = new RingBufferEventRepository(5);

    final BulletinRepository bulletinRepo = new VolatileBulletinRepository();
    flowEngine = new FlowEngine(4, "unit test flow engine");
    extensionManager = new DirectInjectionExtensionManager();

    extensionManager.injectExtensionType(FlowFileRepository.class, WriteAheadFlowFileRepository.class);
    extensionManager.injectExtensionType(ContentRepository.class, FileSystemRepository.class);
    extensionManager.injectExtensionType(ProvenanceRepository.class, WriteAheadProvenanceRepository.class);
    extensionManager.injectExtensionType(StateProvider.class, WriteAheadLocalStateProvider.class);
    extensionManager.injectExtensionType(ComponentStatusRepository.class, VolatileComponentStatusRepository.class);
    extensionManager.injectExtensionType(FlowFileSwapManager.class, FileSystemSwapManager.class);

    extensionManager.injectExtensionType(Processor.class, BiConsumerProcessor.class);
    extensionManager.injectExtensionType(Processor.class, GenerateProcessor.class);
    extensionManager.injectExtensionType(Processor.class, TerminateOnce.class);
    extensionManager.injectExtensionType(Processor.class, TerminateAll.class);
    extensionManager.injectExtensionType(Processor.class, NopProcessor.class);
    extensionManager.injectExtensionType(Processor.class, UsernamePasswordProcessor.class);

    injectExtensionTypes(extensionManager);
    systemBundle = SystemBundle.create(nifiProperties);
    extensionManager.discoverExtensions(systemBundle, Collections.emptySet());

    final StringEncryptor encryptor = StringEncryptor.createEncryptor("PBEWITHMD5AND256BITAES-CBC-OPENSSL", "BC", "unit-test");
    final Authorizer authorizer = new AlwaysAuthorizedAuthorizer();
    final AuditService auditService = new NopAuditService();

    if (isClusteredTest()) {
        final File zookeeperDir = new File("target/state/zookeeper");
        final File version2Dir =  new File(zookeeperDir, "version-2");

        if (!version2Dir.exists()) {
            assertTrue(version2Dir.mkdirs());
        }

        final File[] children = version2Dir.listFiles();
        if (children != null) {
            for (final File file : children) {
                FileUtils.deleteFile(file, true);
            }
        }

        clusterCoordinator = Mockito.mock(ClusterCoordinator.class);
        final HeartbeatMonitor heartbeatMonitor = Mockito.mock(HeartbeatMonitor.class);
        final NodeProtocolSender protocolSender = Mockito.mock(NodeProtocolSender.class);
        final LeaderElectionManager leaderElectionManager = new CuratorLeaderElectionManager(2, nifiProperties);

        final NodeIdentifier localNodeId = new NodeIdentifier(UUID.randomUUID().toString(), "localhost", 8111, "localhost", 8081,
            "localhost", 8082, "localhost", 8083, 8084, false, Collections.emptySet());
        final NodeIdentifier node2Id = new NodeIdentifier(UUID.randomUUID().toString(), "localhost", 8222, "localhost", 8081,
            "localhost", 8082, "localhost", 8083, 8084, false, Collections.emptySet());

        final Set<NodeIdentifier> nodeIdentifiers = new HashSet<>();
        nodeIdentifiers.add(localNodeId);
        nodeIdentifiers.add(node2Id);
        Mockito.when(clusterCoordinator.getNodeIdentifiers()).thenReturn(nodeIdentifiers);
        Mockito.when(clusterCoordinator.getLocalNodeIdentifier()).thenReturn(localNodeId);

        flowController = FlowController.createClusteredInstance(flowFileEventRepository, nifiProperties, authorizer, auditService, encryptor, protocolSender, bulletinRepo, clusterCoordinator,
            heartbeatMonitor, leaderElectionManager, VariableRegistry.ENVIRONMENT_SYSTEM_REGISTRY, flowRegistryClient, extensionManager);

        flowController.setClustered(true, UUID.randomUUID().toString());
        flowController.setNodeId(localNodeId);

        flowController.setConnectionStatus(new NodeConnectionStatus(localNodeId, NodeConnectionState.CONNECTED));
    } else {
        flowController = FlowController.createStandaloneInstance(flowFileEventRepository, nifiProperties, authorizer, auditService, encryptor, bulletinRepo,
            VariableRegistry.ENVIRONMENT_SYSTEM_REGISTRY, flowRegistryClient, extensionManager);
    }

    processScheduler = new StandardProcessScheduler(flowEngine, flowController, encryptor, flowController.getStateManagerProvider(), nifiProperties);

    final RepositoryContextFactory repositoryContextFactory = flowController.getRepositoryContextFactory();
    final SchedulingAgent timerDrivenSchedulingAgent = new TimerDrivenSchedulingAgent(flowController, flowEngine, repositoryContextFactory, encryptor, nifiProperties);
    processScheduler.setSchedulingAgent(SchedulingStrategy.TIMER_DRIVEN, timerDrivenSchedulingAgent);

    flowFileSwapManager = flowController.createSwapManager();
    resourceClaimManager = flowController.getResourceClaimManager();
}
 

public ProvenanceRepository getProvenanceRepository() {
    return provenanceRepo;
}
 

/**
 * 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);
    }
}
 

@Override
public ProvenanceRepository getProvenanceRepository() {
    return flowController.getProvenanceRepository();
}
 

/**
 * 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());
}
 

private void test(TestConfiguration tc) throws Exception {
    final ReportLineageToAtlas reportingTask = new ReportLineageToAtlas();
    final MockComponentLog logger = new MockComponentLog("reporting-task-id", reportingTask);

    final ReportingInitializationContext initializationContext = mock(ReportingInitializationContext.class);
    when(initializationContext.getLogger()).thenReturn(logger);
    final ConfigurationContext configurationContext = new MockConfigurationContext(tc.properties, null);
    final ValidationContext validationContext = mock(ValidationContext.class);
    when(validationContext.getProperty(any())).then(invocation -> new MockPropertyValue(tc.properties.get(invocation.getArguments()[0])));
    final ReportingContext reportingContext = mock(ReportingContext.class);
    final MockStateManager stateManager = new MockStateManager(reportingTask);
    final EventAccess eventAccess = mock(EventAccess.class);
    when(reportingContext.getProperties()).thenReturn(tc.properties);
    when(reportingContext.getProperty(any())).then(invocation -> new MockPropertyValue(tc.properties.get(invocation.getArguments()[0])));
    when(reportingContext.getStateManager()).thenReturn(stateManager);
    when(reportingContext.getEventAccess()).thenReturn(eventAccess);
    when(eventAccess.getGroupStatus(eq("root"))).thenReturn(tc.rootPgStatus);

    final ProvenanceRepository provenanceRepository = mock(ProvenanceRepository.class);
    when(eventAccess.getControllerStatus()).thenReturn(tc.rootPgStatus);
    when(eventAccess.getProvenanceRepository()).thenReturn(provenanceRepository);
    when(eventAccess.getProvenanceEvents(eq(-1L), anyInt())).thenReturn(tc.provenanceRecords);
    when(provenanceRepository.getMaxEventId()).thenReturn((long) tc.provenanceRecords.size() - 1);
    when(provenanceRepository.getEvent(anyLong())).then(invocation -> tc.provenanceRecords.get(((Long) invocation.getArguments()[0]).intValue()));

    // To mock this async method invocations, keep the requested event ids in a stack.
    final ComputeLineageSubmission lineageComputationSubmission = mock(ComputeLineageSubmission.class);
    when(provenanceRepository.submitLineageComputation(anyLong(), any())).thenAnswer(invocation -> {
        requestedLineageComputationIds.push((Long) invocation.getArguments()[0]);
        return lineageComputationSubmission;
    });
    when(lineageComputationSubmission.getResult()).then(invocation -> tc.lineageResults.get(requestedLineageComputationIds.pop()));

    final ComputeLineageSubmission expandParentsSubmission = mock(ComputeLineageSubmission.class);
    when(provenanceRepository.submitExpandParents(anyLong(), any())).thenAnswer(invocation -> {
        requestedExpandParentsIds.push(((Long) invocation.getArguments()[0]));
        return expandParentsSubmission;
    });
    when(expandParentsSubmission.getResult()).then(invocation -> tc.parentLineageResults.get(requestedExpandParentsIds.pop()));

    tc.properties.put(ATLAS_NIFI_URL, "http://localhost:8080/nifi");
    tc.properties.put(ATLAS_URLS, TARGET_ATLAS_URL);
    tc.properties.put(ATLAS_USER, "admin");
    tc.properties.put(ATLAS_PASSWORD, "admin");
    tc.properties.put(new PropertyDescriptor.Builder().name("hostnamePattern.example").dynamic(true).build(), ".*");


    reportingTask.initialize(initializationContext);
    reportingTask.validate(validationContext);
    reportingTask.setup(configurationContext);
    reportingTask.onTrigger(reportingContext);
    reportingTask.onUnscheduled();
    reportingTask.onStopped();
}
 

public StandardAnalysisContext(NiFiFlow nifiFlow, NamespaceResolver namespaceResolver,
                               ProvenanceRepository provenanceRepository) {
    this.nifiFlow = nifiFlow;
    this.namespaceResolver = namespaceResolver;
    this.provenanceRepository = provenanceRepository;
}
 

@Override
public ProvenanceRepository getProvenanceRepository() {
    return provenanceRepository;
}
 

/**
 * 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());
}