Examples with NiFiRestClient com.thinkbiganalytics.nifi.rest.client.NiFiRestClient used on opensource projects

Example 1 with NiFiRestClient

use of com.thinkbiganalytics.nifi.rest.client.NiFiRestClient in project kylo by Teradata.

the class TemplateCreationHelperTest method updateControllerServiceReferencesWithRecursive.

/**
 * Verify recursively enabling controller services when updating processor properties.
 */
@Test
public void updateControllerServiceReferencesWithRecursive() {
    final List<ControllerServiceDTO> updatedControllerServices = new ArrayList<>();
    final List<NifiProperty> updatedProperties = new ArrayList<>();
    // Mock NiFi client
    final NiFiControllerServicesRestClient controllerServicesRestClient = Mockito.mock(NiFiControllerServicesRestClient.class);
    Mockito.when(controllerServicesRestClient.update(Mockito.any())).thenAnswer(answer -> {
        final ControllerServiceDTO controllerService = answer.getArgumentAt(0, ControllerServiceDTO.class);
        updatedControllerServices.add(controllerService);
        return controllerService;
    });
    final NiFiRestClient restClient = Mockito.mock(NiFiRestClient.class);
    Mockito.when(restClient.controllerServices()).thenReturn(controllerServicesRestClient);
    // Mock Legacy NiFi client
    final LegacyNifiRestClient legacyRestClient = Mockito.mock(LegacyNifiRestClient.class);
    Mockito.when(legacyRestClient.enableControllerServiceAndSetProperties(Mockito.any(), Mockito.any())).thenReturn(new ControllerServiceDTO());
    Mockito.when(legacyRestClient.getNiFiRestClient()).thenReturn(restClient);
    Mockito.when(legacyRestClient.getPropertyDescriptorTransform()).thenReturn(new MockNiFiPropertyDescriptorTransform());
    Mockito.doAnswer(invocation -> {
        updatedProperties.add(invocation.getArgumentAt(2, NifiProperty.class));
        return null;
    }).when(legacyRestClient).updateProcessorProperty(Mockito.isNull(String.class), Mockito.eq("P1"), Mockito.any());
    final ControllerServiceDTO service1 = new ControllerServiceDTO();
    service1.setDescriptors(Collections.singletonMap("service", newPropertyDescriptor("service", "com.example.Service2", "S2")));
    service1.setId("S1");
    service1.setName("Service1");
    service1.setProperties(newHashMap("service", "invalid"));
    service1.setState("DISABLED");
    final ControllerServiceDTO service2 = new ControllerServiceDTO();
    service2.setId("S2");
    service2.setName("Service2");
    service2.setProperties(new HashMap<>());
    service2.setState("ENABLED");
    Mockito.when(legacyRestClient.getControllerServices()).thenReturn(ImmutableSet.of(service1, service2));
    // Mock processors
    final ProcessorConfigDTO config = new ProcessorConfigDTO();
    config.setDescriptors(Collections.singletonMap("service", newPropertyDescriptor("service", "com.example.Service1", "S1")));
    config.setProperties(Collections.singletonMap("service", "invalid"));
    final ProcessorDTO processor = new ProcessorDTO();
    processor.setId("P1");
    processor.setName("Processor1");
    processor.setConfig(config);
    // Update processors
    final TemplateCreationHelper helper = new TemplateCreationHelper(legacyRestClient);
    helper.snapshotControllerServiceReferences();
    helper.identifyNewlyCreatedControllerServiceReferences();
    helper.updateControllerServiceReferences(Collections.singletonList(processor));
    // Verify updated properties
    Assert.assertEquals("Property 'Service' not set on controller service 'Server1'.", 1, updatedControllerServices.size());
    Assert.assertEquals("S2", updatedControllerServices.get(0).getProperties().get("service"));
    Assert.assertEquals("Property 'Service' not set on processor 'Processor1'.", 1, updatedProperties.size());
    Assert.assertEquals("S1", updatedProperties.get(0).getValue());
}

Example 2 with NiFiRestClient

use of com.thinkbiganalytics.nifi.rest.client.NiFiRestClient in project kylo by Teradata.

