com.couchbase.client.java.document.BinaryDocument Java Examples

@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);
}
 

protected void handleGetAsyncBinaryDocument(BinaryDocument binaryDocument, CompletableFuture<V> future, String key) {
  ByteBuf buffer = binaryDocument.content();
  try {
    byte[] bytes;
    if (buffer.hasArray() && buffer.arrayOffset() == 0 && buffer.readableBytes() == buffer.array().length) {
      bytes = buffer.array();
    } else {
      bytes = new byte[buffer.readableBytes()];
      buffer.readBytes(bytes);
    }
    future.complete(valueSerde.fromBytes(bytes));
  } catch (Exception e) {
    future.completeExceptionally(
        new SamzaException(String.format("Failed to deserialize value of key %s with given serde", key), e));
  } finally {
    ReferenceCountUtil.release(buffer);
  }
}
 

@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);
}
 

@Test
public void testBinaryDocumentToAttribute() 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.setProperty(PUT_VALUE_TO_ATTRIBUTE, "targetAttribute");
    testRunner.run();

    testRunner.assertTransferCount(REL_SUCCESS, 1);
    testRunner.assertTransferCount(REL_ORIGINAL, 0);
    testRunner.assertTransferCount(REL_RETRY, 0);
    testRunner.assertTransferCount(REL_FAILURE, 0);
    MockFlowFile outFile = testRunner.getFlowFilesForRelationship(REL_SUCCESS).get(0);
    outFile.assertContentEquals(inFileDataStr);
    outFile.assertAttributeEquals("targetAttribute", "binary");
}
 

@Test
public void testGet() throws Exception {
    final CouchbaseMapCacheClient client = new CouchbaseMapCacheClient();
    final CouchbaseClusterControllerService couchbaseService = mock(CouchbaseClusterControllerService.class);
    final Bucket bucket = mock(Bucket.class);

    final MockControllerServiceInitializationContext serviceInitializationContext
            = new MockControllerServiceInitializationContext(couchbaseService, "couchbaseService");
    final Map<PropertyDescriptor, String> properties = new HashMap<>();
    properties.put(COUCHBASE_CLUSTER_SERVICE, "couchbaseService");
    properties.put(BUCKET_NAME, "bucketA");

    final ByteBuf contents = Unpooled.copiedBuffer("value".getBytes(StandardCharsets.UTF_8));
    final BinaryDocument doc = BinaryDocument.create("key", contents);
    when(couchbaseService.openBucket(eq("bucketA"))).thenReturn(bucket);
    when(bucket.get(any(BinaryDocument.class))).thenReturn(doc);

    final MockConfigurationContext context = new MockConfigurationContext(properties, serviceInitializationContext);
    client.configure(context);
    final String cacheEntry = client.get("key", stringSerializer, stringDeserializer);

    assertEquals("value", cacheEntry);
}
 

@Override
public CompletableFuture<Void> putAsync(String key, V record) {
  Preconditions.checkArgument(StringUtils.isNotBlank(key), "key must not be null, empty or blank");
  Preconditions.checkArgument(!key.contains(" "), String.format("key should not contain spaces: %s", key));
  Preconditions.checkNotNull(record);
  Document<?> document = record instanceof JsonObject
      ? JsonDocument.create(key, (int) ttl.getSeconds(), (JsonObject) record)
      : BinaryDocument.create(key, (int) ttl.getSeconds(), Unpooled.copiedBuffer(valueSerde.toBytes(record)));
  return asyncWriteHelper(
      bucket.async().upsert(document, timeout.toMillis(), TimeUnit.MILLISECONDS),
      String.format("Failed to insert key %s into bucket %s", key, bucketName));
}
 

@Override
public CompletableFuture<V> getAsync(String key) {
  Preconditions.checkArgument(StringUtils.isNotBlank(key), "key must not be null, empty or blank");
  CompletableFuture<V> future = new CompletableFuture<>();
  Single<? extends Document<?>> singleObservable =
      bucket.async().get(key, documentType, timeout.toMillis(), TimeUnit.MILLISECONDS).toSingle();
  singleObservable.subscribe(new SingleSubscriber<Document<?>>() {
    @Override
    public void onSuccess(Document<?> document) {
      if (document != null) {
        if (document instanceof BinaryDocument) {
          handleGetAsyncBinaryDocument((BinaryDocument) document, future, key);
        } else {
          // V is of type JsonObject
          future.complete((V) document.content());
        }
      } else {
        // The Couchbase async client should not return null
        future.completeExceptionally(new SamzaException(String.format("Got unexpected null value from key %s", key)));
      }
    }

    @Override
    public void onError(Throwable throwable) {
      if (throwable instanceof NoSuchElementException) {
        // There is no element returned by the observable, meaning the key doesn't exist.
        future.complete(null);
      } else {
        future.completeExceptionally(new SamzaException(String.format("Failed to get key %s", key), throwable));
      }
    }
  });
  return future;
}
 

@Test
public void testGetAsyncFailedToDeserialize() {
  String key = "key";
  Bucket bucket = mock(Bucket.class);
  AsyncBucket asyncBucket = mock(AsyncBucket.class);
  CouchbaseTableReadFunction readFunction = createAndInit(String.class, new StringSerde(), bucket, asyncBucket);
  when(asyncBucket.get(eq(key), anyObject(), anyLong(), any(TimeUnit.class))).thenReturn(
      Observable.just(BinaryDocument.create(key, null)));
  assertTrue(readFunction.getAsync(key).isCompletedExceptionally());
}
 

