org.apache.nifi.controller.repository.ContentRepository 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 ContentRepository createContentRepository(final NiFiProperties properties) throws InstantiationException, IllegalAccessException, ClassNotFoundException {
    final String implementationClassName = properties.getProperty(NiFiProperties.CONTENT_REPOSITORY_IMPLEMENTATION, DEFAULT_CONTENT_REPO_IMPLEMENTATION);
    if (implementationClassName == null) {
        throw new RuntimeException("Cannot create Content Repository because the NiFi Properties is missing the following property: "
                + NiFiProperties.CONTENT_REPOSITORY_IMPLEMENTATION);
    }

    try {
        final ContentRepository contentRepo = NarThreadContextClassLoader.createInstance(implementationClassName, ContentRepository.class, properties);
        synchronized (contentRepo) {
            contentRepo.initialize(resourceClaimManager);
        }
        return contentRepo;
    } catch (final Exception e) {
        throw new RuntimeException(e);
    }
}
 

@Test
public void testContentNotFoundPropagated() throws IOException {
    final ContentRepository contentRepo = mock(ContentRepository.class);

    final ResourceClaimManager claimManager = new StandardResourceClaimManager();
    final ResourceClaim resourceClaim = new StandardResourceClaim(claimManager, "container", "section", "id", false);
    final ContentClaim contentClaim = new StandardContentClaim(resourceClaim, 5L);

    final FlowFileRecord flowFile = mock(FlowFileRecord.class);
    when(flowFile.getContentClaim()).thenReturn(contentClaim);

    final ContentNotFoundException cnfe = new ContentNotFoundException(contentClaim);
    when(contentRepo.read(contentClaim)).thenThrow(cnfe);

    final ContentRepositoryFlowFileAccess flowAccess = new ContentRepositoryFlowFileAccess(contentRepo);

    try {
        flowAccess.read(flowFile);
        Assert.fail("Expected ContentNotFoundException but it did not happen");
    } catch (final ContentNotFoundException thrown) {
        // expected
        thrown.getFlowFile().orElseThrow(() -> new AssertionError("Expected FlowFile to be provided"));
    }
}
 

@Test
public void testInputStreamFromContentRepo() throws IOException {
    final ContentRepository contentRepo = mock(ContentRepository.class);

    final ResourceClaimManager claimManager = new StandardResourceClaimManager();
    final ResourceClaim resourceClaim = new StandardResourceClaim(claimManager, "container", "section", "id", false);
    final ContentClaim contentClaim = new StandardContentClaim(resourceClaim, 5L);

    final FlowFileRecord flowFile = mock(FlowFileRecord.class);
    when(flowFile.getContentClaim()).thenReturn(contentClaim);
    when(flowFile.getSize()).thenReturn(5L);

    final InputStream inputStream = new ByteArrayInputStream("hello".getBytes());
    when(contentRepo.read(contentClaim)).thenReturn(inputStream);

    final ContentRepositoryFlowFileAccess flowAccess = new ContentRepositoryFlowFileAccess(contentRepo);

    final InputStream repoStream = flowAccess.read(flowFile);
    verify(contentRepo, times(1)).read(contentClaim);

    final byte[] buffer = new byte[5];
    StreamUtils.fillBuffer(repoStream, buffer);
    assertEquals(-1, repoStream.read());
    assertArrayEquals("hello".getBytes(), buffer);
}
 

private ContentRepository createContentRepository(final NiFiProperties properties) throws InstantiationException, IllegalAccessException, ClassNotFoundException {
    final String implementationClassName = properties.getProperty(NiFiProperties.CONTENT_REPOSITORY_IMPLEMENTATION, DEFAULT_CONTENT_REPO_IMPLEMENTATION);
    if (implementationClassName == null) {
        throw new RuntimeException("Cannot create Content Repository because the NiFi Properties is missing the following property: "
                + NiFiProperties.CONTENT_REPOSITORY_IMPLEMENTATION);
    }

    try {
        final ContentRepository contentRepo = NarThreadContextClassLoader.createInstance(extensionManager, implementationClassName, ContentRepository.class, properties);
        synchronized (contentRepo) {
            contentRepo.initialize(resourceClaimManager);
        }
        return contentRepo;
    } catch (final Exception e) {
        throw new RuntimeException(e);
    }
}
 

