Examples with ProcessSessionFactory org.apache.nifi.processor.ProcessSessionFactory used on opensource projects

Example 16 with ProcessSessionFactory

use of org.apache.nifi.processor.ProcessSessionFactory in project nifi by apache.

the class FlowFileIngestService method send.

/**
 * Handle receipt of a FlowFileRequest and route it to the appropriate process relationship.
 *
 * @param request          the flowfile request
 * @param responseObserver the mechanism by which to reply to the client
 */
@Override
public void send(final org.apache.nifi.processors.grpc.FlowFileRequest request, final StreamObserver<FlowFileReply> responseObserver) {
    final FlowFileReply.Builder replyBuilder = FlowFileReply.newBuilder();
    final String remoteHost = FlowFileIngestServiceInterceptor.REMOTE_HOST_KEY.get();
    final String remoteDN = FlowFileIngestServiceInterceptor.REMOTE_DN_KEY.get();
    // block until we have a session factory (occurs when processor is triggered)
    ProcessSessionFactory sessionFactory = null;
    while (sessionFactory == null) {
        sessionFactory = sessionFactoryReference.get();
        if (sessionFactory == null) {
            try {
                Thread.sleep(10);
            } catch (final InterruptedException e) {
            }
        }
    }
    final ProcessSession session = sessionFactory.createSession();
    // if there's no space available, reject the request.
    final long n = filesReceived.getAndIncrement() % FILES_BEFORE_CHECKING_DESTINATION_SPACE;
    if (n == 0 || !spaceAvailable.get()) {
        if (context.getAvailableRelationships().isEmpty()) {
            spaceAvailable.set(false);
            final String message = "Received request from " + remoteHost + " but no space available; Indicating Service Unavailable";
            if (logger.isDebugEnabled()) {
                logger.debug(message);
            }
            final FlowFileReply reply = replyBuilder.setResponseCode(FlowFileReply.ResponseCode.ERROR).setBody(message).build();
            responseObserver.onNext(reply);
            responseObserver.onCompleted();
            return;
        } else {
            spaceAvailable.set(true);
        }
    }
    if (logger.isDebugEnabled()) {
        logger.debug("Received request from " + remoteHost);
    }
    final long startNanos = System.nanoTime();
    FlowFile flowFile = session.create();
    // push the attributes provided onto the created flowfile
    final Map<String, String> attributes = Maps.newHashMap();
    attributes.putAll(request.getAttributesMap());
    String sourceSystemFlowFileIdentifier = attributes.get(CoreAttributes.UUID.key());
    if (sourceSystemFlowFileIdentifier != null) {
        sourceSystemFlowFileIdentifier = "urn:nifi:" + sourceSystemFlowFileIdentifier;
        // If we receveied a UUID, we want to give the FlowFile a new UUID and register the sending system's
        // identifier as the SourceSystemFlowFileIdentifier field in the Provenance RECEIVE event
        attributes.put(CoreAttributes.UUID.key(), UUID.randomUUID().toString());
    }
    flowFile = session.putAllAttributes(flowFile, attributes);
    final ByteString content = request.getContent();
    final InputStream contentStream = content.newInput();
    // write the provided content to the flowfile
    flowFile = session.write(flowFile, out -> {
        try (final BufferedOutputStream bos = new BufferedOutputStream(out, 65536)) {
            IOUtils.copy(contentStream, bos);
        }
    });
    final long transferNanos = System.nanoTime() - startNanos;
    final long transferMillis = TimeUnit.MILLISECONDS.convert(transferNanos, TimeUnit.NANOSECONDS);
    session.getProvenanceReporter().receive(flowFile, SERVICE_NAME, sourceSystemFlowFileIdentifier, "Remote DN=" + remoteDN, transferMillis);
    flowFile = session.putAttribute(flowFile, ListenGRPC.REMOTE_HOST, remoteHost);
    flowFile = session.putAttribute(flowFile, ListenGRPC.REMOTE_USER_DN, remoteDN);
    // register success
    session.transfer(flowFile, ListenGRPC.REL_SUCCESS);
    session.commit();
    // reply to client
    final FlowFileReply reply = replyBuilder.setResponseCode(FlowFileReply.ResponseCode.SUCCESS).setBody("FlowFile successfully received.").build();
    responseObserver.onNext(reply);
    responseObserver.onCompleted();
}

Example 17 with ProcessSessionFactory

use of org.apache.nifi.processor.ProcessSessionFactory in project nifi by apache.

the class ITListenGRPC method testSuccessfulRoundTrip.

