Search in sources :

Example 16 with Entity

use of com.google.datastore.v1.Entity in project beam by apache.

the class V1ReadIT method writeEntitiesToDatastore.

// Creates entities and write them to datastore
private static void writeEntitiesToDatastore(V1TestOptions options, String project, String ancestor, long numEntities) throws Exception {
    Datastore datastore = getDatastore(options, project);
    // Write test entities to datastore
    V1TestWriter writer = new V1TestWriter(datastore, new UpsertMutationBuilder());
    Key ancestorKey = makeAncestorKey(options.getNamespace(), options.getKind(), ancestor);
    for (long i = 0; i < numEntities; i++) {
        Entity entity = makeEntity(i, ancestorKey, options.getKind(), options.getNamespace(), 0);
        writer.write(entity);
    }
    writer.close();
}
Also used : Entity(com.google.datastore.v1.Entity) V1TestUtil.makeEntity(org.apache.beam.sdk.io.gcp.datastore.V1TestUtil.makeEntity) UpsertMutationBuilder(org.apache.beam.sdk.io.gcp.datastore.V1TestUtil.UpsertMutationBuilder) V1TestUtil.getDatastore(org.apache.beam.sdk.io.gcp.datastore.V1TestUtil.getDatastore) Datastore(com.google.datastore.v1.client.Datastore) V1TestWriter(org.apache.beam.sdk.io.gcp.datastore.V1TestUtil.V1TestWriter) V1TestUtil.makeAncestorKey(org.apache.beam.sdk.io.gcp.datastore.V1TestUtil.makeAncestorKey) Key(com.google.datastore.v1.Key)

Example 17 with Entity

use of com.google.datastore.v1.Entity in project beam by apache.

the class V1TestUtil method deleteAllEntities.

/**
   * Delete all entities with the given ancestor.
   */
static void deleteAllEntities(V1TestOptions options, String project, String ancestor) throws Exception {
    Datastore datastore = getDatastore(options, project);
    Query query = V1TestUtil.makeAncestorKindQuery(options.getKind(), options.getNamespace(), ancestor);
    V1TestReader reader = new V1TestReader(datastore, query, options.getNamespace());
    V1TestWriter writer = new V1TestWriter(datastore, new DeleteMutationBuilder());
    long numEntities = 0;
    while (reader.advance()) {
        Entity entity = reader.getCurrent();
        numEntities++;
        writer.write(entity);
    }
    writer.close();
    LOG.info("Successfully deleted {} entities", numEntities);
}
Also used : Entity(com.google.datastore.v1.Entity) Datastore(com.google.datastore.v1.client.Datastore) Query(com.google.datastore.v1.Query)

Example 18 with Entity

use of com.google.datastore.v1.Entity in project beam by apache.

the class DatastoreV1Test method makeLatestTimestampResponse.

/** Builds a response of the given timestamp. */
private static RunQueryResponse makeLatestTimestampResponse(long timestamp) {
    RunQueryResponse.Builder timestampResponse = RunQueryResponse.newBuilder();
    Entity.Builder entity = Entity.newBuilder();
    entity.setKey(makeKey("dummyKind", "dummyId"));
    entity.putProperties("timestamp", makeValue(new Date(timestamp * 1000)).build());
    EntityResult.Builder entityResult = EntityResult.newBuilder();
    entityResult.setEntity(entity);
    QueryResultBatch.Builder batch = QueryResultBatch.newBuilder();
    batch.addEntityResults(entityResult);
    timestampResponse.setBatch(batch);
    return timestampResponse.build();
}
Also used : Entity(com.google.datastore.v1.Entity) DeleteEntity(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteEntity) QueryResultBatch(com.google.datastore.v1.QueryResultBatch) RunQueryResponse(com.google.datastore.v1.RunQueryResponse) EntityResult(com.google.datastore.v1.EntityResult) Date(java.util.Date)

Example 19 with Entity

use of com.google.datastore.v1.Entity in project beam by apache.

the class DatastoreV1Test method testDatatoreWriterFnWithLargeEntities.

/**
   * Tests {@link DatastoreWriterFn} with large entities that need to be split into more batches.
   */
