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