Java Code Examples for org.apache.nifi.util.Tuple#getValue()

public void releaseWriteLock(final File file) {
    final String key = getMapKey(file);

    boolean updated = false;
    while (!updated) {
        final Tuple<ReadWriteLock, Integer> tuple = lockMap.get(key);
        if (tuple == null) {
            throw new IllegalMonitorStateException("Lock is not owned");
        }

        // If this is the only reference to the lock, remove it from the map and then unlock.
        if (tuple.getValue() <= 1) {
            updated = lockMap.remove(key, tuple);
            if (updated) {
                tuple.getKey().writeLock().unlock();
            }
        } else {
            final Tuple<ReadWriteLock, Integer> updatedTuple = new Tuple<>(tuple.getKey(), tuple.getValue() - 1);
            updated = lockMap.replace(key, tuple, updatedTuple);
            if (updated) {
                tuple.getKey().writeLock().unlock();
            }
        }
    }
}
 

@Override
public ValidationResult validate(final String subject, final String input, final ValidationContext validationContext) {
    final Tuple<String, ValidationResult> lastResult = this.cachedResult;
    if (lastResult != null && lastResult.getKey().equals(input)) {
        return lastResult.getValue();
    } else {
        String error = null;
        final File stylesheet = new File(input);
        final TransformerFactory tFactory = new net.sf.saxon.TransformerFactoryImpl();
        final StreamSource styleSource = new StreamSource(stylesheet);

        try {
            tFactory.newTransformer(styleSource);
        } catch (final Exception e) {
            error = e.toString();
        }

        this.cachedResult = new Tuple<>(input, new ValidationResult.Builder()
                .input(input)
                .subject(subject)
                .valid(error == null)
                .explanation(error)
                .build());
        return this.cachedResult.getValue();
    }
}
 

public void releaseWriteLock(final File file) {
    final String key = getMapKey(file);

    boolean updated = false;
    while (!updated) {
        final Tuple<ReadWriteLock, Integer> tuple = lockMap.get(key);
        if (tuple == null) {
            throw new IllegalMonitorStateException("Lock is not owned");
        }

        // If this is the only reference to the lock, remove it from the map and then unlock.
        if (tuple.getValue() <= 1) {
            updated = lockMap.remove(key, tuple);
            if (updated) {
                tuple.getKey().writeLock().unlock();
            }
        } else {
            final Tuple<ReadWriteLock, Integer> updatedTuple = new Tuple<>(tuple.getKey(), tuple.getValue() - 1);
            updated = lockMap.replace(key, tuple, updatedTuple);
            if (updated) {
                tuple.getKey().writeLock().unlock();
            }
        }
    }
}
 

private void writeRecord(final Record record, final Relationship relationship, final Map<Relationship, Tuple<FlowFile, RecordSetWriter>> writers, final ProcessSession session,
                         final FlowFile original, final Map<String, String> originalAttributes, final RecordSetWriterFactory writerFactory) throws IOException, SchemaNotFoundException {
    final RecordSetWriter recordSetWriter;
    Tuple<FlowFile, RecordSetWriter> tuple = writers.get(relationship);

    if (tuple == null) {
        final FlowFile outFlowFile = session.create(original);
        final OutputStream out = session.write(outFlowFile);

        final RecordSchema recordWriteSchema = writerFactory.getSchema(originalAttributes, record.getSchema());
        recordSetWriter = writerFactory.createWriter(getLogger(), recordWriteSchema, out, outFlowFile);
        recordSetWriter.beginRecordSet();

        tuple = new Tuple<>(outFlowFile, recordSetWriter);
        writers.put(relationship, tuple);
    } else {
        recordSetWriter = tuple.getValue();
    }

    recordSetWriter.write(record);
}
 

@Override
public ValidationResult validate(final String subject, final String input, final ValidationContext validationContext) {
    final Tuple<String, ValidationResult> lastResult = this.cachedResult;
    if (lastResult != null && lastResult.getKey().equals(input)) {
        return lastResult.getValue();
    } else {
        String error = null;
        final File stylesheet = new File(input);
        final TransformerFactory tFactory = new net.sf.saxon.TransformerFactoryImpl();
        final StreamSource styleSource = new StreamSource(stylesheet);

        try {
            tFactory.newTransformer(styleSource);
        } catch (final Exception e) {
            error = e.toString();
        }

        this.cachedResult = new Tuple<>(input, new ValidationResult.Builder()
                .input(input)
                .subject(subject)
                .valid(error == null)
                .explanation(error)
                .build());
        return this.cachedResult.getValue();
    }
}
 