public RemoteQueuePartition(final NodeIdentifier nodeId, final SwappablePriorityQueue priorityQueue, final TransferFailureDestination failureDestination,
                            final FlowFileRepository flowFileRepo, final ProvenanceEventRepository provRepo, final ContentRepository contentRepository,
                            final AsyncLoadBalanceClientRegistry clientRegistry, final LoadBalancedFlowFileQueue flowFileQueue) {

    this.nodeIdentifier = nodeId;
    this.priorityQueue = priorityQueue;
    this.flowFileQueue = flowFileQueue;
    this.failureDestination = failureDestination;
    this.flowFileRepo = flowFileRepo;
    this.provRepo = provRepo;
    this.contentRepo = contentRepository;
    this.clientRegistry = clientRegistry;
    this.description = "RemoteQueuePartition[queueId=" + flowFileQueue.getIdentifier() + ", nodeId=" + nodeIdentifier + "]";
}
 

@Test
public void testEOFExceptionIfNotEnoughData() throws IOException {
    final ContentRepository contentRepo = mock(ContentRepository.class);

    final ResourceClaimManager claimManager = new StandardResourceClaimManager();
    final ResourceClaim resourceClaim = new StandardResourceClaim(claimManager, "container", "section", "id", false);
    final ContentClaim contentClaim = new StandardContentClaim(resourceClaim, 5L);

    final FlowFileRecord flowFile = mock(FlowFileRecord.class);
    when(flowFile.getContentClaim()).thenReturn(contentClaim);
    when(flowFile.getSize()).thenReturn(100L);

    final InputStream inputStream = new ByteArrayInputStream("hello".getBytes());
    when(contentRepo.read(contentClaim)).thenReturn(inputStream);

    final ContentRepositoryFlowFileAccess flowAccess = new ContentRepositoryFlowFileAccess(contentRepo);

    final InputStream repoStream = flowAccess.read(flowFile);
    verify(contentRepo, times(1)).read(contentClaim);

    final byte[] buffer = new byte[5];
    StreamUtils.fillBuffer(repoStream, buffer);

    try {
        repoStream.read();
        Assert.fail("Expected EOFException because not enough bytes were in the InputStream for the FlowFile");
    } catch (final EOFException eof) {
        // expected
    }
}
 

@Before
public void setup() throws IOException {
    repo = mock(ContentRepository.class);
    contentClaim = mock(ContentClaim.class);

    closed.set(false);
    Mockito.when(repo.read(contentClaim)).thenAnswer(invocation -> new ByteArrayInputStream("hello".getBytes()) {
        @Override
        public void close() throws IOException {
            super.close();
            closed.set(true);
        }
    });
}
 

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

    details.add("Bytes Read (Last 5 mins): " + FormatUtils.formatDataSize(status.getBytesRead()));
    details.add("Bytes Written (Last 5 mins): " + FormatUtils.formatDataSize(status.getBytesWritten()));
}
 

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;
}
 

public ProcessContextFactory(final ContentRepository contentRepository, final FlowFileRepository flowFileRepository,
        final FlowFileEventRepository flowFileEventRepository, final CounterRepository counterRepository,
        final ProvenanceEventRepository provenanceRepository) {

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

public ContentClaimInputStream(final ContentRepository contentRepository, final ContentClaim contentClaim, final long claimOffset) {
    this.contentRepository = contentRepository;
    this.contentClaim = contentClaim;
    this.claimOffset = claimOffset;

    this.currentOffset = claimOffset;
}
 

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;
}
 

public SocketLoadBalancedFlowFileQueue(final String identifier, final ConnectionEventListener eventListener, final ProcessScheduler scheduler, final FlowFileRepository flowFileRepo,
                                       final ProvenanceEventRepository provRepo, final ContentRepository contentRepo, final ResourceClaimManager resourceClaimManager,
                                       final ClusterCoordinator clusterCoordinator, final AsyncLoadBalanceClientRegistry clientRegistry, final FlowFileSwapManager swapManager,
                                       final int swapThreshold, final EventReporter eventReporter) {

    super(identifier, scheduler, flowFileRepo, provRepo, resourceClaimManager);
    this.eventListener = eventListener;
    this.eventReporter = eventReporter;
    this.swapManager = swapManager;
    this.flowFileRepo = flowFileRepo;
    this.provRepo = provRepo;
    this.contentRepo = contentRepo;
    this.clusterCoordinator = clusterCoordinator;
    this.clientRegistry = clientRegistry;

    localPartition = new SwappablePriorityQueueLocalPartition(swapManager, swapThreshold, eventReporter, this, this::drop);
    rebalancingPartition = new StandardRebalancingPartition(swapManager, swapThreshold, eventReporter, this, this::drop);

    // Create a RemoteQueuePartition for each node
    nodeIdentifiers = clusterCoordinator == null ? Collections.emptySet() : clusterCoordinator.getNodeIdentifiers();

    final List<NodeIdentifier> sortedNodeIdentifiers = new ArrayList<>(nodeIdentifiers);
    sortedNodeIdentifiers.sort(Comparator.comparing(NodeIdentifier::getApiAddress));

    if (sortedNodeIdentifiers.isEmpty()) {
        // No Node Identifiers are known yet. Just create the partitions using the local partition.
        queuePartitions = new QueuePartition[] { localPartition };
    } else {
        // The node identifiers are known. Create the partitions using the local partition and 1 Remote Partition for each node
        // that is not the local node identifier. If the Local Node Identifier is not yet known, that's okay. When it becomes known,
        // the queuePartitions array will be recreated with the appropriate partitions.
        final List<QueuePartition> partitionList = new ArrayList<>();

        final NodeIdentifier localNodeId = clusterCoordinator.getLocalNodeIdentifier();
        for (final NodeIdentifier nodeId : sortedNodeIdentifiers) {
            if (nodeId.equals(localNodeId)) {
                partitionList.add(localPartition);
            } else {
                partitionList.add(createRemotePartition(nodeId));
            }
        }

        // Ensure that our list of queue partitions always contains the local partition.
        if (!partitionList.contains(localPartition)) {
            partitionList.add(localPartition);
        }

        queuePartitions = partitionList.toArray(new QueuePartition[0]);
    }

    partitioner = new LocalPartitionPartitioner();

    if (clusterCoordinator != null) {
        clusterCoordinator.registerEventListener(new ClusterEventListener());
    }

    rebalancingPartition.start(partitioner);
}
 

