Java Code Examples for org.apache.nifi.registry.flow.VersionedFlowSnapshot#setExternalControllerServices()

@Test
public void testSerializeDeserializeWithExternalServices() throws SerializationException {
    final VersionedProcessGroup processGroup1 = new VersionedProcessGroup();
    processGroup1.setIdentifier("pg1");
    processGroup1.setName("My Process Group");

    final ExternalControllerServiceReference serviceReference1 = new ExternalControllerServiceReference();
    serviceReference1.setIdentifier("1");
    serviceReference1.setName("Service 1");

    final ExternalControllerServiceReference serviceReference2 = new ExternalControllerServiceReference();
    serviceReference2.setIdentifier("2");
    serviceReference2.setName("Service 2");

    final Map<String,ExternalControllerServiceReference> serviceReferences = new HashMap<>();
    serviceReferences.put(serviceReference1.getIdentifier(), serviceReference1);
    serviceReferences.put(serviceReference2.getIdentifier(), serviceReference2);

    final VersionedFlowSnapshot snapshot = new VersionedFlowSnapshot();
    snapshot.setFlowContents(processGroup1);
    snapshot.setExternalControllerServices(serviceReferences);

    final FlowContent flowContent = new FlowContent();
    flowContent.setFlowSnapshot(snapshot);

    final ByteArrayOutputStream out = new ByteArrayOutputStream();
    serializer.serializeFlowContent(flowContent, out);

    final ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray());

    // make sure we can read the version from the input stream and it should be the current version
    final Integer version = serializer.readDataModelVersion(in);
    assertEquals(serializer.getCurrentDataModelVersion(), version);

    // make sure we can deserialize back to FlowContent
    final FlowContent deserializedFlowContent = serializer.deserializeFlowContent(version, in);
    assertNotNull(deserializedFlowContent);

    final VersionedFlowSnapshot deserializedSnapshot = deserializedFlowContent.getFlowSnapshot();
    assertNotNull(deserializedSnapshot);

    final VersionedProcessGroup deserializedProcessGroup = deserializedSnapshot.getFlowContents();
    assertEquals(processGroup1.getIdentifier(), deserializedProcessGroup.getIdentifier());
    assertEquals(processGroup1.getName(), deserializedProcessGroup.getName());

    final Map<String,ExternalControllerServiceReference> deserializedServiceReferences = deserializedSnapshot.getExternalControllerServices();
    assertNotNull(deserializedServiceReferences);
    assertEquals(2, deserializedServiceReferences.size());

    final ExternalControllerServiceReference deserializedServiceReference1 = deserializedServiceReferences.get(serviceReference1.getIdentifier());
    assertNotNull(deserializedServiceReference1);
    assertEquals(serviceReference1.getIdentifier(), deserializedServiceReference1.getIdentifier());
    assertEquals(serviceReference1.getName(), deserializedServiceReference1.getName());
}
 

