Examples with ByteArrayOutputStream org.apache.nifi.stream.io.ByteArrayOutputStream used on opensource projects

Example 46 with ByteArrayOutputStream

use of org.apache.nifi.stream.io.ByteArrayOutputStream in project nifi by apache.

the class HttpClientTransaction method readTransactionResponse.

@Override
protected Response readTransactionResponse() throws IOException {
    HttpCommunicationsSession commSession = (HttpCommunicationsSession) peer.getCommunicationsSession();
    ByteArrayOutputStream bos = new ByteArrayOutputStream();
    DataOutputStream dos = new DataOutputStream(bos);
    if (TransferDirection.RECEIVE.equals(direction)) {
        switch(state) {
            case TRANSACTION_STARTED:
            case DATA_EXCHANGED:
                logger.debug("{} {} readTransactionResponse. checksum={}", this, peer, commSession.getChecksum());
                if (StringUtils.isEmpty(commSession.getChecksum())) {
                    // We don't know if there's more data to receive, so just continue it.
                    ResponseCode.CONTINUE_TRANSACTION.writeResponse(dos);
                } else {
                    // We got a checksum to send to server.
                    if (TransactionState.TRANSACTION_STARTED.equals(state)) {
                        logger.debug("{} {} There's no transaction to confirm.", this, peer);
                        ResponseCode.CONFIRM_TRANSACTION.writeResponse(dos, "");
                    } else {
                        TransactionResultEntity transactionResult = apiClient.commitReceivingFlowFiles(transactionUrl, ResponseCode.CONFIRM_TRANSACTION, commSession.getChecksum());
                        ResponseCode responseCode = ResponseCode.fromCode(transactionResult.getResponseCode());
                        if (responseCode.containsMessage()) {
                            String message = transactionResult.getMessage();
                            responseCode.writeResponse(dos, message == null ? "" : message);
                        } else {
                            responseCode.writeResponse(dos);
                        }
                    }
                }
                break;
        }
    } else {
        switch(state) {
            case DATA_EXCHANGED:
                // Some flow files have been sent via stream, finish transferring.
                apiClient.finishTransferFlowFiles(commSession);
                ResponseCode.CONFIRM_TRANSACTION.writeResponse(dos, commSession.getChecksum());
                break;
            case TRANSACTION_CONFIRMED:
                TransactionResultEntity resultEntity = apiClient.commitTransferFlowFiles(transactionUrl, ResponseCode.CONFIRM_TRANSACTION);
                ResponseCode responseCode = ResponseCode.fromCode(resultEntity.getResponseCode());
                if (responseCode.containsMessage()) {
                    responseCode.writeResponse(dos, resultEntity.getMessage());
                } else {
                    responseCode.writeResponse(dos);
                }
                break;
        }
    }
    ByteArrayInputStream bis = new ByteArrayInputStream(bos.toByteArray());
    return Response.read(new DataInputStream(bis));
}

Example 47 with ByteArrayOutputStream

use of org.apache.nifi.stream.io.ByteArrayOutputStream in project nifi by apache.

the class TestHttpClient method consumeDataPacket.

private static void consumeDataPacket(DataPacket packet) throws IOException {
    final ByteArrayOutputStream bos = new ByteArrayOutputStream();
    StreamUtils.copy(packet.getData(), bos);
    String contents = new String(bos.toByteArray());
    logger.info("received: {}, {}", contents, packet.getAttributes());
}

Example 48 with ByteArrayOutputStream

use of org.apache.nifi.stream.io.ByteArrayOutputStream in project nifi by apache.

the class TestSiteToSiteClient method testSerialization.

@Test
public void testSerialization() {
    final SiteToSiteClientConfig clientConfig = new SiteToSiteClient.Builder().url("http://localhost:8080/nifi").portName("input").buildConfig();
    final Kryo kryo = new Kryo();
    final ByteArrayOutputStream out = new ByteArrayOutputStream();
    final Output output = new Output(out);
    try {
        kryo.writeObject(output, clientConfig);
    } finally {
        output.close();
    }
    final ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray());
    final Input input = new Input(in);
    try {
        SiteToSiteClientConfig clientConfig2 = kryo.readObject(input, SiteToSiteClient.StandardSiteToSiteClientConfig.class);
        Assert.assertEquals(clientConfig.getUrls(), clientConfig2.getUrls());
    } finally {
        input.close();
    }
}