private Tuple<ProcessGroupStatus, String> getNextProcessGroupStatus() {
    Tuple<Iterator<ProcessGroupStatus>, String> top = iteratorBreadcrumb.peek();
    if (top == null) {
        return null;
    }
    Iterator<ProcessGroupStatus> i = top.getKey();
    String parentId = top.getValue();

    if (i.hasNext()) {
        ProcessGroupStatus nextPG = i.next();
        iteratorBreadcrumb.push(new Tuple<>(nextPG.getProcessGroupStatus().iterator(), nextPG.getId()));
        return new Tuple<>(nextPG, parentId);
    } else {
        // No more child PGs, remove it from the breadcrumb trail and try again
        Tuple<Iterator<ProcessGroupStatus>, String> pop = iteratorBreadcrumb.pop();
        return getNextProcessGroupStatus();
    }
}
 

public static Set<ConfigurableComponent> getRestrictedComponents(final VersionedProcessGroup group, final NiFiServiceFacade serviceFacade) {
    final Set<ConfigurableComponent> restrictedComponents = new HashSet<>();

    final Set<Tuple<String, BundleCoordinate>> componentTypes = new HashSet<>();
    populateComponentTypes(group, componentTypes);

    for (final Tuple<String, BundleCoordinate> tuple : componentTypes) {
        final ConfigurableComponent component = serviceFacade.getTempComponent(tuple.getKey(), tuple.getValue());
        if (component == null) {
            throw new NiFiCoreException("Could not create an instance of component " + tuple.getKey() + " using bundle coordinates " + tuple.getValue());
        }

        final boolean isRestricted = component.getClass().isAnnotationPresent(Restricted.class);
        if (isRestricted) {
            restrictedComponents.add(component);
        }
    }

    return restrictedComponents;
}
 

/**
 * Tests attempting to copy/paste a restricted processor.
 *
 * @throws Exception ex
 */
@Test
public void testCopyPasteRestrictedProcessor() throws Exception {
    final String copyUrl = helper.getBaseUrl() + "/process-groups/root/snippet-instance";
    final Tuple<ProcessorEntity, SnippetEntity> tuple = createSnippetWithRestrictedComponent();
    final SnippetEntity snippetEntity = tuple.getValue();

    // build the copy/paste request
    final CopySnippetRequestEntity copyRequest = new CopySnippetRequestEntity();
    copyRequest.setSnippetId(snippetEntity.getSnippet().getId());
    copyRequest.setOriginX(0.0);
    copyRequest.setOriginY(0.0);

    // create the snippet
    ClientResponse response = helper.getReadWriteUser().testPost(copyUrl, copyRequest);

    // ensure the request failed... need privileged users since snippet comprised of the restricted components
    assertEquals(403, response.getStatus());

    // create the snippet
    response = helper.getPrivilegedUser().testPost(copyUrl, copyRequest);

    // ensure the request is successful
    assertEquals(201, response.getStatus());

    final FlowEntity flowEntity = response.getEntity(FlowEntity.class);

    // remove the restricted processors
    deleteRestrictedComponent(tuple.getKey());
    deleteRestrictedComponent(flowEntity.getFlow().getProcessors().stream().findFirst().orElse(null));
}
 

