Examples with Write com.google.firestore.v1.Write used on opensource projects

Example 11 with Write

use of com.google.firestore.v1.Write in project beam by apache.

the class DatastoreV1Test method testWriteDisplayData.

@Test
public void testWriteDisplayData() {
    Write write = DatastoreIO.v1().write().withProjectId(PROJECT_ID);
    DisplayData displayData = DisplayData.from(write);
    assertThat(displayData, hasDisplayItem("projectId", PROJECT_ID));
}

Example 12 with Write

use of com.google.firestore.v1.Write in project beam by apache.

the class DatastoreV1Test method testBuildWrite.

/**
 * Test building a Write using builder methods.
 */
@Test
public void testBuildWrite() throws Exception {
    DatastoreV1.Write write = DatastoreIO.v1().write().withProjectId(PROJECT_ID);
    assertEquals(PROJECT_ID, write.getProjectId());
}

Example 13 with Write

use of com.google.firestore.v1.Write in project beam by apache.

the class FirestoreV1IT method batchWrite_partialFailureOutputsToDeadLetterQueue.

@Test
public void batchWrite_partialFailureOutputsToDeadLetterQueue() throws InterruptedException, ExecutionException, TimeoutException {
    String collectionId = "a";
    String docId = helper.docId();
    Write validWrite = Write.newBuilder().setUpdate(Document.newBuilder().setName(docPath(helper.getBaseDocumentPath(), collectionId, docId)).putFields("foo", Value.newBuilder().setStringValue(docId).build())).build();
    long millis = System.currentTimeMillis();
    Timestamp timestamp = Timestamp.newBuilder().setSeconds(millis / 1000).setNanos((int) ((millis % 1000) * 1000000)).build();
    String docId2 = helper.docId();
    helper.getBaseDocument().collection(collectionId).document(docId2).create(ImmutableMap.of("foo", "baz")).get(10, TimeUnit.SECONDS);
    // this will fail because we're setting a updateTime precondition to before it was created
    Write conditionalUpdate = Write.newBuilder().setUpdate(Document.newBuilder().setName(docPath(helper.getBaseDocumentPath(), collectionId, docId2)).putFields("foo", Value.newBuilder().setStringValue(docId).build())).setCurrentDocument(Precondition.newBuilder().setUpdateTime(timestamp)).build();
    List<Write> writes = newArrayList(validWrite, conditionalUpdate);
    RpcQosOptions options = BaseFirestoreIT.rpcQosOptions.toBuilder().withBatchMaxCount(2).build();
    PCollection<WriteFailure> writeFailurePCollection = testPipeline.apply(Create.of(writes)).apply(FirestoreIO.v1().write().batchWrite().withDeadLetterQueue().withRpcQosOptions(options).build());
    PAssert.that(writeFailurePCollection).satisfies((writeFailures) -> {
        Iterator<WriteFailure> iterator = writeFailures.iterator();
        assertTrue(iterator.hasNext());
        WriteFailure failure = iterator.next();
        assertEquals(Code.FAILED_PRECONDITION, Code.forNumber(failure.getStatus().getCode()));
        assertNotNull(failure.getWriteResult());
        assertFalse(failure.getWriteResult().hasUpdateTime());
        assertEquals(conditionalUpdate, failure.getWrite());
        assertFalse(iterator.hasNext());
        return null;
    });
    testPipeline.run(this.options);
    ApiFuture<QuerySnapshot> actualDocsQuery = helper.getBaseDocument().collection(collectionId).orderBy("__name__").get();
    QuerySnapshot querySnapshot = actualDocsQuery.get(10, TimeUnit.SECONDS);
    List<QueryDocumentSnapshot> documents = querySnapshot.getDocuments();
    List<KV<String, String>> actualDocumentIds = documents.stream().map(doc -> KV.of(doc.getId(), doc.getString("foo"))).collect(Collectors.toList());
    List<KV<String, String>> expected = newArrayList(KV.of(docId, docId), KV.of(docId2, "baz"));
    assertEquals(expected, actualDocumentIds);
}

Example 14 with Write

use of com.google.firestore.v1.Write in project beam by apache.

the class BaseFirestoreV1WriteFnTest method queuedWritesMaintainPriorityIfNotFlushed.

