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

protected <T> CompletableFuture<T>  asyncWriteHelper(Observable<? extends Document> observable, String errorMessage,
    boolean isVoid) {
  CompletableFuture<T> future = new CompletableFuture<>();
  observable.toSingle().subscribe(new SingleSubscriber<Document>() {
    @Override
    public void onSuccess(Document document) {
      if (isVoid) {
        future.complete(null);
      } else {
        future.complete((T) document.content());
      }
    }

    @Override
    public void onError(Throwable error) {
      future.completeExceptionally(new SamzaException(errorMessage, error));
    }
  });
  return future;
}
 

@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 testPutAsyncException() {
  String key = "throwExceptionKey";
  JsonObject value = JsonObject.create();
  Bucket bucket = mock(Bucket.class);
  AsyncBucket asyncBucket = mock(AsyncBucket.class);
  CouchbaseTableWriteFunction<JsonObject> writeFunction = createAndInit(bucket, asyncBucket);
  when(asyncBucket.upsert(any(Document.class), anyLong(), any(TimeUnit.class))).thenReturn(
      Observable.error(new CouchbaseException()));
  assertTrue(writeFunction.putAsync(key, value).isCompletedExceptionally());
}
 

@Test
public void testPutAsyncJsonObjectValue() throws Exception {
  String key = "key";
  JsonObject value = JsonObject.fromJson("{\"id\": 1}");
  Bucket bucket = mock(Bucket.class);
  AsyncBucket asyncBucket = mock(AsyncBucket.class);
  CouchbaseTableWriteFunction<JsonObject> writeFunction = createAndInit(bucket, asyncBucket);
  when(asyncBucket.upsert(any(Document.class), anyLong(), any(TimeUnit.class))).thenReturn(Observable.just(null));
  assertNull(writeFunction.putAsync(key, value).get());
}
 

@Test
public void testPutAsyncStringValue() throws Exception {
  String key = "key";
  String value = "value";
  Bucket bucket = mock(Bucket.class);
  AsyncBucket asyncBucket = mock(AsyncBucket.class);
  CouchbaseTableWriteFunction<String> writeFunction =
      createAndInit(String.class, new StringSerde(), bucket, asyncBucket);
  when(asyncBucket.upsert(any(Document.class), anyLong(), any(TimeUnit.class))).thenReturn(Observable.just(null));
  assertNull(writeFunction.putAsync(key, value).get());
}
 

@Override
public <K, V> boolean putIfAbsent(K key, V value, Serializer<K> keySerializer, Serializer<V> valueSerializer) throws IOException {
    final String docId = toDocumentId(key, keySerializer);
    final Document doc = toDocument(docId, value, valueSerializer);
    try {
        bucket.insert(doc);
        return true;
    } catch (DocumentAlreadyExistsException e) {
        return false;
    }
}
 

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

protected CompletableFuture<Void>  asyncWriteHelper(Observable<? extends Document> observable, String errorMessage) {
  return asyncWriteHelper(observable, errorMessage, true);
}
 

@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 (context.getProperty(DOC_ID).isSet()) {
        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: {
                doc = ByteArrayDocument.create(docId, content);
                break;
            }
        }

        final PersistTo persistTo = PersistTo.valueOf(context.getProperty(PERSIST_TO).getValue());
        final ReplicateTo replicateTo = ReplicateTo.valueOf(context.getProperty(REPLICATE_TO).getValue());
        final Bucket bucket = openBucket(context);
        doc = bucket.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(), bucket.name());
        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(bucket, 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) throws IOException {
    return toDocument(docId, value, valueSerializer, 0);
}
 

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> void put(K key, V value, Serializer<K> keySerializer, Serializer<V> valueSerializer) throws IOException {
    final String docId = toDocumentId(key, keySerializer);
    final Document doc = toDocument(docId, value, valueSerializer);
    bucket.upsert(doc);
}