private SchemaVersionInfo getSchemaVersionInfo(final SchemaRegistryClient client, final SchemaVersionKey key) throws org.apache.nifi.schema.access.SchemaNotFoundException {
    try {
        // Try to fetch the SchemaVersionInfo from the cache.
        final Tuple<SchemaVersionInfo, Long> timestampedVersionInfo = schemaVersionByKeyCache.get(key);

        // Determine if the timestampedVersionInfo is expired
        boolean fetch = false;
        if (timestampedVersionInfo == null) {
            fetch = true;
        } else {
            final long minTimestamp = System.nanoTime() - versionInfoCacheNanos;
            fetch = timestampedVersionInfo.getValue() < minTimestamp;
        }

        // If not expired, use what we got from the cache
        if (!fetch) {
            return timestampedVersionInfo.getKey();
        }

        // schema version info was expired or not found in cache. Fetch from schema registry
        final SchemaVersionInfo versionInfo = client.getSchemaVersionInfo(key);
        if (versionInfo == null) {
            throw new org.apache.nifi.schema.access.SchemaNotFoundException("Could not find schema with name '" + key.getSchemaName() + "' and version " + key.getVersion());
        }

        // Store new version in cache.
        final Tuple<SchemaVersionInfo, Long> tuple = new Tuple<>(versionInfo, System.nanoTime());
        schemaVersionByKeyCache.put(key, tuple);
        return versionInfo;
    } catch (final SchemaNotFoundException e) {
        throw new org.apache.nifi.schema.access.SchemaNotFoundException(e);
    }
}
 

/**
 * Tests attempting to copy/paste a restricted processor.
 *
 * @throws Exception ex
 */
@Test
public void testCopyPasteRestrictedProcessor() throws Exception {
    final String copyUrl = helper.getBaseUrl() + "/process-groups/root/snippet-instance";
    final Tuple<ProcessorEntity, SnippetEntity> tuple = createSnippetWithRestrictedComponent(RestrictedProcessor.class.getName(), helper.getPrivilegedUser());
    final SnippetEntity snippetEntity = tuple.getValue();

    // build the copy/paste request
    final CopySnippetRequestEntity copyRequest = new CopySnippetRequestEntity();
    copyRequest.setSnippetId(snippetEntity.getSnippet().getId());
    copyRequest.setOriginX(0.0);
    copyRequest.setOriginY(0.0);

    // create the snippet
    Response response = helper.getReadWriteUser().testPost(copyUrl, copyRequest);

    // ensure the request failed... need privileged users since snippet comprised of the restricted components
    assertEquals(403, response.getStatus());

    // perform the request as a user with read/write and only execute code restricted access
    response = helper.getExecuteCodeUser().testPost(copyUrl, copyRequest);

    // ensure the request is successful
    assertEquals(403, response.getStatus());

    // create the snippet
    response = helper.getPrivilegedUser().testPost(copyUrl, copyRequest);

    // ensure the request is successful
    assertEquals(201, response.getStatus());

    final FlowEntity flowEntity = response.readEntity(FlowEntity.class);

    // remove the restricted processors
    deleteRestrictedComponent(tuple.getKey(), helper.getPrivilegedUser());
    deleteRestrictedComponent(flowEntity.getFlow().getProcessors().stream().findFirst().orElse(null), helper.getPrivilegedUser());
}
 

private void copyPasteExecuteCodeRestrictedProcessor(final NiFiTestUser user) throws Exception {
    final String copyUrl = helper.getBaseUrl() + "/process-groups/root/snippet-instance";
    final Tuple<ProcessorEntity, SnippetEntity> tuple = createSnippetWithRestrictedComponent(ExecuteCodeRestrictedProcessor.class.getName(), user);
    final SnippetEntity snippetEntity = tuple.getValue();

    // build the copy/paste request
    final CopySnippetRequestEntity copyRequest = new CopySnippetRequestEntity();
    copyRequest.setSnippetId(snippetEntity.getSnippet().getId());
    copyRequest.setOriginX(0.0);
    copyRequest.setOriginY(0.0);

    // create the snippet
    Response response = helper.getReadWriteUser().testPost(copyUrl, copyRequest);

    // ensure the request failed... need privileged users since snippet comprised of the restricted components
    assertEquals(403, response.getStatus());

    // perform the request as a user with read/write and only execute code restricted access
    response = user.testPost(copyUrl, copyRequest);

    // ensure the request is successful
    assertEquals(201, response.getStatus());

    final FlowEntity flowEntity = response.readEntity(FlowEntity.class);

    // remove the restricted processors
    deleteRestrictedComponent(tuple.getKey(), user);
    deleteRestrictedComponent(flowEntity.getFlow().getProcessors().stream().findFirst().orElse(null), user);
}
 

/**
 * Tests attempting to use a template with a restricted processor.
 *
 * @throws Exception ex
 */