@Test
public void testDatatoreWriterFnWithLargeEntities() throws Exception {
    List<Mutation> mutations = new ArrayList<>();
    int propertySize = 900_000;
    for (int i = 0; i < 12; ++i) {
        Entity.Builder entity = Entity.newBuilder().setKey(makeKey("key" + i, i + 1));
        entity.putProperties("long", makeValue(new String(new char[propertySize])).setExcludeFromIndexes(true).build());
        mutations.add(makeUpsert(entity.build()).build());
    }
    DatastoreWriterFn datastoreWriter = new DatastoreWriterFn(StaticValueProvider.of(PROJECT_ID), null, mockDatastoreFactory);
    DoFnTester<Mutation, Void> doFnTester = DoFnTester.of(datastoreWriter);
    doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
    doFnTester.processBundle(mutations);
    // This test is over-specific currently; it requires that we split the 12 entity writes into 3
    // requests, but we only need each CommitRequest to be less than 10MB in size.
    int propertiesPerRpc = DATASTORE_BATCH_UPDATE_BYTES_LIMIT / propertySize;
    int start = 0;
    while (start < mutations.size()) {
        int end = Math.min(mutations.size(), start + propertiesPerRpc);
        CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
        commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
        commitRequest.addAllMutations(mutations.subList(start, end));
        // Verify all the batch requests were made with the expected mutations.
        verify(mockDatastore).commit(commitRequest.build());
        start = end;
    }
}
Also used : CommitRequest(com.google.datastore.v1.CommitRequest) Entity(com.google.datastore.v1.Entity) DeleteEntity(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteEntity) ArrayList(java.util.ArrayList) Mutation(com.google.datastore.v1.Mutation) DatastoreWriterFn(org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DatastoreWriterFn) Test(org.junit.Test)

Example 20 with Entity

use of com.google.datastore.v1.Entity in project beam by apache.

the class V1ReadIT method testE2EV1Read.

/**
   * An end-to-end test for {@link DatastoreV1.Read#withQuery(Query)}
   *
   * <p>Write some test entities to datastore and then run a pipeline that
   * reads and counts the total number of entities. Verify that the count matches the
   * number of entities written.
   */
@Test
public void testE2EV1Read() throws Exception {
    // Read from datastore
    Query query = V1TestUtil.makeAncestorKindQuery(options.getKind(), options.getNamespace(), ancestor);
    DatastoreV1.Read read = DatastoreIO.v1().read().withProjectId(project).withQuery(query).withNamespace(options.getNamespace());
    // Count the total number of entities
    Pipeline p = Pipeline.create(options);
    PCollection<Long> count = p.apply(read).apply(Count.<Entity>globally());
    PAssert.thatSingleton(count).isEqualTo(numEntities);
    p.run();
}
Also used : Query(com.google.datastore.v1.Query) TestPipeline(org.apache.beam.sdk.testing.TestPipeline) Pipeline(org.apache.beam.sdk.Pipeline) Test(org.junit.Test)

Aggregations

Entity (com.google.datastore.v1.Entity)9 Test (org.junit.Test)9 DeleteEntity (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteEntity)7 Key (com.google.datastore.v1.Key)5 RunQueryResponse (com.google.datastore.v1.RunQueryResponse)5 QueryResultBatch (com.google.datastore.v1.QueryResultBatch)4 RunQueryRequest (com.google.datastore.v1.RunQueryRequest)4 DatastoreHelper.makeKey (com.google.datastore.v1.client.DatastoreHelper.makeKey)4 DeleteKey (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteKey)4 DatastoreV1.isValidKey (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.isValidKey)4 Mutation (com.google.datastore.v1.Mutation)3 Query (com.google.datastore.v1.Query)3 ArrayList (java.util.ArrayList)3 Function (com.google.common.base.Function)2 Optional (com.google.common.base.Optional)2 Datastore (com.google.datastore.v1.client.Datastore)2 DatastoreException (com.google.datastore.v1.client.DatastoreException)2 Status (com.yahoo.ycsb.Status)2 StringByteIterator (com.yahoo.ycsb.StringByteIterator)2 Date (java.util.Date)2