Example 49 with ByteArrayOutputStream

use of org.apache.nifi.stream.io.ByteArrayOutputStream in project nifi by apache.

the class TestHttpClientTransaction method testReceiveOneFlowFile.

@Test
public void testReceiveOneFlowFile() throws IOException {
    SiteToSiteRestApiClient apiClient = mock(SiteToSiteRestApiClient.class);
    final String transactionUrl = "http://www.example.com/data-transfer/input-ports/portId/transactions/transactionId";
    doReturn(true).when(apiClient).openConnectionForReceive(eq(transactionUrl), any(Peer.class));
    TransactionResultEntity resultEntity = new TransactionResultEntity();
    resultEntity.setResponseCode(CONFIRM_TRANSACTION.getCode());
    doReturn(resultEntity).when(apiClient).commitReceivingFlowFiles(eq(transactionUrl), eq(CONFIRM_TRANSACTION), eq("3680976076"));
    ByteArrayOutputStream serverResponseBos = new ByteArrayOutputStream();
    codec.encode(createDataPacket("contents on server 1"), serverResponseBos);
    ByteArrayInputStream serverResponse = new ByteArrayInputStream(serverResponseBos.toByteArray());
    ByteArrayOutputStream clientRequest = new ByteArrayOutputStream();
    HttpClientTransaction transaction = getClientTransaction(serverResponse, clientRequest, apiClient, TransferDirection.RECEIVE, transactionUrl);
    execReceiveOneFlowFile(transaction);
    assertEquals("Client sends nothing as payload to receive flow files.", 0, clientRequest.toByteArray().length);
    verify(apiClient).commitReceivingFlowFiles(transactionUrl, CONFIRM_TRANSACTION, "3680976076");
}

Example 50 with ByteArrayOutputStream

use of org.apache.nifi.stream.io.ByteArrayOutputStream in project nifi by apache.

the class TestHttpClientTransaction method testReceiveWithInvalidChecksum.

@Test
public void testReceiveWithInvalidChecksum() throws IOException {
    SiteToSiteRestApiClient apiClient = mock(SiteToSiteRestApiClient.class);
    final String transactionUrl = "http://www.example.com/data-transfer/input-ports/portId/transactions/transactionId";
    doReturn(true).when(apiClient).openConnectionForReceive(eq(transactionUrl), any(Peer.class));
    // The checksum is correct, but here we simulate as if it's wrong, BAD_CHECKSUM.
    TransactionResultEntity resultEntity = new TransactionResultEntity();
    resultEntity.setResponseCode(ResponseCode.BAD_CHECKSUM.getCode());
    doReturn(resultEntity).when(apiClient).commitReceivingFlowFiles(eq(transactionUrl), eq(CONFIRM_TRANSACTION), eq("2969091230"));
    ByteArrayOutputStream serverResponseBos = new ByteArrayOutputStream();
    codec.encode(createDataPacket("contents on server 1"), serverResponseBos);
    codec.encode(createDataPacket("contents on server 2"), serverResponseBos);
    ByteArrayInputStream serverResponse = new ByteArrayInputStream(serverResponseBos.toByteArray());
    ByteArrayOutputStream clientRequest = new ByteArrayOutputStream();
    HttpClientTransaction transaction = getClientTransaction(serverResponse, clientRequest, apiClient, TransferDirection.RECEIVE, transactionUrl);
    execReceiveWithInvalidChecksum(transaction);
    assertEquals("Client sends nothing as payload to receive flow files.", 0, clientRequest.toByteArray().length);
    verify(apiClient).commitReceivingFlowFiles(transactionUrl, CONFIRM_TRANSACTION, "2969091230");
}