@Test
public void testGetAsyncStringValue() throws Exception {
  String key = "key";
  String value = "value";
  StringSerde stringSerde = new StringSerde();
  Bucket bucket = mock(Bucket.class);
  AsyncBucket asyncBucket = mock(AsyncBucket.class);
  CouchbaseTableReadFunction readFunction = createAndInit(String.class, stringSerde, bucket, asyncBucket);
  when(asyncBucket.get(eq(key), anyObject(), anyLong(), any(TimeUnit.class))).thenReturn(
      Observable.just(BinaryDocument.create(key, Unpooled.wrappedBuffer(stringSerde.toBytes(value)))));
  assertEquals(value, readFunction.getAsync(key).get());
}
 

@Override
public <K, V> AtomicCacheEntry<K, V, Long> fetch(K key, Serializer<K> keySerializer, Deserializer<V> valueDeserializer) throws IOException {
    final String docId = toDocumentId(key, keySerializer);
    final BinaryDocument doc = bucket.get(BinaryDocument.create(docId));
    if (doc == null) {
        return null;
    }
    final V value = deserialize(doc, valueDeserializer);
    return new AtomicCacheEntry<>(key, value, doc.cas());
}
 

private <V> V deserialize(BinaryDocument doc, Deserializer<V> valueDeserializer) throws IOException {
    if (doc == null) {
        return null;
    }
    final ByteBuf byteBuf = doc.content();
    final byte[] bytes = new byte[byteBuf.readableBytes()];
    byteBuf.readBytes(bytes);
    byteBuf.release();
    return valueDeserializer.deserialize(bytes);
}
 

@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);
    }
}
 

private <V> Document toDocument(String docId, V value, Serializer<V> valueSerializer, long revision) throws IOException {
    ByteArrayOutputStream bos = new ByteArrayOutputStream();
    valueSerializer.serialize(value, bos);
    final ByteBuf byteBuf = Unpooled.wrappedBuffer(bos.toByteArray());
    return BinaryDocument.create(docId, byteBuf, revision);
}
 

@Override
public <K, V> V get(K key, Serializer<K> keySerializer, Deserializer<V> valueDeserializer) throws IOException {
    final String docId = toDocumentId(key, keySerializer);
    final BinaryDocument doc = bucket.get(BinaryDocument.create(docId));
    return deserialize(doc, valueDeserializer);
}
 

@Ignore("This test method requires a live Couchbase Server instance")
@Test
public void testDocumentTypesAndStringConversion() {
    final CouchbaseCluster cluster = CouchbaseCluster.fromConnectionString("couchbase://192.168.99.100:8091");
    final Bucket bucket = cluster.openBucket("b1", "b1password");

    bucket.upsert(JsonDocument.create("JsonDocument", JsonObject.create().put("one", 1)));
    bucket.upsert(JsonArrayDocument.create("JsonArray", JsonArray.create().add(1).add(2).add(3)));
    bucket.upsert(JsonDoubleDocument.create("JsonDouble", 0.123));
    bucket.upsert(JsonStringDocument.create("JsonString", "value"));
    bucket.upsert(JsonBooleanDocument.create("JsonBoolean", true));
    bucket.upsert(JsonLongDocument.create("JsonLong", 123L));

    bucket.upsert(RawJsonDocument.create("RawJsonDocument", "value"));
    bucket.upsert(StringDocument.create("StringDocument", "value"));

    bucket.upsert(BinaryDocument.create("BinaryDocument", Unpooled.copiedBuffer("value".getBytes(StandardCharsets.UTF_8))));
    bucket.upsert(ByteArrayDocument.create("ByteArrayDocument", "value".getBytes(StandardCharsets.UTF_8)));

    final String[][] expectations = {
            {"JsonDocument", "String", "{\"one\":1}"},
            {"JsonArray", "String", "[1,2,3]"},
            {"JsonDouble", "String", "0.123"},
            {"JsonString", "String", "\"value\""},
            {"JsonBoolean", "String", "true"},
            {"JsonLong", "String", "123"},
            {"RawJsonDocument", "String", "value"},
            {"StringDocument", "String", "value"},
            {"BinaryDocument", "byte[]", "value"},
            {"ByteArrayDocument", "byte[]", "value"},
    };

    for (String[] expectation : expectations) {
        final LegacyDocument document = bucket.get(LegacyDocument.create(expectation[0]));
        assertEquals(expectation[1], document.content().getClass().getSimpleName());
        assertEquals(expectation[2], CouchbaseUtils.getStringContent(document.content()));
    }

    final BinaryDocument binaryDocument = bucket.get(BinaryDocument.create("BinaryDocument"));
    final String stringFromByteBuff = CouchbaseUtils.getStringContent(binaryDocument.content());
    assertEquals("value", stringFromByteBuff);

    try {
        bucket.get(BinaryDocument.create("JsonDocument"));
        fail("Getting a JSON document as a BinaryDocument fails");
    } catch (TranscodingException e) {
        assertTrue(e.getMessage().contains("Flags (0x2000000) indicate non-binary document for id JsonDocument"));
    }

}
 

/**
 * Construct an instance of {@link CouchbaseTableReadFunction}.
 * @param bucketName Name of the couchbase bucket
 * @param clusterNodes Some Hosts of the Couchbase cluster. Recommended to provide more than one nodes so that if
 *                     the first node could not be connected, other nodes can be tried.
 * @param valueClass Type of values
 */
public CouchbaseTableReadFunction(String bucketName, Class<V> valueClass, String... clusterNodes) {
  super(bucketName, valueClass, clusterNodes);
  documentType = JsonObject.class.isAssignableFrom(valueClass) ? JsonDocument.class : BinaryDocument.class;
}