public ContentRepositoryFlowFileAccess(final ContentRepository contentRepository) {
    this.contentRepository = contentRepository;
}
 

@Override
public DiagnosticsDumpElement captureDump(final boolean verbose) {
    if (!verbose) {
        // This task is very expensive, as it must scan the contents of the Content Repository. As such, it will not
        // run at all if verbose output is disabled.
        return null;
    }

    final ContentRepository contentRepository = flowController.getRepositoryContextFactory().getContentRepository();
    if (!contentRepository.isActiveResourceClaimsSupported()) {
        return new StandardDiagnosticsDumpElement("Content Repository Scan", Collections.singletonList("Current Content Repository does not support scanning for in-use content"));
    }

    final FlowFileRepository flowFileRepository = flowController.getRepositoryContextFactory().getFlowFileRepository();
    final ResourceClaimManager resourceClaimManager = flowController.getResourceClaimManager();
    final FlowFileSwapManager swapManager = flowController.createSwapManager();

    final List<String> details = new ArrayList<>();

    for (final String containerName : contentRepository.getContainerNames()) {
        try {
            final Set<ResourceClaim> resourceClaims = contentRepository.getActiveResourceClaims(containerName);

            final Map<ResourceClaim, Set<ResourceClaimReference>> referenceMap = flowFileRepository.findResourceClaimReferences(resourceClaims, swapManager);

            for (final ResourceClaim resourceClaim : resourceClaims) {
                final int claimCount = resourceClaimManager.getClaimantCount(resourceClaim);
                final boolean inUse = resourceClaim.isInUse();
                final boolean destructable = resourceClaimManager.isDestructable(resourceClaim);

                final Set<ResourceClaimReference> references = referenceMap == null ? Collections.emptySet() : referenceMap.getOrDefault(resourceClaim, Collections.emptySet());

                final String path = resourceClaim.getContainer() + "/" + resourceClaim.getSection() + "/" + resourceClaim.getId();
                details.add(String.format("%1$s, Claimant Count = %2$d, In Use = %3$b, Awaiting Destruction = %4$b, References (%5$d) = %6$s",
                    path, claimCount, inUse, destructable,  references.size(), references.toString()));
            }
        } catch (final Exception e) {
            logger.error("Failed to obtain listing of Active Resource Claims for container {}", containerName, e);
            details.add("Failed to obtain listing of Active Resource Claims in container " + containerName);
        }
    }


    return new StandardDiagnosticsDumpElement("Content Repository Scan", details);
}
 

public ContentClaimWriteCache(final ContentRepository contentRepo, final int bufferSize) {
    this.contentRepo = contentRepo;
    this.bufferSize = bufferSize;
}
 

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

protected final ContentRepository getContentRepository() {
    return getRepositoryContext().getContentRepository();
}
 

public ContentClaimWriteCache(final ContentRepository contentRepo) {
    this(contentRepo, 8192);
}
 

public ContentRepository getContentRepository() {
    return contentRepo;
}
 

public ContentClaimWriteCache(final ContentRepository contentRepo, final int bufferSize) {
    this.contentRepo = contentRepo;
    this.bufferSize = bufferSize;
}
 

public ContentClaimWriteCache(final ContentRepository contentRepo) {
    this(contentRepo, 8192);
}