gobblin.couchbase.writer.CouchbaseWriter.java Source code

Java tutorial

Introduction

Here is the source code for gobblin.couchbase.writer.CouchbaseWriter.java

Source

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package gobblin.couchbase.writer;

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;

import org.apache.commons.math3.util.Pair;

import com.couchbase.client.core.lang.Tuple;
import com.couchbase.client.core.lang.Tuple2;
import com.couchbase.client.core.message.ResponseStatus;
import com.couchbase.client.core.message.kv.MutationToken;
import com.couchbase.client.deps.io.netty.buffer.ByteBuf;
import com.couchbase.client.java.Bucket;
import com.couchbase.client.java.Cluster;
import com.couchbase.client.java.CouchbaseCluster;
import com.couchbase.client.java.document.AbstractDocument;
import com.couchbase.client.java.document.Document;
import com.couchbase.client.java.document.RawJsonDocument;
import com.couchbase.client.java.env.CouchbaseEnvironment;
import com.couchbase.client.java.transcoder.Transcoder;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.typesafe.config.Config;

import lombok.extern.slf4j.Slf4j;
import rx.Observable;
import rx.Subscriber;

import gobblin.couchbase.common.TupleDocument;
import gobblin.util.ConfigUtils;
import gobblin.writer.AsyncDataWriter;
import gobblin.writer.GenericWriteResponse;
import gobblin.writer.GenericWriteResponseWrapper;
import gobblin.writer.SyncDataWriter;
import gobblin.writer.WriteCallback;
import gobblin.writer.WriteResponse;
import gobblin.writer.WriteResponseFuture;
import gobblin.writer.WriteResponseMapper;

/**
 * A single bucket Couchbase writer.
 */
@Slf4j
public class CouchbaseWriter<D extends AbstractDocument> implements AsyncDataWriter<D>, SyncDataWriter<D> {

    private final Cluster _cluster;
    private final Bucket _bucket;
    private final long _operationTimeout;
    private final TimeUnit _operationTimeunit;
    private final WriteResponseMapper<D> _defaultWriteResponseMapper;

    // A basic transcoder that just passes through the embedded binary content.
    private final Transcoder<TupleDocument, Tuple2<ByteBuf, Integer>> _tupleDocumentTranscoder = new Transcoder<TupleDocument, Tuple2<ByteBuf, Integer>>() {
        @Override
        public TupleDocument decode(String id, ByteBuf content, long cas, int expiry, int flags,
                ResponseStatus status) {
            return newDocument(id, expiry, Tuple.create(content, flags), cas);
        }

        @Override
        public Tuple2<ByteBuf, Integer> encode(TupleDocument document) {
            return document.content();
        }

        @Override
        public TupleDocument newDocument(String id, int expiry, Tuple2<ByteBuf, Integer> content, long cas) {
            return new TupleDocument(id, expiry, content, cas);
        }

        @Override
        public TupleDocument newDocument(String id, int expiry, Tuple2<ByteBuf, Integer> content, long cas,
                MutationToken mutationToken) {
            return new TupleDocument(id, expiry, content, cas);
        }

        @Override
        public Class<TupleDocument> documentType() {
            return TupleDocument.class;
        }
    };

    public CouchbaseWriter(CouchbaseEnvironment couchbaseEnvironment, Config config) {

        List<String> hosts = ConfigUtils.getStringList(config, CouchbaseWriterConfigurationKeys.BOOTSTRAP_SERVERS);

        _cluster = CouchbaseCluster.create(couchbaseEnvironment, hosts);

        String bucketName = ConfigUtils.getString(config, CouchbaseWriterConfigurationKeys.BUCKET,
                CouchbaseWriterConfigurationKeys.BUCKET_DEFAULT);

        String password = ConfigUtils.getString(config, CouchbaseWriterConfigurationKeys.PASSWORD, "");

        _bucket = _cluster.openBucket(bucketName, password,
                Collections.<Transcoder<? extends Document, ?>>singletonList(_tupleDocumentTranscoder));

        _operationTimeout = ConfigUtils.getLong(config, CouchbaseWriterConfigurationKeys.OPERATION_TIMEOUT_MILLIS,
                CouchbaseWriterConfigurationKeys.OPERATION_TIMEOUT_DEFAULT);
        _operationTimeunit = TimeUnit.MILLISECONDS;

        _defaultWriteResponseMapper = new GenericWriteResponseWrapper<>();

        log.info("Couchbase writer configured with: hosts: {}, bucketName: {}, operationTimeoutInMillis: {}", hosts,
                bucketName, _operationTimeout);
    }

