Make the large entities test neater.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/de95c7f5 Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/de95c7f5 Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/de95c7f5 Branch: refs/heads/master Commit: de95c7f5bf3b8391486b898b0ffb1a0a05338725 Parents: c4d14f8 Author: Colin Phipps <fi...@google.com> Authored: Tue May 9 15:07:51 2017 +0000 Committer: Eugene Kirpichov <kirpic...@google.com> Committed: Fri May 19 13:11:17 2017 -0700 ---------------------------------------------------------------------- .../org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java | 1 + .../apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java | 8 +++++--- 2 files changed, 6 insertions(+), 3 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/beam/blob/de95c7f5/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java index 4cfb801..b198a6f 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1.java @@ -212,6 +212,7 @@ public class DatastoreV1 { * exceeds this limit. This is set lower than the 10MB limit on the RPC, as this only accounts for * the mutations themselves and not the CommitRequest wrapper around them. */ + @VisibleForTesting static final int DATASTORE_BATCH_UPDATE_BYTES_LIMIT = 5_000_000; /** http://git-wip-us.apache.org/repos/asf/beam/blob/de95c7f5/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java ---------------------------------------------------------------------- diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java index 3597b54..460049e 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/DatastoreV1Test.java @@ -26,6 +26,7 @@ import static com.google.datastore.v1.client.DatastoreHelper.makeKey; import static com.google.datastore.v1.client.DatastoreHelper.makeOrder; import static com.google.datastore.v1.client.DatastoreHelper.makeUpsert; import static com.google.datastore.v1.client.DatastoreHelper.makeValue; +import static org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DATASTORE_BATCH_UPDATE_BYTES_LIMIT; import static org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DATASTORE_BATCH_UPDATE_LIMIT; import static org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read.DEFAULT_BUNDLE_SIZE_BYTES; import static org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read.QUERY_BATCH_LIMIT; @@ -60,7 +61,6 @@ import com.google.datastore.v1.Query; import com.google.datastore.v1.QueryResultBatch; import com.google.datastore.v1.RunQueryRequest; import com.google.datastore.v1.RunQueryResponse; -import com.google.datastore.v1.Value; import com.google.datastore.v1.client.Datastore; import com.google.datastore.v1.client.DatastoreException; import com.google.datastore.v1.client.QuerySplitter; @@ -651,9 +651,10 @@ public class DatastoreV1Test { @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", Value.newBuilder().setStringValue(new String(new char[1_000_000]) + entity.putProperties("long", makeValue(new String(new char[propertySize]) ).setExcludeFromIndexes(true).build()); mutations.add(makeUpsert(entity.build()).build()); } @@ -666,9 +667,10 @@ public class DatastoreV1Test { // 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 + 4); + 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));