@Test
public void testTemplateWithRestrictedProcessor() throws Exception {
    final String createTemplateUrl = helper.getBaseUrl() + "/process-groups/root/templates";
    final String instantiateTemplateUrl = helper.getBaseUrl() + "/process-groups/root/template-instance";
    final Tuple<ProcessorEntity, SnippetEntity> tuple = createSnippetWithRestrictedComponent();
    final SnippetEntity snippetEntity = tuple.getValue();

    // create the template
    final CreateTemplateRequestEntity createTemplateRequest = new CreateTemplateRequestEntity();
    createTemplateRequest.setSnippetId(snippetEntity.getSnippet().getId());
    createTemplateRequest.setName("test");

    // create the snippet
    ClientResponse response = helper.getWriteUser().testPost(createTemplateUrl, createTemplateRequest);

    // ensure the request failed... need read perms to the components in the snippet
    assertEquals(403, response.getStatus());

    response = helper.getReadWriteUser().testPost(createTemplateUrl, createTemplateRequest);

    // ensure the request is successfull
    assertEquals(201, response.getStatus());

    final TemplateEntity templateEntity = response.getEntity(TemplateEntity.class);

    // build the template request
    final InstantiateTemplateRequestEntity instantiateTemplateRequest = new InstantiateTemplateRequestEntity();
    instantiateTemplateRequest.setTemplateId(templateEntity.getTemplate().getId());
    instantiateTemplateRequest.setOriginX(0.0);
    instantiateTemplateRequest.setOriginY(0.0);

    // create the snippet
    response = helper.getReadWriteUser().testPost(instantiateTemplateUrl, instantiateTemplateRequest);

    // ensure the request failed... need privileged user since the template is comprised of restricted components
    assertEquals(403, response.getStatus());

    // create the snippet
    response = helper.getPrivilegedUser().testPost(instantiateTemplateUrl, instantiateTemplateRequest);

    // ensure the request is successful
    assertEquals(201, response.getStatus());

    final FlowEntity flowEntity = response.getEntity(FlowEntity.class);

    // clean up the resources created during this test
    deleteTemplate(templateEntity);
    deleteRestrictedComponent(tuple.getKey());
    deleteRestrictedComponent(flowEntity.getFlow().getProcessors().stream().findFirst().orElse(null));
}
 

private SchemaVersionInfo getLatestSchemaVersionInfo(final SchemaRegistryClient client, final String schemaName, final String branchName)
        throws org.apache.nifi.schema.access.SchemaNotFoundException {
    try {
        // Try to fetch the SchemaVersionInfo from the cache.
        final Tuple<String,String> nameAndBranch = new Tuple<>(schemaName, branchName);
        final Tuple<SchemaVersionInfo, Long> timestampedVersionInfo = schemaVersionByNameCache.get(nameAndBranch);

        // Determine if the timestampedVersionInfo is expired
        boolean fetch = false;
        if (timestampedVersionInfo == null) {
            fetch = true;
        } else {
            final long minTimestamp = System.nanoTime() - versionInfoCacheNanos;
            fetch = timestampedVersionInfo.getValue() < minTimestamp;
        }

        // If not expired, use what we got from the cache
        if (!fetch) {
            return timestampedVersionInfo.getKey();
        }

        // schema version info was expired or not found in cache. Fetch from schema registry
        final SchemaVersionInfo versionInfo;
        if (StringUtils.isBlank(branchName)) {
            versionInfo = client.getLatestSchemaVersionInfo(schemaName);
        } else {
            versionInfo = client.getLatestSchemaVersionInfo(branchName, schemaName);
        }

        if (versionInfo == null) {
            throw new org.apache.nifi.schema.access.SchemaNotFoundException("Could not find schema with name '" + schemaName + "'");
        }

        // Store new version in cache.
        final Tuple<SchemaVersionInfo, Long> tuple = new Tuple<>(versionInfo, System.nanoTime());
        schemaVersionByNameCache.put(nameAndBranch, tuple);
        return versionInfo;
    } catch (final SchemaNotFoundException e) {
        throw new org.apache.nifi.schema.access.SchemaNotFoundException(e);
    }
}
 

