Search in sources :

Example 1 with ByteBuf

use of com.couchbase.client.deps.io.netty.buffer.ByteBuf in project nifi by apache.

the class PutCouchbaseKey method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    final ComponentLog logger = getLogger();
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }
    final byte[] content = new byte[(int) flowFile.getSize()];
    session.read(flowFile, new InputStreamCallback() {

        @Override
        public void process(final InputStream in) throws IOException {
            StreamUtils.fillBuffer(in, content, true);
        }
    });
    String docId = flowFile.getAttribute(CoreAttributes.UUID.key());
    if (!StringUtils.isEmpty(context.getProperty(DOC_ID).getValue())) {
        docId = context.getProperty(DOC_ID).evaluateAttributeExpressions(flowFile).getValue();
    }
    try {
        Document<?> doc = null;
        final DocumentType documentType = DocumentType.valueOf(context.getProperty(DOCUMENT_TYPE).getValue());
        switch(documentType) {
            case Json:
                {
                    doc = RawJsonDocument.create(docId, new String(content, StandardCharsets.UTF_8));
                    break;
                }
            case Binary:
                {
                    final ByteBuf buf = Unpooled.copiedBuffer(content);
                    doc = BinaryDocument.create(docId, buf);
                    break;
                }
        }
        final PersistTo persistTo = PersistTo.valueOf(context.getProperty(PERSIST_TO).getValue());
        final ReplicateTo replicateTo = ReplicateTo.valueOf(context.getProperty(REPLICATE_TO).getValue());
        doc = openBucket(context).upsert(doc, persistTo, replicateTo);
        final Map<String, String> updatedAttrs = new HashMap<>();
        updatedAttrs.put(CouchbaseAttributes.Cluster.key(), context.getProperty(COUCHBASE_CLUSTER_SERVICE).getValue());
        updatedAttrs.put(CouchbaseAttributes.Bucket.key(), context.getProperty(BUCKET_NAME).getValue());
        updatedAttrs.put(CouchbaseAttributes.DocId.key(), docId);
        updatedAttrs.put(CouchbaseAttributes.Cas.key(), String.valueOf(doc.cas()));
        updatedAttrs.put(CouchbaseAttributes.Expiry.key(), String.valueOf(doc.expiry()));
        flowFile = session.putAllAttributes(flowFile, updatedAttrs);
        session.getProvenanceReporter().send(flowFile, getTransitUrl(context, docId));
        session.transfer(flowFile, REL_SUCCESS);
    } catch (final CouchbaseException e) {
        String errMsg = String.format("Writing document %s to Couchbase Server using %s failed due to %s", docId, flowFile, e);
        handleCouchbaseException(context, session, logger, flowFile, e, errMsg);
    }
}
Also used : FlowFile(org.apache.nifi.flowfile.FlowFile) HashMap(java.util.HashMap) InputStream(java.io.InputStream) IOException(java.io.IOException) ByteBuf(com.couchbase.client.deps.io.netty.buffer.ByteBuf) ComponentLog(org.apache.nifi.logging.ComponentLog) ReplicateTo(com.couchbase.client.java.ReplicateTo) CouchbaseException(com.couchbase.client.core.CouchbaseException) PersistTo(com.couchbase.client.java.PersistTo) InputStreamCallback(org.apache.nifi.processor.io.InputStreamCallback)

Example 2 with ByteBuf

use of com.couchbase.client.deps.io.netty.buffer.ByteBuf in project components by Talend.

the class CouchbaseStreamingConnection method stopStreaming.

public void stopStreaming() {
    if (resultsQueue != null) {
        client.stopStreaming(partitionsToStream()).await();
        BlockingQueue<ByteBuf> queue = resultsQueue;
        resultsQueue = null;
        List<ByteBuf> drained = new ArrayList<ByteBuf>();
        queue.drainTo(drained);
        for (ByteBuf byteBuf : drained) {
            byteBuf.release();
        }
        client.disconnect();
    }
}
Also used : ArrayList(java.util.ArrayList) ByteBuf(com.couchbase.client.deps.io.netty.buffer.ByteBuf)

Example 3 with ByteBuf

use of com.couchbase.client.deps.io.netty.buffer.ByteBuf in project components by Talend.

the class CouchbaseEventGenericRecordConverterTest method testConvertToAvroDeletion.

@Test
public void testConvertToAvroDeletion() {
    ByteBuf buffer = Mockito.mock(ByteBuf.class);
    // Mocking key object
    Mockito.when(buffer.getByte(0)).thenReturn(MessageUtil.MAGIC_REQ);
    Mockito.when(buffer.getByte(1)).thenReturn(MessageUtil.DCP_DELETION_OPCODE);
    Mockito.when(buffer.getByte(4)).thenReturn(OFFSET);
    Mockito.when(buffer.getShort(2)).thenReturn(LENGTH);
    ByteBuf key = Mockito.mock(ByteBuf.class);
    Mockito.when(key.readableBytes()).thenReturn(4);
    Mockito.when(buffer.slice(29, 10)).thenReturn(key);
    IndexedRecord record = converter.convertToAvro(buffer);
    assertIndexedRecordResult(record, "deletion", null);
}
Also used : IndexedRecord(org.apache.avro.generic.IndexedRecord) ByteBuf(com.couchbase.client.deps.io.netty.buffer.ByteBuf) Test(org.junit.Test)

