Examples with DocumentGenerator org.apache.beam.sdk.io.gcp.firestore.it.FirestoreTestingHelper.DocumentGenerator used on opensource projects

Example 1 with DocumentGenerator

use of org.apache.beam.sdk.io.gcp.firestore.it.FirestoreTestingHelper.DocumentGenerator in project beam by apache.

the class BaseFirestoreIT method partitionQuery.

@Test
public final void partitionQuery() throws Exception {
    String collectionGroupId = UUID.randomUUID().toString();
    // currently firestore will only generate a partition every 128 documents, so generate enough
    // documents to get 2 cursors returned, resulting in 3 partitions
    int partitionCount = 3;
    int documentCount = (partitionCount * 128) - 1;
    // create some documents for listing and asserting in the test
    DocumentGenerator documentGenerator = helper.documentGenerator(documentCount, collectionGroupId);
    documentGenerator.generateDocuments().get(10, TimeUnit.SECONDS);
    PCollection<String> listDocumentPaths = testPipeline.apply(Create.of(collectionGroupId)).apply(getPartitionQueryPTransform(testName.getMethodName(), partitionCount)).apply(FirestoreIO.v1().read().partitionQuery().withNameOnlyQuery().build()).apply(FirestoreIO.v1().read().runQuery().build()).apply(ParDo.of(new RunQueryResponseToDocument())).apply(ParDo.of(new DocumentToName()));
    PAssert.that(listDocumentPaths).containsInAnyOrder(documentGenerator.expectedDocumentPaths());
    testPipeline.run(options);
}

Example 2 with DocumentGenerator

use of org.apache.beam.sdk.io.gcp.firestore.it.FirestoreTestingHelper.DocumentGenerator in project beam by apache.

the class BaseFirestoreIT method listDocuments.

@Test
public final void listDocuments() throws Exception {
    DocumentGenerator documentGenerator = helper.documentGenerator(NUM_ITEMS_TO_GENERATE, "a");
    documentGenerator.generateDocuments().get(10, TimeUnit.SECONDS);
    PCollection<String> listDocumentPaths = testPipeline.apply(Create.of("a")).apply(getListDocumentsPTransform(testName.getMethodName())).apply(FirestoreIO.v1().read().listDocuments().withRpcQosOptions(rpcQosOptions).build()).apply(ParDo.of(new DocumentToName()));
    PAssert.that(listDocumentPaths).containsInAnyOrder(documentGenerator.expectedDocumentPaths());
    testPipeline.run(options);
}

Example 3 with DocumentGenerator

use of org.apache.beam.sdk.io.gcp.firestore.it.FirestoreTestingHelper.DocumentGenerator in project beam by apache.

the class BaseFirestoreIT method batchGet.

@Test
public final void batchGet() throws Exception {
    String collectionId = "a";
    DocumentGenerator documentGenerator = helper.documentGenerator(NUM_ITEMS_TO_GENERATE, collectionId);
    documentGenerator.generateDocuments().get(10, TimeUnit.SECONDS);
    PCollection<String> listDocumentPaths = testPipeline.apply(Create.of(Collections.singletonList(documentGenerator.getDocumentIds()))).apply(getBatchGetDocumentsPTransform(testName.getMethodName(), collectionId)).apply(FirestoreIO.v1().read().batchGetDocuments().withRpcQosOptions(rpcQosOptions).build()).apply(Filter.by(BatchGetDocumentsResponse::hasFound)).apply(ParDo.of(new BatchGetDocumentsResponseToDocument())).apply(ParDo.of(new DocumentToName()));
    PAssert.that(listDocumentPaths).containsInAnyOrder(documentGenerator.expectedDocumentPaths());
    testPipeline.run(options);
}

Example 4 with DocumentGenerator

use of org.apache.beam.sdk.io.gcp.firestore.it.FirestoreTestingHelper.DocumentGenerator in project beam by apache.

the class BaseFirestoreIT method runQuery.

@Test
public final void runQuery() throws Exception {
    String collectionId = "a";
    DocumentGenerator documentGenerator = helper.documentGenerator(NUM_ITEMS_TO_GENERATE, collectionId, /* addBazDoc = */
    true);
    documentGenerator.generateDocuments().get(10, TimeUnit.SECONDS);
    PCollection<String> listDocumentPaths = testPipeline.apply(Create.of(collectionId)).apply(getRunQueryPTransform(testName.getMethodName())).apply(FirestoreIO.v1().read().runQuery().withRpcQosOptions(rpcQosOptions).build()).apply(ParDo.of(new RunQueryResponseToDocument())).apply(ParDo.of(new DocumentToName()));
    PAssert.that(listDocumentPaths).containsInAnyOrder(documentGenerator.expectedDocumentPaths());
    testPipeline.run(options);
}