the class TemplateCreationHelper method instantiateFlowFromTemplate.

/**
 * Instantiates the specified template in the specified process group.
 *
 * <p>Controller services that are created under the specified process group will be moved to the root process group. This side-effect may be removed in the future.</p>
 *
 * @param processGroupId the process group id
 * @param templateId     the template id
 * @return the instantiated flow
 * @throws NifiComponentNotFoundException if the process group or template does not exist
 */
@Nonnull
public TemplateInstance instantiateFlowFromTemplate(@Nonnull final String processGroupId, @Nonnull final String templateId) throws NifiComponentNotFoundException {
    // Instantiate template
    final NiFiRestClient nifiClient = restClient.getNiFiRestClient();
    final FlowSnippetDTO templateFlow = nifiClient.processGroups().instantiateTemplate(processGroupId, templateId);
    TemplateInstance instance = new TemplateInstance(templateFlow);
    // Move controller services to root process group (NiFi >= v1.0)
    final Set<ControllerServiceDTO> groupControllerServices = nifiClient.processGroups().getControllerServices(processGroupId);
    final Map<String, String> idMap = new HashMap<>(groupControllerServices.size());
    groupControllerServices.stream().filter(controllerService -> controllerService.getParentGroupId().equals(processGroupId)).forEach(groupControllerService -> {
        // Delete scoped service
        final String oldId = groupControllerService.getId();
        nifiClient.controllerServices().delete(groupControllerService.getId());
        // Create root service
        final ControllerServiceDTO rootControllerService = new ControllerServiceDTO();
        rootControllerService.setComments(groupControllerService.getComments());
        rootControllerService.setName(groupControllerService.getName());
        rootControllerService.setType(groupControllerService.getType());
        ControllerServiceDTO newRootService = nifiClient.processGroups().createControllerService("root", rootControllerService);
        final String rootId = newRootService.getId();
        // Map old ID to new ID
        idMap.put(oldId, rootId);
        instance.movedScopedControllerService(groupControllerService, newRootService);
    });
    // Set properties on root controller services
    groupControllerServices.stream().filter(controllerService -> controllerService.getParentGroupId().equals(processGroupId)).forEach(groupControllerService -> {
        final Map<String, String> properties = groupControllerService.getProperties();
        groupControllerService.getDescriptors().values().stream().filter(descriptor -> StringUtils.isNotBlank(descriptor.getIdentifiesControllerService())).forEach(descriptor -> {
            final String name = descriptor.getName();
            final String oldId = properties.get(name);
            properties.put(name, idMap.get(oldId));
        });
        final ControllerServiceDTO rootControllerService = new ControllerServiceDTO();
        rootControllerService.setId(idMap.get(groupControllerService.getId()));
        rootControllerService.setProperties(properties);
        nifiClient.controllerServices().update(rootControllerService);
    });
    // Return flow
    return instance;
}

Example 3 with NiFiRestClient

use of com.thinkbiganalytics.nifi.rest.client.NiFiRestClient in project kylo by Teradata.

the class NifiRestTest method setupRestClient.

@Before
public void setupRestClient() {
    restClient = new LegacyNifiRestClient();
    NifiRestClientConfig clientConfig = new NifiRestClientConfig();
    // clientConfig.setHost("localhost");
    clientConfig.setHost("34.208.236.190");
    clientConfig.setPort(8079);
    NiFiRestClient c = new NiFiRestClientV1(clientConfig);
    restClient.setClient(c);
    nifiFlowCache = new NifiFlowCacheImpl();
    propertyDescriptorTransform = new NiFiPropertyDescriptorTransformV1();
    createFeedBuilderCache = new NiFiObjectCache();
    createFeedBuilderCache.setRestClient(restClient);
    templateConnectionUtil = new TemplateConnectionUtil();
    templateConnectionUtil.setRestClient(restClient);
    templateConnectionUtil.setNifiFlowCache(nifiFlowCache);
    templateConnectionUtil.setNiFiObjectCache(createFeedBuilderCache);
    templateConnectionUtil.setPropertyDescriptorTransform(propertyDescriptorTransform);
}