@Override
public void processEvent(AnalysisContext analysisContext, NiFiFlow nifiFlow, ProvenanceEventRecord event) {
    if (!ProvenanceEventType.DROP.equals(event.getEventType())) {
        return;
    }
    final ComputeLineageResult lineage = analysisContext.queryLineage(event.getEventId());

    // Construct a tree model to traverse backwards.
    final Map<String, List<LineageNode>> lineageTree = new HashMap<>();
    analyzeLineageTree(lineage, lineageTree);

    final LineagePath lineagePath = new LineagePath();
    extractLineagePaths(analysisContext, lineageTree, lineagePath, event);

    analyzeLineagePath(analysisContext, lineagePath);

    // Input and output data set are both required to report lineage.
    List<Tuple<NiFiFlowPath, DataSetRefs>> createdFlowPaths = new ArrayList<>();
    if (lineagePath.isComplete()) {
        createCompleteFlowPath(nifiFlow, lineagePath, createdFlowPaths);
        for (Tuple<NiFiFlowPath, DataSetRefs> createdFlowPath : createdFlowPaths) {
            final NiFiFlowPath flowPath = createdFlowPath.getKey();
            // NOTE 1: FlowPath creation and DataSet references should be reported separately
            // ------------------------------------------------------------------------------
            // For example, with following provenance event inputs:
            //   CREATE(F1), FORK (F1 -> F2, F3), DROP(F1), SEND (F2), SEND(F3), DROP(F2), DROP(F3),
            // there is no guarantee that DROP(F2) and DROP(F3) are processed within the same cycle.
            // If DROP(F3) is processed in different cycle, it needs to be added to the existing FlowPath
            // that contains F1 -> F2, to be F1 -> F2, F3.
            // Execution cycle 1: Path1 (source of F1 -> ForkA), ForkA_queue (F1 -> F2), Path2 (ForkA -> dest of F2)
            // Execution cycle 2: Path1 (source of F1 -> ForkB), ForkB_queue (F1 -> F3), Path3 (ForkB -> dest of F3)

            // NOTE 2: Both FlowPath creation and FlowPath update messages are required
            // ------------------------------------------------------------------------
            // For the 1st time when a lineage is found, nifi_flow_path and referred DataSets are created.
            // If we notify these entities by a create 3 entities message (Path1, DataSet1, DataSet2)
            // followed by 1 partial update message to add lineage (DataSet1 -> Path1 -> DataSet2), then
            // the update message may arrive at Atlas earlier than the create message gets processed.
            // If that happens, lineage among these entities will be missed.
            // But as written in NOTE1, there is a case where existing nifi_flow_paths need to be updated.
            // Also, we don't know if this is the 1st time or 2nd or later.
            // So, we need to notify entity creation and also partial update messages.

            // Create flow path entity with DataSet refs.
            final Referenceable flowPathRef = toReferenceable(flowPath, nifiFlow);
            final DataSetRefs dataSetRefs = createdFlowPath.getValue();
            addDataSetRefs(dataSetRefs.getInputs(), flowPathRef, ATTR_INPUTS);
            addDataSetRefs(dataSetRefs.getOutputs(), flowPathRef, ATTR_OUTPUTS);
            createEntity(flowPathRef);
            // Also, sending partial update message to update existing flow_path.
            addDataSetRefs(nifiFlow, Collections.singleton(flowPath), createdFlowPath.getValue());

        }
        createdFlowPaths.clear();
    }
}
 

/**
 * Tests attempting to use a template with a restricted processor.
 *
 * @throws Exception ex
 */
