Examples with WriteResponse org.apache.gobblin.writer.WriteResponse used on opensource projects

Example 6 with WriteResponse

use of org.apache.gobblin.writer.WriteResponse in project incubator-gobblin by apache.

the class EventhubDataWriterTest method testSingleBatch.

@Test
public void testSingleBatch() {
    // mock eventhub data writer
    Properties props = new Properties();
    EventhubDataWriter eventhubDataWriter = Mockito.spy(new EventhubDataWriter(props, mockHttpClient));
    Mockito.doNothing().when(eventhubDataWriter).refreshSignature();
    List<String> records = new LinkedList<>();
    for (int i = 0; i < 50; ++i) records.add(new String("abcdefgh"));
    Batch<String> batch = mock(Batch.class);
    WriteCallback callback = mock(WriteCallback.class);
    Mockito.when(batch.getRecords()).thenReturn(records);
    Future<WriteResponse> future = eventhubDataWriter.write(batch, callback);
    verify(callback, times(1)).onSuccess(isA(WriteResponse.class));
    verify(callback, never()).onFailure(isA(Exception.class));
    Assert.assertTrue(future.isDone(), "Future should be done");
}

Example 7 with WriteResponse

use of org.apache.gobblin.writer.WriteResponse in project incubator-gobblin by apache.

the class CouchbaseWriter method write.

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

Example 8 with WriteResponse

use of org.apache.gobblin.writer.WriteResponse in project incubator-gobblin by apache.

the class StreamingKafkaSpecExecutorTest method testUpdateSpec.

@Test(dependsOnMethods = "testAddSpec")
public void testUpdateSpec() throws Exception {
    // update is only treated as an update for existing job specs
    String updatedSpecUriString = "/foo/bar/addedSpec";
    Spec spec = initJobSpec(updatedSpecUriString);
    WriteResponse writeResponse = (WriteResponse) _seip.updateSpec(spec).get();
    log.info("WriteResponse: " + writeResponse);
    List<Pair<SpecExecutor.Verb, Spec>> consumedEvent = _seic.changedSpecs().get();
    Assert.assertTrue(consumedEvent.size() == 1, "Consumption did not match production");
    Map.Entry<SpecExecutor.Verb, Spec> consumedSpecAction = consumedEvent.get(0);
    Assert.assertTrue(consumedSpecAction.getKey().equals(SpecExecutor.Verb.UPDATE), "Verb did not match");
    Assert.assertTrue(consumedSpecAction.getValue().getUri().toString().equals(updatedSpecUriString), "Expected URI did not match");
    Assert.assertTrue(consumedSpecAction.getValue() instanceof JobSpec, "Expected JobSpec");
}

Example 9 with WriteResponse

use of org.apache.gobblin.writer.WriteResponse in project incubator-gobblin by apache.

the class SimpleKafkaSpecExecutorTest method testAddSpec.

@Test
public void testAddSpec() throws Exception {
    _closer = Closer.create();
    _properties = new Properties();
    // Properties for Producer
    _properties.setProperty(KafkaWriterConfigurationKeys.KAFKA_TOPIC, TOPIC);
    _properties.setProperty(KafkaWriterConfigurationKeys.KAFKA_PRODUCER_CONFIG_PREFIX + "bootstrap.servers", _kafkaBrokers);
    _properties.setProperty(KafkaWriterConfigurationKeys.KAFKA_PRODUCER_CONFIG_PREFIX + "value.serializer", "org.apache.kafka.common.serialization.ByteArraySerializer");
    // Properties for Consumer
    _properties.setProperty(ConfigurationKeys.KAFKA_BROKERS, _kafkaBrokers);
    _properties.setProperty(SimpleKafkaSpecExecutor.SPEC_KAFKA_TOPICS_KEY, TOPIC);
    // SEI Producer
    _seip = _closer.register(new SimpleKafkaSpecProducer(ConfigUtils.propertiesToConfig(_properties)));
    String addedSpecUriString = "/foo/bar/addedSpec";
    Spec spec = initJobSpec(addedSpecUriString);
    WriteResponse writeResponse = (WriteResponse) _seip.addSpec(spec).get();
    log.info("WriteResponse: " + writeResponse);
    try {
        Thread.sleep(1000);
    } catch (InterruptedException ex) {
        Thread.currentThread().interrupt();
    }
    _seic = _closer.register(new SimpleKafkaSpecConsumer(ConfigUtils.propertiesToConfig(_properties)));
    List<Pair<SpecExecutor.Verb, Spec>> consumedEvent = _seic.changedSpecs().get();
    Assert.assertTrue(consumedEvent.size() == 1, "Consumption did not match production");
    Map.Entry<SpecExecutor.Verb, Spec> consumedSpecAction = consumedEvent.get(0);
    Assert.assertTrue(consumedSpecAction.getKey().equals(SpecExecutor.Verb.ADD), "Verb did not match");
    Assert.assertTrue(consumedSpecAction.getValue().getUri().toString().equals(addedSpecUriString), "Expected URI did not match");
    Assert.assertTrue(consumedSpecAction.getValue() instanceof JobSpec, "Expected JobSpec");
}