Examples with DataPacket org.apache.nifi.remote.protocol.DataPacket used on opensource projects

Example 1 with DataPacket

use of org.apache.nifi.remote.protocol.DataPacket in project nifi by apache.

the class TestStandardRemoteGroupPort method testReceiveRaw.

@Test
public void testReceiveRaw() throws Exception {
    setupMock(SiteToSiteTransportProtocol.RAW, TransferDirection.RECEIVE);
    setupMockProcessSession();
    final String peerUrl = "nifi://node1.example.com:9090";
    final PeerDescription peerDescription = new PeerDescription("node1.example.com", 9090, true);
    try (final SocketChannel socketChannel = SocketChannel.open()) {
        final CommunicationsSession commsSession = new SocketChannelCommunicationsSession(socketChannel);
        commsSession.setUserDn("nifi.node1.example.com");
        final Peer peer = new Peer(peerDescription, commsSession, peerUrl, REMOTE_CLUSTER_URL);
        doReturn(peer).when(transaction).getCommunicant();
        final String sourceFlowFileUuid = "flowfile-uuid";
        final Map<String, String> attributes = new HashMap<>();
        attributes.put(CoreAttributes.UUID.key(), sourceFlowFileUuid);
        final byte[] dataPacketContents = "DataPacket Contents".getBytes();
        final ByteArrayInputStream dataPacketInputStream = new ByteArrayInputStream(dataPacketContents);
        final DataPacket dataPacket = new StandardDataPacket(attributes, dataPacketInputStream, dataPacketContents.length);
        // Return null when it gets called second time.
        doReturn(dataPacket).doReturn(null).when(this.transaction).receive();
        port.onTrigger(processContext, processSession);
        // Assert provenance.
        final List<ProvenanceEventRecord> provenanceEvents = sessionState.getProvenanceEvents();
        assertEquals(1, provenanceEvents.size());
        final ProvenanceEventRecord provenanceEvent = provenanceEvents.get(0);
        assertEquals(ProvenanceEventType.RECEIVE, provenanceEvent.getEventType());
        assertEquals(peerUrl + "/" + sourceFlowFileUuid, provenanceEvent.getTransitUri());
        assertEquals("Remote DN=nifi.node1.example.com", provenanceEvent.getDetails());
        // Assert received flow files.
        processSession.assertAllFlowFilesTransferred(Relationship.ANONYMOUS);
        final List<MockFlowFile> flowFiles = processSession.getFlowFilesForRelationship(Relationship.ANONYMOUS);
        assertEquals(1, flowFiles.size());
        final MockFlowFile flowFile = flowFiles.get(0);
        flowFile.assertAttributeEquals(SiteToSiteAttributes.S2S_HOST.key(), peer.getHost());
        flowFile.assertAttributeEquals(SiteToSiteAttributes.S2S_ADDRESS.key(), peer.getHost() + ":" + peer.getPort());
    }
}

Example 2 with DataPacket

use of org.apache.nifi.remote.protocol.DataPacket in project nifi by apache.

the class TestHttpFlowFileServerProtocol method receiveFlowFiles.

private void receiveFlowFiles(final HttpFlowFileServerProtocol serverProtocol, final String transactionId, final Peer peer, final DataPacket... dataPackets) throws IOException {
    final HttpRemoteSiteListener remoteSiteListener = HttpRemoteSiteListener.getInstance(NiFiProperties.createBasicNiFiProperties(null, null));
    final HttpServerCommunicationsSession commsSession = (HttpServerCommunicationsSession) peer.getCommunicationsSession();
    serverProtocol.handshake(peer);
    assertTrue(serverProtocol.isHandshakeSuccessful());
    setupMockProcessSession();
    // Emulate dataPackets sent from a Site-to-Site client.
    final FlowFileCodec negotiatedCoded = serverProtocol.negotiateCodec(peer);
    final ByteArrayOutputStream testDataOs = new ByteArrayOutputStream();
    for (final DataPacket dataPacket : dataPackets) {
        negotiatedCoded.encode(dataPacket, testDataOs);
    }
    final InputStream httpInputStream = new ByteArrayInputStream(testDataOs.toByteArray());
    ((HttpInput) commsSession.getInput()).setInputStream(httpInputStream);
    // Execute test using mock
    final int flowFileReceived = serverProtocol.receiveFlowFiles(peer, processContext, processSession, negotiatedCoded);
    assertEquals(dataPackets.length, flowFileReceived);
    assertTrue(remoteSiteListener.isTransactionActive(transactionId));
}