@Test
public void testFlowSnapshotsWithParameterContextAndEncodingVersion() throws IOException, NiFiRegistryException {
    final RevisionInfo initialRevision = new RevisionInfo(null, 0L);

    // Create a bucket
    final Bucket bucket = new Bucket();
    bucket.setName("Test Bucket");
    bucket.setRevision(initialRevision);

    final Bucket createdBucket = client.getBucketClient().create(bucket);
    assertNotNull(createdBucket);

    // Create the flow
    final VersionedFlow flow = new VersionedFlow();
    flow.setName("My Flow");
    flow.setBucketIdentifier(createdBucket.getIdentifier());
    flow.setRevision(initialRevision);

    final VersionedFlow createdFlow = client.getFlowClient().create(flow);
    assertNotNull(createdFlow);

    // Create a param context
    final VersionedParameter param1 = new VersionedParameter();
    param1.setName("Param 1");
    param1.setValue("Param 1 Value");
    param1.setDescription("Description");

    final VersionedParameter param2 = new VersionedParameter();
    param2.setName("Param 2");
    param2.setValue("Param 2 Value");
    param2.setDescription("Description");
    param2.setSensitive(true);

    final VersionedParameterContext context1 = new VersionedParameterContext();
    context1.setName("Parameter Context 1");
    context1.setParameters(new HashSet<>(Arrays.asList(param1, param2)));

    final Map<String,VersionedParameterContext> contexts = new HashMap<>();
    contexts.put(context1.getName(), context1);

    // Create an external controller service reference
    final ExternalControllerServiceReference serviceReference = new ExternalControllerServiceReference();
    serviceReference.setName("External Service 1");
    serviceReference.setIdentifier(UUID.randomUUID().toString());

    final Map<String,ExternalControllerServiceReference> serviceReferences = new HashMap<>();
    serviceReferences.put(serviceReference.getIdentifier(), serviceReference);

    // Create the snapshot
    final VersionedFlowSnapshot snapshot = buildSnapshot(createdFlow, 1);
    snapshot.setFlowEncodingVersion("2.0.0");
    snapshot.setParameterContexts(contexts);
    snapshot.setExternalControllerServices(serviceReferences);

    final VersionedFlowSnapshot createdSnapshot = client.getFlowSnapshotClient().create(snapshot);
    assertNotNull(createdSnapshot);
    assertNotNull(createdSnapshot.getFlowEncodingVersion());
    assertNotNull(createdSnapshot.getParameterContexts());
    assertNotNull(createdSnapshot.getExternalControllerServices());
    assertEquals(snapshot.getFlowEncodingVersion(), createdSnapshot.getFlowEncodingVersion());
    assertEquals(1, createdSnapshot.getParameterContexts().size());
    assertEquals(1, createdSnapshot.getExternalControllerServices().size());

    // Retrieve the snapshot
    final VersionedFlowSnapshot retrievedSnapshot = client.getFlowSnapshotClient().get(
            createdSnapshot.getSnapshotMetadata().getFlowIdentifier(),
            createdSnapshot.getSnapshotMetadata().getVersion());
    assertNotNull(retrievedSnapshot);
    assertNotNull(retrievedSnapshot.getFlowEncodingVersion());
    assertNotNull(retrievedSnapshot.getParameterContexts());
    assertNotNull(retrievedSnapshot.getExternalControllerServices());
    assertEquals(snapshot.getFlowEncodingVersion(), retrievedSnapshot.getFlowEncodingVersion());
    assertEquals(1, retrievedSnapshot.getParameterContexts().size());
    assertEquals(1, retrievedSnapshot.getExternalControllerServices().size());
}
 

@Override
public StringResult doExecute(final NiFiRegistryClient client, final Properties properties)
        throws ParseException, IOException, NiFiRegistryException {
    final String flowId = getRequiredArg(properties, CommandOption.FLOW_ID);
    final String inputFile = getRequiredArg(properties, CommandOption.INPUT_SOURCE);

    final String contents = getInputSourceContent(inputFile);

    final FlowClient flowClient = client.getFlowClient();
    final FlowSnapshotClient snapshotClient = client.getFlowSnapshotClient();

    final ObjectMapper objectMapper = JacksonUtils.getObjectMapper();
    final VersionedFlowSnapshot deserializedSnapshot = objectMapper.readValue(contents, VersionedFlowSnapshot.class);
    if (deserializedSnapshot == null) {
        throw new IOException("Unable to deserialize flow version from " + inputFile);
    }

    final VersionedFlow versionedFlow = flowClient.get(flowId);

    // determine the latest existing version in the destination system
    Integer version;
    try {
        final VersionedFlowSnapshotMetadata latestMetadata = snapshotClient.getLatestMetadata(flowId);
        version = latestMetadata.getVersion() + 1;
    } catch (NiFiRegistryException e) {
        // when there are no versions it produces a 404 not found
        version = 1;
    }

    // create new metadata using the passed in bucket and flow in the target registry, keep comments
    final VersionedFlowSnapshotMetadata metadata = new VersionedFlowSnapshotMetadata();
    metadata.setBucketIdentifier(versionedFlow.getBucketIdentifier());
    metadata.setFlowIdentifier(flowId);
    metadata.setVersion(version);

    final VersionedFlowSnapshotMetadata deserializedSnapshotMetadata = deserializedSnapshot.getSnapshotMetadata();
    if (deserializedSnapshotMetadata != null) {
        metadata.setComments(deserializedSnapshotMetadata.getComments());
    }

    // create a new snapshot using the new metadata and the contents from the deserialized snapshot
    final VersionedFlowSnapshot snapshot = new VersionedFlowSnapshot();
    snapshot.setSnapshotMetadata(metadata);
    snapshot.setFlowContents(deserializedSnapshot.getFlowContents());
    snapshot.setExternalControllerServices(deserializedSnapshot.getExternalControllerServices());
    snapshot.setParameterContexts(deserializedSnapshot.getParameterContexts());
    snapshot.setFlowEncodingVersion(deserializedSnapshot.getFlowEncodingVersion());

    final VersionedFlowSnapshot createdSnapshot = snapshotClient.create(snapshot);
    final VersionedFlowSnapshotMetadata createdMetadata = createdSnapshot.getSnapshotMetadata();

    return new StringResult(String.valueOf(createdMetadata.getVersion()), getContext().isInteractive());
}