@Test
public void testTemplateWithRestrictedProcessor() throws Exception {
    final String createTemplateUrl = helper.getBaseUrl() + "/process-groups/root/templates";
    final String instantiateTemplateUrl = helper.getBaseUrl() + "/process-groups/root/template-instance";
    final Tuple<ProcessorEntity, SnippetEntity> tuple = createSnippetWithRestrictedComponent(RestrictedProcessor.class.getName(), helper.getPrivilegedUser());
    final SnippetEntity snippetEntity = tuple.getValue();

    // create the template
    final CreateTemplateRequestEntity createTemplateRequest = new CreateTemplateRequestEntity();
    createTemplateRequest.setSnippetId(snippetEntity.getSnippet().getId());
    createTemplateRequest.setName("test");

    // create the snippet
    Response response = helper.getWriteUser().testPost(createTemplateUrl, createTemplateRequest);

    // ensure the request failed... need read perms to the components in the snippet
    assertEquals(403, response.getStatus());

    response = helper.getReadWriteUser().testPost(createTemplateUrl, createTemplateRequest);

    // ensure the request is successful
    assertEquals(201, response.getStatus());

    final TemplateEntity templateEntity = response.readEntity(TemplateEntity.class);

    // build the template request
    final InstantiateTemplateRequestEntity instantiateTemplateRequest = new InstantiateTemplateRequestEntity();
    instantiateTemplateRequest.setTemplateId(templateEntity.getTemplate().getId());
    instantiateTemplateRequest.setOriginX(0.0);
    instantiateTemplateRequest.setOriginY(0.0);

    // create the snippet
    response = helper.getReadWriteUser().testPost(instantiateTemplateUrl, instantiateTemplateRequest);

    // ensure the request failed... need privileged user since the template is comprised of restricted components
    assertEquals(403, response.getStatus());

    // create the snippet
    response = helper.getExecuteCodeUser().testPost(instantiateTemplateUrl, instantiateTemplateRequest);

    // ensure the request failed... need privileged user since the template is comprised of restricted components
    assertEquals(403, response.getStatus());

    // create the snippet
    response = helper.getPrivilegedUser().testPost(instantiateTemplateUrl, instantiateTemplateRequest);

    // ensure the request is successful
    assertEquals(201, response.getStatus());

    final FlowEntity flowEntity = response.readEntity(FlowEntity.class);

    // clean up the resources created during this test
    deleteTemplate(templateEntity);
    deleteRestrictedComponent(tuple.getKey(), helper.getPrivilegedUser());
    deleteRestrictedComponent(flowEntity.getFlow().getProcessors().stream().findFirst().orElse(null), helper.getPrivilegedUser());
}
 

private void templateWithExecuteCodeRestrictedProcessor(final NiFiTestUser user) throws Exception {
    final String createTemplateUrl = helper.getBaseUrl() + "/process-groups/root/templates";
    final String instantiateTemplateUrl = helper.getBaseUrl() + "/process-groups/root/template-instance";
    final Tuple<ProcessorEntity, SnippetEntity> tuple = createSnippetWithRestrictedComponent(ExecuteCodeRestrictedProcessor.class.getName(), helper.getPrivilegedUser());
    final SnippetEntity snippetEntity = tuple.getValue();

    // create the template
    final CreateTemplateRequestEntity createTemplateRequest = new CreateTemplateRequestEntity();
    createTemplateRequest.setSnippetId(snippetEntity.getSnippet().getId());
    createTemplateRequest.setName("test");

    // create the snippet
    Response response = helper.getWriteUser().testPost(createTemplateUrl, createTemplateRequest);

    // ensure the request failed... need read perms to the components in the snippet
    assertEquals(403, response.getStatus());

    response = helper.getReadWriteUser().testPost(createTemplateUrl, createTemplateRequest);

    // ensure the request is successful
    assertEquals(201, response.getStatus());

    final TemplateEntity templateEntity = response.readEntity(TemplateEntity.class);

    // build the template request
    final InstantiateTemplateRequestEntity instantiateTemplateRequest = new InstantiateTemplateRequestEntity();
    instantiateTemplateRequest.setTemplateId(templateEntity.getTemplate().getId());
    instantiateTemplateRequest.setOriginX(0.0);
    instantiateTemplateRequest.setOriginY(0.0);

    // create the snippet
    response = helper.getReadWriteUser().testPost(instantiateTemplateUrl, instantiateTemplateRequest);

    // ensure the request failed... need privileged user since the template is comprised of restricted components
    assertEquals(403, response.getStatus());

    // create the snippet
    response = user.testPost(instantiateTemplateUrl, instantiateTemplateRequest);

    // ensure the request is successful
    assertEquals(201, response.getStatus());

    final FlowEntity flowEntity = response.readEntity(FlowEntity.class);

    // clean up the resources created during this test
    deleteTemplate(templateEntity);
    deleteRestrictedComponent(tuple.getKey(), user);
    deleteRestrictedComponent(flowEntity.getFlow().getProcessors().stream().findFirst().orElse(null), user);
}