Example 3 with DataPacket

use of org.apache.nifi.remote.protocol.DataPacket in project nifi by apache.

the class DataPacketDto method getDataPacketStream.

public static Stream<DataPacket> getDataPacketStream(InputStream inputStream) throws IOException {
    JsonParser jsonParser = new JsonFactory().createParser(inputStream);
    if (jsonParser.nextToken() != JsonToken.START_ARRAY) {
        throw new IOException("Expecting start array token to begin object array.");
    }
    jsonParser.setCodec(new ObjectMapper());
    return StreamSupport.stream(Spliterators.spliteratorUnknownSize(new Iterator<DataPacket>() {

        DataPacket next = getNext();

        @Override
        public boolean hasNext() {
            return next != null;
        }

        @Override
        public DataPacket next() {
            DataPacket next = this.next;
            this.next = getNext();
            return next;
        }

        DataPacket getNext() throws RuntimeException {
            try {
                if (jsonParser.nextToken() == JsonToken.END_ARRAY) {
                    return null;
                }
                DataPacketDto dataPacketDto = jsonParser.readValueAs(DATA_PACKET_DTO_TYPE_REFERENCE);
                return dataPacketDto.toDataPacket();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }, Spliterator.ORDERED), false);
}

Example 4 with DataPacket

use of org.apache.nifi.remote.protocol.DataPacket in project nifi by apache.

the class SiteToSiteReceiver method receiveFiles.

public TransactionCompletion receiveFiles() throws IOException {
    Transaction transaction = siteToSiteClient.createTransaction(TransferDirection.RECEIVE);
    JsonGenerator jsonGenerator = new JsonFactory().createJsonGenerator(output);
    jsonGenerator.writeStartArray();
    DataPacket dataPacket;
    while ((dataPacket = transaction.receive()) != null) {
        jsonGenerator.writeStartObject();
        jsonGenerator.writeFieldName("attributes");
        jsonGenerator.writeStartObject();
        Map<String, String> attributes = dataPacket.getAttributes();
        if (attributes != null) {
            for (Map.Entry<String, String> stringStringEntry : attributes.entrySet()) {
                jsonGenerator.writeStringField(stringStringEntry.getKey(), stringStringEntry.getValue());
            }
        }
        jsonGenerator.writeEndObject();
        InputStream data = dataPacket.getData();
        if (data != null) {
            jsonGenerator.writeBinaryField("data", IOUtils.toByteArray(data));
        }
        jsonGenerator.writeEndObject();
    }
    jsonGenerator.writeEndArray();
    jsonGenerator.close();
    transaction.confirm();
    return transaction.complete();
}

Example 5 with DataPacket

use of org.apache.nifi.remote.protocol.DataPacket in project nifi by apache.

the class TestHttpClient method testReceiveSlowClientSuccess.

@Test
public void testReceiveSlowClientSuccess() throws Exception {
    try (SiteToSiteClient client = getDefaultBuilder().portName("output-running").build()) {
        final Transaction transaction = client.createTransaction(TransferDirection.RECEIVE);
        assertNotNull(transaction);
        DataPacket packet;
        while ((packet = transaction.receive()) != null) {
            consumeDataPacket(packet);
            Thread.sleep(500);
        }
        transaction.confirm();
        transaction.complete();
    }
}