Examples with ProvenanceRepository org.apache.nifi.provenance.ProvenanceRepository used on opensource projects

Example 6 with ProvenanceRepository

use of org.apache.nifi.provenance.ProvenanceRepository in project nifi by apache.

the class ITReportLineageToAtlas method test.

private void test(TestConfiguration tc) throws InitializationException, IOException {
    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();
}

Example 7 with ProvenanceRepository

use of org.apache.nifi.provenance.ProvenanceRepository in project nifi by apache.

the class ControllerFacade method deleteProvenanceQuery.

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

Example 8 with ProvenanceRepository

use of org.apache.nifi.provenance.ProvenanceRepository in project nifi by apache.

the class ControllerFacade method submitLineage.

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