@Test
public void testSuccessfulRoundTrip() throws UnrecoverableKeyException, CertificateException, NoSuchAlgorithmException, KeyStoreException, IOException {
    final int randPort = TestGRPCClient.randomPort();
    final ManagedChannel channel = TestGRPCClient.buildChannel(HOST, randPort);
    final FlowFileServiceGrpc.FlowFileServiceBlockingStub stub = FlowFileServiceGrpc.newBlockingStub(channel);
    final ListenGRPC listenGRPC = new ListenGRPC();
    final TestRunner runner = TestRunners.newTestRunner(listenGRPC);
    runner.setProperty(ListenGRPC.PROP_SERVICE_PORT, String.valueOf(randPort));
    final ProcessContext processContext = runner.getProcessContext();
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
    try {
        // start the server. The order of the following statements shouldn't matter, because the
        // startServer() method waits for a processSessionFactory to be available to it.
        listenGRPC.startServer(processContext);
        listenGRPC.onTrigger(processContext, processSessionFactory);
        final FlowFileRequest ingestFile = FlowFileRequest.newBuilder().putAttributes("FOO", "BAR").putAttributes(CoreAttributes.UUID.key(), SOURCE_SYSTEM_UUID).setContent(ByteString.copyFrom("content".getBytes())).build();
        final FlowFileReply reply = stub.send(ingestFile);
        assertThat(reply.getResponseCode(), equalTo(FlowFileReply.ResponseCode.SUCCESS));
        assertThat(reply.getBody(), equalTo("FlowFile successfully received."));
        runner.assertTransferCount(ListenGRPC.REL_SUCCESS, 1);
        final List<MockFlowFile> successFiles = runner.getFlowFilesForRelationship(ListenGRPC.REL_SUCCESS);
        assertThat(successFiles.size(), equalTo(1));
        final MockFlowFile mockFlowFile = successFiles.get(0);
        assertThat(mockFlowFile.getAttribute("FOO"), equalTo("BAR"));
        assertThat(mockFlowFile.getAttribute(ListenGRPC.REMOTE_HOST), equalTo("127.0.0.1"));
        assertThat(mockFlowFile.getAttribute(ListenGRPC.REMOTE_USER_DN), equalTo(FlowFileIngestServiceInterceptor.DEFAULT_FOUND_SUBJECT));
    } finally {
        // stop the server
        listenGRPC.stopServer(processContext);
        channel.shutdown();
    }
}

Example 18 with ProcessSessionFactory

use of org.apache.nifi.processor.ProcessSessionFactory in project nifi by apache.

the class ITListenGRPC method testOutOfSpaceRoundTrip.

@Test
public void testOutOfSpaceRoundTrip() throws UnrecoverableKeyException, CertificateException, NoSuchAlgorithmException, KeyStoreException, IOException {
    final int randPort = TestGRPCClient.randomPort();
    final ManagedChannel channel = TestGRPCClient.buildChannel(HOST, randPort);
    final FlowFileServiceGrpc.FlowFileServiceBlockingStub stub = FlowFileServiceGrpc.newBlockingStub(channel);
    final ListenGRPC listenGRPC = new ListenGRPC();
    final TestRunner runner = TestRunners.newTestRunner(listenGRPC);
    runner.setProperty(ListenGRPC.PROP_SERVICE_PORT, String.valueOf(randPort));
    final ProcessContext processContext = spy(runner.getProcessContext());
    // force the context to return that space isn't available, prompting an error message to be returned.
    when(processContext.getAvailableRelationships()).thenReturn(Sets.newHashSet());
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
    try {
        // start the server. The order of the following statements shouldn't matter, because the
        // startServer() method waits for a processSessionFactory to be available to it.
        listenGRPC.startServer(processContext);
        listenGRPC.onTrigger(processContext, processSessionFactory);
        final FlowFileRequest ingestFile = FlowFileRequest.newBuilder().putAttributes("FOO", "BAR").setContent(ByteString.copyFrom("content".getBytes())).build();
        final FlowFileReply reply = stub.send(ingestFile);
        assertThat(reply.getResponseCode(), equalTo(FlowFileReply.ResponseCode.ERROR));
        assertThat(reply.getBody(), containsString("but no space available; Indicating Service Unavailable"));
        runner.assertTransferCount(ListenGRPC.REL_SUCCESS, 0);
    } finally {
        // stop the server
        listenGRPC.stopServer(processContext);
        channel.shutdown();
    }
}

Example 19 with ProcessSessionFactory

use of org.apache.nifi.processor.ProcessSessionFactory in project nifi by apache.

the class ITListenGRPC method testSecureTwoWaySSLPassAuthorizedDNCheck.