    @VisibleForTesting
    Bucket getBucket() {
        return _bucket;
    }

    private void assertRecordWritable(D record) {
        boolean recordIsTupleDocument = (record instanceof TupleDocument);
        boolean recordIsJsonDocument = (record instanceof RawJsonDocument);
        Preconditions.checkArgument(recordIsTupleDocument || recordIsJsonDocument,
                "This writer only supports TupleDocument or RawJsonDocument. Found " + record.getClass().getName());
    }

    @Override
    public Future<WriteResponse> write(final D record, final WriteCallback callback) {
        assertRecordWritable(record);
        if (record instanceof TupleDocument) {
            ((TupleDocument) record).content().value1().retain();
        }
        Observable<D> observable = _bucket.async().upsert(record);
        if (callback == null) {
            return new WriteResponseFuture<>(
                    observable.timeout(_operationTimeout, _operationTimeunit).toBlocking().toFuture(),
                    _defaultWriteResponseMapper);
        } else {

            final AtomicBoolean callbackFired = new AtomicBoolean(false);
            final BlockingQueue<Pair<WriteResponse, Throwable>> writeResponseQueue = new ArrayBlockingQueue<>(1);

            final Future<WriteResponse> writeResponseFuture = new Future<WriteResponse>() {
                @Override
                public boolean cancel(boolean mayInterruptIfRunning) {
                    return false;
                }

                @Override
                public boolean isCancelled() {
                    return false;
                }

                @Override
                public boolean isDone() {
                    return callbackFired.get();
                }

                @Override
                public WriteResponse get() throws InterruptedException, ExecutionException {
                    Pair<WriteResponse, Throwable> writeResponseThrowablePair = writeResponseQueue.take();
                    return getWriteResponseorThrow(writeResponseThrowablePair);
                }

                @Override
                public WriteResponse get(long timeout, TimeUnit unit)
                        throws InterruptedException, ExecutionException, TimeoutException {
                    Pair<WriteResponse, Throwable> writeResponseThrowablePair = writeResponseQueue.poll(timeout,
                            unit);
                    if (writeResponseThrowablePair == null) {
                        throw new TimeoutException("Timeout exceeded while waiting for future to be done");
                    } else {
                        return getWriteResponseorThrow(writeResponseThrowablePair);
                    }
                }
            };

            observable.timeout(_operationTimeout, _operationTimeunit).subscribe(new Subscriber<D>() {
                @Override
                public void onCompleted() {
                }

                @Override
                public void onError(Throwable e) {
                    callbackFired.set(true);
                    writeResponseQueue.add(new Pair<WriteResponse, Throwable>(null, e));
                    callback.onFailure(e);
                }

                @Override
                public void onNext(D doc) {
                    try {
                        callbackFired.set(true);
                        WriteResponse writeResponse = new GenericWriteResponse<D>(doc);
                        writeResponseQueue.add(new Pair<WriteResponse, Throwable>(writeResponse, null));
                        callback.onSuccess(writeResponse);
                    } finally {
                        if (doc instanceof TupleDocument) {
                            ((TupleDocument) doc).content().value1().release();
                        }
                    }
                }
            });
            return writeResponseFuture;
        }
    }

    @Override
    public void flush() throws IOException {

    }

    private WriteResponse getWriteResponseorThrow(Pair<WriteResponse, Throwable> writeResponseThrowablePair)
            throws ExecutionException {
        if (writeResponseThrowablePair.getFirst() != null) {
            return writeResponseThrowablePair.getFirst();
        } else if (writeResponseThrowablePair.getSecond() != null) {
            throw new ExecutionException(writeResponseThrowablePair.getSecond());
        } else {
            throw new ExecutionException(new RuntimeException("Could not find non-null WriteResponse pair"));
        }
    }

    @Override
    public void cleanup() throws IOException {

    }

    @Override
    public WriteResponse write(D record) throws IOException {

        try {
            D doc = _bucket.upsert(record);
            return new GenericWriteResponse(doc);
        } catch (Exception e) {
            throw new IOException("Failed to write to Couchbase cluster", e);
        }
    }

    @Override
    public void close() {
        if (!_bucket.isClosed()) {
            try {
                _bucket.close();
            } catch (Exception e) {
                log.warn("Failed to close bucket", e);
            }
        }
        try {
            _cluster.disconnect();
        } catch (Exception e) {
            log.warn("Failed to disconnect from cluster", e);
        }
    }
}