Example 4 with ByteBuf

use of com.couchbase.client.deps.io.netty.buffer.ByteBuf in project components by Talend.

the class CouchbaseStreamingConnectionTest method testStartStreaming.

@Test
public void testStartStreaming() throws InterruptedException {
    Mockito.when(client.initializeState(StreamFrom.BEGINNING, StreamTo.NOW)).thenReturn(Completable.complete());
    Mockito.when(client.startStreaming(Mockito.<Short[]>anyVararg())).thenReturn(Completable.complete());
    SessionState sessionState = Mockito.mock(SessionState.class);
    Mockito.when(sessionState.isAtEnd()).thenReturn(false, false, true);
    Mockito.when(client.sessionState()).thenReturn(sessionState);
    BlockingQueue<ByteBuf> resultsQueue = new ArrayBlockingQueue<>(3);
    streamingConnection.startStreaming(resultsQueue);
    Assert.assertTrue(streamingConnection.isStreaming());
    Thread.sleep(2000);
    Mockito.verify(client, Mockito.times(3)).sessionState();
}
Also used : SessionState(com.couchbase.client.dcp.state.SessionState) ArrayBlockingQueue(java.util.concurrent.ArrayBlockingQueue) ByteBuf(com.couchbase.client.deps.io.netty.buffer.ByteBuf) Test(org.junit.Test) PrepareForTest(org.powermock.core.classloader.annotations.PrepareForTest)

Example 5 with ByteBuf

use of com.couchbase.client.deps.io.netty.buffer.ByteBuf in project nifi by apache.

the class TestGetCouchbaseKey method testBinaryDocument.

@Test
public void testBinaryDocument() throws Exception {
    Bucket bucket = mock(Bucket.class);
    String inFileDataStr = "doc-in";
    String content = "binary";
    ByteBuf buf = Unpooled.copiedBuffer(content.getBytes(StandardCharsets.UTF_8));
    when(bucket.get(inFileDataStr, BinaryDocument.class)).thenReturn(BinaryDocument.create(inFileDataStr, buf));
    setupMockBucket(bucket);
    byte[] inFileData = inFileDataStr.getBytes(StandardCharsets.UTF_8);
    testRunner.enqueue(inFileData);
    testRunner.setProperty(DOCUMENT_TYPE, DocumentType.Binary.toString());
    testRunner.run();
    testRunner.assertTransferCount(REL_SUCCESS, 1);
    testRunner.assertTransferCount(REL_ORIGINAL, 1);
    testRunner.assertTransferCount(REL_RETRY, 0);
    testRunner.assertTransferCount(REL_FAILURE, 0);
    MockFlowFile outFile = testRunner.getFlowFilesForRelationship(REL_SUCCESS).get(0);
    outFile.assertContentEquals(content);
    MockFlowFile orgFile = testRunner.getFlowFilesForRelationship(REL_ORIGINAL).get(0);
    orgFile.assertContentEquals(inFileDataStr);
}
Also used : MockFlowFile(org.apache.nifi.util.MockFlowFile) Bucket(com.couchbase.client.java.Bucket) Matchers.anyString(org.mockito.Matchers.anyString) ByteBuf(com.couchbase.client.deps.io.netty.buffer.ByteBuf) Test(org.junit.Test)

Aggregations

ByteBuf (com.couchbase.client.deps.io.netty.buffer.ByteBuf)12 Test (org.junit.Test)8 IndexedRecord (org.apache.avro.generic.IndexedRecord)5 SessionState (com.couchbase.client.dcp.state.SessionState)2 ArrayBlockingQueue (java.util.concurrent.ArrayBlockingQueue)2 PrepareForTest (org.powermock.core.classloader.annotations.PrepareForTest)2 CouchbaseException (com.couchbase.client.core.CouchbaseException)1 Bucket (com.couchbase.client.java.Bucket)1 PersistTo (com.couchbase.client.java.PersistTo)1 ReplicateTo (com.couchbase.client.java.ReplicateTo)1 IOException (java.io.IOException)1 InputStream (java.io.InputStream)1 ByteBuffer (java.nio.ByteBuffer)1 ArrayList (java.util.ArrayList)1 HashMap (java.util.HashMap)1 GenericRecord (org.apache.avro.generic.GenericRecord)1 SingleRecordIterable (org.apache.gobblin.converter.SingleRecordIterable)1 TupleDocument (org.apache.gobblin.couchbase.common.TupleDocument)1 FlowFile (org.apache.nifi.flowfile.FlowFile)1 ComponentLog (org.apache.nifi.logging.ComponentLog)1