@Test
public void testSecureTwoWaySSLPassAuthorizedDNCheck() throws UnrecoverableKeyException, CertificateException, NoSuchAlgorithmException, KeyStoreException, IOException {
    final int randPort = TestGRPCClient.randomPort();
    final Map<String, String> sslProperties = getKeystoreProperties();
    sslProperties.putAll(getTruststoreProperties());
    final ManagedChannel channel = TestGRPCClient.buildChannel(HOST, randPort, sslProperties);
    final FlowFileServiceGrpc.FlowFileServiceBlockingStub stub = FlowFileServiceGrpc.newBlockingStub(channel);
    final ListenGRPC listenGRPC = new ListenGRPC();
    final TestRunner runner = TestRunners.newTestRunner(listenGRPC);
    runner.setProperty(ListenGRPC.PROP_SERVICE_PORT, String.valueOf(randPort));
    runner.setProperty(ListenGRPC.PROP_USE_SECURE, "true");
    runner.setProperty(ListenGRPC.PROP_AUTHORIZED_DN_PATTERN, "CN=localhost.*");
    useSSLContextService(runner, sslProperties);
    final ProcessContext processContext = runner.getProcessContext();
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
    try {
        // start the server. The order of the following statements shouldn't matter, because the
        // startServer() method waits for a processSessionFactory to be available to it.
        listenGRPC.startServer(processContext);
        listenGRPC.onTrigger(processContext, processSessionFactory);
        final FlowFileRequest ingestFile = FlowFileRequest.newBuilder().putAttributes("FOO", "BAR").setContent(ByteString.copyFrom("content".getBytes())).build();
        final FlowFileReply reply = stub.send(ingestFile);
        assertThat(reply.getResponseCode(), equalTo(FlowFileReply.ResponseCode.SUCCESS));
        assertThat(reply.getBody(), equalTo("FlowFile successfully received."));
        runner.assertTransferCount(ListenGRPC.REL_SUCCESS, 1);
        final List<MockFlowFile> successFiles = runner.getFlowFilesForRelationship(ListenGRPC.REL_SUCCESS);
        assertThat(successFiles.size(), equalTo(1));
        final MockFlowFile mockFlowFile = successFiles.get(0);
        assertThat(mockFlowFile.getAttribute("FOO"), equalTo("BAR"));
        assertThat(mockFlowFile.getAttribute(ListenGRPC.REMOTE_HOST), equalTo("127.0.0.1"));
        assertThat(mockFlowFile.getAttribute(ListenGRPC.REMOTE_USER_DN), equalTo(CERT_DN));
    } finally {
        // stop the server
        listenGRPC.stopServer(processContext);
        channel.shutdown();
    }
}

Example 20 with ProcessSessionFactory

use of org.apache.nifi.processor.ProcessSessionFactory in project nifi by apache.

the class ITListenGRPC method testExceedMaxMessageSize.

@Test(expected = io.grpc.StatusRuntimeException.class)
public void testExceedMaxMessageSize() throws UnrecoverableKeyException, CertificateException, NoSuchAlgorithmException, KeyStoreException, IOException {
    final int randPort = TestGRPCClient.randomPort();
    final ManagedChannel channel = TestGRPCClient.buildChannel(HOST, randPort);
    final FlowFileServiceGrpc.FlowFileServiceBlockingStub stub = FlowFileServiceGrpc.newBlockingStub(channel);
    final ListenGRPC listenGRPC = new ListenGRPC();
    final TestRunner runner = TestRunners.newTestRunner(listenGRPC);
    runner.setProperty(ListenGRPC.PROP_SERVICE_PORT, String.valueOf(randPort));
    // set max message size to 1 byte to force exception to be thrown.
    runner.setProperty(ListenGRPC.PROP_MAX_MESSAGE_SIZE, "1B");
    final ProcessContext processContext = runner.getProcessContext();
    final ProcessSessionFactory processSessionFactory = runner.getProcessSessionFactory();
    try {
        // start the server. The order of the following statements shouldn't matter, because the
        // startServer() method waits for a processSessionFactory to be available to it.
        listenGRPC.startServer(processContext);
        listenGRPC.onTrigger(processContext, processSessionFactory);
        final FlowFileRequest ingestFile = FlowFileRequest.newBuilder().putAttributes("FOO", "BAR").putAttributes(CoreAttributes.UUID.key(), SOURCE_SYSTEM_UUID).setContent(ByteString.copyFrom("content".getBytes())).build();
        // this should throw a runtime exception
        final FlowFileReply reply = stub.send(ingestFile);
        assertThat(reply.getResponseCode(), equalTo(FlowFileReply.ResponseCode.SUCCESS));
        assertThat(reply.getBody(), equalTo("FlowFile successfully received."));
        runner.assertTransferCount(ListenGRPC.REL_SUCCESS, 1);
        final List<MockFlowFile> successFiles = runner.getFlowFilesForRelationship(ListenGRPC.REL_SUCCESS);
        assertThat(successFiles.size(), equalTo(1));
        final MockFlowFile mockFlowFile = successFiles.get(0);
        assertThat(mockFlowFile.getAttribute("FOO"), equalTo("BAR"));
        assertThat(mockFlowFile.getAttribute(ListenGRPC.REMOTE_HOST), equalTo("127.0.0.1"));
        assertThat(mockFlowFile.getAttribute(ListenGRPC.REMOTE_USER_DN), equalTo(FlowFileIngestServiceInterceptor.DEFAULT_FOUND_SUBJECT));
    } finally {
        // stop the server
        listenGRPC.stopServer(processContext);
        channel.shutdown();
    }
}