@Test
public final void queuedWritesMaintainPriorityIfNotFlushed() throws Exception {
    RpcQosOptions options = rpcQosOptions.toBuilder().withMaxAttempts(1).build();
    Write write0 = newWrite(0);
    Write write1 = newWrite(1);
    Write write2 = newWrite(2);
    Write write3 = newWrite(3);
    Write write4 = newWrite(4);
    Instant write4Start = Instant.ofEpochMilli(4);
    when(processContext.element()).thenReturn(write0, write1, write2, write3, write4).thenThrow(new IllegalStateException("too many calls"));
    when(rpcQos.newWriteAttempt(any())).thenReturn(attempt);
    when(attempt.awaitSafeToProceed(any())).thenReturn(true);
    when(attempt.<Write, Element<Write>>newFlushBuffer(any())).thenAnswer(invocation -> newFlushBuffer(options));
    FnT fn = getFn(clock, ff, options, CounterFactory.DEFAULT, DistributionFactory.DEFAULT);
    fn.setup();
    fn.startBundle(startBundleContext);
    // write0
    fn.processElement(processContext, window);
    // write1
    fn.processElement(processContext, window);
    // write2
    fn.processElement(processContext, window);
    // write3
    fn.processElement(processContext, window);
    // write4
    fn.processElement(processContext, window);
    List<WriteElement> expectedWrites = newArrayList(new WriteElement(0, write0, window), new WriteElement(1, write1, window), new WriteElement(2, write2, window), new WriteElement(3, write3, window), new WriteElement(4, write4, window));
    List<WriteElement> actualWrites = new ArrayList<>(fn.writes);
    assertEquals(expectedWrites, actualWrites);
    assertEquals(5, fn.queueNextEntryPriority);
    verify(attempt, times(1)).newFlushBuffer(write4Start);
    verifyNoMoreInteractions(callable);
}

Example 15 with Write

use of com.google.firestore.v1.Write in project beam by apache.

the class BaseFirestoreV1WriteFnTest method endToEnd_awaitSafeToProceed_falseIsTerminalForAttempt.

@Test
public final void endToEnd_awaitSafeToProceed_falseIsTerminalForAttempt() throws Exception {
    RpcQosOptions options = rpcQosOptions.toBuilder().withBatchMaxCount(2).build();
    Instant rpc1Start = Instant.ofEpochMilli(3);
    Instant rpc1End = Instant.ofEpochMilli(4);
    ArgumentCaptor<BatchWriteRequest> requestCaptor = ArgumentCaptor.forClass(BatchWriteRequest.class);
    Write write = newWrite();
    BatchWriteRequest expectedRequest = BatchWriteRequest.newBuilder().setDatabase("projects/testing-project/databases/(default)").addWrites(write).build();
    BatchWriteResponse response = BatchWriteResponse.newBuilder().addStatus(STATUS_OK).build();
    when(processContext.element()).thenReturn(write);
    // process element attempt 1
    when(attempt.awaitSafeToProceed(any())).thenReturn(false).thenThrow(new IllegalStateException("too many attempt1#awaitSafeToProceed"));
    // process element attempt 2
    when(attempt2.awaitSafeToProceed(any())).thenReturn(true).thenThrow(new IllegalStateException("too many attempt2#awaitSafeToProceed"));
    when(attempt2.<Write, Element<Write>>newFlushBuffer(any())).thenAnswer(invocation -> newFlushBuffer(options));
    // finish bundle attempt
    RpcQos.RpcWriteAttempt finishBundleAttempt = mock(RpcWriteAttempt.class);
    when(finishBundleAttempt.awaitSafeToProceed(any())).thenReturn(true, true).thenThrow(new IllegalStateException("too many finishBundleAttempt#awaitSafeToProceed"));
    when(finishBundleAttempt.<Write, Element<Write>>newFlushBuffer(any())).thenAnswer(invocation -> newFlushBuffer(options));
    when(rpcQos.newWriteAttempt(any())).thenReturn(attempt, attempt2, finishBundleAttempt);
    when(callable.call(requestCaptor.capture())).thenReturn(response);
    FnT fn = getFn(clock, ff, options, CounterFactory.DEFAULT, DistributionFactory.DEFAULT);
    runFunction(fn);
    assertEquals(expectedRequest, requestCaptor.getValue());
    verify(attempt, times(1)).awaitSafeToProceed(any());
    verifyNoMoreInteractions(attempt);
    verify(attempt2, times(1)).awaitSafeToProceed(any());
    verify(attempt2, times(1)).newFlushBuffer(any());
    verifyNoMoreInteractions(attempt2);
    verify(finishBundleAttempt, times(1)).recordRequestStart(rpc1Start, 1);
    verify(finishBundleAttempt, times(1)).recordWriteCounts(rpc1End, 1, 0);
    verify(finishBundleAttempt, times(1)).completeSuccess();
    verify(finishBundleAttempt, never()).checkCanRetry(any(), any());
}