You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by jo...@apache.org on 2023/12/19 18:27:10 UTC
(beam) branch master updated: Support DatabaseID in Datastore beam connector (#29815)
This is an automated email from the ASF dual-hosted git repository.
johncasey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new f841098c54d Support DatabaseID in Datastore beam connector (#29815)
f841098c54d is described below
commit f841098c54df1a42bb3c71dd208fbae4402073b1
Author: Jing <ls...@gmail.com>
AuthorDate: Tue Dec 19 10:26:55 2023 -0800
Support DatabaseID in Datastore beam connector (#29815)
* Fix routing header issue in Firestore Beam connector
* Apply the lint
* Support multi-db in datastore beam
* Support multi-db in datastore beam
* Remove databaseRef in WriteFn
* retrigger checks
* resolve comments
* resolve comments
* resolve comments
* resolve comments
* lint
* Remove the PipelineOption reference
* Update comment
* Update ValueProviders usage
* Update IT test to refer default database
* Remove unused import
---------
Co-authored-by: Sichen Liu <si...@google.com>
---
.../beam/sdk/io/gcp/datastore/DatastoreV1.java | 349 ++++++++++++++++++---
.../beam/sdk/io/gcp/datastore/DatastoreV1Test.java | 94 +++++-
.../beam/sdk/io/gcp/datastore/SplitQueryFnIT.java | 17 +-
.../apache/beam/sdk/io/gcp/datastore/V1ReadIT.java | 26 +-
.../beam/sdk/io/gcp/datastore/V1TestUtil.java | 59 +++-
.../beam/sdk/io/gcp/datastore/V1WriteIT.java | 11 +-
6 files changed, 458 insertions(+), 98 deletions(-)
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 cdd003abae2..b94186b9893 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
@@ -126,8 +126,9 @@ import org.slf4j.LoggerFactory;
* <p>To read a {@link PCollection} from a query to Cloud Datastore, use {@link DatastoreV1#read}
* and its methods {@link DatastoreV1.Read#withProjectId} and {@link DatastoreV1.Read#withQuery} to
* specify the project to query and the query to read from. You can optionally provide a namespace
- * to query within using {@link DatastoreV1.Read#withNamespace}. You could also optionally specify
- * how many splits you want for the query using {@link DatastoreV1.Read#withNumQuerySplits}.
+ * to query within using {@link DatastoreV1.Read#withDatabase} or {@link
+ * DatastoreV1.Read#withNamespace}. You could also optionally specify how many splits you want for
+ * the query using {@link DatastoreV1.Read#withNumQuerySplits}.
*
* <p>For example:
*
@@ -135,12 +136,14 @@ import org.slf4j.LoggerFactory;
* // Read a query from Datastore
* PipelineOptions options = PipelineOptionsFactory.fromArgs(args).create();
* Query query = ...;
+ * String databaseId = "...";
* String projectId = "...";
*
* Pipeline p = Pipeline.create(options);
* PCollection<Entity> entities = p.apply(
* DatastoreIO.v1().read()
* .withProjectId(projectId)
+ * .withDatabaseId(databaseId)
* .withQuery(query));
* }</pre>
*
@@ -156,7 +159,7 @@ import org.slf4j.LoggerFactory;
*
* <pre>{@code
* PCollection<Entity> entities = ...;
- * entities.apply(DatastoreIO.v1().write().withProjectId(projectId));
+ * entities.apply(DatastoreIO.v1().write().withProjectId(projectId).withDatabaseId(databaseId));
* p.run();
* }</pre>
*
@@ -165,7 +168,7 @@ import org.slf4j.LoggerFactory;
*
* <pre>{@code
* PCollection<Entity> entities = ...;
- * entities.apply(DatastoreIO.v1().deleteEntity().withProjectId(projectId));
+ * entities.apply(DatastoreIO.v1().deleteEntity().withProjectId(projectId).withDatabaseId(databaseId));
* p.run();
* }</pre>
*
@@ -174,7 +177,7 @@ import org.slf4j.LoggerFactory;
*
* <pre>{@code
* PCollection<Entity> entities = ...;
- * entities.apply(DatastoreIO.v1().deleteKey().withProjectId(projectId));
+ * entities.apply(DatastoreIO.v1().deleteKey().withProjectId(projectId).withDatabaseId(databaseId));
* p.run();
* }</pre>
*
@@ -275,6 +278,9 @@ public class DatastoreV1 {
Code.PERMISSION_DENIED,
Code.UNAUTHENTICATED);
+ /** Database ID for the default database. */
+ private static final String DEFAULT_DATABASE = "";
+
/**
* Returns an empty {@link DatastoreV1.Read} builder. Configure the source {@code projectId},
* {@code query}, and optionally {@code namespace} and {@code numQuerySplits} using {@link
@@ -313,6 +319,8 @@ public class DatastoreV1 {
public abstract @Nullable ValueProvider<String> getProjectId();
+ public abstract @Nullable ValueProvider<String> getDatabaseId();
+
public abstract @Nullable Query getQuery();
public abstract @Nullable ValueProvider<String> getLiteralGqlQuery();
@@ -334,6 +342,8 @@ public class DatastoreV1 {
abstract static class Builder {
abstract Builder setProjectId(ValueProvider<String> projectId);
+ abstract Builder setDatabaseId(ValueProvider<String> databaseId);
+
abstract Builder setQuery(Query query);
abstract Builder setLiteralGqlQuery(ValueProvider<String> literalGqlQuery);
@@ -354,10 +364,16 @@ public class DatastoreV1 {
* size from Cloud Datastore.
*/
static int getEstimatedNumSplits(
- Datastore datastore, Query query, @Nullable String namespace, @Nullable Instant readTime) {
+ Datastore datastore,
+ String projectId,
+ String databaseId,
+ Query query,
+ @Nullable String namespace,
+ @Nullable Instant readTime) {
int numSplits;
try {
- long estimatedSizeBytes = getEstimatedSizeBytes(datastore, query, namespace, readTime);
+ long estimatedSizeBytes =
+ getEstimatedSizeBytes(datastore, projectId, databaseId, query, namespace, readTime);
LOG.info("Estimated size bytes for the query is: {}", estimatedSizeBytes);
numSplits =
(int)
@@ -378,7 +394,11 @@ public class DatastoreV1 {
* table.
*/
private static long queryLatestStatisticsTimestamp(
- Datastore datastore, @Nullable String namespace, @Nullable Instant readTime)
+ Datastore datastore,
+ String projectId,
+ String databaseId,
+ @Nullable String namespace,
+ @Nullable Instant readTime)
throws DatastoreException {
Query.Builder query = Query.newBuilder();
// Note: namespace either being null or empty represents the default namespace, in which
@@ -390,7 +410,8 @@ public class DatastoreV1 {
}
query.addOrder(makeOrder("timestamp", DESCENDING));
query.setLimit(Int32Value.newBuilder().setValue(1));
- RunQueryRequest request = makeRequest(query.build(), namespace, readTime);
+ RunQueryRequest request =
+ makeRequest(projectId, databaseId, query.build(), namespace, readTime);
RunQueryResponse response = datastore.runQuery(request);
QueryResultBatch batch = response.getBatch();
@@ -406,9 +427,15 @@ public class DatastoreV1 {
* readTime specified, the latest statistics at or before readTime is retrieved.
*/
private static Entity getLatestTableStats(
- String ourKind, @Nullable String namespace, Datastore datastore, @Nullable Instant readTime)
+ String projectId,
+ String databaseId,
+ String ourKind,
+ @Nullable String namespace,
+ Datastore datastore,
+ @Nullable Instant readTime)
throws DatastoreException {
- long latestTimestamp = queryLatestStatisticsTimestamp(datastore, namespace, readTime);
+ long latestTimestamp =
+ queryLatestStatisticsTimestamp(datastore, projectId, databaseId, namespace, readTime);
LOG.info("Latest stats timestamp for kind {} is {}", ourKind, latestTimestamp);
Query.Builder queryBuilder = Query.newBuilder();
@@ -423,7 +450,8 @@ public class DatastoreV1 {
makeFilter("kind_name", EQUAL, makeValue(ourKind).build()).build(),
makeFilter("timestamp", EQUAL, makeValue(latestTimestamp).build()).build()));
- RunQueryRequest request = makeRequest(queryBuilder.build(), namespace, readTime);
+ RunQueryRequest request =
+ makeRequest(projectId, databaseId, queryBuilder.build(), namespace, readTime);
long now = System.currentTimeMillis();
RunQueryResponse response = datastore.runQuery(request);
@@ -447,10 +475,16 @@ public class DatastoreV1 {
* <p>See https://cloud.google.com/datastore/docs/concepts/stats.
*/
static long getEstimatedSizeBytes(
- Datastore datastore, Query query, @Nullable String namespace, @Nullable Instant readTime)
+ Datastore datastore,
+ String projectId,
+ String databaseId,
+ Query query,
+ @Nullable String namespace,
+ @Nullable Instant readTime)
throws DatastoreException {
String ourKind = query.getKind(0).getName();
- Entity entity = getLatestTableStats(ourKind, namespace, datastore, readTime);
+ Entity entity =
+ getLatestTableStats(projectId, databaseId, ourKind, namespace, datastore, readTime);
return entity.getPropertiesOrThrow("entity_bytes").getIntegerValue();
}
@@ -470,9 +504,18 @@ public class DatastoreV1 {
* the requested {@code readTime}.
*/
static RunQueryRequest makeRequest(
- Query query, @Nullable String namespace, @Nullable Instant readTime) {
+ String projectId,
+ String databaseId,
+ Query query,
+ @Nullable String namespace,
+ @Nullable Instant readTime) {
RunQueryRequest.Builder request =
- RunQueryRequest.newBuilder().setQuery(query).setPartitionId(forNamespace(namespace));
+ RunQueryRequest.newBuilder()
+ .setProjectId(projectId)
+ .setDatabaseId(databaseId)
+ .setQuery(query)
+ .setPartitionId(
+ forNamespace(namespace).setProjectId(projectId).setDatabaseId(databaseId));
if (readTime != null) {
Timestamp readTimeProto = Timestamps.fromMillis(readTime.getMillis());
request.setReadOptions(ReadOptions.newBuilder().setReadTime(readTimeProto).build());
@@ -486,11 +529,18 @@ public class DatastoreV1 {
* at the requested {@code readTime}.
*/
static RunQueryRequest makeRequest(
- GqlQuery gqlQuery, @Nullable String namespace, @Nullable Instant readTime) {
+ String projectId,
+ String databaseId,
+ GqlQuery gqlQuery,
+ @Nullable String namespace,
+ @Nullable Instant readTime) {
RunQueryRequest.Builder request =
RunQueryRequest.newBuilder()
+ .setProjectId(projectId)
+ .setDatabaseId(databaseId)
.setGqlQuery(gqlQuery)
- .setPartitionId(forNamespace(namespace));
+ .setPartitionId(
+ forNamespace(namespace).setProjectId(projectId).setDatabaseId(databaseId));
if (readTime != null) {
Timestamp readTimeProto = Timestamps.fromMillis(readTime.getMillis());
request.setReadOptions(ReadOptions.newBuilder().setReadTime(readTimeProto).build());
@@ -504,6 +554,8 @@ public class DatastoreV1 {
* namespace}.
*/
private static List<Query> splitQuery(
+ String projectId,
+ String databaseId,
Query query,
@Nullable String namespace,
Datastore datastore,
@@ -512,7 +564,8 @@ public class DatastoreV1 {
@Nullable Instant readTime)
throws DatastoreException {
// If namespace is set, include it in the split request so splits are calculated accordingly.
- PartitionId partitionId = forNamespace(namespace).build();
+ PartitionId partitionId =
+ forNamespace(namespace).setProjectId(projectId).setDatabaseId(databaseId).build();
if (readTime != null) {
Timestamp readTimeProto = Timestamps.fromMillis(readTime.getMillis());
return querySplitter.getSplits(query, partitionId, numSplits, datastore, readTimeProto);
@@ -535,12 +588,18 @@ public class DatastoreV1 {
*/
@VisibleForTesting
static Query translateGqlQueryWithLimitCheck(
- String gql, Datastore datastore, String namespace, @Nullable Instant readTime)
+ String gql,
+ Datastore datastore,
+ String projectId,
+ String databaseId,
+ String namespace,
+ @Nullable Instant readTime)
throws DatastoreException {
String gqlQueryWithZeroLimit = gql + " LIMIT 0";
try {
Query translatedQuery =
- translateGqlQuery(gqlQueryWithZeroLimit, datastore, namespace, readTime);
+ translateGqlQuery(
+ gqlQueryWithZeroLimit, datastore, projectId, databaseId, namespace, readTime);
// Clear the limit that we set.
return translatedQuery.toBuilder().clearLimit().build();
} catch (DatastoreException e) {
@@ -551,7 +610,7 @@ public class DatastoreV1 {
LOG.warn("Failed to translate Gql query '{}': {}", gqlQueryWithZeroLimit, e.getMessage());
LOG.warn("User query might have a limit already set, so trying without zero limit");
// Retry without the zero limit.
- return translateGqlQuery(gql, datastore, namespace, readTime);
+ return translateGqlQuery(gql, datastore, projectId, databaseId, namespace, readTime);
} else {
throw e;
}
@@ -560,13 +619,27 @@ public class DatastoreV1 {
/** Translates a gql query string to {@link Query}. */
private static Query translateGqlQuery(
- String gql, Datastore datastore, String namespace, @Nullable Instant readTime)
+ String gql,
+ Datastore datastore,
+ String projectId,
+ String databaseId,
+ String namespace,
+ @Nullable Instant readTime)
throws DatastoreException {
GqlQuery gqlQuery = GqlQuery.newBuilder().setQueryString(gql).setAllowLiterals(true).build();
- RunQueryRequest req = makeRequest(gqlQuery, namespace, readTime);
+ RunQueryRequest req = makeRequest(projectId, databaseId, gqlQuery, namespace, readTime);
return datastore.runQuery(req).getQuery();
}
+ /**
+ * Returns a new {@link DatastoreV1.Read} that reads from the Cloud Datastore for the specified
+ * database.
+ */
+ public DatastoreV1.Read withDatabaseId(String databaseId) {
+ checkArgument(databaseId != null, "databaseId can not be null");
+ return toBuilder().setDatabaseId(StaticValueProvider.of(databaseId)).build();
+ }
+
/**
* Returns a new {@link DatastoreV1.Read} that reads from the Cloud Datastore for the specified
* project.
@@ -677,13 +750,24 @@ public class DatastoreV1 {
public long getNumEntities(
PipelineOptions options, String ourKind, @Nullable String namespace) {
try {
- V1Options v1Options = V1Options.from(getProjectId(), getNamespace(), getLocalhost());
+ V1Options v1Options =
+ V1Options.from(getProjectId(), getDatabaseId(), getNamespace(), getLocalhost());
V1DatastoreFactory datastoreFactory = new V1DatastoreFactory();
Datastore datastore =
datastoreFactory.getDatastore(
- options, v1Options.getProjectId(), v1Options.getLocalhost());
-
- Entity entity = getLatestTableStats(ourKind, namespace, datastore, getReadTime());
+ options,
+ v1Options.getProjectId(),
+ v1Options.getDatabaseId(),
+ v1Options.getLocalhost());
+
+ Entity entity =
+ getLatestTableStats(
+ v1Options.getProjectId(),
+ v1Options.getDatabaseId(),
+ ourKind,
+ namespace,
+ datastore,
+ getReadTime());
return entity.getPropertiesOrThrow("count").getIntegerValue();
} catch (Exception e) {
return -1;
@@ -704,7 +788,8 @@ public class DatastoreV1 {
getQuery() == null || getLiteralGqlQuery() == null,
"withQuery() and withLiteralGqlQuery() are exclusive");
- V1Options v1Options = V1Options.from(getProjectId(), getNamespace(), getLocalhost());
+ V1Options v1Options =
+ V1Options.from(getProjectId(), getDatabaseId(), getNamespace(), getLocalhost());
/*
* This composite transform involves the following steps:
@@ -748,6 +833,7 @@ public class DatastoreV1 {
String query = getQuery() == null ? null : getQuery().toString();
builder
.addIfNotNull(DisplayData.item("projectId", getProjectId()).withLabel("ProjectId"))
+ .addIfNotNull(DisplayData.item("databaseId", getDatabaseId()).withLabel("DatabaseId"))
.addIfNotNull(DisplayData.item("namespace", getNamespace()).withLabel("Namespace"))
.addIfNotNull(DisplayData.item("query", query).withLabel("Query"))
.addIfNotNull(DisplayData.item("gqlQuery", getLiteralGqlQuery()).withLabel("GqlQuery"))
@@ -757,30 +843,46 @@ public class DatastoreV1 {
@VisibleForTesting
static class V1Options implements HasDisplayData, Serializable {
private final ValueProvider<String> project;
+ private final ValueProvider<String> database;
private final @Nullable ValueProvider<String> namespace;
private final @Nullable String localhost;
private V1Options(
- ValueProvider<String> project, ValueProvider<String> namespace, String localhost) {
+ ValueProvider<String> project,
+ ValueProvider<String> database,
+ ValueProvider<String> namespace,
+ String localhost) {
this.project = project;
+ this.database = database;
this.namespace = namespace;
this.localhost = localhost;
}
- public static V1Options from(String projectId, String namespace, String localhost) {
+ public static V1Options from(
+ String projectId, ValueProvider<String> databaseId, String namespace, String localhost) {
return from(
- StaticValueProvider.of(projectId), StaticValueProvider.of(namespace), localhost);
+ StaticValueProvider.of(projectId),
+ databaseId,
+ StaticValueProvider.of(namespace),
+ localhost);
}
public static V1Options from(
- ValueProvider<String> project, ValueProvider<String> namespace, String localhost) {
- return new V1Options(project, namespace, localhost);
+ ValueProvider<String> project,
+ ValueProvider<String> databaseId,
+ ValueProvider<String> namespace,
+ String localhost) {
+ return new V1Options(project, databaseId, namespace, localhost);
}
public String getProjectId() {
return project.get();
}
+ public String getDatabaseId() {
+ return database == null ? DEFAULT_DATABASE : database.get();
+ }
+
public @Nullable String getNamespace() {
return namespace == null ? null : namespace.get();
}
@@ -789,6 +891,10 @@ public class DatastoreV1 {
return project;
}
+ public ValueProvider<String> getDatabaseValueProvider() {
+ return database;
+ }
+
public @Nullable ValueProvider<String> getNamespaceValueProvider() {
return namespace;
}
@@ -802,6 +908,8 @@ public class DatastoreV1 {
builder
.addIfNotNull(
DisplayData.item("projectId", getProjectValueProvider()).withLabel("ProjectId"))
+ .addIfNotNull(
+ DisplayData.item("databaseId", getDatabaseValueProvider()).withLabel("DatabaseId"))
.addIfNotNull(
DisplayData.item("namespace", getNamespaceValueProvider()).withLabel("Namespace"));
}
@@ -833,7 +941,10 @@ public class DatastoreV1 {
public void startBundle(StartBundleContext c) throws Exception {
datastore =
datastoreFactory.getDatastore(
- c.getPipelineOptions(), v1Options.getProjectId(), v1Options.getLocalhost());
+ c.getPipelineOptions(),
+ v1Options.getProjectId(),
+ v1Options.getDatabaseId(),
+ v1Options.getLocalhost());
}
@ProcessElement
@@ -842,7 +953,12 @@ public class DatastoreV1 {
LOG.info("User query: '{}'", gqlQuery);
Query query =
translateGqlQueryWithLimitCheck(
- gqlQuery, datastore, v1Options.getNamespace(), readTime);
+ gqlQuery,
+ datastore,
+ v1Options.getProjectId(),
+ v1Options.getDatabaseId(),
+ v1Options.getNamespace(),
+ readTime);
LOG.info("User gql query translated to Query({})", query);
c.output(query);
}
@@ -890,7 +1006,10 @@ public class DatastoreV1 {
public void startBundle(StartBundleContext c) throws Exception {
datastore =
datastoreFactory.getDatastore(
- c.getPipelineOptions(), options.getProjectId(), options.getLocalhost());
+ c.getPipelineOptions(),
+ options.getProjectId(),
+ options.getDatabaseId(),
+ options.getLocalhost());
querySplitter = datastoreFactory.getQuerySplitter();
}
@@ -908,7 +1027,13 @@ public class DatastoreV1 {
// Compute the estimated numSplits if numSplits is not specified by the user.
if (numSplits <= 0) {
estimatedNumSplits =
- getEstimatedNumSplits(datastore, query, options.getNamespace(), readTime);
+ getEstimatedNumSplits(
+ datastore,
+ options.getProjectId(),
+ options.getDatabaseId(),
+ query,
+ options.getNamespace(),
+ readTime);
} else {
estimatedNumSplits = numSplits;
}
@@ -918,6 +1043,8 @@ public class DatastoreV1 {
try {
querySplits =
splitQuery(
+ options.getProjectId(),
+ options.getDatabaseId(),
query,
options.getNamespace(),
datastore,
@@ -985,7 +1112,10 @@ public class DatastoreV1 {
public void startBundle(StartBundleContext c) throws Exception {
datastore =
datastoreFactory.getDatastore(
- c.getPipelineOptions(), options.getProjectId(), options.getLocalhost());
+ c.getPipelineOptions(),
+ options.getProjectId(),
+ options.getDatabaseId(),
+ options.getLocalhost());
}
private RunQueryResponse runQueryWithRetries(RunQueryRequest request) throws Exception {
@@ -1045,7 +1175,13 @@ public class DatastoreV1 {
queryBuilder.setStartCursor(currentBatch.getEndCursor());
}
- RunQueryRequest request = makeRequest(queryBuilder.build(), namespace, readTime);
+ RunQueryRequest request =
+ makeRequest(
+ options.getProjectId(),
+ options.getDatabaseId(),
+ queryBuilder.build(),
+ namespace,
+ readTime);
RunQueryResponse response = runQueryWithRetries(request);
currentBatch = response.getBatch();
@@ -1128,21 +1264,42 @@ public class DatastoreV1 {
@Nullable String localhost,
boolean throttleRampup,
ValueProvider<Integer> hintNumWorkers) {
- super(projectId, localhost, new UpsertFn(), throttleRampup, hintNumWorkers);
+ super(projectId, null, localhost, new UpsertFn(), throttleRampup, hintNumWorkers);
+ }
+
+ Write(
+ @Nullable ValueProvider<String> projectId,
+ @Nullable ValueProvider<String> databaseId,
+ @Nullable String localhost,
+ boolean throttleRampup,
+ ValueProvider<Integer> hintNumWorkers) {
+ super(projectId, databaseId, localhost, new UpsertFn(), throttleRampup, hintNumWorkers);
}
- /** Returns a new {@link Write} that writes to the Cloud Datastore for the specified project. */
+ /** Returns a new {@link Write} that writes to the Cloud Datastore for the default database. */
public Write withProjectId(String projectId) {
checkArgument(projectId != null, "projectId can not be null");
return withProjectId(StaticValueProvider.of(projectId));
}
+ /** Returns a new {@link Write} that writes to the Cloud Datastore for the database id. */
+ public Write withDatabaseId(String databaseId) {
+ checkArgument(databaseId != null, "databaseId can not be null");
+ return withDatabaseId(StaticValueProvider.of(databaseId));
+ }
+
/** Same as {@link Write#withProjectId(String)} but with a {@link ValueProvider}. */
public Write withProjectId(ValueProvider<String> projectId) {
checkArgument(projectId != null, "projectId can not be null");
return new Write(projectId, localhost, throttleRampup, hintNumWorkers);
}
+ /** Same as {@link Write#withDatabaseId(String)} but with a {@link ValueProvider}. */
+ public Write withDatabaseId(ValueProvider<String> databaseId) {
+ checkArgument(databaseId != null, "databaseId can not be null");
+ return new Write(projectId, databaseId, localhost, throttleRampup, hintNumWorkers);
+ }
+
/**
* Returns a new {@link Write} that writes to the Cloud Datastore Emulator running locally on
* the specified host port.
@@ -1188,7 +1345,16 @@ public class DatastoreV1 {
@Nullable String localhost,
boolean throttleRampup,
ValueProvider<Integer> hintNumWorkers) {
- super(projectId, localhost, new DeleteEntityFn(), throttleRampup, hintNumWorkers);
+ super(projectId, null, localhost, new DeleteEntityFn(), throttleRampup, hintNumWorkers);
+ }
+
+ DeleteEntity(
+ @Nullable ValueProvider<String> projectId,
+ @Nullable ValueProvider<String> databaseId,
+ @Nullable String localhost,
+ boolean throttleRampup,
+ ValueProvider<Integer> hintNumWorkers) {
+ super(projectId, databaseId, localhost, new DeleteEntityFn(), throttleRampup, hintNumWorkers);
}
/**
@@ -1200,12 +1366,27 @@ public class DatastoreV1 {
return withProjectId(StaticValueProvider.of(projectId));
}
+ /**
+ * Returns a new {@link DeleteEntity} that deletes entities from the Cloud Datastore for the
+ * specified database.
+ */
+ public DeleteEntity withDatabaseId(String databaseId) {
+ checkArgument(databaseId != null, "databaseId can not be null");
+ return withDatabaseId(StaticValueProvider.of(databaseId));
+ }
+
/** Same as {@link DeleteEntity#withProjectId(String)} but with a {@link ValueProvider}. */
public DeleteEntity withProjectId(ValueProvider<String> projectId) {
checkArgument(projectId != null, "projectId can not be null");
return new DeleteEntity(projectId, localhost, throttleRampup, hintNumWorkers);
}
+ /** Same as {@link DeleteEntity#withDatabaseId(String)} but with a {@link ValueProvider}. */
+ public DeleteEntity withDatabaseId(ValueProvider<String> databaseId) {
+ checkArgument(databaseId != null, "databaseId can not be null");
+ return new DeleteEntity(projectId, databaseId, localhost, throttleRampup, hintNumWorkers);
+ }
+
/**
* Returns a new {@link DeleteEntity} that deletes entities from the Cloud Datastore Emulator
* running locally on the specified host port.
@@ -1253,7 +1434,16 @@ public class DatastoreV1 {
@Nullable String localhost,
boolean throttleRampup,
ValueProvider<Integer> hintNumWorkers) {
- super(projectId, localhost, new DeleteKeyFn(), throttleRampup, hintNumWorkers);
+ super(projectId, null, localhost, new DeleteKeyFn(), throttleRampup, hintNumWorkers);
+ }
+
+ DeleteKey(
+ @Nullable ValueProvider<String> projectId,
+ @Nullable ValueProvider<String> databaseId,
+ @Nullable String localhost,
+ boolean throttleRampup,
+ ValueProvider<Integer> hintNumWorkers) {
+ super(projectId, databaseId, localhost, new DeleteKeyFn(), throttleRampup, hintNumWorkers);
}
/**
@@ -1265,6 +1455,15 @@ public class DatastoreV1 {
return withProjectId(StaticValueProvider.of(projectId));
}
+ /**
+ * Returns a new {@link DeleteKey} that deletes entities from the Cloud Datastore for the
+ * specified database.
+ */
+ public DeleteKey withDatabaseId(String databaseId) {
+ checkArgument(databaseId != null, "databaseId can not be null");
+ return withDatabaseId(StaticValueProvider.of(databaseId));
+ }
+
/**
* Returns a new {@link DeleteKey} that deletes entities from the Cloud Datastore Emulator
* running locally on the specified host port.
@@ -1280,6 +1479,12 @@ public class DatastoreV1 {
return new DeleteKey(projectId, localhost, throttleRampup, hintNumWorkers);
}
+ /** Same as {@link DeleteKey#withDatabaseId(String)} but with a {@link ValueProvider}. */
+ public DeleteKey withDatabaseId(ValueProvider<String> databaseId) {
+ checkArgument(databaseId != null, "databaseId can not be null");
+ return new DeleteKey(projectId, databaseId, localhost, throttleRampup, hintNumWorkers);
+ }
+
/** Returns a new {@link DeleteKey} that does not throttle during ramp-up. */
public DeleteKey withRampupThrottlingDisabled() {
return new DeleteKey(projectId, localhost, false, hintNumWorkers);
@@ -1312,6 +1517,7 @@ public class DatastoreV1 {
private abstract static class Mutate<T> extends PTransform<PCollection<T>, PDone> {
protected ValueProvider<String> projectId;
+ protected ValueProvider<String> databaseId;
protected @Nullable String localhost;
protected boolean throttleRampup;
protected ValueProvider<Integer> hintNumWorkers;
@@ -1326,11 +1532,13 @@ public class DatastoreV1 {
*/
Mutate(
@Nullable ValueProvider<String> projectId,
+ @Nullable ValueProvider<String> databaseId,
@Nullable String localhost,
SimpleFunction<T, Mutation> mutationFn,
boolean throttleRampup,
ValueProvider<Integer> hintNumWorkers) {
this.projectId = projectId;
+ this.databaseId = databaseId;
this.localhost = localhost;
this.throttleRampup = throttleRampup;
this.hintNumWorkers = hintNumWorkers;
@@ -1372,7 +1580,14 @@ public class DatastoreV1 {
ParDo.of(rampupThrottlingFn).withSideInputs(startTimestampView));
}
intermediateOutput.apply(
- "Write Mutation to Datastore", ParDo.of(new DatastoreWriterFn(projectId, localhost)));
+ "Write Mutation to Datastore",
+ ParDo.of(
+ new DatastoreWriterFn(
+ projectId,
+ databaseId,
+ localhost,
+ new V1DatastoreFactory(),
+ new WriteBatcherImpl())));
return PDone.in(input.getPipeline());
}
@@ -1390,6 +1605,7 @@ public class DatastoreV1 {
super.populateDisplayData(builder);
builder
.addIfNotNull(DisplayData.item("projectId", projectId).withLabel("Output Project"))
+ .addIfNotNull(DisplayData.item("databaseId", databaseId).withLabel("Output Database"))
.include("mutationFn", mutationFn);
if (rampupThrottlingFn != null) {
builder.include("rampupThrottlingFn", rampupThrottlingFn);
@@ -1399,6 +1615,10 @@ public class DatastoreV1 {
public String getProjectId() {
return projectId.get();
}
+
+ public String getDatabaseId() {
+ return databaseId.get();
+ }
}
/** Determines batch sizes for commit RPCs. */
@@ -1480,6 +1700,7 @@ public class DatastoreV1 {
private static final Logger LOG = LoggerFactory.getLogger(DatastoreWriterFn.class);
private final ValueProvider<String> projectId;
+ private final ValueProvider<String> databaseId;
private final @Nullable String localhost;
private transient Datastore datastore;
private final V1DatastoreFactory datastoreFactory;
@@ -1511,13 +1732,14 @@ public class DatastoreV1 {
DatastoreWriterFn(String projectId, @Nullable String localhost) {
this(
StaticValueProvider.of(projectId),
+ null,
localhost,
new V1DatastoreFactory(),
new WriteBatcherImpl());
}
DatastoreWriterFn(ValueProvider<String> projectId, @Nullable String localhost) {
- this(projectId, localhost, new V1DatastoreFactory(), new WriteBatcherImpl());
+ this(projectId, null, localhost, new V1DatastoreFactory(), new WriteBatcherImpl());
}
@VisibleForTesting
@@ -1526,7 +1748,18 @@ public class DatastoreV1 {
@Nullable String localhost,
V1DatastoreFactory datastoreFactory,
WriteBatcher writeBatcher) {
+ this(projectId, null, localhost, datastoreFactory, writeBatcher);
+ }
+
+ @VisibleForTesting
+ DatastoreWriterFn(
+ ValueProvider<String> projectId,
+ ValueProvider<String> databaseId,
+ @Nullable String localhost,
+ V1DatastoreFactory datastoreFactory,
+ WriteBatcher writeBatcher) {
this.projectId = checkNotNull(projectId, "projectId");
+ this.databaseId = databaseId;
this.localhost = localhost;
this.datastoreFactory = datastoreFactory;
this.writeBatcher = writeBatcher;
@@ -1534,7 +1767,10 @@ public class DatastoreV1 {
@StartBundle
public void startBundle(StartBundleContext c) {
- datastore = datastoreFactory.getDatastore(c.getPipelineOptions(), projectId.get(), localhost);
+ String databaseIdOrDefaultDatabase = databaseId == null ? DEFAULT_DATABASE : databaseId.get();
+ datastore =
+ datastoreFactory.getDatastore(
+ c.getPipelineOptions(), projectId.get(), databaseIdOrDefaultDatabase, localhost);
writeBatcher.start();
if (adaptiveThrottler == null) {
// Initialize throttler at first use, because it is not serializable.
@@ -1602,11 +1838,14 @@ public class DatastoreV1 {
batchSize.update(mutations.size());
+ String databaseIdOrDefaultDatabase = databaseId == null ? DEFAULT_DATABASE : databaseId.get();
while (true) {
// Batch upsert entities.
CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
commitRequest.addAllMutations(mutations);
commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
+ commitRequest.setProjectId(projectId.get());
+ commitRequest.setDatabaseId(databaseIdOrDefaultDatabase);
long startTime = System.currentTimeMillis(), endTime;
if (adaptiveThrottler.throttleRequest(startTime)) {
@@ -1628,6 +1867,7 @@ public class DatastoreV1 {
ServiceCallMetric serviceCallMetric =
new ServiceCallMetric(MonitoringInfoConstants.Urns.API_REQUEST_COUNT, baseLabels);
try {
+
datastore.commit(commitRequest.build());
endTime = System.currentTimeMillis();
serviceCallMetric.call("ok");
@@ -1768,8 +2008,9 @@ public class DatastoreV1 {
static class V1DatastoreFactory implements Serializable {
/** Builds a Cloud Datastore client for the given pipeline options and project. */
- public Datastore getDatastore(PipelineOptions pipelineOptions, String projectId) {
- return getDatastore(pipelineOptions, projectId, null);
+ public Datastore getDatastore(
+ PipelineOptions pipelineOptions, String projectId, String databaseId) {
+ return getDatastore(pipelineOptions, projectId, databaseId, null);
}
/**
@@ -1777,7 +2018,10 @@ public class DatastoreV1 {
* locahost.
*/
public Datastore getDatastore(
- PipelineOptions pipelineOptions, String projectId, @Nullable String localhost) {
+ PipelineOptions pipelineOptions,
+ String projectId,
+ String databaseId,
+ @Nullable String localhost) {
Credentials credential = pipelineOptions.as(GcpOptions.class).getGcpCredential();
// Add Beam version to user agent header.
@@ -1797,7 +2041,10 @@ public class DatastoreV1 {
}
DatastoreOptions.Builder builder =
- new DatastoreOptions.Builder().projectId(projectId).initializer(initializer);
+ new DatastoreOptions.Builder()
+ .projectId(projectId)
+ .databaseId(databaseId)
+ .initializer(initializer);
if (localhost != null) {
builder.localHost(localhost);
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 3280e17998e..ab6f5567dd9 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
@@ -125,6 +125,7 @@ import org.mockito.MockitoAnnotations;
@RunWith(Enclosed.class)
public class DatastoreV1Test {
private static final String PROJECT_ID = "testProject";
+ private static final String DATABASE_ID = "";
private static final String NAMESPACE = "testNamespace";
private static final String KIND = "testKind";
private static final Query QUERY;
@@ -138,7 +139,7 @@ public class DatastoreV1Test {
Query.Builder q = Query.newBuilder();
q.addKindBuilder().setName(KIND);
QUERY = q.build();
- V_1_OPTIONS = V1Options.from(PROJECT_ID, NAMESPACE, null);
+ V_1_OPTIONS = V1Options.from(PROJECT_ID, StaticValueProvider.of(DATABASE_ID), NAMESPACE, null);
}
@Mock protected Datastore mockDatastore;
@@ -153,6 +154,9 @@ public class DatastoreV1Test {
when(mockDatastoreFactory.getDatastore(any(PipelineOptions.class), any(String.class), any()))
.thenReturn(mockDatastore);
+ when(mockDatastoreFactory.getDatastore(
+ any(PipelineOptions.class), any(String.class), any(String.class), any()))
+ .thenReturn(mockDatastore);
when(mockDatastoreFactory.getQuerySplitter()).thenReturn(mockQuerySplitter);
// Setup the ProcessWideContainer for testing metrics are set.
MetricsContainerImpl container = new MetricsContainerImpl(null);
@@ -167,10 +171,12 @@ public class DatastoreV1Test {
DatastoreIO.v1()
.read()
.withProjectId(PROJECT_ID)
+ .withDatabaseId(DATABASE_ID)
.withQuery(QUERY)
.withNamespace(NAMESPACE);
assertEquals(QUERY, read.getQuery());
assertEquals(PROJECT_ID, read.getProjectId().get());
+ assertEquals(DATABASE_ID, read.getDatabaseId().get());
assertEquals(NAMESPACE, read.getNamespace().get());
}
@@ -180,11 +186,13 @@ public class DatastoreV1Test {
DatastoreIO.v1()
.read()
.withProjectId(PROJECT_ID)
+ .withDatabaseId(DATABASE_ID)
.withQuery(QUERY)
.withReadTime(TIMESTAMP);
assertEquals(TIMESTAMP, read.getReadTime());
assertEquals(QUERY, read.getQuery());
assertEquals(PROJECT_ID, read.getProjectId().get());
+ assertEquals(DATABASE_ID, read.getDatabaseId().get());
}
@Test
@@ -193,10 +201,12 @@ public class DatastoreV1Test {
DatastoreIO.v1()
.read()
.withProjectId(PROJECT_ID)
+ .withDatabaseId(DATABASE_ID)
.withLiteralGqlQuery(GQL_QUERY)
.withNamespace(NAMESPACE);
assertEquals(GQL_QUERY, read.getLiteralGqlQuery().get());
assertEquals(PROJECT_ID, read.getProjectId().get());
+ assertEquals(DATABASE_ID, read.getDatabaseId().get());
assertEquals(NAMESPACE, read.getNamespace().get());
}
@@ -210,10 +220,12 @@ public class DatastoreV1Test {
.withQuery(QUERY)
.withNamespace(NAMESPACE)
.withProjectId(PROJECT_ID)
+ .withDatabaseId(DATABASE_ID)
.withLocalhost(LOCALHOST);
assertEquals(TIMESTAMP, read.getReadTime());
assertEquals(QUERY, read.getQuery());
assertEquals(PROJECT_ID, read.getProjectId().get());
+ assertEquals(DATABASE_ID, read.getDatabaseId().get());
assertEquals(NAMESPACE, read.getNamespace().get());
assertEquals(LOCALHOST, read.getLocalhost());
}
@@ -224,6 +236,7 @@ public class DatastoreV1Test {
DatastoreIO.v1()
.read()
.withProjectId(PROJECT_ID)
+ .withDatabaseId(DATABASE_ID)
.withLiteralGqlQuery(GQL_QUERY)
.withQuery(QUERY);
@@ -257,6 +270,7 @@ public class DatastoreV1Test {
DatastoreIO.v1()
.read()
.withProjectId(PROJECT_ID)
+ .withDatabaseId(DATABASE_ID)
.withQuery(QUERY)
.withNamespace(NAMESPACE)
.withReadTime(TIMESTAMP);
@@ -264,6 +278,7 @@ public class DatastoreV1Test {
DisplayData displayData = DisplayData.from(read);
assertThat(displayData, hasDisplayItem("projectId", PROJECT_ID));
+ assertThat(displayData, hasDisplayItem("databaseId", DATABASE_ID));
assertThat(displayData, hasDisplayItem("query", QUERY.toString()));
assertThat(displayData, hasDisplayItem("namespace", NAMESPACE));
assertThat(displayData, hasDisplayItem("readTime", TIMESTAMP));
@@ -275,6 +290,7 @@ public class DatastoreV1Test {
DatastoreIO.v1()
.read()
.withProjectId(PROJECT_ID)
+ .withDatabaseId(DATABASE_ID)
.withLiteralGqlQuery(GQL_QUERY)
.withNamespace(NAMESPACE)
.withReadTime(TIMESTAMP);
@@ -295,6 +311,7 @@ public class DatastoreV1Test {
DatastoreIO.v1()
.read()
.withProjectId(PROJECT_ID)
+ .withDatabaseId(DATABASE_ID)
.withQuery(Query.newBuilder().build())
.withNumQuerySplits(numSplits);
@@ -313,29 +330,34 @@ public class DatastoreV1Test {
@Test
public void testWriteDisplayData() {
- Write write = DatastoreIO.v1().write().withProjectId(PROJECT_ID);
+ Write write = DatastoreIO.v1().write().withProjectId(PROJECT_ID).withDatabaseId(DATABASE_ID);
DisplayData displayData = DisplayData.from(write);
assertThat(displayData, hasDisplayItem("projectId", PROJECT_ID));
+ assertThat(displayData, hasDisplayItem("databaseId", DATABASE_ID));
}
@Test
public void testDeleteEntityDisplayData() {
- DeleteEntity deleteEntity = DatastoreIO.v1().deleteEntity().withProjectId(PROJECT_ID);
+ DeleteEntity deleteEntity =
+ DatastoreIO.v1().deleteEntity().withProjectId(PROJECT_ID).withDatabaseId(DATABASE_ID);
DisplayData displayData = DisplayData.from(deleteEntity);
assertThat(displayData, hasDisplayItem("projectId", PROJECT_ID));
+ assertThat(displayData, hasDisplayItem("databaseId", DATABASE_ID));
}
@Test
public void testDeleteKeyDisplayData() {
- DeleteKey deleteKey = DatastoreIO.v1().deleteKey().withProjectId(PROJECT_ID);
+ DeleteKey deleteKey =
+ DatastoreIO.v1().deleteKey().withProjectId(PROJECT_ID).withDatabaseId(DATABASE_ID);
DisplayData displayData = DisplayData.from(deleteKey);
assertThat(displayData, hasDisplayItem("projectId", PROJECT_ID));
+ assertThat(displayData, hasDisplayItem("databaseId", DATABASE_ID));
}
@Test
@@ -600,6 +622,8 @@ public class DatastoreV1Test {
while (start < numMutations) {
int end = Math.min(numMutations, start + DatastoreV1.DATASTORE_BATCH_UPDATE_ENTITIES_START);
CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
+ commitRequest.setProjectId(PROJECT_ID);
+ commitRequest.setDatabaseId(DATABASE_ID);
commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
commitRequest.addAllMutations(mutations.subList(start, end));
// Verify all the batch requests were made with the expected mutations.
@@ -644,6 +668,8 @@ public class DatastoreV1Test {
while (start < mutations.size()) {
int end = Math.min(mutations.size(), start + entitiesPerRpc);
CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
+ commitRequest.setProjectId(PROJECT_ID);
+ commitRequest.setDatabaseId(DATABASE_ID);
commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
commitRequest.addAllMutations(mutations.subList(start, end));
// Verify all the batch requests were made with the expected mutations.
@@ -681,12 +707,16 @@ public class DatastoreV1Test {
CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
commitRequest.addAllMutations(mutations.subList(0, 2));
+ commitRequest.setProjectId(PROJECT_ID);
+ commitRequest.setDatabaseId(DATABASE_ID);
verify(mockDatastore, times(1)).commit(commitRequest.build());
// second invocation has key [0, 2] because the second 0 triggered a flush batch
commitRequest = CommitRequest.newBuilder();
commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
commitRequest.addAllMutations(mutations.subList(2, 4));
+ commitRequest.setProjectId(PROJECT_ID);
+ commitRequest.setDatabaseId(DATABASE_ID);
verify(mockDatastore, times(1)).commit(commitRequest.build());
verifyMetricWasSet("BatchDatastoreWrite", "ok", "", 2);
}
@@ -841,17 +871,26 @@ public class DatastoreV1Test {
long timestamp = 1234L;
RunQueryRequest latestTimestampRequest =
- makeRequest(makeLatestTimestampQuery(NAMESPACE), NAMESPACE, readTime);
+ makeRequest(
+ PROJECT_ID, DATABASE_ID, makeLatestTimestampQuery(NAMESPACE), NAMESPACE, readTime);
RunQueryResponse latestTimestampResponse = makeLatestTimestampResponse(timestamp);
// Per Kind statistics request and response
RunQueryRequest statRequest =
- makeRequest(makeStatKindQuery(NAMESPACE, timestamp), NAMESPACE, readTime);
+ makeRequest(
+ PROJECT_ID,
+ DATABASE_ID,
+ makeStatKindQuery(NAMESPACE, timestamp),
+ NAMESPACE,
+ readTime);
RunQueryResponse statResponse = makeStatKindResponse(entityBytes);
when(mockDatastore.runQuery(latestTimestampRequest)).thenReturn(latestTimestampResponse);
when(mockDatastore.runQuery(statRequest)).thenReturn(statResponse);
- assertEquals(entityBytes, getEstimatedSizeBytes(mockDatastore, QUERY, NAMESPACE, readTime));
+ assertEquals(
+ entityBytes,
+ getEstimatedSizeBytes(
+ mockDatastore, PROJECT_ID, DATABASE_ID, QUERY, NAMESPACE, readTime));
verify(mockDatastore, times(1)).runQuery(latestTimestampRequest);
verify(mockDatastore, times(1)).runQuery(statRequest);
}
@@ -917,12 +956,18 @@ public class DatastoreV1Test {
long timestamp = 1234L;
RunQueryRequest latestTimestampRequest =
- makeRequest(makeLatestTimestampQuery(NAMESPACE), NAMESPACE, readTime);
+ makeRequest(
+ PROJECT_ID, DATABASE_ID, makeLatestTimestampQuery(NAMESPACE), NAMESPACE, readTime);
RunQueryResponse latestTimestampResponse = makeLatestTimestampResponse(timestamp);
// Per Kind statistics request and response
RunQueryRequest statRequest =
- makeRequest(makeStatKindQuery(NAMESPACE, timestamp), NAMESPACE, readTime);
+ makeRequest(
+ PROJECT_ID,
+ DATABASE_ID,
+ makeStatKindQuery(NAMESPACE, timestamp),
+ NAMESPACE,
+ readTime);
RunQueryResponse statResponse = makeStatKindResponse(entityBytes);
when(mockDatastore.runQuery(latestTimestampRequest)).thenReturn(latestTimestampResponse);
@@ -1030,9 +1075,11 @@ public class DatastoreV1Test {
GqlQuery gqlQueryWithZeroLimit =
GqlQuery.newBuilder().setQueryString(gqlWithZeroLimit).setAllowLiterals(true).build();
- RunQueryRequest gqlRequest = makeRequest(gqlQuery, V_1_OPTIONS.getNamespace(), readTime);
+ RunQueryRequest gqlRequest =
+ makeRequest(PROJECT_ID, DATABASE_ID, gqlQuery, V_1_OPTIONS.getNamespace(), readTime);
RunQueryRequest gqlRequestWithZeroLimit =
- makeRequest(gqlQueryWithZeroLimit, V_1_OPTIONS.getNamespace(), readTime);
+ makeRequest(
+ PROJECT_ID, DATABASE_ID, gqlQueryWithZeroLimit, V_1_OPTIONS.getNamespace(), readTime);
when(mockDatastore.runQuery(gqlRequestWithZeroLimit))
.thenThrow(
new DatastoreException(
@@ -1044,7 +1091,8 @@ public class DatastoreV1Test {
when(mockDatastore.runQuery(gqlRequest))
.thenReturn(RunQueryResponse.newBuilder().setQuery(QUERY).build());
assertEquals(
- translateGqlQueryWithLimitCheck(gql, mockDatastore, V_1_OPTIONS.getNamespace(), readTime),
+ translateGqlQueryWithLimitCheck(
+ gql, mockDatastore, PROJECT_ID, DATABASE_ID, V_1_OPTIONS.getNamespace(), readTime),
QUERY);
verify(mockDatastore, times(1)).runQuery(gqlRequest);
verify(mockDatastore, times(1)).runQuery(gqlRequestWithZeroLimit);
@@ -1058,11 +1106,18 @@ public class DatastoreV1Test {
GqlQuery.newBuilder().setQueryString(gqlWithZeroLimit).setAllowLiterals(true).build();
RunQueryRequest gqlRequestWithZeroLimit =
- makeRequest(gqlQueryWithZeroLimit, V_1_OPTIONS.getNamespace(), readTime);
+ makeRequest(
+ PROJECT_ID, DATABASE_ID, gqlQueryWithZeroLimit, V_1_OPTIONS.getNamespace(), readTime);
when(mockDatastore.runQuery(gqlRequestWithZeroLimit))
.thenReturn(RunQueryResponse.newBuilder().setQuery(QUERY).build());
assertEquals(
- translateGqlQueryWithLimitCheck(gql, mockDatastore, V_1_OPTIONS.getNamespace(), readTime),
+ translateGqlQueryWithLimitCheck(
+ gql,
+ mockDatastore,
+ V_1_OPTIONS.getProjectId(),
+ V_1_OPTIONS.getDatabaseId(),
+ V_1_OPTIONS.getNamespace(),
+ readTime),
QUERY);
verify(mockDatastore, times(1)).runQuery(gqlRequestWithZeroLimit);
}
@@ -1074,13 +1129,20 @@ public class DatastoreV1Test {
GqlQuery gqlQueryWithZeroLimit =
GqlQuery.newBuilder().setQueryString(gqlWithZeroLimit).setAllowLiterals(true).build();
RunQueryRequest gqlRequestWithZeroLimit =
- makeRequest(gqlQueryWithZeroLimit, V_1_OPTIONS.getNamespace(), readTime);
+ makeRequest(
+ PROJECT_ID, DATABASE_ID, gqlQueryWithZeroLimit, V_1_OPTIONS.getNamespace(), readTime);
when(mockDatastore.runQuery(gqlRequestWithZeroLimit))
.thenThrow(new RuntimeException("TestException"));
thrown.expect(RuntimeException.class);
thrown.expectMessage("TestException");
- translateGqlQueryWithLimitCheck(gql, mockDatastore, V_1_OPTIONS.getNamespace(), readTime);
+ translateGqlQueryWithLimitCheck(
+ gql,
+ mockDatastore,
+ V_1_OPTIONS.getProjectId(),
+ V_1_OPTIONS.getDatabaseId(),
+ V_1_OPTIONS.getNamespace(),
+ readTime);
}
}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/SplitQueryFnIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/SplitQueryFnIT.java
index ea00821f360..913e05c8612 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/SplitQueryFnIT.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/SplitQueryFnIT.java
@@ -24,6 +24,7 @@ import com.google.datastore.v1.Query;
import java.util.List;
import org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read.SplitQueryFn;
import org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.Read.V1Options;
+import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
import org.apache.beam.sdk.transforms.DoFnTester;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.Duration;
@@ -58,29 +59,32 @@ public class SplitQueryFnIT {
@Test
public void testSplitQueryFnWithLargeDataset() throws Exception {
String projectId = "apache-beam-testing";
+ String databaseId = "";
String kind = "sort_1G";
String namespace = null;
// Num splits is computed based on the entity_bytes size of the input_sort_1G kind reported by
// Datastore stats.
int expectedNumSplits = 32;
- testSplitQueryFn(projectId, kind, namespace, expectedNumSplits, null);
- testSplitQueryFn(projectId, kind, namespace, expectedNumSplits, readTime);
+ testSplitQueryFn(projectId, databaseId, kind, namespace, expectedNumSplits, null);
+ testSplitQueryFn(projectId, databaseId, kind, namespace, expectedNumSplits, readTime);
}
/** Tests {@link SplitQueryFn} to fallback to NUM_QUERY_SPLITS_MIN for a small dataset. */
@Test
public void testSplitQueryFnWithSmallDataset() throws Exception {
String projectId = "apache-beam-testing";
+ String databaseId = "";
String kind = "shakespeare";
String namespace = null;
int expectedNumSplits = NUM_QUERY_SPLITS_MIN;
- testSplitQueryFn(projectId, kind, namespace, expectedNumSplits, null);
- testSplitQueryFn(projectId, kind, namespace, expectedNumSplits, readTime);
+ testSplitQueryFn(projectId, databaseId, kind, namespace, expectedNumSplits, null);
+ testSplitQueryFn(projectId, databaseId, kind, namespace, expectedNumSplits, readTime);
}
/** A helper method to test {@link SplitQueryFn} to generate the expected number of splits. */
private void testSplitQueryFn(
String projectId,
+ String databaseId,
String kind,
@Nullable String namespace,
int expectedNumSplits,
@@ -90,7 +94,10 @@ public class SplitQueryFnIT {
query.addKindBuilder().setName(kind);
SplitQueryFn splitQueryFn =
- new SplitQueryFn(V1Options.from(projectId, namespace, null), 0, readTime);
+ new SplitQueryFn(
+ V1Options.from(projectId, StaticValueProvider.of(databaseId), namespace, null),
+ 0,
+ readTime);
DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn);
List<Query> queries = doFnTester.processBundle(query.build());
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/V1ReadIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/V1ReadIT.java
index 249cadd4886..c9530ba43cb 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/V1ReadIT.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/V1ReadIT.java
@@ -48,6 +48,7 @@ import org.junit.runners.JUnit4;
public class V1ReadIT {
private V1TestOptions options;
private String project;
+ private String database;
private String ancestor;
private final long numEntitiesBeforeReadTime = 600;
private final long totalNumEntities = 1000;
@@ -58,9 +59,12 @@ public class V1ReadIT {
PipelineOptionsFactory.register(V1TestOptions.class);
options = TestPipeline.testingPipelineOptions().as(V1TestOptions.class);
project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject();
+ // The default database.
+ database = "";
+
ancestor = UUID.randomUUID().toString();
// Create entities and write them to datastore
- writeEntitiesToDatastore(options, project, ancestor, 0, numEntitiesBeforeReadTime);
+ writeEntitiesToDatastore(options, project, database, ancestor, 0, numEntitiesBeforeReadTime);
Thread.sleep(1000);
readTime = Instant.now();
@@ -68,12 +72,12 @@ public class V1ReadIT {
long moreEntitiesToWrite = totalNumEntities - numEntitiesBeforeReadTime;
writeEntitiesToDatastore(
- options, project, ancestor, numEntitiesBeforeReadTime, moreEntitiesToWrite);
+ options, project, database, ancestor, numEntitiesBeforeReadTime, moreEntitiesToWrite);
}
@After
public void tearDown() throws Exception {
- deleteAllEntities(options, project, ancestor);
+ deleteAllEntities(options, project, database, ancestor);
}
/**
@@ -93,6 +97,7 @@ public class V1ReadIT {
DatastoreIO.v1()
.read()
.withProjectId(project)
+ .withDatabaseId(database)
.withQuery(query)
.withNamespace(options.getNamespace());
@@ -108,6 +113,7 @@ public class V1ReadIT {
DatastoreIO.v1()
.read()
.withProjectId(project)
+ .withDatabaseId(database)
.withQuery(query)
.withNamespace(options.getNamespace())
.withReadTime(readTime);
@@ -152,6 +158,7 @@ public class V1ReadIT {
DatastoreIO.v1()
.read()
.withProjectId(project)
+ .withDatabaseId(database)
.withLiteralGqlQuery(gqlQuery)
.withNamespace(options.getNamespace());
@@ -167,6 +174,7 @@ public class V1ReadIT {
DatastoreIO.v1()
.read()
.withProjectId(project)
+ .withDatabaseId(database)
.withLiteralGqlQuery(gqlQuery)
.withNamespace(options.getNamespace())
.withReadTime(readTime);
@@ -181,11 +189,17 @@ public class V1ReadIT {
// Creates entities and write them to datastore
private static void writeEntitiesToDatastore(
- V1TestOptions options, String project, String ancestor, long valueOffset, long numEntities)
+ V1TestOptions options,
+ String project,
+ String database,
+ String ancestor,
+ long valueOffset,
+ long numEntities)
throws Exception {
- Datastore datastore = getDatastore(options, project);
+ Datastore datastore = getDatastore(options, project, database);
// Write test entities to datastore
- V1TestWriter writer = new V1TestWriter(datastore, new UpsertMutationBuilder());
+ V1TestWriter writer =
+ new V1TestWriter(datastore, project, database, new UpsertMutationBuilder());
Key ancestorKey = makeAncestorKey(options.getNamespace(), options.getKind(), ancestor);
for (long i = 0; i < numEntities; i++) {
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/V1TestUtil.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/V1TestUtil.java
index e3057c4b90c..0cb3ffac322 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/V1TestUtil.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/V1TestUtil.java
@@ -137,10 +137,12 @@ class V1TestUtil {
}
/** Build a new datastore client. */
- static Datastore getDatastore(PipelineOptions pipelineOptions, String projectId) {
+ static Datastore getDatastore(
+ PipelineOptions pipelineOptions, String projectId, String databaseId) {
Credentials credential = pipelineOptions.as(GcpOptions.class).getGcpCredential();
HttpRequestInitializer initializer;
if (credential != null) {
+
initializer =
new ChainingHttpRequestInitializer(
new HttpCredentialsAdapter(credential), new RetryHttpRequestInitializer());
@@ -149,14 +151,22 @@ class V1TestUtil {
}
DatastoreOptions.Builder builder =
- new DatastoreOptions.Builder().projectId(projectId).initializer(initializer);
+ new DatastoreOptions.Builder()
+ .projectId(projectId)
+ .databaseId(databaseId)
+ .initializer(initializer);
return DatastoreFactory.get().create(builder.build());
}
/** Build a datastore query request. */
- private static RunQueryRequest makeRequest(Query query, @Nullable String namespace) {
- RunQueryRequest.Builder requestBuilder = RunQueryRequest.newBuilder().setQuery(query);
+ private static RunQueryRequest makeRequest(
+ String projectId, String databaseId, Query query, @Nullable String namespace) {
+ RunQueryRequest.Builder requestBuilder =
+ RunQueryRequest.newBuilder()
+ .setQuery(query)
+ .setProjectId(projectId)
+ .setDatabaseId(databaseId);
if (namespace != null) {
requestBuilder.getPartitionIdBuilder().setNamespaceId(namespace);
}
@@ -164,14 +174,16 @@ class V1TestUtil {
}
/** Delete all entities with the given ancestor. */
- static void deleteAllEntities(V1TestOptions options, String project, String ancestor)
- throws Exception {
- Datastore datastore = getDatastore(options, project);
+ static void deleteAllEntities(
+ V1TestOptions options, String project, String database, String ancestor) throws Exception {
+ Datastore datastore = getDatastore(options, project, database);
Query query =
V1TestUtil.makeAncestorKindQuery(options.getKind(), options.getNamespace(), ancestor);
- V1TestReader reader = new V1TestReader(datastore, query, options.getNamespace());
- V1TestWriter writer = new V1TestWriter(datastore, new DeleteMutationBuilder());
+ V1TestReader reader =
+ new V1TestReader(datastore, project, database, query, options.getNamespace());
+ V1TestWriter writer =
+ new V1TestWriter(datastore, project, database, new DeleteMutationBuilder());
long numEntities = 0;
while (reader.advance()) {
@@ -185,14 +197,15 @@ class V1TestUtil {
}
/** Returns the total number of entities for the given datastore. */
- static long countEntities(V1TestOptions options, String project, String ancestor)
+ static long countEntities(V1TestOptions options, String project, String database, String ancestor)
throws Exception {
// Read from datastore.
- Datastore datastore = V1TestUtil.getDatastore(options, project);
+ Datastore datastore = V1TestUtil.getDatastore(options, project, database);
Query query =
V1TestUtil.makeAncestorKindQuery(options.getKind(), options.getNamespace(), ancestor);
- V1TestReader reader = new V1TestReader(datastore, query, options.getNamespace());
+ V1TestReader reader =
+ new V1TestReader(datastore, project, database, query, options.getNamespace());
long numEntitiesRead = 0;
while (reader.advance()) {
@@ -247,12 +260,17 @@ class V1TestUtil {
return (lastElement.getId() != 0 || !lastElement.getName().isEmpty());
}
+ private final String projectId;
+ private final String databaseId;
private final Datastore datastore;
private final MutationBuilder mutationBuilder;
private final List<Entity> entities = new ArrayList<>();
- V1TestWriter(Datastore datastore, MutationBuilder mutationBuilder) {
+ V1TestWriter(
+ Datastore datastore, String projectId, String databaseId, MutationBuilder mutationBuilder) {
this.datastore = datastore;
+ this.projectId = projectId;
+ this.databaseId = databaseId;
this.mutationBuilder = mutationBuilder;
}
@@ -295,6 +313,8 @@ class V1TestUtil {
commitRequest.addMutations(mutationBuilder.apply(entity));
}
commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
+ commitRequest.setProjectId(projectId);
+ commitRequest.setDatabaseId(databaseId);
datastore.commit(commitRequest.build());
// Break if the commit threw no exception.
break;
@@ -318,6 +338,8 @@ class V1TestUtil {
static class V1TestReader {
private static final int QUERY_BATCH_LIMIT = 500;
private final Datastore datastore;
+ private final String projectId;
+ private final String databaseId;
private final Query query;
private final @Nullable String namespace;
private boolean moreResults;
@@ -326,8 +348,15 @@ class V1TestUtil {
private QueryResultBatch currentBatch;
private Entity currentEntity;
- V1TestReader(Datastore datastore, Query query, @Nullable String namespace) {
+ V1TestReader(
+ Datastore datastore,
+ String projectId,
+ String databaseId,
+ Query query,
+ @Nullable String namespace) {
this.datastore = datastore;
+ this.projectId = projectId;
+ this.databaseId = databaseId;
this.query = query;
this.namespace = namespace;
}
@@ -362,7 +391,7 @@ class V1TestUtil {
query.setStartCursor(currentBatch.getEndCursor());
}
- RunQueryRequest request = makeRequest(query.build(), namespace);
+ RunQueryRequest request = makeRequest(projectId, databaseId, query.build(), namespace);
RunQueryResponse response = datastore.runQuery(request);
currentBatch = response.getBatch();
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/V1WriteIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/V1WriteIT.java
index 4f794a03e58..0062208630f 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/V1WriteIT.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/V1WriteIT.java
@@ -56,6 +56,7 @@ import org.junit.runners.JUnit4;
public class V1WriteIT {
private V1TestOptions options;
private String project;
+ private String database = "";
private String ancestor;
private final long numEntities = 1000;
@@ -83,12 +84,12 @@ public class V1WriteIT {
ParDo.of(
new V1TestUtil.CreateEntityFn(
options.getKind(), options.getNamespace(), ancestor, 0)))
- .apply(DatastoreIO.v1().write().withProjectId(project));
+ .apply(DatastoreIO.v1().write().withProjectId(project).withDatabaseId(database));
p.run();
// Count number of entities written to datastore.
- long numEntitiesWritten = countEntities(options, project, ancestor);
+ long numEntitiesWritten = countEntities(options, project, database, ancestor);
assertEquals(numEntities, numEntitiesWritten);
}
@@ -185,18 +186,18 @@ public class V1WriteIT {
ParDo.of(
new V1TestUtil.CreateEntityFn(
options.getKind(), options.getNamespace(), ancestor, rawPropertySize)))
- .apply(DatastoreIO.v1().write().withProjectId(project));
+ .apply(DatastoreIO.v1().write().withProjectId(project).withDatabaseId(database));
p.run();
// Count number of entities written to datastore.
- long numEntitiesWritten = countEntities(options, project, ancestor);
+ long numEntitiesWritten = countEntities(options, project, database, ancestor);
assertEquals(numLargeEntities, numEntitiesWritten);
}
@After
public void tearDown() throws Exception {
- deleteAllEntities(options, project, ancestor);
+ deleteAllEntities(options, project, database, ancestor);
}
}