diff --git a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/DatastoreEmulator.java b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/DatastoreEmulator.java
index 2fb11ebe8..dcca102a1 100644
--- a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/DatastoreEmulator.java
+++ b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/DatastoreEmulator.java
@@ -262,6 +262,7 @@ public synchronized File getProjectDirectory() {
static class StartupMonitor extends Thread {
private final InputStream inputStream;
private volatile boolean success = false;
+
/** This latch will reach 0 once server startup has completed. */
private final CountDownLatch startupCompleteLatch = new CountDownLatch(1);
diff --git a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/DatastoreException.java b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/DatastoreException.java
index 400e63308..b283603fb 100644
--- a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/DatastoreException.java
+++ b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/DatastoreException.java
@@ -28,12 +28,16 @@ public DatastoreException(String methodName, Code code, String message, Throwabl
this.code = code;
}
- /** @return the canonical error code */
+ /**
+ * @return the canonical error code
+ */
public Code getCode() {
return code;
}
- /** @return the datastore method name */
+ /**
+ * @return the datastore method name
+ */
public String getMethodName() {
return methodName;
}
diff --git a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/DatastoreHelper.java b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/DatastoreHelper.java
index 23ac315b5..239673061 100644
--- a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/DatastoreHelper.java
+++ b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/DatastoreHelper.java
@@ -336,7 +336,9 @@ private static void setProjectEndpointFromEnv(DatastoreOptions.Builder options)
return;
}
- /** @see #getOptionsFromEnv() */
+ /**
+ * @see #getOptionsFromEnv()
+ */
public static Datastore getDatastoreFromEnv() throws GeneralSecurityException, IOException {
return DatastoreFactory.get().create(getOptionsFromEnv().build());
}
diff --git a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/EndToEndChecksumHandler.java b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/EndToEndChecksumHandler.java
index 3f840a8ba..339fe4878 100644
--- a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/EndToEndChecksumHandler.java
+++ b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/EndToEndChecksumHandler.java
@@ -24,6 +24,7 @@
class EndToEndChecksumHandler {
/** The checksum http header on http requests */
static final String HTTP_REQUEST_CHECKSUM_HEADER = "x-request-checksum-348659783";
+
/** The checksum http header on http responses */
static final String HTTP_RESPONSE_CHECKSUM_HEADER = "x-response-checksum-348659783";
diff --git a/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/QuerySplitterTest.java b/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/QuerySplitterTest.java
index b064e137a..0802a62aa 100644
--- a/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/QuerySplitterTest.java
+++ b/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/QuerySplitterTest.java
@@ -101,8 +101,7 @@ public void disallowsSortOrder() {
public void disallowsMultipleKinds() {
Datastore datastore = factory.create(options.build());
Query queryWithMultipleKinds =
- query
- .toBuilder()
+ query.toBuilder()
.addKind(KindExpression.newBuilder().setName("another-kind").build())
.build();
IllegalArgumentException exception =
@@ -129,8 +128,7 @@ public void disallowsKindlessQuery() {
public void disallowsInequalityFilter() {
Datastore datastore = factory.create(options.build());
Query queryWithInequality =
- query
- .toBuilder()
+ query.toBuilder()
.setFilter(makeFilter("foo", Operator.GREATER_THAN, makeValue("value")))
.build();
IllegalArgumentException exception =
@@ -177,16 +175,13 @@ public void getSplits() throws Exception {
assertThat(splittedQueries)
.containsExactly(
- query
- .toBuilder()
+ query.toBuilder()
.setFilter(makeFilterWithKeyRange(propertyFilter, null, splitKey1))
.build(),
- query
- .toBuilder()
+ query.toBuilder()
.setFilter(makeFilterWithKeyRange(propertyFilter, splitKey1, splitKey3))
.build(),
- query
- .toBuilder()
+ query.toBuilder()
.setFilter(makeFilterWithKeyRange(propertyFilter, splitKey3, null))
.build());
@@ -229,16 +224,13 @@ public void getSplitsWithDatabaseId() throws Exception {
assertThat(splitQueries)
.containsExactly(
- query
- .toBuilder()
+ query.toBuilder()
.setFilter(makeFilterWithKeyRange(propertyFilter, null, splitKey1))
.build(),
- query
- .toBuilder()
+ query.toBuilder()
.setFilter(makeFilterWithKeyRange(propertyFilter, splitKey1, splitKey3))
.build(),
- query
- .toBuilder()
+ query.toBuilder()
.setFilter(makeFilterWithKeyRange(propertyFilter, splitKey3, null))
.build());
@@ -277,12 +269,10 @@ public void notEnoughSplits() throws Exception {
assertThat(splittedQueries)
.containsExactly(
- query
- .toBuilder()
+ query.toBuilder()
.setFilter(makeFilterWithKeyRange(propertyFilter, null, splitKey0))
.build(),
- query
- .toBuilder()
+ query.toBuilder()
.setFilter(makeFilterWithKeyRange(propertyFilter, splitKey0, null))
.build());
@@ -325,16 +315,13 @@ public void getSplits_withReadTime() throws Exception {
assertThat(splittedQueries)
.containsExactly(
- query
- .toBuilder()
+ query.toBuilder()
.setFilter(makeFilterWithKeyRange(propertyFilter, null, splitKey1))
.build(),
- query
- .toBuilder()
+ query.toBuilder()
.setFilter(makeFilterWithKeyRange(propertyFilter, splitKey1, splitKey3))
.build(),
- query
- .toBuilder()
+ query.toBuilder()
.setFilter(makeFilterWithKeyRange(propertyFilter, splitKey3, null))
.build());
diff --git a/generation_config.yaml b/generation_config.yaml
new file mode 100644
index 000000000..c7f6ff188
--- /dev/null
+++ b/generation_config.yaml
@@ -0,0 +1,25 @@
+gapic_generator_version: 2.56.2
+googleapis_commitish: ce291b3bc967923f89e0e54ed33d18802672b171
+libraries_bom_version: 26.59.0
+libraries:
+ - api_shortname: datastore
+ name_pretty: Cloud Datastore
+ product_documentation: https://cloud.google.com/datastore
+ client_documentation: https://cloud.google.com/java/docs/reference/google-cloud-datastore/latest/history
+ issue_tracker: https://issuetracker.google.com/savedsearches/559768
+ release_level: stable
+ language: java
+ repo: googleapis/java-datastore
+ repo_short: java-datastore
+ distribution_name: com.google.cloud:google-cloud-datastore
+ codeowner_team: '@googleapis/cloud-native-db-dpes @googleapis/api-datastore-sdk @googleapis/api-firestore-partners'
+ api_id: datastore.googleapis.com
+ library_type: GAPIC_COMBO
+ api_description: is a fully managed, schemaless database for\nstoring non-relational data. Cloud Datastore automatically scales with\nyour users and supports ACID transactions, high availability of reads and\nwrites, strong consistency for reads and ancestor queries, and eventual\nconsistency for all other queries.
+ excluded_dependencies: grpc-google-cloud-datastore-v1
+ extra_versioned_modules: datastore-v1-proto-client
+ excluded_poms: grpc-google-cloud-datastore-v1
+ recommended_package: com.google.cloud.datastore
+ GAPICs:
+ - proto_path: google/datastore/v1
+ - proto_path: google/datastore/admin/v1
diff --git a/google-cloud-datastore-bom/pom.xml b/google-cloud-datastore-bom/pom.xml
index 2651713e3..6842d693e 100644
--- a/google-cloud-datastore-bom/pom.xml
+++ b/google-cloud-datastore-bom/pom.xml
@@ -3,12 +3,12 @@
4.0.0
com.google.cloud
google-cloud-datastore-bom
- 2.19.2-SNAPSHOT
+ 2.28.2-SNAPSHOT
pom
com.google.cloud
sdk-platform-java-config
- 3.29.0
+ 3.47.0
Google Cloud datastore BOM
@@ -52,22 +52,22 @@
com.google.cloud
google-cloud-datastore
- 2.19.2-SNAPSHOT
+ 2.28.2-SNAPSHOT
com.google.api.grpc
grpc-google-cloud-datastore-admin-v1
- 2.19.2-SNAPSHOT
+ 2.28.2-SNAPSHOT
com.google.api.grpc
proto-google-cloud-datastore-v1
- 0.110.2-SNAPSHOT
+ 0.119.2-SNAPSHOT
com.google.api.grpc
proto-google-cloud-datastore-admin-v1
- 2.19.2-SNAPSHOT
+ 2.28.2-SNAPSHOT
diff --git a/google-cloud-datastore-utils/pom.xml b/google-cloud-datastore-utils/pom.xml
new file mode 100644
index 000000000..0360a44a6
--- /dev/null
+++ b/google-cloud-datastore-utils/pom.xml
@@ -0,0 +1,111 @@
+
+
+ 4.0.0
+ com.google.cloud
+ google-cloud-datastore-utils
+ 2.28.2-SNAPSHOT
+ jar
+ Google Cloud Datastore Utilities
+ https://github.com/googleapis/java-datastore
+
+ Java datastore client utility library.
+
+
+ com.google.cloud
+ google-cloud-datastore-parent
+ 2.28.2-SNAPSHOT
+
+
+ google-cloud-datastore-utils
+
+
+
+ com.google.api-client
+ google-api-client
+
+
+ com.google.http-client
+ google-http-client-protobuf
+
+
+ com.google.http-client
+ google-http-client-gson
+
+
+ com.google.api.grpc
+ proto-google-cloud-datastore-v1
+
+
+ com.google.api
+ api-common
+
+
+ com.google.protobuf
+ protobuf-java
+
+
+ com.google.guava
+ guava
+
+
+ com.google.api.grpc
+ proto-google-common-protos
+
+
+ com.google.http-client
+ google-http-client
+
+
+ com.google.http-client
+ google-http-client-jackson2
+
+
+ com.google.oauth-client
+ google-oauth-client
+
+
+ com.google.code.findbugs
+ jsr305
+
+
+
+ junit
+ junit
+ test
+
+
+ com.google.truth
+ truth
+ 1.4.2
+ test
+
+
+ org.checkerframework
+ checker-qual
+
+
+
+
+
+
+
+
+ org.apache.maven.plugins
+ maven-surefire-plugin
+
+ -Xmx2048m
+
+
+
+
+
+
+
+ native
+
+
+ true
+
+
+
+
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/Datastore.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/Datastore.java
new file mode 100644
index 000000000..d66e9ce60
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/Datastore.java
@@ -0,0 +1,136 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.datastore.utils;
+
+import com.google.datastore.v1.*;
+import com.google.rpc.Code;
+import java.io.IOException;
+import java.io.InputStream;
+
+/**
+ * Provides access to Cloud Datastore.
+ *
+ * This class is thread-safe.
+ */
+public class Datastore {
+
+ final RemoteRpc remoteRpc;
+
+ Datastore(RemoteRpc remoteRpc) {
+ this.remoteRpc = remoteRpc;
+ }
+
+ /** Reset the RPC count. */
+ public void resetRpcCount() {
+ remoteRpc.resetRpcCount();
+ }
+
+ /**
+ * Returns the number of RPC calls made since the client was created or {@link #resetRpcCount} was
+ * called.
+ */
+ public int getRpcCount() {
+ return remoteRpc.getRpcCount();
+ }
+
+ private com.google.datastore.utils.DatastoreException invalidResponseException(
+ String method, IOException exception) {
+ return RemoteRpc.makeException(
+ remoteRpc.getUrl(), method, Code.UNAVAILABLE, "Invalid response", exception);
+ }
+
+ public AllocateIdsResponse allocateIds(AllocateIdsRequest request)
+ throws com.google.datastore.utils.DatastoreException {
+ try (InputStream is =
+ remoteRpc.call("allocateIds", request, request.getProjectId(), request.getDatabaseId())) {
+ return AllocateIdsResponse.parseFrom(is);
+ } catch (IOException exception) {
+ throw invalidResponseException("allocateIds", exception);
+ }
+ }
+
+ public BeginTransactionResponse beginTransaction(BeginTransactionRequest request)
+ throws com.google.datastore.utils.DatastoreException {
+ try (InputStream is =
+ remoteRpc.call(
+ "beginTransaction", request, request.getProjectId(), request.getDatabaseId())) {
+ return BeginTransactionResponse.parseFrom(is);
+ } catch (IOException exception) {
+ throw invalidResponseException("beginTransaction", exception);
+ }
+ }
+
+ public CommitResponse commit(CommitRequest request)
+ throws com.google.datastore.utils.DatastoreException {
+ try (InputStream is =
+ remoteRpc.call("commit", request, request.getProjectId(), request.getDatabaseId())) {
+ return CommitResponse.parseFrom(is);
+ } catch (IOException exception) {
+ throw invalidResponseException("commit", exception);
+ }
+ }
+
+ public LookupResponse lookup(LookupRequest request)
+ throws com.google.datastore.utils.DatastoreException {
+ try (InputStream is =
+ remoteRpc.call("lookup", request, request.getProjectId(), request.getDatabaseId())) {
+ return LookupResponse.parseFrom(is);
+ } catch (IOException exception) {
+ throw invalidResponseException("lookup", exception);
+ }
+ }
+
+ public ReserveIdsResponse reserveIds(ReserveIdsRequest request)
+ throws com.google.datastore.utils.DatastoreException {
+ try (InputStream is =
+ remoteRpc.call("reserveIds", request, request.getProjectId(), request.getDatabaseId())) {
+ return ReserveIdsResponse.parseFrom(is);
+ } catch (IOException exception) {
+ throw invalidResponseException("reserveIds", exception);
+ }
+ }
+
+ public RollbackResponse rollback(RollbackRequest request)
+ throws com.google.datastore.utils.DatastoreException {
+ try (InputStream is =
+ remoteRpc.call("rollback", request, request.getProjectId(), request.getDatabaseId())) {
+ return RollbackResponse.parseFrom(is);
+ } catch (IOException exception) {
+ throw invalidResponseException("rollback", exception);
+ }
+ }
+
+ public RunQueryResponse runQuery(RunQueryRequest request)
+ throws com.google.datastore.utils.DatastoreException {
+ try (InputStream is =
+ remoteRpc.call("runQuery", request, request.getProjectId(), request.getDatabaseId())) {
+ return RunQueryResponse.parseFrom(is);
+ } catch (IOException exception) {
+ throw invalidResponseException("runQuery", exception);
+ }
+ }
+
+ public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request)
+ throws DatastoreException {
+ try (InputStream is =
+ remoteRpc.call(
+ "runAggregationQuery", request, request.getProjectId(), request.getDatabaseId())) {
+ return RunAggregationQueryResponse.parseFrom(is);
+ } catch (IOException exception) {
+ throw invalidResponseException("runAggregationQuery", exception);
+ }
+ }
+}
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreException.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreException.java
new file mode 100644
index 000000000..67dbc1f9a
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreException.java
@@ -0,0 +1,49 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.datastore.utils;
+
+import com.google.rpc.Code;
+
+/** Indicates an error in a {@link Datastore} call. */
+public class DatastoreException extends Exception {
+ private final String methodName;
+ private final Code code;
+
+ public DatastoreException(String methodName, Code code, String message, Throwable cause) {
+ super(message, cause);
+ this.methodName = methodName;
+ this.code = code;
+ }
+
+ /**
+ * @return the canonical error code
+ */
+ public Code getCode() {
+ return code;
+ }
+
+ /**
+ * @return the datastore method name
+ */
+ public String getMethodName() {
+ return methodName;
+ }
+
+ @Override
+ public String toString() {
+ return String.format("%s, code=%s", super.toString(), code);
+ }
+}
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreFactory.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreFactory.java
new file mode 100644
index 000000000..2befe276e
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreFactory.java
@@ -0,0 +1,127 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.datastore.utils;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.api.client.auth.oauth2.Credential;
+import com.google.api.client.http.HttpRequestFactory;
+import com.google.api.client.http.HttpTransport;
+import com.google.api.client.http.javanet.NetHttpTransport;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Arrays;
+import java.util.logging.*;
+
+/** Client factory for {@link com.google.datastore.utils.Datastore}. */
+public class DatastoreFactory {
+
+ // Lazy load this because we might be running inside App Engine and this
+ // class isn't on the whitelist.
+ private static ConsoleHandler methodHandler;
+
+ /** API version. */
+ public static final String VERSION = "v1";
+
+ public static final String DEFAULT_HOST = "https://datastore.googleapis.com";
+
+ /** Singleton factory instance. */
+ private static final DatastoreFactory INSTANCE = new DatastoreFactory();
+
+ public static DatastoreFactory get() {
+ return INSTANCE;
+ }
+
+ /**
+ * Provides access to a datastore using the provided options. Logs into the application using the
+ * credentials available via these options.
+ *
+ * @throws IllegalArgumentException if the server or credentials weren't provided.
+ */
+ public com.google.datastore.utils.Datastore create(
+ com.google.datastore.utils.DatastoreOptions options) {
+ return new com.google.datastore.utils.Datastore(newRemoteRpc(options));
+ }
+
+ /** Constructs a Google APIs HTTP client with the associated credentials. */
+ public HttpRequestFactory makeClient(com.google.datastore.utils.DatastoreOptions options) {
+ Credential credential = options.getCredential();
+ HttpTransport transport = options.getTransport();
+ if (transport == null) {
+ transport = credential == null ? new NetHttpTransport() : credential.getTransport();
+ transport = transport == null ? new NetHttpTransport() : transport;
+ }
+ return transport.createRequestFactory(credential);
+ }
+
+ /** Starts logging datastore method calls to the console. (Useful within tests.) */
+ public static void logMethodCalls() {
+ Logger logger = Logger.getLogger(Datastore.class.getName());
+ logger.setLevel(Level.FINE);
+ if (!Arrays.asList(logger.getHandlers()).contains(getStreamHandler())) {
+ logger.addHandler(getStreamHandler());
+ }
+ }
+
+ /** Build a valid datastore URL. */
+ String buildProjectEndpoint(com.google.datastore.utils.DatastoreOptions options) {
+ if (options.getProjectEndpoint() != null) {
+ return options.getProjectEndpoint();
+ }
+ // DatastoreOptions ensures either project endpoint or project ID is set.
+ String projectId = checkNotNull(options.getProjectId());
+ if (options.getHost() != null) {
+ return validateUrl(
+ String.format("https://%s/%s/projects/%s", options.getHost(), VERSION, projectId));
+ } else if (options.getLocalHost() != null) {
+ return validateUrl(
+ String.format("http://%s/%s/projects/%s", options.getLocalHost(), VERSION, projectId));
+ }
+ return validateUrl(String.format("%s/%s/projects/%s", DEFAULT_HOST, VERSION, projectId));
+ }
+
+ protected com.google.datastore.utils.RemoteRpc newRemoteRpc(DatastoreOptions options) {
+ checkNotNull(options);
+ HttpRequestFactory client = makeClient(options);
+ return new com.google.datastore.utils.RemoteRpc(
+ client, options.getInitializer(), buildProjectEndpoint(options));
+ }
+
+ private static String validateUrl(String url) {
+ try {
+ return new URI(url).toString();
+ } catch (URISyntaxException e) {
+ throw new IllegalArgumentException(e);
+ }
+ }
+
+ // TODO: Support something other than console handler for when we're
+ // running in App Engine
+ private static synchronized StreamHandler getStreamHandler() {
+ if (methodHandler == null) {
+ methodHandler = new ConsoleHandler();
+ methodHandler.setFormatter(
+ new Formatter() {
+ @Override
+ public String format(LogRecord record) {
+ return record.getMessage() + "\n";
+ }
+ });
+ methodHandler.setLevel(Level.FINE);
+ }
+ return methodHandler;
+ }
+}
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreHelper.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreHelper.java
new file mode 100644
index 000000000..a937dffad
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreHelper.java
@@ -0,0 +1,731 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.datastore.utils;
+
+import com.google.api.client.auth.oauth2.Credential;
+import com.google.api.client.googleapis.auth.oauth2.GoogleCredential;
+import com.google.api.client.googleapis.javanet.GoogleNetHttpTransport;
+import com.google.api.client.http.GenericUrl;
+import com.google.api.client.http.HttpRequest;
+import com.google.api.client.http.HttpTransport;
+import com.google.api.client.json.JsonFactory;
+import com.google.api.client.json.gson.GsonFactory;
+import com.google.datastore.v1.ArrayValue;
+import com.google.datastore.v1.CompositeFilter;
+import com.google.datastore.v1.Entity;
+import com.google.datastore.v1.Filter;
+import com.google.datastore.v1.Key;
+import com.google.datastore.v1.Key.PathElement;
+import com.google.datastore.v1.Key.PathElement.IdTypeCase;
+import com.google.datastore.v1.Mutation;
+import com.google.datastore.v1.PartitionId;
+import com.google.datastore.v1.PropertyFilter;
+import com.google.datastore.v1.PropertyOrder;
+import com.google.datastore.v1.PropertyReference;
+import com.google.datastore.v1.Value;
+import com.google.datastore.v1.Value.ValueTypeCase;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Timestamp;
+import com.google.protobuf.TimestampOrBuilder;
+import com.google.type.LatLng;
+import java.io.File;
+import java.io.IOException;
+import java.security.GeneralSecurityException;
+import java.security.PrivateKey;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.Date;
+import java.util.Iterator;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import javax.annotation.Nullable;
+
+/** Helper methods for {@link Datastore}. */
+// TODO: Accept OrBuilders when possible.
+public final class DatastoreHelper {
+ private static final Logger logger =
+ Logger.getLogger(com.google.datastore.utils.DatastoreHelper.class.getName());
+
+ private static final int MICROSECONDS_PER_SECOND = 1000 * 1000;
+ private static final int NANOSECONDS_PER_MICROSECOND = 1000;
+
+ /** The property used in the Datastore to give us a random distribution. * */
+ public static final String SCATTER_PROPERTY_NAME = "__scatter__";
+
+ /** The property used in the Datastore to get the key of the entity. * */
+ public static final String KEY_PROPERTY_NAME = "__key__";
+
+ /** Name of the environment variable used to set the project ID. */
+ public static final String PROJECT_ID_ENV_VAR = "DATASTORE_PROJECT_ID";
+
+ /** Name of the environment variable used to set the local host. */
+ public static final String LOCAL_HOST_ENV_VAR = "DATASTORE_EMULATOR_HOST";
+
+ /** Name of the environment variable used to set the service account. */
+ public static final String SERVICE_ACCOUNT_ENV_VAR = "DATASTORE_SERVICE_ACCOUNT";
+
+ /** Name of the environment variable used to set the private key file. */
+ public static final String PRIVATE_KEY_FILE_ENV_VAR = "DATASTORE_PRIVATE_KEY_FILE";
+
+ private static final String URL_OVERRIDE_ENV_VAR = "__DATASTORE_URL_OVERRIDE";
+
+ private static final AtomicReference projectIdFromComputeEngine = new AtomicReference<>();
+
+ /** Comparator for Keys */
+ private static final class KeyComparator implements Comparator {
+
+ static final com.google.datastore.utils.DatastoreHelper.KeyComparator INSTANCE =
+ new com.google.datastore.utils.DatastoreHelper.KeyComparator();
+
+ private int comparePathElement(PathElement thisElement, PathElement otherElement) {
+ int result = thisElement.getKind().compareTo(otherElement.getKind());
+ if (result != 0) {
+ return result;
+ }
+ if (thisElement.getIdTypeCase() == IdTypeCase.ID) {
+ if (otherElement.getIdTypeCase() != IdTypeCase.ID) {
+ return -1;
+ }
+ return Long.valueOf(thisElement.getId()).compareTo(otherElement.getId());
+ }
+ if (otherElement.getIdTypeCase() == IdTypeCase.ID) {
+ return 1;
+ }
+
+ return thisElement.getName().compareTo(otherElement.getName());
+ }
+
+ @Override
+ public int compare(Key thisKey, Key otherKey) {
+ if (!thisKey.getPartitionId().equals(otherKey.getPartitionId())) {
+ throw new IllegalArgumentException("Cannot compare keys with different partition ids.");
+ }
+
+ Iterator thisPath = thisKey.getPathList().iterator();
+ Iterator otherPath = otherKey.getPathList().iterator();
+ while (thisPath.hasNext()) {
+ if (!otherPath.hasNext()) {
+ return 1;
+ }
+ int result = comparePathElement(thisPath.next(), otherPath.next());
+ if (result != 0) {
+ return result;
+ }
+ }
+
+ return otherPath.hasNext() ? -1 : 0;
+ }
+ }
+
+ private DatastoreHelper() {}
+
+ private static HttpTransport newTransport() throws GeneralSecurityException, IOException {
+ return GoogleNetHttpTransport.newTrustedTransport();
+ }
+
+ static JsonFactory newJsonFactory() {
+ return new GsonFactory();
+ }
+
+ /**
+ * Constructs credentials for the given account and key.
+ *
+ * @param serviceAccountId service account ID (typically an e-mail address).
+ * @param privateKeyFile the file name from which to get the private key.
+ * @return valid credentials or {@code null}
+ */
+ public static Credential getServiceAccountCredential(
+ String serviceAccountId, String privateKeyFile) throws GeneralSecurityException, IOException {
+ return getServiceAccountCredential(serviceAccountId, privateKeyFile, DatastoreOptions.SCOPES);
+ }
+
+ /**
+ * Constructs credentials for the given account and key file.
+ *
+ * @param serviceAccountId service account ID (typically an e-mail address).
+ * @param privateKeyFile the file name from which to get the private key.
+ * @param serviceAccountScopes Collection of OAuth scopes to use with the the service account flow
+ * or {@code null} if not.
+ * @return valid credentials or {@code null}
+ */
+ public static Credential getServiceAccountCredential(
+ String serviceAccountId, String privateKeyFile, Collection serviceAccountScopes)
+ throws GeneralSecurityException, IOException {
+ return getCredentialBuilderWithoutPrivateKey(serviceAccountId, serviceAccountScopes)
+ .setServiceAccountPrivateKeyFromP12File(new File(privateKeyFile))
+ .build();
+ }
+
+ /**
+ * Constructs credentials for the given account and key.
+ *
+ * @param serviceAccountId service account ID (typically an e-mail address).
+ * @param privateKey the private key for the given account.
+ * @param serviceAccountScopes Collection of OAuth scopes to use with the the service account flow
+ * or {@code null} if not.
+ * @return valid credentials or {@code null}
+ */
+ public static Credential getServiceAccountCredential(
+ String serviceAccountId, PrivateKey privateKey, Collection serviceAccountScopes)
+ throws GeneralSecurityException, IOException {
+ return getCredentialBuilderWithoutPrivateKey(serviceAccountId, serviceAccountScopes)
+ .setServiceAccountPrivateKey(privateKey)
+ .build();
+ }
+
+ private static GoogleCredential.Builder getCredentialBuilderWithoutPrivateKey(
+ String serviceAccountId, Collection serviceAccountScopes)
+ throws GeneralSecurityException, IOException {
+ HttpTransport transport = newTransport();
+ JsonFactory jsonFactory = newJsonFactory();
+ return new GoogleCredential.Builder()
+ .setTransport(transport)
+ .setJsonFactory(jsonFactory)
+ .setServiceAccountId(serviceAccountId)
+ .setServiceAccountScopes(serviceAccountScopes);
+ }
+
+ /**
+ * Constructs a {@link Datastore} from environment variables and/or the Compute Engine metadata
+ * server.
+ *
+ * The project ID is determined from, in order of preference:
+ *
+ *
+ * - DATASTORE_PROJECT_ID environment variable
+ *
- Compute Engine
+ *
+ *
+ * Credentials are taken from, in order of preference:
+ *
+ *
+ * - No credentials (if the DATASTORE_EMULATOR_HOST environment variable is set)
+ *
- Service Account specified by the DATASTORE_SERVICE_ACCOUNT and DATASTORE_PRIVATE_KEY_FILE
+ * environment variables
+ *
- Google Application Default as described here.
+ *
+ */
+ public static DatastoreOptions.Builder getOptionsFromEnv()
+ throws GeneralSecurityException, IOException {
+ DatastoreOptions.Builder options = new DatastoreOptions.Builder();
+ setProjectEndpointFromEnv(options);
+ options.credential(getCredentialFromEnv());
+ return options;
+ }
+
+ private static Credential getCredentialFromEnv() throws GeneralSecurityException, IOException {
+ if (System.getenv(LOCAL_HOST_ENV_VAR) != null) {
+ logger.log(
+ Level.INFO,
+ "{0} environment variable was set. Not using credentials.",
+ new Object[] {LOCAL_HOST_ENV_VAR});
+ return null;
+ }
+ String serviceAccount = System.getenv(SERVICE_ACCOUNT_ENV_VAR);
+ String privateKeyFile = System.getenv(PRIVATE_KEY_FILE_ENV_VAR);
+ if (serviceAccount != null && privateKeyFile != null) {
+ logger.log(
+ Level.INFO,
+ "{0} and {1} environment variables were set. " + "Using service account credential.",
+ new Object[] {SERVICE_ACCOUNT_ENV_VAR, PRIVATE_KEY_FILE_ENV_VAR});
+ return getServiceAccountCredential(serviceAccount, privateKeyFile);
+ }
+ return GoogleCredential.getApplicationDefault().createScoped(DatastoreOptions.SCOPES);
+ }
+
+ /**
+ * Determines the project id from the environment. Uses the following sources in order of
+ * preference:
+ *
+ *
+ * - Value of the DATASTORE_PROJECT_ID environment variable
+ *
- Compute Engine
+ *
+ *
+ * @throws IllegalStateException if the project ID cannot be determined
+ */
+ private static String getProjectIdFromEnv() {
+ if (System.getenv(PROJECT_ID_ENV_VAR) != null) {
+ return System.getenv(PROJECT_ID_ENV_VAR);
+ }
+ String projectIdFromComputeEngine = getProjectIdFromComputeEngine();
+ if (projectIdFromComputeEngine != null) {
+ return projectIdFromComputeEngine;
+ }
+ throw new IllegalStateException(
+ String.format(
+ "Could not determine project ID."
+ + " If you are not running on Compute Engine, set the"
+ + " %s environment variable.",
+ PROJECT_ID_ENV_VAR));
+ }
+
+ /**
+ * Gets the project ID from the Compute Engine metadata server. Returns {@code null} if the
+ * project ID cannot be determined (because, for instance, the code is not running on Compute
+ * Engine).
+ */
+ @Nullable
+ public static String getProjectIdFromComputeEngine() {
+ String cachedProjectId = projectIdFromComputeEngine.get();
+ return cachedProjectId != null ? cachedProjectId : queryProjectIdFromComputeEngine();
+ }
+
+ @Nullable
+ private static String queryProjectIdFromComputeEngine() {
+ HttpTransport transport;
+
+ try {
+ transport = newTransport();
+ } catch (GeneralSecurityException | IOException e) {
+ logger.log(Level.WARNING, "Failed to create HttpTransport.", e);
+ return null;
+ }
+
+ try {
+ GenericUrl projectIdUrl =
+ new GenericUrl("http://metadata/computeMetadata/v1/project/project-id");
+ HttpRequest request = transport.createRequestFactory().buildGetRequest(projectIdUrl);
+ request.getHeaders().set("Metadata-Flavor", "Google");
+ String result = request.execute().parseAsString();
+ projectIdFromComputeEngine.set(result);
+ return result;
+ } catch (IOException e) {
+ logger.log(Level.INFO, "Could not determine project ID from Compute Engine.", e);
+ return null;
+ }
+ }
+
+ private static void setProjectEndpointFromEnv(DatastoreOptions.Builder options) {
+ // DATASTORE_HOST is deprecated.
+ if (System.getenv("DATASTORE_HOST") != null) {
+ logger.warning(
+ String.format(
+ "Ignoring value of environment variable DATASTORE_HOST. "
+ + "To point datastore to a host running locally, use "
+ + "the environment variable %s.",
+ LOCAL_HOST_ENV_VAR));
+ }
+ String projectId = getProjectIdFromEnv();
+ if (System.getenv(URL_OVERRIDE_ENV_VAR) != null) {
+ options.projectEndpoint(
+ String.format("%s/projects/%s", System.getenv(URL_OVERRIDE_ENV_VAR), projectId));
+ return;
+ }
+ if (System.getenv(LOCAL_HOST_ENV_VAR) != null) {
+ options.projectId(projectId);
+ options.localHost(System.getenv(LOCAL_HOST_ENV_VAR));
+ return;
+ }
+ options.projectId(projectId);
+ return;
+ }
+
+ /**
+ * @see #getOptionsFromEnv()
+ */
+ public static Datastore getDatastoreFromEnv() throws GeneralSecurityException, IOException {
+ return DatastoreFactory.get().create(getOptionsFromEnv().build());
+ }
+
+ /**
+ * Gets a {@link com.google.datastore.utils.QuerySplitter}.
+ *
+ * The returned {@link com.google.datastore.utils.QuerySplitter#getSplits} cannot accept a
+ * query that contains inequality filters, a sort filter, or a missing kind.
+ */
+ public static QuerySplitter getQuerySplitter() {
+ return com.google.datastore.utils.QuerySplitterImpl.INSTANCE;
+ }
+
+ public static Comparator getKeyComparator() {
+ return com.google.datastore.utils.DatastoreHelper.KeyComparator.INSTANCE;
+ }
+
+ /** Make a sort order for use in a query. */
+ public static PropertyOrder.Builder makeOrder(
+ String property, PropertyOrder.Direction direction) {
+ return PropertyOrder.newBuilder()
+ .setProperty(makePropertyReference(property))
+ .setDirection(direction);
+ }
+
+ /** Makes an ancestor filter. */
+ public static Filter.Builder makeAncestorFilter(Key ancestor) {
+ return makeFilter(
+ com.google.datastore.utils.DatastoreHelper.KEY_PROPERTY_NAME,
+ PropertyFilter.Operator.HAS_ANCESTOR,
+ makeValue(ancestor));
+ }
+
+ /** Make a filter on a property for use in a query. */
+ public static Filter.Builder makeFilter(
+ String property, PropertyFilter.Operator operator, Value value) {
+ return Filter.newBuilder()
+ .setPropertyFilter(
+ PropertyFilter.newBuilder()
+ .setProperty(makePropertyReference(property))
+ .setOp(operator)
+ .setValue(value));
+ }
+
+ /** Make a filter on a property for use in a query. */
+ public static Filter.Builder makeFilter(
+ String property, PropertyFilter.Operator operator, Value.Builder value) {
+ return makeFilter(property, operator, value.build());
+ }
+
+ /** Make a composite filter from the given sub-filters using AND to combine filters. */
+ public static Filter.Builder makeAndFilter(Filter... subfilters) {
+ return makeAndFilter(Arrays.asList(subfilters));
+ }
+
+ /** Make a composite filter from the given sub-filters using AND to combine filters. */
+ public static Filter.Builder makeAndFilter(Iterable subfilters) {
+ return Filter.newBuilder()
+ .setCompositeFilter(
+ CompositeFilter.newBuilder()
+ .addAllFilters(subfilters)
+ .setOp(CompositeFilter.Operator.AND));
+ }
+
+ /** Make a property reference for use in a query. */
+ public static PropertyReference.Builder makePropertyReference(String propertyName) {
+ return PropertyReference.newBuilder().setName(propertyName);
+ }
+
+ /** Make an array value containing the specified values. */
+ public static Value.Builder makeValue(Iterable values) {
+ return Value.newBuilder().setArrayValue(ArrayValue.newBuilder().addAllValues(values));
+ }
+
+ /** Make a list value containing the specified values. */
+ public static Value.Builder makeValue(Value value1, Value value2, Value... rest) {
+ ArrayValue.Builder arrayValue = ArrayValue.newBuilder();
+ arrayValue.addValues(value1);
+ arrayValue.addValues(value2);
+ arrayValue.addAllValues(Arrays.asList(rest));
+ return Value.newBuilder().setArrayValue(arrayValue);
+ }
+
+ /** Make an array value containing the specified values. */
+ public static Value.Builder makeValue(
+ Value.Builder value1, Value.Builder value2, Value.Builder... rest) {
+ ArrayValue.Builder arrayValue = ArrayValue.newBuilder();
+ arrayValue.addValues(value1);
+ arrayValue.addValues(value2);
+ for (Value.Builder builder : rest) {
+ arrayValue.addValues(builder);
+ }
+ return Value.newBuilder().setArrayValue(arrayValue);
+ }
+
+ /** Make a key value. */
+ public static Value.Builder makeValue(Key key) {
+ return Value.newBuilder().setKeyValue(key);
+ }
+
+ /** Make a key value. */
+ public static Value.Builder makeValue(Key.Builder key) {
+ return makeValue(key.build());
+ }
+
+ /** Make an integer value. */
+ public static Value.Builder makeValue(long key) {
+ return Value.newBuilder().setIntegerValue(key);
+ }
+
+ /** Make a floating point value. */
+ public static Value.Builder makeValue(double value) {
+ return Value.newBuilder().setDoubleValue(value);
+ }
+
+ /** Make a boolean value. */
+ public static Value.Builder makeValue(boolean value) {
+ return Value.newBuilder().setBooleanValue(value);
+ }
+
+ /** Make a string value. */
+ public static Value.Builder makeValue(String value) {
+ return Value.newBuilder().setStringValue(value);
+ }
+
+ /** Make an entity value. */
+ public static Value.Builder makeValue(Entity entity) {
+ return Value.newBuilder().setEntityValue(entity);
+ }
+
+ /** Make a entity value. */
+ public static Value.Builder makeValue(Entity.Builder entity) {
+ return makeValue(entity.build());
+ }
+
+ /** Make a ByteString value. */
+ public static Value.Builder makeValue(ByteString blob) {
+ return Value.newBuilder().setBlobValue(blob);
+ }
+
+ /** Make a timestamp value given a date. */
+ public static Value.Builder makeValue(Date date) {
+ return Value.newBuilder().setTimestampValue(toTimestamp(date.getTime() * 1000L));
+ }
+
+ /** Makes a GeoPoint value. */
+ public static Value.Builder makeValue(LatLng value) {
+ return Value.newBuilder().setGeoPointValue(value);
+ }
+
+ /** Makes a GeoPoint value. */
+ public static Value.Builder makeValue(LatLng.Builder value) {
+ return makeValue(value.build());
+ }
+
+ private static Timestamp.Builder toTimestamp(long microseconds) {
+ long seconds = microseconds / MICROSECONDS_PER_SECOND;
+ long microsecondsRemainder = microseconds % MICROSECONDS_PER_SECOND;
+ if (microsecondsRemainder < 0) {
+ // Nanos must be positive even if microseconds is negative.
+ // Java modulus doesn't take care of this for us.
+ microsecondsRemainder += MICROSECONDS_PER_SECOND;
+ seconds -= 1;
+ }
+ return Timestamp.newBuilder()
+ .setSeconds(seconds)
+ .setNanos((int) microsecondsRemainder * NANOSECONDS_PER_MICROSECOND);
+ }
+
+ /**
+ * Make a key from the specified path of kind/id-or-name pairs and/or Keys.
+ *
+ * The id-or-name values must be either String, Long, Integer or Short.
+ *
+ *
The last id-or-name value may be omitted, in which case an entity without an id is created
+ * (for use with automatic id allocation).
+ *
+ *
The PartitionIds of all Keys in the path must be equal. The returned Key.Builder will use
+ * this PartitionId.
+ */
+ public static Key.Builder makeKey(Object... elements) {
+ Key.Builder key = Key.newBuilder();
+ PartitionId partitionId = null;
+ for (int pathIndex = 0; pathIndex < elements.length; pathIndex += 2) {
+ PathElement.Builder pathElement = PathElement.newBuilder();
+ Object element = elements[pathIndex];
+ if (element instanceof Key) {
+ Key subKey = (Key) element;
+ if (partitionId == null) {
+ partitionId = subKey.getPartitionId();
+ } else if (!partitionId.equals(subKey.getPartitionId())) {
+ throw new IllegalArgumentException(
+ "Partition IDs did not match, found: "
+ + partitionId
+ + " and "
+ + subKey.getPartitionId());
+ }
+ key.addAllPath(((Key) element).getPathList());
+ // We increment by 2, but since we got a Key argument we're only consuming 1 element in this
+ // iteration of the loop. Decrement the index so that when we jump by 2 we end up in the
+ // right spot.
+ pathIndex--;
+ } else {
+ String kind;
+ try {
+ kind = (String) element;
+ } catch (ClassCastException e) {
+ throw new IllegalArgumentException(
+ "Expected string or Key, got: " + element.getClass(), e);
+ }
+ pathElement.setKind(kind);
+ if (pathIndex + 1 < elements.length) {
+ Object value = elements[pathIndex + 1];
+ if (value instanceof String) {
+ pathElement.setName((String) value);
+ } else if (value instanceof Long) {
+ pathElement.setId((Long) value);
+ } else if (value instanceof Integer) {
+ pathElement.setId((Integer) value);
+ } else if (value instanceof Short) {
+ pathElement.setId((Short) value);
+ } else {
+ throw new IllegalArgumentException(
+ "Expected string or integer, got: " + value.getClass());
+ }
+ }
+ key.addPath(pathElement);
+ }
+ }
+ if (partitionId != null && !partitionId.equals(PartitionId.getDefaultInstance())) {
+ key.setPartitionId(partitionId);
+ }
+ return key;
+ }
+
+ /**
+ * @return the double contained in value
+ * @throws IllegalArgumentException if the value does not contain a double.
+ */
+ public static double getDouble(Value value) {
+ if (value.getValueTypeCase() != ValueTypeCase.DOUBLE_VALUE) {
+ throw new IllegalArgumentException("Value does not contain a double.");
+ }
+ return value.getDoubleValue();
+ }
+
+ /**
+ * @return the key contained in value
+ * @throws IllegalArgumentException if the value does not contain a key.
+ */
+ public static Key getKey(Value value) {
+ if (value.getValueTypeCase() != ValueTypeCase.KEY_VALUE) {
+ throw new IllegalArgumentException("Value does not contain a key.");
+ }
+ return value.getKeyValue();
+ }
+
+ /**
+ * @return the blob contained in value
+ * @throws IllegalArgumentException if the value does not contain a blob.
+ */
+ public static ByteString getByteString(Value value) {
+ if (value.getMeaning() == 18 && value.getValueTypeCase() == ValueTypeCase.STRING_VALUE) {
+ return value.getStringValueBytes();
+ } else if (value.getValueTypeCase() == ValueTypeCase.BLOB_VALUE) {
+ return value.getBlobValue();
+ }
+ throw new IllegalArgumentException("Value does not contain a blob.");
+ }
+
+ /**
+ * @return the entity contained in value
+ * @throws IllegalArgumentException if the value does not contain an entity.
+ */
+ public static Entity getEntity(Value value) {
+ if (value.getValueTypeCase() != ValueTypeCase.ENTITY_VALUE) {
+ throw new IllegalArgumentException("Value does not contain an Entity.");
+ }
+ return value.getEntityValue();
+ }
+
+ /**
+ * @return the string contained in value
+ * @throws IllegalArgumentException if the value does not contain a string.
+ */
+ public static String getString(Value value) {
+ if (value.getValueTypeCase() != ValueTypeCase.STRING_VALUE) {
+ throw new IllegalArgumentException("Value does not contain a string.");
+ }
+ return value.getStringValue();
+ }
+
+ /**
+ * @return the boolean contained in value
+ * @throws IllegalArgumentException if the value does not contain a boolean.
+ */
+ public static boolean getBoolean(Value value) {
+ if (value.getValueTypeCase() != ValueTypeCase.BOOLEAN_VALUE) {
+ throw new IllegalArgumentException("Value does not contain a boolean.");
+ }
+ return value.getBooleanValue();
+ }
+
+ /**
+ * @return the long contained in value
+ * @throws IllegalArgumentException if the value does not contain a long.
+ */
+ public static long getLong(Value value) {
+ if (value.getValueTypeCase() != ValueTypeCase.INTEGER_VALUE) {
+ throw new IllegalArgumentException("Value does not contain an integer.");
+ }
+ return value.getIntegerValue();
+ }
+
+ /**
+ * @return the timestamp in microseconds contained in value
+ * @throws IllegalArgumentException if the value does not contain a timestamp.
+ */
+ public static long getTimestamp(Value value) {
+ if (value.getMeaning() == 18 && value.getValueTypeCase() == ValueTypeCase.INTEGER_VALUE) {
+ return value.getIntegerValue();
+ } else if (value.getValueTypeCase() == ValueTypeCase.TIMESTAMP_VALUE) {
+ return toMicroseconds(value.getTimestampValue());
+ }
+ throw new IllegalArgumentException("Value does not contain a timestamp.");
+ }
+
+ private static long toMicroseconds(TimestampOrBuilder timestamp) {
+ // Nanosecond precision is lost.
+ return timestamp.getSeconds() * MICROSECONDS_PER_SECOND
+ + timestamp.getNanos() / NANOSECONDS_PER_MICROSECOND;
+ }
+
+ /**
+ * @return the array contained in value as a list.
+ * @throws IllegalArgumentException if the value does not contain an array.
+ */
+ public static List getList(Value value) {
+ if (value.getValueTypeCase() != ValueTypeCase.ARRAY_VALUE) {
+ throw new IllegalArgumentException("Value does not contain an array.");
+ }
+ return value.getArrayValue().getValuesList();
+ }
+
+ /**
+ * Convert a timestamp value into a {@link Date} clipping off the microseconds.
+ *
+ * @param value a timestamp value to convert
+ * @return the resulting {@link Date}
+ * @throws IllegalArgumentException if the value does not contain a timestamp.
+ */
+ public static Date toDate(Value value) {
+ return new Date(getTimestamp(value) / 1000);
+ }
+
+ /**
+ * @param entity the entity to insert
+ * @return a mutation that will insert an entity
+ */
+ public static Mutation.Builder makeInsert(Entity entity) {
+ return Mutation.newBuilder().setInsert(entity);
+ }
+
+ /**
+ * @param entity the entity to update
+ * @return a mutation that will update an entity
+ */
+ public static Mutation.Builder makeUpdate(Entity entity) {
+ return Mutation.newBuilder().setUpdate(entity);
+ }
+
+ /**
+ * @param entity the entity to upsert
+ * @return a mutation that will upsert an entity
+ */
+ public static Mutation.Builder makeUpsert(Entity entity) {
+ return Mutation.newBuilder().setUpsert(entity);
+ }
+
+ /**
+ * @param key the key of the entity to delete
+ * @return a mutation that will delete an entity
+ */
+ public static Mutation.Builder makeDelete(Key key) {
+ return Mutation.newBuilder().setDelete(key);
+ }
+}
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreOptions.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreOptions.java
new file mode 100644
index 000000000..f6e91a41a
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/DatastoreOptions.java
@@ -0,0 +1,204 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.datastore.utils;
+
+import static com.google.common.base.Preconditions.checkArgument;
+
+import com.google.api.client.auth.oauth2.Credential;
+import com.google.api.client.http.HttpRequestInitializer;
+import com.google.api.client.http.HttpTransport;
+import java.util.Arrays;
+import java.util.List;
+
+/**
+ * An immutable object containing settings for the datastore.
+ *
+ * Example for connecting to a datastore:
+ *
+ *
+ * DatastoreOptions options = new DatastoreOptions.Builder()
+ * .projectId("my-project-id")
+ * .credential(DatastoreHelper.getComputeEngineCredential())
+ * .build();
+ * DatastoreFactory.get().create(options);
+ *
+ *
+ * The options should be passed to {@link DatastoreFactory#create}.
+ */
+public class DatastoreOptions {
+ private final String projectId;
+
+ private final String projectEndpoint;
+ private final String host;
+ private final String localHost;
+
+ private final HttpRequestInitializer initializer;
+
+ private final Credential credential;
+ private final HttpTransport transport;
+ public static final List SCOPES =
+ Arrays.asList("https://www.googleapis.com/auth/datastore");
+
+ DatastoreOptions(Builder b) {
+ checkArgument(
+ b.projectId != null || b.projectEndpoint != null,
+ "Either project ID or project endpoint must be provided.");
+ this.projectId = b.projectId;
+ this.projectEndpoint = b.projectEndpoint;
+ this.host = b.host;
+ this.localHost = b.localHost;
+ this.initializer = b.initializer;
+ this.credential = b.credential;
+ this.transport = b.transport;
+ }
+
+ /** Builder for {@link DatastoreOptions}. */
+ public static class Builder {
+ private static final String PROJECT_ENDPOINT_AND_PROJECT_ID_ERROR =
+ "Cannot set both project endpoint and project ID.";
+ private static final String PROJECT_ENDPOINT_AND_HOST_ERROR =
+ "Can set at most one of project endpoint, host, and local host.";
+
+ private String projectId;
+
+ private String projectEndpoint;
+ private String host;
+ private String localHost;
+ private HttpRequestInitializer initializer;
+ private Credential credential;
+ private HttpTransport transport;
+
+ public Builder() {}
+
+ public Builder(DatastoreOptions options) {
+ this.projectId = options.projectId;
+ this.projectEndpoint = options.projectEndpoint;
+ this.host = options.host;
+ this.localHost = options.localHost;
+ this.initializer = options.initializer;
+ this.credential = options.credential;
+ this.transport = options.transport;
+ }
+
+ public DatastoreOptions build() {
+ return new DatastoreOptions(this);
+ }
+
+ /** Sets the project ID used to access Cloud Datastore. */
+ public Builder projectId(String projectId) {
+ checkArgument(projectEndpoint == null, PROJECT_ENDPOINT_AND_PROJECT_ID_ERROR);
+ this.projectId = projectId;
+ return this;
+ }
+
+ /**
+ * Sets the host used to access Cloud Datastore. To connect to the Cloud Datastore Emulator, use
+ * {@link #localHost} instead.
+ */
+ public Builder host(String host) {
+ checkArgument(projectEndpoint == null && localHost == null, PROJECT_ENDPOINT_AND_HOST_ERROR);
+ if (includesScheme(host)) {
+ throw new IllegalArgumentException(
+ String.format("Host \"%s\" must not include scheme.", host));
+ }
+ this.host = host;
+ return this;
+ }
+
+ /**
+ * Configures the client to access Cloud Datastore on a local host (typically a Cloud Datastore
+ * Emulator instance). Call this method also configures the client not to attach credentials to
+ * requests.
+ */
+ public Builder localHost(String localHost) {
+ checkArgument(projectEndpoint == null && host == null, PROJECT_ENDPOINT_AND_HOST_ERROR);
+ if (includesScheme(localHost)) {
+ throw new IllegalArgumentException(
+ String.format("Local host \"%s\" must not include scheme.", localHost));
+ }
+ this.localHost = localHost;
+ return this;
+ }
+
+ /**
+ * Sets the project endpoint used to access Cloud Datastore. Prefer using {@link #projectId}
+ * and/or {@link #host}/{@link #localHost} when possible.
+ *
+ * @deprecated Use {@link #projectId} and/or {@link #host}/{@link #localHost} instead.
+ */
+ @Deprecated
+ public Builder projectEndpoint(String projectEndpoint) {
+ checkArgument(projectId == null, PROJECT_ENDPOINT_AND_PROJECT_ID_ERROR);
+ checkArgument(localHost == null && host == null, PROJECT_ENDPOINT_AND_HOST_ERROR);
+ if (!includesScheme(projectEndpoint)) {
+ throw new IllegalArgumentException(
+ String.format("Project endpoint \"%s\" must include scheme.", projectEndpoint));
+ }
+ this.projectEndpoint = projectEndpoint;
+ return this;
+ }
+
+ /** Sets the (optional) initializer to run on HTTP requests to Cloud Datastore. */
+ public Builder initializer(HttpRequestInitializer initializer) {
+ this.initializer = initializer;
+ return this;
+ }
+
+ /** Sets the Google APIs {@link Credential} used to access Cloud Datastore. */
+ public Builder credential(Credential credential) {
+ this.credential = credential;
+ return this;
+ }
+
+ /** Sets the transport used to access Cloud Datastore. */
+ public Builder transport(HttpTransport transport) {
+ this.transport = transport;
+ return this;
+ }
+
+ private static boolean includesScheme(String url) {
+ return url.startsWith("http://") || url.startsWith("https://");
+ }
+ }
+
+ public String getProjectId() {
+ return projectId;
+ }
+
+ public String getProjectEndpoint() {
+ return projectEndpoint;
+ }
+
+ public String getHost() {
+ return host;
+ }
+
+ public String getLocalHost() {
+ return localHost;
+ }
+
+ public HttpRequestInitializer getInitializer() {
+ return initializer;
+ }
+
+ public Credential getCredential() {
+ return credential;
+ }
+
+ public HttpTransport getTransport() {
+ return transport;
+ }
+}
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/QuerySplitter.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/QuerySplitter.java
new file mode 100644
index 000000000..31d1fd7d5
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/QuerySplitter.java
@@ -0,0 +1,56 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.datastore.utils;
+
+import com.google.api.core.BetaApi;
+import com.google.datastore.v1.PartitionId;
+import com.google.datastore.v1.Query;
+import com.google.protobuf.Timestamp;
+import java.util.List;
+
+/** Provides the ability to split a query into multiple shards. */
+public interface QuerySplitter {
+
+ /**
+ * Returns a list of sharded {@link Query}s for the given query.
+ *
+ * This will create up to the desired number of splits, however it may return less splits if
+ * the desired number of splits is unavailable. This will happen if the number of split points
+ * provided by the underlying Datastore is less than the desired number, which will occur if the
+ * number of results for the query is too small.
+ *
+ * @param query the query to split.
+ * @param partition the partition the query is running in.
+ * @param numSplits the desired number of splits.
+ * @param datastore the datastore to run on.
+ * @throws DatastoreException if there was a datastore error while generating query splits.
+ * @throws IllegalArgumentException if the given query or numSplits was invalid.
+ */
+ List getSplits(Query query, PartitionId partition, int numSplits, Datastore datastore)
+ throws DatastoreException;
+
+ /**
+ * Same as {@link #getSplits(Query, PartitionId, int, Datastore)} but the splits are based on
+ * {@code readTime}, and the returned sharded {@link Query}s should also be executed with {@code
+ * readTime}. Reading from a timestamp is currently a private preview feature in Datastore.
+ */
+ @BetaApi
+ default List getSplits(
+ Query query, PartitionId partition, int numSplits, Datastore datastore, Timestamp readTime)
+ throws DatastoreException {
+ throw new UnsupportedOperationException("Not implemented.");
+ }
+}
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/QuerySplitterImpl.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/QuerySplitterImpl.java
new file mode 100644
index 000000000..ac2a6557e
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/QuerySplitterImpl.java
@@ -0,0 +1,309 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.datastore.utils;
+
+import static com.google.datastore.utils.DatastoreHelper.makeAndFilter;
+
+import com.google.api.core.BetaApi;
+import com.google.datastore.v1.EntityResult;
+import com.google.datastore.v1.Filter;
+import com.google.datastore.v1.Key;
+import com.google.datastore.v1.PartitionId;
+import com.google.datastore.v1.Projection;
+import com.google.datastore.v1.PropertyFilter;
+import com.google.datastore.v1.PropertyFilter.Operator;
+import com.google.datastore.v1.PropertyOrder.Direction;
+import com.google.datastore.v1.PropertyReference;
+import com.google.datastore.v1.Query;
+import com.google.datastore.v1.QueryResultBatch;
+import com.google.datastore.v1.QueryResultBatch.MoreResultsType;
+import com.google.datastore.v1.ReadOptions;
+import com.google.datastore.v1.RunQueryRequest;
+import com.google.protobuf.Timestamp;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.EnumSet;
+import java.util.List;
+import javax.annotation.Nullable;
+
+/**
+ * Provides the ability to split a query into multiple shards using Cloud Datastore.
+ *
+ * This implementation of the QuerySplitter uses the __scatter__ property to gather random split
+ * points for a query.
+ */
+final class QuerySplitterImpl implements QuerySplitter {
+
+ /** The number of keys to sample for each split. * */
+ private static final int KEYS_PER_SPLIT = 32;
+
+ private static final EnumSet UNSUPPORTED_OPERATORS =
+ EnumSet.of(
+ Operator.LESS_THAN,
+ Operator.LESS_THAN_OR_EQUAL,
+ Operator.GREATER_THAN,
+ Operator.GREATER_THAN_OR_EQUAL);
+
+ static final QuerySplitter INSTANCE = new QuerySplitterImpl();
+
+ private QuerySplitterImpl() {
+ // No initialization required.
+ }
+
+ @Override
+ public List getSplits(
+ Query query, PartitionId partition, int numSplits, Datastore datastore)
+ throws DatastoreException, IllegalArgumentException {
+ return getSplitsInternal(query, partition, numSplits, datastore, null);
+ }
+
+ @BetaApi
+ @Override
+ public List getSplits(
+ Query query, PartitionId partition, int numSplits, Datastore datastore, Timestamp readTime)
+ throws DatastoreException, IllegalArgumentException {
+ return getSplitsInternal(query, partition, numSplits, datastore, readTime);
+ }
+
+ private List getSplitsInternal(
+ Query query,
+ PartitionId partition,
+ int numSplits,
+ Datastore datastore,
+ @Nullable Timestamp readTime)
+ throws DatastoreException, IllegalArgumentException {
+ List splits = new ArrayList(numSplits);
+ if (numSplits == 1) {
+ splits.add(query);
+ return splits;
+ }
+ validateQuery(query);
+ validateSplitSize(numSplits);
+
+ List scatterKeys = getScatterKeys(numSplits, query, partition, datastore, readTime);
+ Key lastKey = null;
+ for (Key nextKey : getSplitKey(scatterKeys, numSplits)) {
+ splits.add(createSplit(lastKey, nextKey, query));
+ lastKey = nextKey;
+ }
+ splits.add(createSplit(lastKey, null, query));
+ return splits;
+ }
+
+ /**
+ * Verify that the given number of splits is not out of bounds.
+ *
+ * @param numSplits the number of splits.
+ * @throws IllegalArgumentException if the split size is invalid.
+ */
+ private void validateSplitSize(int numSplits) throws IllegalArgumentException {
+ if (numSplits < 1) {
+ throw new IllegalArgumentException("The number of splits must be greater than 0.");
+ }
+ }
+
+ /**
+ * Validates that we only have allowable filters.
+ *
+ * Note that equality and ancestor filters are allowed, however they may result in inefficient
+ * sharding.
+ */
+ private void validateFilter(Filter filter) throws IllegalArgumentException {
+ switch (filter.getFilterTypeCase()) {
+ case COMPOSITE_FILTER:
+ for (Filter subFilter : filter.getCompositeFilter().getFiltersList()) {
+ validateFilter(subFilter);
+ }
+ break;
+ case PROPERTY_FILTER:
+ if (UNSUPPORTED_OPERATORS.contains(filter.getPropertyFilter().getOp())) {
+ throw new IllegalArgumentException("Query cannot have any inequality filters.");
+ }
+ break;
+ default:
+ throw new IllegalArgumentException(
+ "Unsupported filter type: " + filter.getFilterTypeCase());
+ }
+ }
+
+ /**
+ * Verifies that the given query can be properly scattered.
+ *
+ * @param query the query to verify
+ * @throws IllegalArgumentException if the query is invalid.
+ */
+ private void validateQuery(Query query) throws IllegalArgumentException {
+ if (query.getKindCount() != 1) {
+ throw new IllegalArgumentException("Query must have exactly one kind.");
+ }
+ if (query.getOrderCount() != 0) {
+ throw new IllegalArgumentException("Query cannot have any sort orders.");
+ }
+ if (query.hasFilter()) {
+ validateFilter(query.getFilter());
+ }
+ }
+
+ /**
+ * Create a new {@link Query} given the query and range.
+ *
+ * @param lastKey the previous key. If null then assumed to be the beginning.
+ * @param nextKey the next key. If null then assumed to be the end.
+ * @param query the desired query.
+ */
+ private Query createSplit(Key lastKey, Key nextKey, Query query) {
+ if (lastKey == null && nextKey == null) {
+ return query;
+ }
+ List keyFilters = new ArrayList();
+ if (query.hasFilter()) {
+ keyFilters.add(query.getFilter());
+ }
+ if (lastKey != null) {
+ Filter lowerBound =
+ DatastoreHelper.makeFilter(
+ DatastoreHelper.KEY_PROPERTY_NAME,
+ PropertyFilter.Operator.GREATER_THAN_OR_EQUAL,
+ DatastoreHelper.makeValue(lastKey))
+ .build();
+ keyFilters.add(lowerBound);
+ }
+ if (nextKey != null) {
+ Filter upperBound =
+ DatastoreHelper.makeFilter(
+ DatastoreHelper.KEY_PROPERTY_NAME,
+ PropertyFilter.Operator.LESS_THAN,
+ DatastoreHelper.makeValue(nextKey))
+ .build();
+ keyFilters.add(upperBound);
+ }
+ return Query.newBuilder(query).setFilter(makeAndFilter(keyFilters)).build();
+ }
+
+ /**
+ * Gets a list of split keys given a desired number of splits.
+ *
+ * This list will contain multiple split keys for each split. Only a single split key will be
+ * chosen as the split point, however providing multiple keys allows for more uniform sharding.
+ *
+ * @param numSplits the number of desired splits.
+ * @param query the user query.
+ * @param partition the partition to run the query in.
+ * @param datastore the datastore containing the data.
+ * @param readTime read time at which to get the split keys from the datastore.
+ * @throws com.google.datastore.utils.DatastoreException if there was an error when executing the
+ * datastore query.
+ */
+ private List getScatterKeys(
+ int numSplits,
+ Query query,
+ PartitionId partition,
+ Datastore datastore,
+ @Nullable Timestamp readTime)
+ throws DatastoreException {
+ Query.Builder scatterPointQuery = createScatterQuery(query, numSplits);
+
+ List keySplits = new ArrayList();
+
+ QueryResultBatch batch;
+ do {
+ RunQueryRequest.Builder scatterRequest =
+ RunQueryRequest.newBuilder().setPartitionId(partition).setQuery(scatterPointQuery);
+ scatterRequest.setProjectId(partition.getProjectId());
+ scatterRequest.setDatabaseId(partition.getDatabaseId());
+ if (readTime != null) {
+ scatterRequest.setReadOptions(ReadOptions.newBuilder().setReadTime(readTime).build());
+ }
+ batch = datastore.runQuery(scatterRequest.build()).getBatch();
+ for (EntityResult result : batch.getEntityResultsList()) {
+ keySplits.add(result.getEntity().getKey());
+ }
+ scatterPointQuery.setStartCursor(batch.getEndCursor());
+ scatterPointQuery
+ .getLimitBuilder()
+ .setValue(scatterPointQuery.getLimit().getValue() - batch.getEntityResultsCount());
+ } while (batch.getMoreResults() == MoreResultsType.NOT_FINISHED);
+ Collections.sort(keySplits, DatastoreHelper.getKeyComparator());
+ return keySplits;
+ }
+
+ /**
+ * Creates a scatter query from the given user query
+ *
+ * @param query the user's query.
+ * @param numSplits the number of splits to create.
+ */
+ private Query.Builder createScatterQuery(Query query, int numSplits) {
+ // TODO(pcostello): We can potentially support better splits with equality filters in our query
+ // if there exists a composite index on property, __scatter__, __key__. Until an API for
+ // metadata exists, this isn't possible. Note that ancestor and inequality queries fall into
+ // the same category.
+ Query.Builder scatterPointQuery = Query.newBuilder();
+ scatterPointQuery.addAllKind(query.getKindList());
+ scatterPointQuery.addOrder(
+ DatastoreHelper.makeOrder(DatastoreHelper.SCATTER_PROPERTY_NAME, Direction.ASCENDING));
+ // There is a split containing entities before and after each scatter entity:
+ // ||---*------*------*------*------*------*------*---|| = scatter entity
+ // If we represent each split as a region before a scatter entity, there is an extra region
+ // following the last scatter point. Thus, we do not need the scatter entities for the last
+ // region.
+ scatterPointQuery.getLimitBuilder().setValue((numSplits - 1) * KEYS_PER_SPLIT);
+ scatterPointQuery.addProjection(
+ Projection.newBuilder().setProperty(PropertyReference.newBuilder().setName("__key__")));
+ return scatterPointQuery;
+ }
+
+ /**
+ * Given a list of keys and a number of splits find the keys to split on.
+ *
+ * @param keys the list of keys.
+ * @param numSplits the number of splits.
+ */
+ private Iterable getSplitKey(List keys, int numSplits) {
+ // If the number of keys is less than the number of splits, we are limited in the number of
+ // splits we can make.
+ if (keys.size() < numSplits - 1) {
+ return keys;
+ }
+
+ // Calculate the number of keys per split. This should be KEYS_PER_SPLIT, but may
+ // be less if there are not KEYS_PER_SPLIT * (numSplits - 1) scatter entities.
+ //
+ // Consider the following dataset, where - represents an entity and * represents an entity
+ // that is returned as a scatter entity:
+ // ||---*-----*----*-----*-----*------*----*----||
+ // If we want 4 splits in this data, the optimal split would look like:
+ // ||---*-----*----*-----*-----*------*----*----||
+ // | | |
+ // The scatter keys in the last region are not useful to us, so we never request them:
+ // ||---*-----*----*-----*-----*------*---------||
+ // | | |
+ // With 6 scatter keys we want to set scatter points at indexes: 1, 3, 5.
+ //
+ // We keep this as a double so that any "fractional" keys per split get distributed throughout
+ // the splits and don't make the last split significantly larger than the rest.
+ double numKeysPerSplit = Math.max(1.0, ((double) keys.size()) / (numSplits - 1));
+
+ List keysList = new ArrayList(numSplits - 1);
+ // Grab the last sample for each split, otherwise the first split will be too small.
+ for (int i = 1; i < numSplits; i++) {
+ int splitIndex = (int) Math.round(i * numKeysPerSplit) - 1;
+ keysList.add(keys.get(splitIndex));
+ }
+
+ return keysList;
+ }
+}
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/RemoteRpc.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/RemoteRpc.java
new file mode 100644
index 000000000..492936e15
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/RemoteRpc.java
@@ -0,0 +1,239 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.datastore.utils;
+
+import com.google.api.client.http.*;
+import com.google.api.client.http.protobuf.ProtoHttpContent;
+import com.google.api.client.util.IOUtils;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Strings;
+import com.google.protobuf.MessageLite;
+import com.google.rpc.Code;
+import com.google.rpc.Status;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.net.SocketTimeoutException;
+import java.nio.charset.Charset;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.logging.Logger;
+
+/**
+ * An RPC transport that sends protocol buffers over HTTP.
+ *
+ * This class is thread-safe.
+ */
+class RemoteRpc {
+ private static final Logger logger = Logger.getLogger(RemoteRpc.class.getName());
+
+ @VisibleForTesting static final String API_FORMAT_VERSION_HEADER = "X-Goog-Api-Format-Version";
+ private static final String API_FORMAT_VERSION = "2";
+
+ @VisibleForTesting static final String X_GOOG_REQUEST_PARAMS_HEADER = "x-goog-request-params";
+
+ private final HttpRequestFactory client;
+ private final HttpRequestInitializer initializer;
+ private final String url;
+ private final AtomicInteger rpcCount = new AtomicInteger(0);
+ // Not final - so it can be set/reset in Unittests
+ private static boolean enableE2EChecksum =
+ Boolean.parseBoolean(System.getenv("GOOGLE_CLOUD_DATASTORE_HTTP_ENABLE_E2E_CHECKSUM"));
+
+ RemoteRpc(HttpRequestFactory client, HttpRequestInitializer initializer, String url) {
+ this.client = client;
+ this.initializer = initializer;
+ this.url = url;
+ try {
+ resolveURL("dummyRpc");
+ } catch (Exception e) {
+ throw new IllegalArgumentException(
+ "Unable to construct RemoteRpc due to unsupported url: <" + url + ">", e);
+ }
+ }
+
+ /**
+ * Makes an RPC call using the client. Logs how long it took and any exceptions.
+ *
+ *
NOTE: The request could be an InputStream too, but the http client will need to find its
+ * length, which will require buffering it anyways.
+ *
+ * @throws com.google.datastore.utils.DatastoreException if the RPC fails.
+ */
+ public InputStream call(
+ String methodName, MessageLite request, String projectId, String databaseId)
+ throws com.google.datastore.utils.DatastoreException {
+ logger.fine("remote datastore call " + methodName);
+
+ long startTime = System.currentTimeMillis();
+ try {
+ HttpResponse httpResponse;
+ try {
+ rpcCount.incrementAndGet();
+ ProtoHttpContent payload = new ProtoHttpContent(request);
+ HttpRequest httpRequest = client.buildPostRequest(resolveURL(methodName), payload);
+ setHeaders(request, httpRequest, projectId, databaseId);
+ // Don't throw an HTTPResponseException on error. It converts the response to a String and
+ // throws away the original, whereas we need the raw bytes to parse it as a proto.
+ httpRequest.setThrowExceptionOnExecuteError(false);
+ // Datastore requests typically time out after 60s; set the read timeout to slightly longer
+ // than that by default (can be overridden via the HttpRequestInitializer).
+ httpRequest.setReadTimeout(65 * 1000);
+ if (initializer != null) {
+ initializer.initialize(httpRequest);
+ }
+ httpResponse = httpRequest.execute();
+ if (!httpResponse.isSuccessStatusCode()) {
+ try (InputStream content = httpResponse.getContent()) {
+ throw makeException(
+ url,
+ methodName,
+ content,
+ httpResponse.getContentType(),
+ httpResponse.getContentCharset(),
+ null,
+ httpResponse.getStatusCode());
+ }
+ }
+ InputStream inputStream = httpResponse.getContent();
+ return inputStream;
+ } catch (SocketTimeoutException e) {
+ throw makeException(url, methodName, Code.DEADLINE_EXCEEDED, "Deadline exceeded", e);
+ } catch (IOException e) {
+ throw makeException(url, methodName, Code.UNAVAILABLE, "I/O error", e);
+ }
+ } finally {
+ long elapsedTime = System.currentTimeMillis() - startTime;
+ logger.fine("remote datastore call " + methodName + " took " + elapsedTime + " ms");
+ }
+ }
+
+ @VisibleForTesting
+ void setHeaders(
+ MessageLite request, HttpRequest httpRequest, String projectId, String databaseId) {
+ httpRequest.getHeaders().put(API_FORMAT_VERSION_HEADER, API_FORMAT_VERSION);
+ StringBuilder builder = new StringBuilder("project_id=");
+ builder.append(projectId);
+ if (!Strings.isNullOrEmpty(databaseId)) {
+ builder.append("&database_id=");
+ builder.append(databaseId);
+ }
+ httpRequest.getHeaders().put(X_GOOG_REQUEST_PARAMS_HEADER, builder.toString());
+ }
+
+ @VisibleForTesting
+ HttpRequestFactory getClient() {
+ return client;
+ }
+
+ @VisibleForTesting
+ static void setSystemEnvE2EChecksum(boolean enableE2EChecksum) {
+ RemoteRpc.enableE2EChecksum = enableE2EChecksum;
+ }
+
+ void resetRpcCount() {
+ rpcCount.set(0);
+ }
+
+ int getRpcCount() {
+ return rpcCount.get();
+ }
+
+ public String getUrl() {
+ return url;
+ }
+
+ GenericUrl resolveURL(String path) {
+ return new GenericUrl(url + ":" + path);
+ }
+
+ HttpRequestFactory getHttpRequestFactory() {
+ return client;
+ }
+
+ public static com.google.datastore.utils.DatastoreException makeException(
+ String url, String methodName, Code code, String message, Throwable cause) {
+ logger.fine("remote datastore call " + methodName + " against " + url + " failed: " + message);
+ return new com.google.datastore.utils.DatastoreException(methodName, code, message, cause);
+ }
+
+ static DatastoreException makeException(
+ String url,
+ String methodName,
+ InputStream content,
+ String contentType,
+ Charset contentCharset,
+ Throwable cause,
+ int httpStatusCode) {
+ if (!contentType.equals("application/x-protobuf")) {
+ String responseContent;
+ try {
+ ByteArrayOutputStream out = new ByteArrayOutputStream();
+ IOUtils.copy(content, out, false);
+ responseContent = out.toString(contentCharset.name());
+ } catch (IOException e) {
+ responseContent = "";
+ }
+ return makeException(
+ url,
+ methodName,
+ Code.INTERNAL,
+ String.format(
+ "Non-protobuf error: %s. HTTP status code was %d.", responseContent, httpStatusCode),
+ cause);
+ }
+
+ Status rpcStatus;
+ try {
+ rpcStatus = Status.parseFrom(content);
+ } catch (IOException e) {
+ return makeException(
+ url,
+ methodName,
+ Code.INTERNAL,
+ String.format(
+ "Unable to parse Status protocol buffer: HTTP status code was %s.", httpStatusCode),
+ e);
+ }
+
+ Code code = Code.forNumber(rpcStatus.getCode());
+ if (code == null) {
+ return makeException(
+ url,
+ methodName,
+ Code.INTERNAL,
+ String.format(
+ "Invalid error code: %d. Message: %s.", rpcStatus.getCode(), rpcStatus.getMessage()),
+ cause);
+ } else if (code == Code.OK) {
+ // We can end up here because there was no response body (and we successfully parsed an
+ // empty Status message). This may happen for 401s in particular due to special handling
+ // in low-level HTTP libraries.
+ if (httpStatusCode == HttpStatusCodes.STATUS_CODE_UNAUTHORIZED) {
+ return makeException(url, methodName, Code.UNAUTHENTICATED, "Unauthenticated.", cause);
+ }
+ return makeException(
+ url,
+ methodName,
+ Code.INTERNAL,
+ String.format(
+ "Unexpected OK error code with HTTP status code of %d. Message: %s.",
+ httpStatusCode, rpcStatus.getMessage()),
+ cause);
+ }
+
+ return makeException(url, methodName, code, rpcStatus.getMessage(), cause);
+ }
+}
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/testing/MockCredential.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/testing/MockCredential.java
new file mode 100644
index 000000000..d5d16bb65
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/testing/MockCredential.java
@@ -0,0 +1,36 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.datastore.utils.testing;
+
+import com.google.api.client.auth.oauth2.Credential;
+import com.google.api.client.http.HttpRequest;
+import java.io.IOException;
+
+/** Fake credential used for testing purpose. */
+public class MockCredential extends Credential {
+ public MockCredential() {
+ super(
+ new AccessMethod() {
+ @Override
+ public void intercept(HttpRequest request, String accessToken) throws IOException {}
+
+ @Override
+ public String getAccessTokenFromRequest(HttpRequest request) {
+ return "MockAccessToken";
+ }
+ });
+ }
+}
diff --git a/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/testing/MockDatastoreFactory.java b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/testing/MockDatastoreFactory.java
new file mode 100644
index 000000000..d4dd5caef
--- /dev/null
+++ b/google-cloud-datastore-utils/src/main/java/com/google/datastore/utils/testing/MockDatastoreFactory.java
@@ -0,0 +1,132 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.datastore.utils.testing;
+
+import static com.google.common.base.Preconditions.checkState;
+
+import com.google.api.client.auth.oauth2.Credential;
+import com.google.api.client.http.*;
+import com.google.api.client.testing.http.MockHttpTransport;
+import com.google.api.client.testing.http.MockLowLevelHttpRequest;
+import com.google.api.client.testing.http.MockLowLevelHttpResponse;
+import com.google.api.client.testing.util.TestableByteArrayInputStream;
+import com.google.common.collect.Iterables;
+import com.google.datastore.utils.DatastoreFactory;
+import com.google.datastore.utils.DatastoreOptions;
+import com.google.protobuf.Message;
+import com.google.rpc.Code;
+import com.google.rpc.Status;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.util.List;
+
+/** Fake Datastore factory used for testing purposes when a true Datastore service is not needed. */
+public class MockDatastoreFactory extends DatastoreFactory {
+ private int nextStatus;
+ private Message nextResponse;
+ private Status nextError;
+ private IOException nextException;
+
+ private String lastPath;
+ private String lastMimeType;
+ private byte[] lastBody;
+ private List lastCookies;
+ private String lastApiFormatHeaderValue;
+
+ public void setNextResponse(Message response) {
+ nextStatus = HttpStatusCodes.STATUS_CODE_OK;
+ nextResponse = response;
+ nextError = null;
+ nextException = null;
+ }
+
+ public void setNextError(int status, Code code, String message) {
+ nextStatus = status;
+ nextResponse = null;
+ nextError = makeErrorContent(message, code);
+ nextException = null;
+ }
+
+ public void setNextException(IOException exception) {
+ nextStatus = 0;
+ nextResponse = null;
+ nextError = null;
+ nextException = exception;
+ }
+
+ @Override
+ public HttpRequestFactory makeClient(DatastoreOptions options) {
+ HttpTransport transport =
+ new MockHttpTransport() {
+ @Override
+ public LowLevelHttpRequest buildRequest(String method, String url) {
+ return new MockLowLevelHttpRequest(url) {
+ @Override
+ public LowLevelHttpResponse execute() throws IOException {
+ lastPath = new GenericUrl(getUrl()).getRawPath();
+ lastMimeType = getContentType();
+ lastCookies = getHeaderValues("Cookie");
+ lastApiFormatHeaderValue =
+ Iterables.getOnlyElement(getHeaderValues("X-Goog-Api-Format-Version"));
+ ByteArrayOutputStream out = new ByteArrayOutputStream();
+ getStreamingContent().writeTo(out);
+ lastBody = out.toByteArray();
+ if (nextException != null) {
+ throw nextException;
+ }
+ MockLowLevelHttpResponse response =
+ new MockLowLevelHttpResponse()
+ .setStatusCode(nextStatus)
+ .setContentType("application/x-protobuf");
+ if (nextError != null) {
+ checkState(nextResponse == null);
+ response.setContent(new TestableByteArrayInputStream(nextError.toByteArray()));
+ } else {
+ response.setContent(new TestableByteArrayInputStream(nextResponse.toByteArray()));
+ }
+ return response;
+ }
+ };
+ }
+ };
+ Credential credential = options.getCredential();
+ return transport.createRequestFactory(credential);
+ }
+
+ public String getLastPath() {
+ return lastPath;
+ }
+
+ public String getLastMimeType() {
+ return lastMimeType;
+ }
+
+ public String getLastApiFormatHeaderValue() {
+ return lastApiFormatHeaderValue;
+ }
+
+ public byte[] getLastBody() {
+ return lastBody;
+ }
+
+ public List getLastCookies() {
+ return lastCookies;
+ }
+
+ private static Status makeErrorContent(String message, Code code) {
+ return Status.newBuilder().setCode(code.getNumber()).setMessage(message).build();
+ }
+}
diff --git a/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreClientTest.java b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreClientTest.java
new file mode 100644
index 000000000..31b0f6440
--- /dev/null
+++ b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreClientTest.java
@@ -0,0 +1,407 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.datastore.utils;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.fail;
+
+import com.google.api.client.http.HttpRequest;
+import com.google.api.client.http.HttpRequestInitializer;
+import com.google.datastore.utils.testing.MockCredential;
+import com.google.datastore.utils.testing.MockDatastoreFactory;
+import com.google.datastore.v1.AllocateIdsRequest;
+import com.google.datastore.v1.AllocateIdsResponse;
+import com.google.datastore.v1.BeginTransactionRequest;
+import com.google.datastore.v1.BeginTransactionResponse;
+import com.google.datastore.v1.CommitRequest;
+import com.google.datastore.v1.CommitResponse;
+import com.google.datastore.v1.EntityResult;
+import com.google.datastore.v1.LookupRequest;
+import com.google.datastore.v1.LookupResponse;
+import com.google.datastore.v1.QueryResultBatch;
+import com.google.datastore.v1.ReserveIdsRequest;
+import com.google.datastore.v1.ReserveIdsResponse;
+import com.google.datastore.v1.RollbackRequest;
+import com.google.datastore.v1.RollbackResponse;
+import com.google.datastore.v1.RunAggregationQueryRequest;
+import com.google.datastore.v1.RunAggregationQueryResponse;
+import com.google.datastore.v1.RunQueryRequest;
+import com.google.datastore.v1.RunQueryResponse;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Message;
+import com.google.rpc.Code;
+import java.io.IOException;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.net.SocketTimeoutException;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link DatastoreFactory} and {@link Datastore}. */
+@RunWith(JUnit4.class)
+public class DatastoreClientTest {
+ private static final String PROJECT_ID = "project-id";
+
+ private DatastoreFactory factory = new MockDatastoreFactory();
+ private DatastoreOptions.Builder options =
+ new DatastoreOptions.Builder().projectId(PROJECT_ID).credential(new MockCredential());
+
+ @Test
+ public void options_NoProjectIdOrProjectEndpoint() {
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> factory.create(new DatastoreOptions.Builder().build()));
+ assertThat(exception)
+ .hasMessageThat()
+ .contains("Either project ID or project endpoint must be provided");
+ factory.create(options.build());
+ }
+
+ @Test
+ public void options_ProjectIdAndProjectEndpoint() throws Exception {
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ new DatastoreOptions.Builder()
+ .projectId(PROJECT_ID)
+ .projectEndpoint(
+ "http://localhost:1234/datastore/v1beta42/projects/project-id"));
+ assertThat(exception)
+ .hasMessageThat()
+ .contains("Cannot set both project endpoint and project ID");
+ }
+
+ @Test
+ public void options_LocalHostAndProjectEndpoint() throws Exception {
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ new DatastoreOptions.Builder()
+ .localHost("localhost:8080")
+ .projectEndpoint(
+ "http://localhost:1234/datastore/v1beta42/projects/project-id"));
+ assertThat(exception)
+ .hasMessageThat()
+ .contains("Can set at most one of project endpoint, host, and local host");
+ }
+
+ @Test
+ public void options_HostAndProjectEndpoint() throws Exception {
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ new DatastoreOptions.Builder()
+ .host("foo-datastore.googleapis.com")
+ .projectEndpoint(
+ "http://localhost:1234/datastore/v1beta42/projects/project-id"));
+ assertThat(exception)
+ .hasMessageThat()
+ .contains("Can set at most one of project endpoint, host, and local host");
+ }
+
+ @Test
+ public void options_HostAndLocalHost() throws Exception {
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ new DatastoreOptions.Builder()
+ .host("foo-datastore.googleapis.com")
+ .localHost("localhost:8080"));
+ assertThat(exception)
+ .hasMessageThat()
+ .contains("Can set at most one of project endpoint, host, and local host");
+ }
+
+ @Test
+ public void options_InvalidLocalHost() throws Exception {
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ factory.create(
+ new DatastoreOptions.Builder()
+ .projectId(PROJECT_ID)
+ .localHost("!not a valid url!")
+ .build()));
+ assertThat(exception).hasMessageThat().contains("Illegal character");
+ }
+
+ @Test
+ public void options_SchemeInLocalHost() {
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> new DatastoreOptions.Builder().localHost("http://localhost:8080"));
+ assertThat(exception)
+ .hasMessageThat()
+ .contains("Local host \"http://localhost:8080\" must not include scheme");
+ }
+
+ @Test
+ public void options_InvalidHost() {
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ factory.create(
+ new DatastoreOptions.Builder()
+ .projectId(PROJECT_ID)
+ .host("!not a valid url!")
+ .build()));
+ assertThat(exception).hasMessageThat().contains("Illegal character");
+ }
+
+ @Test
+ public void options_SchemeInHost() {
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> new DatastoreOptions.Builder().host("http://foo-datastore.googleapis.com"));
+
+ assertThat(exception)
+ .hasMessageThat()
+ .contains("Host \"http://foo-datastore.googleapis.com\" must not include scheme.");
+ }
+
+ @Test
+ public void create_NullOptions() throws Exception {
+ assertThrows(NullPointerException.class, () -> factory.create(null));
+ }
+
+ @Test
+ public void create_Host() {
+ Datastore datastore =
+ factory.create(
+ new DatastoreOptions.Builder()
+ .projectId(PROJECT_ID)
+ .host("foo-datastore.googleapis.com")
+ .build());
+ assertThat(datastore.remoteRpc.getUrl())
+ .isEqualTo("https://foo-datastore.googleapis.com/v1/projects/project-id");
+ }
+
+ @Test
+ public void create_LocalHost() {
+ Datastore datastore =
+ factory.create(
+ new DatastoreOptions.Builder()
+ .projectId(PROJECT_ID)
+ .localHost("localhost:8080")
+ .build());
+ assertThat(datastore.remoteRpc.getUrl())
+ .isEqualTo("http://localhost:8080/v1/projects/project-id");
+ }
+
+ @Test
+ public void create_LocalHostIp() {
+ Datastore datastore =
+ factory.create(
+ new DatastoreOptions.Builder()
+ .projectId(PROJECT_ID)
+ .localHost("127.0.0.1:8080")
+ .build());
+ assertThat(datastore.remoteRpc.getUrl())
+ .isEqualTo("http://127.0.0.1:8080/v1/projects/project-id");
+ }
+
+ @Test
+ public void create_DefaultHost() {
+ Datastore datastore =
+ factory.create(new DatastoreOptions.Builder().projectId(PROJECT_ID).build());
+ assertThat(datastore.remoteRpc.getUrl())
+ .isEqualTo("https://datastore.googleapis.com/v1/projects/project-id");
+ }
+
+ @Test
+ public void create_ProjectEndpoint() {
+ Datastore datastore =
+ factory.create(
+ new DatastoreOptions.Builder()
+ .projectEndpoint("http://prom-qa/datastore/v1beta42/projects/project-id")
+ .build());
+ assertThat(datastore.remoteRpc.getUrl())
+ .isEqualTo("http://prom-qa/datastore/v1beta42/projects/project-id");
+ }
+
+ @Test
+ public void create_ProjectEndpointNoScheme() {
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ factory.create(
+ new DatastoreOptions.Builder()
+ .projectEndpoint("localhost:1234/datastore/v1beta42/projects/project-id")
+ .build()));
+ assertThat(exception)
+ .hasMessageThat()
+ .contains(
+ "Project endpoint \"localhost:1234/datastore/v1beta42/projects/project-id\" must"
+ + " include scheme.");
+ }
+
+ @Test
+ public void initializer() throws Exception {
+ options.initializer(
+ new HttpRequestInitializer() {
+ @Override
+ public void initialize(HttpRequest request) {
+ request.getHeaders().setCookie("magic");
+ }
+ });
+ Datastore datastore = factory.create(options.build());
+ MockDatastoreFactory mockClient = (MockDatastoreFactory) factory;
+ AllocateIdsRequest request = AllocateIdsRequest.newBuilder().build();
+ AllocateIdsResponse response = AllocateIdsResponse.newBuilder().build();
+ mockClient.setNextResponse(response);
+ assertEquals(response, datastore.allocateIds(request));
+ assertEquals("magic", mockClient.getLastCookies().get(0));
+ }
+
+ @Test
+ public void allocateIds() throws Exception {
+ AllocateIdsRequest.Builder request = AllocateIdsRequest.newBuilder();
+ AllocateIdsResponse.Builder response = AllocateIdsResponse.newBuilder();
+ expectRpc("allocateIds", request.build(), response.build());
+ }
+
+ @Test
+ public void lookup() throws Exception {
+ LookupRequest.Builder request = LookupRequest.newBuilder();
+ LookupResponse.Builder response = LookupResponse.newBuilder();
+ expectRpc("lookup", request.build(), response.build());
+ }
+
+ @Test
+ public void beginTransaction() throws Exception {
+ BeginTransactionRequest.Builder request = BeginTransactionRequest.newBuilder();
+ BeginTransactionResponse.Builder response = BeginTransactionResponse.newBuilder();
+ response.setTransaction(ByteString.copyFromUtf8("project-id"));
+ expectRpc("beginTransaction", request.build(), response.build());
+ }
+
+ @Test
+ public void commit() throws Exception {
+ CommitRequest.Builder request = CommitRequest.newBuilder();
+ request.setTransaction(ByteString.copyFromUtf8("project-id"));
+ CommitResponse.Builder response = CommitResponse.newBuilder();
+ expectRpc("commit", request.build(), response.build());
+ }
+
+ @Test
+ public void reserveIds() throws Exception {
+ ReserveIdsRequest.Builder request = ReserveIdsRequest.newBuilder();
+ ReserveIdsResponse.Builder response = ReserveIdsResponse.newBuilder();
+ expectRpc("reserveIds", request.build(), response.build());
+ }
+
+ @Test
+ public void rollback() throws Exception {
+ RollbackRequest.Builder request = RollbackRequest.newBuilder();
+ request.setTransaction(ByteString.copyFromUtf8("project-id"));
+ RollbackResponse.Builder response = RollbackResponse.newBuilder();
+ expectRpc("rollback", request.build(), response.build());
+ }
+
+ @Test
+ public void runQuery() throws Exception {
+ RunQueryRequest.Builder request = RunQueryRequest.newBuilder();
+ request.getQueryBuilder();
+ RunQueryResponse.Builder response = RunQueryResponse.newBuilder();
+ response
+ .getBatchBuilder()
+ .setEntityResultType(EntityResult.ResultType.FULL)
+ .setMoreResults(QueryResultBatch.MoreResultsType.NOT_FINISHED);
+ expectRpc("runQuery", request.build(), response.build());
+ }
+
+ @Test
+ public void runAggregationQuery() throws Exception {
+ RunAggregationQueryRequest.Builder request = RunAggregationQueryRequest.newBuilder();
+ RunAggregationQueryResponse.Builder response = RunAggregationQueryResponse.newBuilder();
+ expectRpc("runAggregationQuery", request.build(), response.build());
+ }
+
+ private void expectRpc(String methodName, Message request, Message response) throws Exception {
+ Datastore datastore = factory.create(options.build());
+ MockDatastoreFactory mockClient = (MockDatastoreFactory) factory;
+
+ mockClient.setNextResponse(response);
+ @SuppressWarnings("rawtypes")
+ Class[] methodArgs = {request.getClass()};
+ Method call = Datastore.class.getMethod(methodName, methodArgs);
+ Object[] callArgs = {request};
+ assertEquals(response, call.invoke(datastore, callArgs));
+
+ assertEquals("/v1/projects/project-id:" + methodName, mockClient.getLastPath());
+ assertEquals("application/x-protobuf", mockClient.getLastMimeType());
+ assertEquals("2", mockClient.getLastApiFormatHeaderValue());
+ assertArrayEquals(request.toByteArray(), mockClient.getLastBody());
+ assertEquals(1, datastore.getRpcCount());
+
+ datastore.resetRpcCount();
+ assertEquals(0, datastore.getRpcCount());
+
+ mockClient.setNextError(400, Code.INVALID_ARGUMENT, "oops");
+ try {
+ call.invoke(datastore, callArgs);
+ fail();
+ } catch (InvocationTargetException targetException) {
+ DatastoreException exception = (DatastoreException) targetException.getCause();
+ assertEquals(Code.INVALID_ARGUMENT, exception.getCode());
+ assertEquals(methodName, exception.getMethodName());
+ assertEquals("oops", exception.getMessage());
+ }
+
+ SocketTimeoutException socketTimeoutException = new SocketTimeoutException("ste");
+ mockClient.setNextException(socketTimeoutException);
+ try {
+ call.invoke(datastore, callArgs);
+ fail();
+ } catch (InvocationTargetException targetException) {
+ DatastoreException exception = (DatastoreException) targetException.getCause();
+ assertEquals(Code.DEADLINE_EXCEEDED, exception.getCode());
+ assertEquals(methodName, exception.getMethodName());
+ assertEquals("Deadline exceeded", exception.getMessage());
+ assertSame(socketTimeoutException, exception.getCause());
+ }
+
+ IOException ioException = new IOException("ioe");
+ mockClient.setNextException(ioException);
+ try {
+ call.invoke(datastore, callArgs);
+ fail();
+ } catch (InvocationTargetException targetException) {
+ DatastoreException exception = (DatastoreException) targetException.getCause();
+ assertEquals(Code.UNAVAILABLE, exception.getCode());
+ assertEquals(methodName, exception.getMethodName());
+ assertEquals("I/O error", exception.getMessage());
+ assertSame(ioException, exception.getCause());
+ }
+
+ assertEquals(3, datastore.getRpcCount());
+ }
+}
diff --git a/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreFactoryTest.java b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreFactoryTest.java
new file mode 100644
index 000000000..2a3d5a38f
--- /dev/null
+++ b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreFactoryTest.java
@@ -0,0 +1,93 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.datastore.utils;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertTrue;
+
+import com.google.api.client.googleapis.auth.oauth2.GoogleCredential;
+import com.google.api.client.http.HttpRequestFactory;
+import com.google.api.client.http.javanet.NetHttpTransport;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Test for {@link DatastoreFactory}. */
+@RunWith(JUnit4.class)
+public class DatastoreFactoryTest {
+ private static final String PROJECT_ID = "project-id";
+
+ private DatastoreFactory factory = DatastoreFactory.get();
+
+ /**
+ * Without specifying a credential or transport, the factory will create a default transport on
+ * its own.
+ */
+ @Test
+ public void makeClient_Default() {
+ DatastoreOptions options = new DatastoreOptions.Builder().projectId(PROJECT_ID).build();
+ HttpRequestFactory f = factory.makeClient(options);
+ assertNotNull(f.getTransport());
+ assertTrue(f.getTransport() instanceof NetHttpTransport);
+ }
+
+ /**
+ * Specifying a credential, but not a transport, the factory will use the transport from the
+ * credential.
+ */
+ @Test
+ public void makeClient_WithCredential() {
+ NetHttpTransport transport = new NetHttpTransport();
+ GoogleCredential credential = new GoogleCredential.Builder().setTransport(transport).build();
+ DatastoreOptions options =
+ new DatastoreOptions.Builder().projectId(PROJECT_ID).credential(credential).build();
+ HttpRequestFactory f = factory.makeClient(options);
+ assertEquals(transport, f.getTransport());
+ }
+
+ /** Specifying a transport, but not a credential, the factory will use the transport specified. */
+ @Test
+ public void makeClient_WithTransport() {
+ NetHttpTransport transport = new NetHttpTransport();
+ DatastoreOptions options =
+ new DatastoreOptions.Builder().projectId(PROJECT_ID).transport(transport).build();
+ HttpRequestFactory f = factory.makeClient(options);
+ assertEquals(transport, f.getTransport());
+ }
+
+ /**
+ * Specifying both credential and transport, the factory will use the transport specified and not
+ * the one in the credential.
+ */
+ @Test
+ public void makeClient_WithCredentialTransport() {
+ NetHttpTransport credTransport = new NetHttpTransport();
+ NetHttpTransport transport = new NetHttpTransport();
+ GoogleCredential credential =
+ new GoogleCredential.Builder().setTransport(credTransport).build();
+ DatastoreOptions options =
+ new DatastoreOptions.Builder()
+ .projectId(PROJECT_ID)
+ .credential(credential)
+ .transport(transport)
+ .build();
+ HttpRequestFactory f = factory.makeClient(options);
+ assertNotSame(credTransport, f.getTransport());
+ assertEquals(transport, f.getTransport());
+ }
+}
diff --git a/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreHelperTest.java b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreHelperTest.java
new file mode 100644
index 000000000..246202444
--- /dev/null
+++ b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/DatastoreHelperTest.java
@@ -0,0 +1,317 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.datastore.utils;
+
+import static com.google.datastore.utils.DatastoreHelper.*;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+import com.google.datastore.v1.Key;
+import com.google.datastore.v1.PartitionId;
+import com.google.datastore.v1.Value;
+import com.google.datastore.v1.Value.ValueTypeCase;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Timestamp;
+import java.util.Date;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link DatastoreHelper}. */
+@RunWith(JUnit4.class)
+public class DatastoreHelperTest {
+
+ private static final Key PARENT =
+ Key.newBuilder().addPath(Key.PathElement.newBuilder().setKind("Parent").setId(23L)).build();
+ private static final Key GRANDPARENT =
+ Key.newBuilder()
+ .addPath(Key.PathElement.newBuilder().setKind("Grandparent").setId(24L))
+ .build();
+ private static final Key CHILD =
+ Key.newBuilder().addPath(Key.PathElement.newBuilder().setKind("Child").setId(26L)).build();
+
+ @Test
+ public void testMakeKey_BadTypeForKind() {
+ try {
+ DatastoreHelper.makeKey(new Object());
+ fail("Expected IllegalArgumentException");
+ } catch (IllegalArgumentException expected) {
+ }
+ }
+
+ @Test
+ public void testMakeKey_BadTypeForNameId() {
+ try {
+ DatastoreHelper.makeKey("kind", new Object());
+ fail("Expected IllegalArgumentException");
+ } catch (IllegalArgumentException expected) {
+ }
+ }
+
+ @Test
+ public void testMakeKey_Empty() {
+ assertEquals(Key.newBuilder().build(), DatastoreHelper.makeKey().build());
+ }
+
+ @Test
+ public void testMakeKey_Incomplete() {
+ assertEquals(
+ Key.newBuilder().addPath(Key.PathElement.newBuilder().setKind("Foo")).build(),
+ makeKey("Foo").build());
+ }
+
+ @Test
+ public void testMakeKey_IdInt() {
+ assertEquals(
+ Key.newBuilder().addPath(Key.PathElement.newBuilder().setKind("Foo").setId(1)).build(),
+ makeKey("Foo", 1).build());
+ }
+
+ @Test
+ public void testMakeKey_IdLong() {
+ assertEquals(
+ Key.newBuilder().addPath(Key.PathElement.newBuilder().setKind("Foo").setId(1)).build(),
+ makeKey("Foo", 1L).build());
+ }
+
+ @Test
+ public void testMakeKey_IdShort() {
+ assertEquals(
+ Key.newBuilder().addPath(Key.PathElement.newBuilder().setKind("Foo").setId(1)).build(),
+ makeKey("Foo", (short) 1).build());
+ }
+
+ @Test
+ public void testMakeKey_Name() {
+ assertEquals(
+ Key.newBuilder().addPath(Key.PathElement.newBuilder().setKind("Foo").setName("hi")).build(),
+ makeKey("Foo", "hi").build());
+ }
+
+ @Test
+ public void testMakeKey_KindNameKind() {
+ assertEquals(
+ Key.newBuilder()
+ .addPath(Key.PathElement.newBuilder().setKind("Foo").setName("hi"))
+ .addPath(Key.PathElement.newBuilder().setKind("Bar"))
+ .build(),
+ makeKey("Foo", "hi", "Bar").build());
+ }
+
+ @Test
+ public void testMakeKey_KeyKind() {
+ // 1 key at the beginning of the series
+ assertEquals(
+ Key.newBuilder()
+ .addPath(PARENT.getPath(0))
+ .addPath(Key.PathElement.newBuilder().setKind("Child"))
+ .build(),
+ makeKey(PARENT, "Child").build());
+ }
+
+ @Test
+ public void testMakeKey_KindIdKeyKind() {
+ // 1 key in the middle of the series
+ assertEquals(
+ Key.newBuilder()
+ .addPath(Key.PathElement.newBuilder().setKind("Grandparent").setId(24L))
+ .addPath(PARENT.getPath(0))
+ .addPath(Key.PathElement.newBuilder().setKind("Child"))
+ .build(),
+ makeKey("Grandparent", 24L, PARENT, "Child").build());
+ }
+
+ @Test
+ public void testMakeKey_KindIdKey() {
+ // 1 key at the end of the series
+ assertEquals(
+ Key.newBuilder()
+ .addPath(Key.PathElement.newBuilder().setKind("Grandparent").setId(24L))
+ .addPath(PARENT.getPath(0))
+ .build(),
+ makeKey("Grandparent", 24L, PARENT).build());
+ }
+
+ @Test
+ public void testMakeKey_KeyKindIdKey() {
+ // 1 key at the beginning and 1 key at the end of the series
+ assertEquals(
+ Key.newBuilder()
+ .addPath(GRANDPARENT.getPath(0))
+ .addPath(Key.PathElement.newBuilder().setKind("Parent").setId(23L))
+ .addPath(CHILD.getPath(0))
+ .build(),
+ makeKey(GRANDPARENT, "Parent", 23, CHILD).build());
+ }
+
+ @Test
+ public void testMakeKey_Key() {
+ // Just 1 key
+ assertEquals(Key.newBuilder().addPath(CHILD.getPath(0)).build(), makeKey(CHILD).build());
+ }
+
+ @Test
+ public void testMakeKey_KeyKey() {
+ // Just 2 keys
+ assertEquals(
+ Key.newBuilder().addPath(PARENT.getPath(0)).addPath(CHILD.getPath(0)).build(),
+ makeKey(PARENT, CHILD).build());
+ }
+
+ @Test
+ public void testMakeKey_KeyKeyKey() {
+ // Just 3 keys
+ assertEquals(
+ Key.newBuilder()
+ .addPath(GRANDPARENT.getPath(0))
+ .addPath(PARENT.getPath(0))
+ .addPath(CHILD.getPath(0))
+ .build(),
+ makeKey(GRANDPARENT, PARENT, CHILD).build());
+ }
+
+ @Test
+ public void testMakeKey_KeyMultiLevelKey() {
+ // 1 key with 2 elements
+ assertEquals(
+ Key.newBuilder()
+ .addPath(GRANDPARENT.getPath(0))
+ .addPath(PARENT.getPath(0))
+ .addPath(CHILD.getPath(0))
+ .build(),
+ makeKey(GRANDPARENT, makeKey(PARENT, CHILD).build()).build());
+ }
+
+ @Test
+ public void testMakeKey_MultiLevelKeyKey() {
+ // 1 key with 2 elements
+ assertEquals(
+ Key.newBuilder()
+ .addPath(GRANDPARENT.getPath(0))
+ .addPath(PARENT.getPath(0))
+ .addPath(CHILD.getPath(0))
+ .build(),
+ makeKey(makeKey(GRANDPARENT, PARENT).build(), CHILD).build());
+ }
+
+ @Test
+ public void testMakeKey_MultiLevelKey() {
+ // 1 key with 3 elements
+ assertEquals(
+ Key.newBuilder()
+ .addPath(GRANDPARENT.getPath(0))
+ .addPath(PARENT.getPath(0))
+ .addPath(CHILD.getPath(0))
+ .build(),
+ makeKey(makeKey(GRANDPARENT, PARENT, CHILD).build()).build());
+ }
+
+ @Test
+ public void testMakeKey_PartitionId() {
+ PartitionId partitionId = PartitionId.newBuilder().setNamespaceId("namespace-id").build();
+ Key parent = PARENT.toBuilder().setPartitionId(partitionId).build();
+ assertEquals(
+ Key.newBuilder()
+ .setPartitionId(partitionId)
+ .addPath(PARENT.getPath(0))
+ .addPath(Key.PathElement.newBuilder().setKind("Child"))
+ .build(),
+ makeKey(parent, "Child").build());
+ }
+
+ @Test
+ public void testMakeKey_NonMatchingPartitionId2() {
+ PartitionId partitionId1 = PartitionId.newBuilder().setNamespaceId("namespace-id").build();
+ PartitionId partitionId2 =
+ PartitionId.newBuilder().setNamespaceId("another-namespace-id").build();
+ try {
+ makeKey(
+ PARENT.toBuilder().setPartitionId(partitionId1).build(),
+ CHILD.toBuilder().setPartitionId(partitionId2).build());
+ fail("expected IllegalArgumentException");
+ } catch (IllegalArgumentException expected) {
+ }
+ }
+
+ @Test
+ public void testMakeTimestampValue() throws Exception {
+ // Test cases with nanos == 0.
+ assertConversion(-50_000, -50, 0);
+ assertConversion(-1_000, -1, 0);
+ assertConversion(0, 0, 0);
+ assertConversion(1_000, 1, 0);
+ assertConversion(2_000, 2, 0);
+ assertConversion(100_000, 100, 0);
+
+ // Test cases with nanos % 1_000_000 == 0 (no loss of precision).
+ assertConversion(2, 0, 2_000_000);
+ assertConversion(1_003, 1, 3_000_000);
+ assertConversion(2_005, 2, 5_000_000);
+
+ // Timestamp specification requires that nanos >= 0 even if the timestamp
+ // is before the epoch.
+ assertConversion(0, 0, 0);
+ assertConversion(-250, -1, 750_000_000); // 1/4 second before epoch
+ assertConversion(-500, -1, 500_000_000); // 1/2 second before epoch
+ assertConversion(-750, -1, 250_000_000); // 3/4 second before epoch
+
+ // If nanos % 1_000_000 != 0, precision is lost (via truncation) when
+ // converting to milliseconds.
+ assertTimestampToMilliseconds(3_100, 3, 100_000_999);
+ assertMillisecondsToTimestamp(3_100, 3, 100_000_000);
+ assertTimestampToMilliseconds(5_999, 5, 999_999_999);
+ assertMillisecondsToTimestamp(5_999, 5, 999_000_000);
+ assertTimestampToMilliseconds(7_100, 7, 100_000_001);
+ assertMillisecondsToTimestamp(7_100, 7, 100_000_000);
+ }
+
+ private void assertConversion(long millis, long seconds, int nanos) {
+ assertMillisecondsToTimestamp(millis, seconds, nanos);
+ assertTimestampToMilliseconds(millis, seconds, nanos);
+ }
+
+ private void assertMillisecondsToTimestamp(long millis, long seconds, long nanos) {
+ Value timestampValue = makeValue(new Date(millis)).build();
+ assertEquals(ValueTypeCase.TIMESTAMP_VALUE, timestampValue.getValueTypeCase());
+ assertEquals(seconds, timestampValue.getTimestampValue().getSeconds());
+ assertEquals(nanos, timestampValue.getTimestampValue().getNanos());
+ }
+
+ private void assertTimestampToMilliseconds(long millis, long seconds, int nanos) {
+ Value.Builder value =
+ Value.newBuilder()
+ .setTimestampValue(Timestamp.newBuilder().setSeconds(seconds).setNanos(nanos));
+ assertEquals(millis, DatastoreHelper.toDate(value.build()).getTime());
+ }
+
+ @Test
+ public void testProjectionHandling() {
+ assertEquals(
+ ByteString.copyFromUtf8("hi"), getByteString(makeValue("hi").setMeaning(18).build()));
+ try {
+ getByteString(makeValue("hi").build());
+ fail("Expected IllegalArgumentException");
+ } catch (IllegalArgumentException expected) {
+ }
+
+ assertEquals(new Date(1), toDate(makeValue(1000).setMeaning(18).build()));
+ try {
+ toDate(makeValue(1000).build());
+ fail("Expected IllegalArgumentException");
+ } catch (IllegalArgumentException expected) {
+ }
+ }
+}
diff --git a/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/QuerySplitterTest.java b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/QuerySplitterTest.java
new file mode 100644
index 000000000..cad9502ae
--- /dev/null
+++ b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/QuerySplitterTest.java
@@ -0,0 +1,365 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.datastore.utils;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.datastore.utils.DatastoreHelper.*;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertThrows;
+
+import com.google.datastore.utils.testing.MockCredential;
+import com.google.datastore.utils.testing.MockDatastoreFactory;
+import com.google.datastore.v1.*;
+import com.google.datastore.v1.EntityResult.ResultType;
+import com.google.datastore.v1.PropertyFilter.Operator;
+import com.google.datastore.v1.PropertyOrder.Direction;
+import com.google.datastore.v1.QueryResultBatch.MoreResultsType;
+import com.google.protobuf.Int32Value;
+import com.google.protobuf.Timestamp;
+import java.util.List;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link com.google.datastore.utils.QuerySplitterImpl}. */
+@RunWith(JUnit4.class)
+public class QuerySplitterTest {
+ private static final String PROJECT_ID = "project-id";
+ private static final PartitionId PARTITION =
+ PartitionId.newBuilder().setProjectId(PROJECT_ID).build();
+ private static final String KIND = "test-kind";
+
+ private DatastoreFactory factory = new MockDatastoreFactory();
+ private com.google.datastore.utils.DatastoreOptions.Builder options =
+ new DatastoreOptions.Builder().projectId(PROJECT_ID).credential(new MockCredential());
+
+ private Filter propertyFilter = makeFilter("foo", Operator.EQUAL, makeValue("value")).build();
+
+ private Query query =
+ Query.newBuilder()
+ .addKind(KindExpression.newBuilder().setName(KIND).build())
+ .setFilter(propertyFilter)
+ .build();
+
+ private Query splitQuery =
+ Query.newBuilder()
+ .addKind(KindExpression.newBuilder().setName(KIND).build())
+ .addOrder(makeOrder("__scatter__", Direction.ASCENDING))
+ .addProjection(Projection.newBuilder().setProperty(makePropertyReference("__key__")))
+ .build();
+
+ private Key splitKey0 = makeKey(KIND, String.format("%05d", 1)).setPartitionId(PARTITION).build();
+ private Key splitKey1 =
+ makeKey(KIND, String.format("%05d", 101)).setPartitionId(PARTITION).build();
+ private Key splitKey2 =
+ makeKey(KIND, String.format("%05d", 201)).setPartitionId(PARTITION).build();
+ private Key splitKey3 =
+ makeKey(KIND, String.format("%05d", 301)).setPartitionId(PARTITION).build();
+
+ @Test
+ public void disallowsSortOrder() {
+ com.google.datastore.utils.Datastore datastore = factory.create(options.build());
+ Query queryWithOrder =
+ query.toBuilder().addOrder(makeOrder("bar", Direction.ASCENDING)).build();
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ com.google.datastore.utils.QuerySplitterImpl.INSTANCE.getSplits(
+ queryWithOrder, PARTITION, 2, datastore));
+ assertThat(exception).hasMessageThat().contains("Query cannot have any sort orders.");
+ }
+
+ @Test
+ public void disallowsMultipleKinds() {
+ com.google.datastore.utils.Datastore datastore = factory.create(options.build());
+ Query queryWithMultipleKinds =
+ query.toBuilder()
+ .addKind(KindExpression.newBuilder().setName("another-kind").build())
+ .build();
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ com.google.datastore.utils.QuerySplitterImpl.INSTANCE.getSplits(
+ queryWithMultipleKinds, PARTITION, 2, datastore));
+ assertThat(exception).hasMessageThat().contains("Query must have exactly one kind.");
+ }
+
+ @Test
+ public void disallowsKindlessQuery() {
+ com.google.datastore.utils.Datastore datastore = factory.create(options.build());
+ Query kindlessQuery = query.toBuilder().clearKind().build();
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ com.google.datastore.utils.QuerySplitterImpl.INSTANCE.getSplits(
+ kindlessQuery, PARTITION, 2, datastore));
+ assertThat(exception).hasMessageThat().contains("Query must have exactly one kind.");
+ }
+
+ @Test
+ public void disallowsInequalityFilter() {
+ com.google.datastore.utils.Datastore datastore = factory.create(options.build());
+ Query queryWithInequality =
+ query.toBuilder()
+ .setFilter(makeFilter("foo", Operator.GREATER_THAN, makeValue("value")))
+ .build();
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ com.google.datastore.utils.QuerySplitterImpl.INSTANCE.getSplits(
+ queryWithInequality, PARTITION, 2, datastore));
+ assertThat(exception).hasMessageThat().contains("Query cannot have any inequality filters.");
+ }
+
+ @Test
+ public void splitsMustBePositive() {
+ com.google.datastore.utils.Datastore datastore = factory.create(options.build());
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ com.google.datastore.utils.QuerySplitterImpl.INSTANCE.getSplits(
+ query, PARTITION, 0, datastore));
+ assertThat(exception).hasMessageThat().contains("The number of splits must be greater than 0.");
+ }
+
+ @Test
+ public void getSplits() throws Exception {
+ com.google.datastore.utils.Datastore datastore = factory.create(options.build());
+ MockDatastoreFactory mockClient = (MockDatastoreFactory) factory;
+
+ RunQueryResponse splitQueryResponse =
+ RunQueryResponse.newBuilder()
+ .setQuery(splitQuery)
+ .setBatch(
+ QueryResultBatch.newBuilder()
+ .setEntityResultType(ResultType.KEY_ONLY)
+ .setMoreResults(MoreResultsType.NO_MORE_RESULTS)
+ .addEntityResults(makeKeyOnlyEntity(splitKey0))
+ .addEntityResults(makeKeyOnlyEntity(splitKey1))
+ .addEntityResults(makeKeyOnlyEntity(splitKey2))
+ .addEntityResults(makeKeyOnlyEntity(splitKey3))
+ .build())
+ .build();
+
+ mockClient.setNextResponse(splitQueryResponse);
+
+ List splittedQueries =
+ com.google.datastore.utils.QuerySplitterImpl.INSTANCE.getSplits(
+ query, PARTITION, 3, datastore);
+
+ assertThat(splittedQueries)
+ .containsExactly(
+ query.toBuilder()
+ .setFilter(makeFilterWithKeyRange(propertyFilter, null, splitKey1))
+ .build(),
+ query.toBuilder()
+ .setFilter(makeFilterWithKeyRange(propertyFilter, splitKey1, splitKey3))
+ .build(),
+ query.toBuilder()
+ .setFilter(makeFilterWithKeyRange(propertyFilter, splitKey3, null))
+ .build());
+
+ RunQueryRequest expectedSplitQueryRequest =
+ RunQueryRequest.newBuilder()
+ .setPartitionId(PARTITION)
+ .setProjectId(PROJECT_ID)
+ .setQuery(
+ splitQuery.toBuilder().setLimit(Int32Value.newBuilder().setValue(2 * 32).build()))
+ .build();
+
+ assertArrayEquals(expectedSplitQueryRequest.toByteArray(), mockClient.getLastBody());
+ }
+
+ @Test
+ public void getSplitsWithDatabaseId() throws Exception {
+ com.google.datastore.utils.Datastore datastore = factory.create(options.build());
+ MockDatastoreFactory mockClient = (MockDatastoreFactory) factory;
+
+ PartitionId partition =
+ PartitionId.newBuilder().setProjectId(PROJECT_ID).setDatabaseId("test-database").build();
+
+ RunQueryResponse splitQueryResponse =
+ RunQueryResponse.newBuilder()
+ .setQuery(splitQuery)
+ .setBatch(
+ QueryResultBatch.newBuilder()
+ .setEntityResultType(ResultType.KEY_ONLY)
+ .setMoreResults(MoreResultsType.NO_MORE_RESULTS)
+ .addEntityResults(makeKeyOnlyEntity(splitKey0))
+ .addEntityResults(makeKeyOnlyEntity(splitKey1))
+ .addEntityResults(makeKeyOnlyEntity(splitKey2))
+ .addEntityResults(makeKeyOnlyEntity(splitKey3))
+ .build())
+ .build();
+
+ mockClient.setNextResponse(splitQueryResponse);
+
+ List splitQueries =
+ com.google.datastore.utils.QuerySplitterImpl.INSTANCE.getSplits(
+ query, partition, 3, datastore);
+
+ assertThat(splitQueries)
+ .containsExactly(
+ query.toBuilder()
+ .setFilter(makeFilterWithKeyRange(propertyFilter, null, splitKey1))
+ .build(),
+ query.toBuilder()
+ .setFilter(makeFilterWithKeyRange(propertyFilter, splitKey1, splitKey3))
+ .build(),
+ query.toBuilder()
+ .setFilter(makeFilterWithKeyRange(propertyFilter, splitKey3, null))
+ .build());
+
+ RunQueryRequest expectedSplitQueryRequest =
+ RunQueryRequest.newBuilder()
+ .setPartitionId(partition)
+ .setProjectId(PROJECT_ID)
+ .setDatabaseId("test-database")
+ .setQuery(
+ splitQuery.toBuilder().setLimit(Int32Value.newBuilder().setValue(2 * 32).build()))
+ .build();
+
+ assertArrayEquals(expectedSplitQueryRequest.toByteArray(), mockClient.getLastBody());
+ }
+
+ @Test
+ public void notEnoughSplits() throws Exception {
+ com.google.datastore.utils.Datastore datastore = factory.create(options.build());
+ MockDatastoreFactory mockClient = (MockDatastoreFactory) factory;
+
+ RunQueryResponse splitQueryResponse =
+ RunQueryResponse.newBuilder()
+ .setQuery(splitQuery)
+ .setBatch(
+ QueryResultBatch.newBuilder()
+ .setEntityResultType(ResultType.KEY_ONLY)
+ .setMoreResults(MoreResultsType.NO_MORE_RESULTS)
+ .addEntityResults(makeKeyOnlyEntity(splitKey0))
+ .build())
+ .build();
+
+ mockClient.setNextResponse(splitQueryResponse);
+
+ List splittedQueries =
+ com.google.datastore.utils.QuerySplitterImpl.INSTANCE.getSplits(
+ query, PARTITION, 100, datastore);
+
+ assertThat(splittedQueries)
+ .containsExactly(
+ query.toBuilder()
+ .setFilter(makeFilterWithKeyRange(propertyFilter, null, splitKey0))
+ .build(),
+ query.toBuilder()
+ .setFilter(makeFilterWithKeyRange(propertyFilter, splitKey0, null))
+ .build());
+
+ RunQueryRequest expectedSplitQueryRequest =
+ RunQueryRequest.newBuilder()
+ .setPartitionId(PARTITION)
+ .setProjectId(PROJECT_ID)
+ .setQuery(
+ splitQuery.toBuilder().setLimit(Int32Value.newBuilder().setValue(99 * 32).build()))
+ .build();
+
+ assertArrayEquals(expectedSplitQueryRequest.toByteArray(), mockClient.getLastBody());
+ }
+
+ @Test
+ public void getSplits_withReadTime() throws Exception {
+ Datastore datastore = factory.create(options.build());
+ MockDatastoreFactory mockClient = (MockDatastoreFactory) factory;
+
+ RunQueryResponse splitQueryResponse =
+ RunQueryResponse.newBuilder()
+ .setQuery(splitQuery)
+ .setBatch(
+ QueryResultBatch.newBuilder()
+ .setEntityResultType(ResultType.KEY_ONLY)
+ .setMoreResults(MoreResultsType.NO_MORE_RESULTS)
+ .addEntityResults(makeKeyOnlyEntity(splitKey0))
+ .addEntityResults(makeKeyOnlyEntity(splitKey1))
+ .addEntityResults(makeKeyOnlyEntity(splitKey2))
+ .addEntityResults(makeKeyOnlyEntity(splitKey3))
+ .build())
+ .build();
+
+ mockClient.setNextResponse(splitQueryResponse);
+
+ Timestamp readTime = Timestamp.newBuilder().setSeconds(1654651341L).build();
+
+ List splittedQueries =
+ com.google.datastore.utils.QuerySplitterImpl.INSTANCE.getSplits(
+ query, PARTITION, 3, datastore, readTime);
+
+ assertThat(splittedQueries)
+ .containsExactly(
+ query.toBuilder()
+ .setFilter(makeFilterWithKeyRange(propertyFilter, null, splitKey1))
+ .build(),
+ query.toBuilder()
+ .setFilter(makeFilterWithKeyRange(propertyFilter, splitKey1, splitKey3))
+ .build(),
+ query.toBuilder()
+ .setFilter(makeFilterWithKeyRange(propertyFilter, splitKey3, null))
+ .build());
+
+ RunQueryRequest expectedSplitQueryRequest =
+ RunQueryRequest.newBuilder()
+ .setPartitionId(PARTITION)
+ .setProjectId(PROJECT_ID)
+ .setQuery(
+ splitQuery.toBuilder().setLimit(Int32Value.newBuilder().setValue(2 * 32).build()))
+ .setReadOptions(ReadOptions.newBuilder().setReadTime(readTime))
+ .build();
+
+ assertArrayEquals(expectedSplitQueryRequest.toByteArray(), mockClient.getLastBody());
+ }
+
+ private static EntityResult makeKeyOnlyEntity(Key key) {
+ return EntityResult.newBuilder().setEntity(Entity.newBuilder().setKey(key).build()).build();
+ }
+
+ private static Filter makeFilterWithKeyRange(Filter originalFilter, Key startKey, Key endKey) {
+ Filter startKeyFilter =
+ startKey == null
+ ? null
+ : makeFilter("__key__", Operator.GREATER_THAN_OR_EQUAL, makeValue(startKey)).build();
+
+ Filter endKeyFilter =
+ endKey == null
+ ? null
+ : makeFilter("__key__", Operator.LESS_THAN, makeValue(endKey)).build();
+
+ if (startKeyFilter == null && endKeyFilter == null) {
+ throw new IllegalArgumentException();
+ }
+
+ if (startKeyFilter != null && endKeyFilter == null) {
+ return makeAndFilter(originalFilter, startKeyFilter).build();
+ }
+
+ if (startKeyFilter == null && endKeyFilter != null) {
+ return makeAndFilter(originalFilter, endKeyFilter).build();
+ }
+
+ return makeAndFilter(originalFilter, startKeyFilter, endKeyFilter).build();
+ }
+}
diff --git a/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/RemoteRpcTest.java b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/RemoteRpcTest.java
new file mode 100644
index 000000000..ae4d7a23e
--- /dev/null
+++ b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/RemoteRpcTest.java
@@ -0,0 +1,341 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.datastore.utils;
+
+import static org.junit.Assert.*;
+
+import com.google.api.client.http.HttpRequest;
+import com.google.api.client.http.HttpTransport;
+import com.google.api.client.http.LowLevelHttpRequest;
+import com.google.api.client.http.LowLevelHttpResponse;
+import com.google.api.client.http.protobuf.ProtoHttpContent;
+import com.google.api.client.util.Charsets;
+import com.google.datastore.v1.BeginTransactionResponse;
+import com.google.datastore.v1.RollbackRequest;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.MessageLite;
+import com.google.rpc.Code;
+import com.google.rpc.Status;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.zip.GZIPOutputStream;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Test for {@link RemoteRpc}. */
+@RunWith(JUnit4.class)
+public class RemoteRpcTest {
+
+ private static final String METHOD_NAME = "methodName";
+
+ @Test
+ public void testException() {
+ Status statusProto =
+ Status.newBuilder()
+ .setCode(Code.UNAUTHENTICATED_VALUE)
+ .setMessage("The request does not have valid authentication credentials.")
+ .build();
+ DatastoreException exception =
+ RemoteRpc.makeException(
+ "url",
+ METHOD_NAME,
+ new ByteArrayInputStream(statusProto.toByteArray()),
+ "application/x-protobuf",
+ Charsets.UTF_8,
+ new RuntimeException(),
+ 401);
+ assertEquals(Code.UNAUTHENTICATED, exception.getCode());
+ assertEquals(
+ "The request does not have valid authentication credentials.", exception.getMessage());
+ assertEquals(METHOD_NAME, exception.getMethodName());
+ }
+
+ @Test
+ public void testInvalidProtoException() {
+ DatastoreException exception =
+ RemoteRpc.makeException(
+ "url",
+ METHOD_NAME,
+ new ByteArrayInputStream("".getBytes()),
+ "application/x-protobuf",
+ Charsets.UTF_8,
+ new RuntimeException(),
+ 401);
+ assertEquals(Code.INTERNAL, exception.getCode());
+ assertEquals(
+ "Unable to parse Status protocol buffer: HTTP status code was 401.",
+ exception.getMessage());
+ assertEquals(METHOD_NAME, exception.getMethodName());
+ }
+
+ @Test
+ public void testEmptyProtoException() {
+ Status statusProto = Status.newBuilder().build();
+ DatastoreException exception =
+ RemoteRpc.makeException(
+ "url",
+ METHOD_NAME,
+ new ByteArrayInputStream(statusProto.toByteArray()),
+ "application/x-protobuf",
+ Charsets.UTF_8,
+ new RuntimeException(),
+ 404);
+ assertEquals(Code.INTERNAL, exception.getCode());
+ assertEquals(
+ "Unexpected OK error code with HTTP status code of 404. Message: .",
+ exception.getMessage());
+ assertEquals(METHOD_NAME, exception.getMethodName());
+ }
+
+ @Test
+ public void testEmptyProtoExceptionUnauthenticated() {
+ Status statusProto = Status.newBuilder().build();
+ DatastoreException exception =
+ RemoteRpc.makeException(
+ "url",
+ METHOD_NAME,
+ new ByteArrayInputStream(statusProto.toByteArray()),
+ "application/x-protobuf",
+ Charsets.UTF_8,
+ new RuntimeException(),
+ 401);
+ assertEquals(Code.UNAUTHENTICATED, exception.getCode());
+ assertEquals("Unauthenticated.", exception.getMessage());
+ assertEquals(METHOD_NAME, exception.getMethodName());
+ }
+
+ @Test
+ public void testPlainTextException() {
+ DatastoreException exception =
+ RemoteRpc.makeException(
+ "url",
+ METHOD_NAME,
+ new ByteArrayInputStream("Text Error".getBytes()),
+ "text/plain",
+ Charsets.UTF_8,
+ new RuntimeException(),
+ 401);
+ assertEquals(Code.INTERNAL, exception.getCode());
+ assertEquals(
+ "Non-protobuf error: Text Error. HTTP status code was 401.", exception.getMessage());
+ assertEquals(METHOD_NAME, exception.getMethodName());
+ }
+
+ @Test
+ public void testGzip() throws IOException, DatastoreException {
+ BeginTransactionResponse response = newBeginTransactionResponse();
+ InjectedTestValues injectedTestValues =
+ new InjectedTestValues(gzip(response), new byte[1], true);
+ RemoteRpc rpc = newRemoteRpc(injectedTestValues);
+
+ InputStream is =
+ rpc.call("beginTransaction", BeginTransactionResponse.getDefaultInstance(), "", "");
+ BeginTransactionResponse parsedResponse = BeginTransactionResponse.parseFrom(is);
+ is.close();
+
+ assertEquals(response, parsedResponse);
+ // Check that the underlying stream is exhausted.
+ assertEquals(-1, injectedTestValues.inputStream.read());
+ }
+
+ @Test
+ public void testHttpHeaders_apiFormat() throws IOException {
+ String projectId = "project-id";
+ MessageLite request =
+ RollbackRequest.newBuilder().setTransaction(ByteString.copyFromUtf8(projectId)).build();
+ RemoteRpc rpc =
+ newRemoteRpc(
+ new InjectedTestValues(gzip(newBeginTransactionResponse()), new byte[1], true));
+ HttpRequest httpRequest =
+ rpc.getClient().buildPostRequest(rpc.resolveURL("blah"), new ProtoHttpContent(request));
+ rpc.setHeaders(request, httpRequest, projectId, "");
+ assertNotNull(
+ httpRequest.getHeaders().getFirstHeaderStringValue(RemoteRpc.API_FORMAT_VERSION_HEADER));
+ }
+
+ @Test
+ public void testHttpHeaders_prefixHeader() throws IOException {
+ String projectId = "my-project";
+ String databaseId = "my-db";
+ MessageLite request =
+ RollbackRequest.newBuilder()
+ .setTransaction(ByteString.copyFromUtf8(projectId))
+ .setDatabaseId(databaseId)
+ .build();
+ RemoteRpc rpc =
+ newRemoteRpc(
+ new InjectedTestValues(gzip(newBeginTransactionResponse()), new byte[1], true));
+ HttpRequest httpRequest =
+ rpc.getClient().buildPostRequest(rpc.resolveURL("blah"), new ProtoHttpContent(request));
+ rpc.setHeaders(request, httpRequest, projectId, databaseId);
+ assertEquals(
+ "project_id=my-project&database_id=my-db",
+ httpRequest.getHeaders().get(RemoteRpc.X_GOOG_REQUEST_PARAMS_HEADER));
+
+ MessageLite request2 =
+ RollbackRequest.newBuilder().setTransaction(ByteString.copyFromUtf8(projectId)).build();
+ RemoteRpc rpc2 =
+ newRemoteRpc(
+ new InjectedTestValues(gzip(newBeginTransactionResponse()), new byte[1], true));
+ HttpRequest httpRequest2 =
+ rpc2.getClient().buildPostRequest(rpc2.resolveURL("blah"), new ProtoHttpContent(request2));
+ rpc2.setHeaders(request, httpRequest2, projectId, "");
+ assertEquals(
+ "project_id=my-project",
+ httpRequest2.getHeaders().get(RemoteRpc.X_GOOG_REQUEST_PARAMS_HEADER));
+ }
+
+ private static BeginTransactionResponse newBeginTransactionResponse() {
+ return BeginTransactionResponse.newBuilder()
+ .setTransaction(ByteString.copyFromUtf8("blah-blah-blah"))
+ .build();
+ }
+
+ private static RemoteRpc newRemoteRpc(InjectedTestValues injectedTestValues) {
+ return new RemoteRpc(
+ new MyHttpTransport(injectedTestValues).createRequestFactory(),
+ null,
+ "https://www.example.com/v1/projects/p");
+ }
+
+ private byte[] gzip(BeginTransactionResponse response) throws IOException {
+ ByteArrayOutputStream bytesOut = new ByteArrayOutputStream();
+ try (GZIPOutputStream gzipOut = new GZIPOutputStream(bytesOut)) {
+ response.writeTo(gzipOut);
+ }
+ return bytesOut.toByteArray();
+ }
+
+ private static class InjectedTestValues {
+ private final InputStream inputStream;
+ private final int contentLength;
+ private final boolean isGzip;
+
+ public InjectedTestValues(byte[] messageBytes, byte[] additionalBytes, boolean isGzip) {
+ byte[] allBytes = concat(messageBytes, additionalBytes);
+ this.inputStream = new ByteArrayInputStream(allBytes);
+ this.contentLength = allBytes.length;
+ this.isGzip = isGzip;
+ }
+
+ private static byte[] concat(byte[] a, byte[] b) {
+ byte[] c = new byte[a.length + b.length];
+ System.arraycopy(a, 0, c, 0, a.length);
+ System.arraycopy(b, 0, c, a.length, b.length);
+ return c;
+ }
+ }
+
+ /** {@link HttpTransport} that allows injection of the returned {@link LowLevelHttpRequest}. */
+ private static class MyHttpTransport extends HttpTransport {
+
+ private final InjectedTestValues injectedTestValues;
+
+ public MyHttpTransport(InjectedTestValues injectedTestValues) {
+ this.injectedTestValues = injectedTestValues;
+ }
+
+ @Override
+ protected LowLevelHttpRequest buildRequest(String method, String url) throws IOException {
+ return new MyLowLevelHttpRequest(injectedTestValues);
+ }
+ }
+
+ /**
+ * {@link LowLevelHttpRequest} that allows injection of the returned {@link LowLevelHttpResponse}.
+ */
+ private static class MyLowLevelHttpRequest extends LowLevelHttpRequest {
+
+ private final InjectedTestValues injectedTestValues;
+
+ public MyLowLevelHttpRequest(InjectedTestValues injectedTestValues) {
+ this.injectedTestValues = injectedTestValues;
+ }
+
+ @Override
+ public void addHeader(String name, String value) throws IOException {
+ // Do nothing.
+ }
+
+ @Override
+ public LowLevelHttpResponse execute() throws IOException {
+ return new MyLowLevelHttpResponse(injectedTestValues);
+ }
+ }
+
+ /** {@link LowLevelHttpResponse} that allows injected properties. */
+ private static class MyLowLevelHttpResponse extends LowLevelHttpResponse {
+
+ private final InjectedTestValues injectedTestValues;
+
+ public MyLowLevelHttpResponse(InjectedTestValues injectedTestValues) {
+ this.injectedTestValues = injectedTestValues;
+ }
+
+ @Override
+ public InputStream getContent() throws IOException {
+ return injectedTestValues.inputStream;
+ }
+
+ @Override
+ public String getContentEncoding() throws IOException {
+ return injectedTestValues.isGzip ? "gzip" : "";
+ }
+
+ @Override
+ public long getContentLength() throws IOException {
+ return injectedTestValues.contentLength;
+ }
+
+ @Override
+ public String getContentType() throws IOException {
+ return "application/x-protobuf";
+ }
+
+ @Override
+ public String getStatusLine() throws IOException {
+ return null;
+ }
+
+ @Override
+ public int getStatusCode() throws IOException {
+ return 200;
+ }
+
+ @Override
+ public String getReasonPhrase() throws IOException {
+ return null;
+ }
+
+ @Override
+ public int getHeaderCount() throws IOException {
+ return 0;
+ }
+
+ @Override
+ public String getHeaderName(int index) throws IOException {
+ return null;
+ }
+
+ @Override
+ public String getHeaderValue(int index) throws IOException {
+ return null;
+ }
+ }
+}
diff --git a/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java
new file mode 100644
index 000000000..d30c1cbdc
--- /dev/null
+++ b/google-cloud-datastore-utils/src/test/java/com/google/datastore/utils/it/ITDatastoreProtoClientTest.java
@@ -0,0 +1,90 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.datastore.utils.it;
+
+import static com.google.datastore.utils.DatastoreHelper.makeFilter;
+import static com.google.datastore.utils.DatastoreHelper.makeValue;
+
+import com.google.common.truth.Truth;
+import com.google.datastore.utils.Datastore;
+import com.google.datastore.utils.DatastoreException;
+import com.google.datastore.utils.DatastoreHelper;
+import com.google.datastore.v1.*;
+import java.io.IOException;
+import java.security.GeneralSecurityException;
+import java.util.List;
+import org.junit.Before;
+import org.junit.Test;
+
+public class ITDatastoreProtoClientTest {
+
+ private static Datastore DATASTORE;
+
+ private static PartitionId PARTITION;
+
+ private static final String KIND = "test-kind";
+ private static final String PROJECT_ID = System.getenv(DatastoreHelper.PROJECT_ID_ENV_VAR);
+
+ @Before
+ public void setUp() throws GeneralSecurityException, IOException {
+ DATASTORE = DatastoreHelper.getDatastoreFromEnv();
+ }
+
+ @Test
+ public void testQuerySplitterWithDefaultDb() throws DatastoreException {
+ Filter propertyFilter =
+ makeFilter("foo", PropertyFilter.Operator.EQUAL, makeValue("value")).build();
+ Query query =
+ Query.newBuilder()
+ .addKind(KindExpression.newBuilder().setName(KIND).build())
+ .setFilter(propertyFilter)
+ .build();
+
+ PARTITION = PartitionId.newBuilder().setProjectId(PROJECT_ID).build();
+
+ List splits =
+ DatastoreHelper.getQuerySplitter().getSplits(query, PARTITION, 2, DATASTORE);
+ Truth.assertThat(splits).isNotEmpty();
+ splits.forEach(
+ split -> {
+ Truth.assertThat(split.getKind(0).getName()).isEqualTo(KIND);
+ Truth.assertThat(split.getFilter()).isEqualTo(propertyFilter);
+ });
+ }
+
+ @Test
+ public void testQuerySplitterWithDb() throws DatastoreException {
+ Filter propertyFilter =
+ makeFilter("foo", PropertyFilter.Operator.EQUAL, makeValue("value")).build();
+ Query query =
+ Query.newBuilder()
+ .addKind(KindExpression.newBuilder().setName(KIND).build())
+ .setFilter(propertyFilter)
+ .build();
+
+ PARTITION = PartitionId.newBuilder().setProjectId(PROJECT_ID).setDatabaseId("test-db").build();
+
+ List splits =
+ DatastoreHelper.getQuerySplitter().getSplits(query, PARTITION, 2, DATASTORE);
+
+ Truth.assertThat(splits).isNotEmpty();
+ splits.forEach(
+ split -> {
+ Truth.assertThat(split.getKind(0).getName()).isEqualTo(KIND);
+ Truth.assertThat(split.getFilter()).isEqualTo(propertyFilter);
+ });
+ }
+}
diff --git a/google-cloud-datastore/clirr-ignored-differences.xml b/google-cloud-datastore/clirr-ignored-differences.xml
index 1620fd752..d1e3a4ff4 100644
--- a/google-cloud-datastore/clirr-ignored-differences.xml
+++ b/google-cloud-datastore/clirr-ignored-differences.xml
@@ -1,7 +1,50 @@
-
+
+
+ com/google/cloud/datastore/ReadOption$QueryAndReadOptions
+ *
+ 8001
+
+
+ com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer
+ *QueryAndReadOptions*
+ *QueryConfig*
+ 7005
+
+
+
+ com/google/cloud/datastore/DatastoreException
+ 5000
+ com/google/cloud/grpc/BaseGrpcServiceException
+
+
+
+ com/google/cloud/datastore/DatastoreException
+ 5001
+ com/google/cloud/http/BaseHttpServiceException
+
+
+ com/google/cloud/datastore/Datastore
+ void close()
+ 7012
+
+
+ com/google/cloud/datastore/spi/v1/DatastoreRpc
+ void close()
+ 7012
+
+
+ com/google/cloud/datastore/Datastore
+ boolean isClosed()
+ 7012
+
+
+ com/google/cloud/datastore/spi/v1/DatastoreRpc
+ boolean isClosed()
+ 7012
+
com/google/cloud/datastore/Datastore
com.google.cloud.datastore.QueryResults run(com.google.cloud.datastore.Query, com.google.cloud.datastore.models.ExplainOptions, com.google.cloud.datastore.ReadOption[])
@@ -14,7 +57,7 @@
com/google/cloud/datastore/DatastoreReader
- com.google.cloud.datastore.AggregationResults runAggregation(com.google.cloud.datastore.AggregationQuery, com.google.cloud.datastore.models.ExplainOptions)
+ com.google.cloud.datastore.AggregationResults runAggregation(com.google.cloud.datastore.AggregationQuery, com.google.cloud.datastore.models.ExplainOptions)
7012
@@ -28,18 +71,6 @@
7012
-
-
- com/google/cloud/datastore/ReadOption$QueryConfig
- com.google.cloud.datastore.ReadOption$QueryConfig create(com.google.cloud.datastore.Query, java.util.List)
- *com.google.datastore.v1.ExplainOptions*
- 7005
-
-
- com/google/cloud/datastore/ReadOption$QueryConfig
- com.google.cloud.datastore.ReadOption$QueryConfig create(com.google.cloud.datastore.Query)
- 7004
-
com/google/cloud/datastore/execution/AggregationQueryExecutor
com.google.cloud.datastore.AggregationResults execute(com.google.cloud.datastore.AggregationQuery, com.google.cloud.datastore.ReadOption[])
@@ -55,4 +86,21 @@
java.lang.Object execute(com.google.cloud.datastore.Query, com.google.cloud.datastore.ReadOption[])
7004
+
+ com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator
+ RetryAndTraceDatastoreRpcDecorator(com.google.cloud.datastore.spi.v1.DatastoreRpc, com.google.cloud.datastore.TraceUtil, com.google.api.gax.retrying.RetrySettings, com.google.cloud.datastore.DatastoreOptions)
+ RetryAndTraceDatastoreRpcDecorator(com.google.cloud.datastore.spi.v1.DatastoreRpc, com.google.cloud.datastore.telemetry.TraceUtil, com.google.api.gax.retrying.RetrySettings, com.google.cloud.datastore.DatastoreOptions)
+ 7005
+
+
+
+
+ com/google/cloud/datastore/TraceUtil
+ 8001
+
+
+ com/google/cloud/datastore/testing/RemoteDatastoreHelper
+ 8001
+
+
diff --git a/google-cloud-datastore/pom.xml b/google-cloud-datastore/pom.xml
index 8d6b79e8d..ecb5895ae 100644
--- a/google-cloud-datastore/pom.xml
+++ b/google-cloud-datastore/pom.xml
@@ -2,7 +2,7 @@
4.0.0
google-cloud-datastore
- 2.19.2-SNAPSHOT
+ 2.28.2-SNAPSHOT
jar
Google Cloud Datastore
https://github.com/googleapis/java-datastore
@@ -12,12 +12,29 @@
com.google.cloud
google-cloud-datastore-parent
- 2.19.2-SNAPSHOT
+ 2.28.2-SNAPSHOT
google-cloud-datastore
+ 1.42.1
+
+
+
+
+ com.google.cloud
+ gapic-libraries-bom
+ 1.53.0
+ pom
+ import
+
+
+
+
+ com.google.api.grpc
+ grpc-google-cloud-datastore-v1
+
com.google.api.grpc
grpc-google-cloud-datastore-admin-v1
@@ -26,6 +43,10 @@
com.google.cloud
google-cloud-core-http
+
+ com.google.cloud
+ google-cloud-core-grpc
+
com.google.api.grpc
proto-google-cloud-datastore-v1
@@ -38,6 +59,10 @@
com.google.cloud.datastore
datastore-v1-proto-client
+
+ com.google.auth
+ google-auth-library-credentials
+
io.grpc
grpc-api
@@ -111,6 +136,19 @@
jsr305
+
+
+ io.opentelemetry
+ opentelemetry-api
+ ${opentelemetry.version}
+
+
+ io.opentelemetry
+ opentelemetry-context
+ ${opentelemetry.version}
+
+
+
${project.groupId}
@@ -118,11 +156,10 @@
test-jar
test
-
com.google.guava
guava-testlib
- 33.1.0-jre
+ 33.4.0-jre
test
@@ -157,9 +194,65 @@
com.google.truth
truth
- 1.4.2
+ 1.4.4
+ test
+
+
+ com.google.testparameterinjector
+ test-parameter-injector
+ 1.17
+ test
+
+
+
+ io.opentelemetry
+ opentelemetry-sdk
+ ${opentelemetry.version}
+ test
+
+
+ io.opentelemetry
+ opentelemetry-sdk-common
+ ${opentelemetry.version}
+ test
+
+
+ io.opentelemetry
+ opentelemetry-sdk-testing
+ ${opentelemetry.version}
+ test
+
+
+ io.opentelemetry
+ opentelemetry-sdk-trace
+ ${opentelemetry.version}
+ test
+
+
+ io.opentelemetry
+ opentelemetry-semconv
+ 1.1.0-alpha
+ test
+
+
+
+
+ com.google.cloud.opentelemetry
+ exporter-trace
+ 0.15.0
+ test
+
+
+ com.google.cloud
+ google-cloud-trace
+ test
+
+
+ com.google.api.grpc
+ proto-google-cloud-trace-v1
test
+
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Batch.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Batch.java
index eb4abd854..bb162af33 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Batch.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Batch.java
@@ -16,6 +16,7 @@
package com.google.cloud.datastore;
+import com.google.api.core.InternalExtensionOnly;
import java.util.List;
import javax.annotation.concurrent.NotThreadSafe;
@@ -42,6 +43,7 @@
* This class too should not be treated as a thread safe class.
*/
@NotThreadSafe
+@InternalExtensionOnly
public interface Batch extends DatastoreBatchWriter {
interface Response {
@@ -53,9 +55,13 @@ interface Response {
/**
* {@inheritDoc}
*
- * If an entity for {@code entity.getKey()} does not exists, {@code entity} is inserted.
+ *
If an entity for {@code entity.getKey()} does not exist, {@code entity} is inserted.
* Otherwise, {@link #submit()} will throw a {@link DatastoreException} with {@link
* DatastoreException#getReason()} equal to {@code "ALREADY_EXISTS"}.
+ *
+ * @param entity the entity to be added to the datastore
+ * @return The entity that was added
+ * @throws DatastoreException if there was any failure
*/
@Override
Entity add(FullEntity entity);
@@ -67,6 +73,10 @@ interface Response {
* exists, {@link #submit()} will throw a {@link DatastoreException} with {@link
* DatastoreException#getReason()} equal to {@code "ALREADY_EXISTS"}. All entities in {@code
* entities} whose key did not exist are inserted.
+ *
+ * @param entities entities to be added to the datastore
+ * @return A list of entities that have been added
+ * @throws DatastoreException if there was any failure
*/
@Override
List add(FullEntity... entities);
@@ -74,10 +84,15 @@ interface Response {
/**
* Submit the batch to the Datastore.
*
- * @throws DatastoreException if there was any failure or if batch is not longer active
+ * @return Response of the batch submit operation.
+ * @throws DatastoreException if there was any failure or if batch is no longer active
*/
Response submit();
- /** Returns the batch associated {@link Datastore}. */
+ /**
+ * Returns the batch associated {@link Datastore}.
+ *
+ * @return The batch associated datastore
+ */
Datastore getDatastore();
}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java
index 5bd8384a3..0e769a109 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java
@@ -17,6 +17,7 @@
package com.google.cloud.datastore;
import com.google.api.core.BetaApi;
+import com.google.api.core.InternalExtensionOnly;
import com.google.cloud.Service;
import com.google.cloud.datastore.models.ExplainOptions;
import com.google.datastore.v1.TransactionOptions;
@@ -24,7 +25,8 @@
import java.util.List;
/** An interface for Google Cloud Datastore. */
-public interface Datastore extends Service, DatastoreReaderWriter {
+@InternalExtensionOnly
+public interface Datastore extends Service, DatastoreReaderWriter, AutoCloseable {
/**
* Returns a new Datastore transaction.
@@ -51,6 +53,13 @@ public interface Datastore extends Service, DatastoreReaderWri
* @param the type of the return value
*/
interface TransactionCallable {
+ /**
+ * Callback's invoke method for the TransactionCallable.
+ *
+ * @param readerWriter DatastoreReaderWriter associated with the new transaction
+ * @return T The transaction result
+ * @throws Exception upon failure
+ */
T run(DatastoreReaderWriter readerWriter) throws Exception;
}
@@ -481,10 +490,7 @@ interface TransactionCallable {
* @throws DatastoreException upon failure
*/
@BetaApi
- default QueryResults run(
- Query query, ExplainOptions explainOptions, ReadOption... options) {
- throw new UnsupportedOperationException("Not implemented.");
- }
+ QueryResults run(Query query, ExplainOptions explainOptions, ReadOption... options);
/**
* Submits a {@link AggregationQuery} and returns {@link AggregationResults}. {@link ReadOption}s
@@ -529,9 +535,7 @@ default QueryResults run(
* @throws DatastoreException upon failure
* @return {@link AggregationResults}
*/
- default AggregationResults runAggregation(AggregationQuery query, ReadOption... options) {
- throw new UnsupportedOperationException("Not implemented.");
- }
+ AggregationResults runAggregation(AggregationQuery query, ReadOption... options);
/**
* Submits a {@link AggregationQuery} with specified {@link
@@ -557,8 +561,17 @@ default AggregationResults runAggregation(AggregationQuery query, ReadOption...
* @return {@link AggregationResults}
*/
@BetaApi
- default AggregationResults runAggregation(
- AggregationQuery query, ExplainOptions explainOptions, ReadOption... options) {
- throw new UnsupportedOperationException("Not implemented.");
- }
+ AggregationResults runAggregation(
+ AggregationQuery query, ExplainOptions explainOptions, ReadOption... options);
+
+ /**
+ * Closes the gRPC channels associated with this instance and frees up their resources. This
+ * method blocks until all channels are closed. Once this method is called, this Datastore client
+ * is no longer usable.
+ */
+ @Override
+ void close() throws Exception;
+
+ /** Returns true if this background resource has been shut down. */
+ boolean isClosed();
}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreBatchWriter.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreBatchWriter.java
index db4bd3179..28d2569b6 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreBatchWriter.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreBatchWriter.java
@@ -16,6 +16,7 @@
package com.google.cloud.datastore;
+import com.google.api.core.InternalExtensionOnly;
import java.util.List;
import javax.annotation.concurrent.NotThreadSafe;
@@ -31,6 +32,7 @@
* This class too should not be treated as a thread safe class.
*/
@NotThreadSafe
+@InternalExtensionOnly
public interface DatastoreBatchWriter extends DatastoreWriter {
/**
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreException.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreException.java
index 512d0a3dc..44bde2c10 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreException.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreException.java
@@ -16,10 +16,16 @@
package com.google.cloud.datastore;
+import com.google.api.gax.grpc.GrpcStatusCode;
+import com.google.api.gax.rpc.ApiException;
+import com.google.api.gax.rpc.StatusCode;
import com.google.cloud.BaseServiceException;
import com.google.cloud.RetryHelper.RetryHelperException;
import com.google.cloud.http.BaseHttpServiceException;
+import com.google.common.base.Strings;
import com.google.common.collect.ImmutableSet;
+import io.grpc.StatusException;
+import io.grpc.StatusRuntimeException;
import java.io.IOException;
import java.util.Set;
@@ -31,7 +37,7 @@
*/
public final class DatastoreException extends BaseHttpServiceException {
- // see https://cloud.google.com/datastore/docs/concepts/errors#Error_Codes"
+ // see https://cloud.google.com/datastore/docs/concepts/errors#Error_Codes
private static final Set RETRYABLE_ERRORS =
ImmutableSet.of(
new Error(10, "ABORTED", false),
@@ -43,6 +49,10 @@ public DatastoreException(int code, String message, String reason) {
this(code, message, reason, true, null);
}
+ public DatastoreException(int code, String message, Throwable cause) {
+ super(code, message, null, true, RETRYABLE_ERRORS, cause);
+ }
+
public DatastoreException(int code, String message, String reason, Throwable cause) {
super(code, message, reason, true, RETRYABLE_ERRORS, cause);
}
@@ -64,7 +74,76 @@ public DatastoreException(IOException exception) {
*/
static DatastoreException translateAndThrow(RetryHelperException ex) {
BaseServiceException.translate(ex);
- throw new DatastoreException(UNKNOWN_CODE, ex.getMessage(), null, ex.getCause());
+ throw transformThrowable(ex);
+ }
+
+ static BaseServiceException transformThrowable(Throwable t) {
+ if (t instanceof BaseServiceException) {
+ return (BaseServiceException) t;
+ }
+ if (t.getCause() instanceof BaseServiceException) {
+ return (BaseServiceException) t.getCause();
+ }
+ if (t instanceof ApiException) {
+ return asDatastoreException((ApiException) t);
+ }
+ if (t.getCause() instanceof ApiException) {
+ return asDatastoreException((ApiException) t.getCause());
+ }
+ return getDatastoreException(t);
+ }
+
+ private static DatastoreException getDatastoreException(Throwable t) {
+ // unwrap a RetryHelperException if that is what is being translated
+ if (t instanceof RetryHelperException) {
+ return new DatastoreException(UNKNOWN_CODE, t.getMessage(), null, t.getCause());
+ }
+ return new DatastoreException(UNKNOWN_CODE, t.getMessage(), t);
+ }
+
+ static DatastoreException asDatastoreException(ApiException apiEx) {
+ int datastoreStatusCode = 0;
+ StatusCode statusCode = apiEx.getStatusCode();
+ if (statusCode instanceof GrpcStatusCode) {
+ GrpcStatusCode gsc = (GrpcStatusCode) statusCode;
+ datastoreStatusCode =
+ GrpcToDatastoreCodeTranslation.grpcCodeToDatastoreStatusCode(gsc.getTransportCode());
+ }
+
+ // If there is a gRPC exception in our cause, pull its error message up to be our
+ // message otherwise, create a generic error message with the status code.
+ String statusCodeName = statusCode.getCode().name();
+ String statusExceptionMessage = getStatusExceptionMessage(apiEx);
+
+ String message;
+ if (statusExceptionMessage != null) {
+ message = statusCodeName + ": " + statusExceptionMessage;
+ } else {
+ message = "Error: " + statusCodeName;
+ }
+
+ String reason = "";
+ if (Strings.isNullOrEmpty(apiEx.getReason())) {
+ if (apiEx.getStatusCode() != null) {
+ reason = apiEx.getStatusCode().getCode().name();
+ }
+ }
+ // It'd be better to use ExceptionData and BaseServiceException#(ExceptionData) but,
+ // BaseHttpServiceException does not pass that through so we're stuck using this for now.
+ // TODO: When we can break the coupling to BaseHttpServiceException replace this
+ return new DatastoreException(datastoreStatusCode, message, reason, apiEx);
+ }
+
+ private static String getStatusExceptionMessage(Exception apiEx) {
+ if (apiEx.getMessage() != null) {
+ return apiEx.getMessage();
+ } else {
+ Throwable cause = apiEx.getCause();
+ if (cause instanceof StatusRuntimeException || cause instanceof StatusException) {
+ return cause.getMessage();
+ }
+ return null;
+ }
}
/**
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreFactory.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreFactory.java
index 1b443066d..54274e7bb 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreFactory.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreFactory.java
@@ -16,7 +16,9 @@
package com.google.cloud.datastore;
+import com.google.api.core.InternalExtensionOnly;
import com.google.cloud.ServiceFactory;
/** An interface for Datastore factories. */
+@InternalExtensionOnly
public interface DatastoreFactory extends ServiceFactory {}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
index a3bfb3796..c64474fa8 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
@@ -16,6 +16,26 @@
package com.google.cloud.datastore;
+import static com.google.cloud.datastore.telemetry.TraceUtil.ATTRIBUTES_KEY_DEFERRED;
+import static com.google.cloud.datastore.telemetry.TraceUtil.ATTRIBUTES_KEY_DOCUMENT_COUNT;
+import static com.google.cloud.datastore.telemetry.TraceUtil.ATTRIBUTES_KEY_MISSING;
+import static com.google.cloud.datastore.telemetry.TraceUtil.ATTRIBUTES_KEY_MORE_RESULTS;
+import static com.google.cloud.datastore.telemetry.TraceUtil.ATTRIBUTES_KEY_READ_CONSISTENCY;
+import static com.google.cloud.datastore.telemetry.TraceUtil.ATTRIBUTES_KEY_RECEIVED;
+import static com.google.cloud.datastore.telemetry.TraceUtil.ATTRIBUTES_KEY_TRANSACTIONAL;
+import static com.google.cloud.datastore.telemetry.TraceUtil.ATTRIBUTES_KEY_TRANSACTION_ID;
+import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_ALLOCATE_IDS;
+import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_BEGIN_TRANSACTION;
+import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_COMMIT;
+import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP;
+import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RESERVE_IDS;
+import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_ROLLBACK;
+import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY;
+import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_COMMIT;
+import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_LOOKUP;
+import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN;
+import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN_QUERY;
+
import com.google.api.core.BetaApi;
import com.google.api.gax.retrying.RetrySettings;
import com.google.cloud.BaseService;
@@ -25,20 +45,23 @@
import com.google.cloud.ServiceOptions;
import com.google.cloud.datastore.execution.AggregationQueryExecutor;
import com.google.cloud.datastore.spi.v1.DatastoreRpc;
+import com.google.cloud.datastore.telemetry.TraceUtil;
+import com.google.cloud.datastore.telemetry.TraceUtil.Scope;
import com.google.common.base.MoreObjects;
import com.google.common.base.Preconditions;
import com.google.common.collect.AbstractIterator;
import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import com.google.common.collect.Sets;
+import com.google.datastore.v1.CommitResponse;
import com.google.datastore.v1.ExplainOptions;
import com.google.datastore.v1.ReadOptions;
import com.google.datastore.v1.ReserveIdsRequest;
+import com.google.datastore.v1.RunQueryResponse;
import com.google.datastore.v1.TransactionOptions;
import com.google.protobuf.ByteString;
-import io.opencensus.common.Scope;
-import io.opencensus.trace.Span;
-import io.opencensus.trace.Status;
+import io.opentelemetry.context.Context;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
@@ -50,16 +73,22 @@
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.Callable;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import javax.annotation.Nullable;
final class DatastoreImpl extends BaseService implements Datastore {
+ Logger logger = Logger.getLogger(Datastore.class.getName());
private final DatastoreRpc datastoreRpc;
private final RetrySettings retrySettings;
private static final ExceptionHandler TRANSACTION_EXCEPTION_HANDLER =
TransactionExceptionHandler.build();
private static final ExceptionHandler TRANSACTION_OPERATION_EXCEPTION_HANDLER =
TransactionOperationExceptionHandler.build();
- private final TraceUtil traceUtil = TraceUtil.getInstance();
+
+ private final com.google.cloud.datastore.telemetry.TraceUtil otelTraceUtil =
+ getOptions().getTraceUtil();
private final ReadOptionProtoPreparer readOptionProtoPreparer;
private final AggregationQueryExecutor aggregationQueryExecutor;
@@ -73,7 +102,8 @@ final class DatastoreImpl extends BaseService implements Datas
readOptionProtoPreparer = new ReadOptionProtoPreparer();
aggregationQueryExecutor =
new AggregationQueryExecutor(
- new RetryAndTraceDatastoreRpcDecorator(datastoreRpc, traceUtil, retrySettings, options),
+ new RetryAndTraceDatastoreRpcDecorator(
+ datastoreRpc, otelTraceUtil, retrySettings, options),
options);
}
@@ -92,8 +122,82 @@ public Transaction newTransaction() {
return new TransactionImpl(this);
}
- static class ReadWriteTransactionCallable implements Callable {
+ static class TracedReadWriteTransactionCallable implements Callable {
+ private final Datastore datastore;
+ private final TransactionCallable callable;
+ private volatile TransactionOptions options;
+ private volatile Transaction transaction;
+
+ private final TraceUtil.Span parentSpan;
+
+ TracedReadWriteTransactionCallable(
+ Datastore datastore,
+ TransactionCallable callable,
+ TransactionOptions options,
+ @Nullable com.google.cloud.datastore.telemetry.TraceUtil.Span parentSpan) {
+ this.datastore = datastore;
+ this.callable = callable;
+ this.options = options;
+ this.transaction = null;
+ this.parentSpan = parentSpan;
+ }
+
+ Datastore getDatastore() {
+ return datastore;
+ }
+
+ TransactionOptions getOptions() {
+ return options;
+ }
+
+ Transaction getTransaction() {
+ return transaction;
+ }
+
+ void setPrevTransactionId(ByteString transactionId) {
+ TransactionOptions.ReadWrite readWrite =
+ TransactionOptions.ReadWrite.newBuilder().setPreviousTransaction(transactionId).build();
+ options = options.toBuilder().setReadWrite(readWrite).build();
+ }
+ @Override
+ public T call() throws DatastoreException {
+ try (io.opentelemetry.context.Scope ignored =
+ Context.current().with(parentSpan.getSpan()).makeCurrent()) {
+ transaction = datastore.newTransaction(options);
+ T value = callable.run(transaction);
+ transaction.commit();
+ return value;
+ } catch (Exception ex) {
+ transaction.rollback();
+ throw DatastoreException.propagateUserException(ex);
+ } finally {
+ if (transaction.isActive()) {
+ transaction.rollback();
+ }
+ if (options != null
+ && options.getModeCase().equals(TransactionOptions.ModeCase.READ_WRITE)) {
+ setPrevTransactionId(transaction.getTransactionId());
+ }
+ }
+ }
+ }
+
+ @Override
+ public void close() throws Exception {
+ try {
+ datastoreRpc.close();
+ } catch (Exception e) {
+ logger.log(Level.WARNING, "Failed to close channels", e);
+ }
+ }
+
+ @Override
+ public boolean isClosed() {
+ return datastoreRpc.isClosed();
+ }
+
+ static class ReadWriteTransactionCallable implements Callable {
private final Datastore datastore;
private final TransactionCallable callable;
private volatile TransactionOptions options;
@@ -127,8 +231,8 @@ void setPrevTransactionId(ByteString transactionId) {
@Override
public T call() throws DatastoreException {
- transaction = datastore.newTransaction(options);
try {
+ transaction = datastore.newTransaction(options);
T value = callable.run(transaction);
transaction.commit();
return value;
@@ -149,36 +253,47 @@ public T call() throws DatastoreException {
@Override
public T runInTransaction(final TransactionCallable callable) {
- Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION);
- try (Scope scope = traceUtil.getTracer().withSpan(span)) {
+ TraceUtil.Span span = otelTraceUtil.startSpan(SPAN_NAME_TRANSACTION_RUN);
+ Callable transactionCallable =
+ (getOptions().getOpenTelemetryOptions().isEnabled()
+ ? new TracedReadWriteTransactionCallable(
+ this, callable, /* transactionOptions= */ null, span)
+ : new ReadWriteTransactionCallable(this, callable, /* transactionOptions= */ null));
+ try (Scope ignored = span.makeCurrent()) {
return RetryHelper.runWithRetries(
- new ReadWriteTransactionCallable(this, callable, null),
+ transactionCallable,
retrySettings,
TRANSACTION_EXCEPTION_HANDLER,
getOptions().getClock());
} catch (RetryHelperException e) {
- span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
+ span.end(e);
throw DatastoreException.translateAndThrow(e);
} finally {
- span.end(TraceUtil.END_SPAN_OPTIONS);
+ span.end();
}
}
@Override
public T runInTransaction(
final TransactionCallable callable, TransactionOptions transactionOptions) {
- Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION);
- try (Scope scope = traceUtil.getTracer().withSpan(span)) {
+ TraceUtil.Span span = otelTraceUtil.startSpan(SPAN_NAME_TRANSACTION_RUN);
+
+ Callable transactionCallable =
+ (getOptions().getOpenTelemetryOptions().isEnabled()
+ ? new TracedReadWriteTransactionCallable(this, callable, transactionOptions, span)
+ : new ReadWriteTransactionCallable(this, callable, transactionOptions));
+
+ try (Scope ignored = span.makeCurrent()) {
return RetryHelper.runWithRetries(
- new ReadWriteTransactionCallable(this, callable, transactionOptions),
+ transactionCallable,
retrySettings,
TRANSACTION_EXCEPTION_HANDLER,
getOptions().getClock());
} catch (RetryHelperException e) {
- span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
+ span.end(e);
throw DatastoreException.translateAndThrow(e);
} finally {
- span.end(TraceUtil.END_SPAN_OPTIONS);
+ span.end();
}
}
@@ -236,20 +351,39 @@ public AggregationResults runAggregation(
com.google.datastore.v1.RunQueryResponse runQuery(
final com.google.datastore.v1.RunQueryRequest requestPb) {
- Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_RUNQUERY);
- try (Scope scope = traceUtil.getTracer().withSpan(span)) {
- return RetryHelper.runWithRetries(
- () -> datastoreRpc.runQuery(requestPb),
- retrySettings,
- requestPb.getReadOptions().getTransaction().isEmpty()
- ? EXCEPTION_HANDLER
- : TRANSACTION_OPERATION_EXCEPTION_HANDLER,
- getOptions().getClock());
+ ReadOptions readOptions = requestPb.getReadOptions();
+ boolean isTransactional = readOptions.hasTransaction() || readOptions.hasNewTransaction();
+ String spanName = (isTransactional ? SPAN_NAME_TRANSACTION_RUN_QUERY : SPAN_NAME_RUN_QUERY);
+ com.google.cloud.datastore.telemetry.TraceUtil.Span span = otelTraceUtil.startSpan(spanName);
+
+ try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
+ RunQueryResponse response =
+ RetryHelper.runWithRetries(
+ () -> datastoreRpc.runQuery(requestPb),
+ retrySettings,
+ requestPb.getReadOptions().getTransaction().isEmpty()
+ ? EXCEPTION_HANDLER
+ : TRANSACTION_OPERATION_EXCEPTION_HANDLER,
+ getOptions().getClock());
+ span.addEvent(
+ spanName + " complete.",
+ new ImmutableMap.Builder()
+ .put(ATTRIBUTES_KEY_DOCUMENT_COUNT, response.getBatch().getEntityResultsCount())
+ .put(ATTRIBUTES_KEY_TRANSACTIONAL, isTransactional)
+ .put(ATTRIBUTES_KEY_READ_CONSISTENCY, readOptions.getReadConsistency().toString())
+ .put(
+ ATTRIBUTES_KEY_TRANSACTION_ID,
+ (isTransactional
+ ? requestPb.getReadOptions().getTransaction().toStringUtf8()
+ : ""))
+ .put(ATTRIBUTES_KEY_MORE_RESULTS, response.getBatch().getMoreResults().toString())
+ .build());
+ return response;
} catch (RetryHelperException e) {
- span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
+ span.end(e);
throw DatastoreException.translateAndThrow(e);
} finally {
- span.end(TraceUtil.END_SPAN_OPTIONS);
+ span.end();
}
}
@@ -291,8 +425,9 @@ public List allocateId(IncompleteKey... keys) {
private com.google.datastore.v1.AllocateIdsResponse allocateIds(
final com.google.datastore.v1.AllocateIdsRequest requestPb) {
- Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_ALLOCATEIDS);
- try (Scope scope = traceUtil.getTracer().withSpan(span)) {
+ com.google.cloud.datastore.telemetry.TraceUtil.Span span =
+ otelTraceUtil.startSpan(SPAN_NAME_ALLOCATE_IDS);
+ try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
return RetryHelper.runWithRetries(
new Callable() {
@Override
@@ -304,10 +439,10 @@ public com.google.datastore.v1.AllocateIdsResponse call() throws DatastoreExcept
EXCEPTION_HANDLER,
getOptions().getClock());
} catch (RetryHelperException e) {
- span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
+ span.end(e);
throw DatastoreException.translateAndThrow(e);
} finally {
- span.end(TraceUtil.END_SPAN_OPTIONS);
+ span.end();
}
}
@@ -450,14 +585,27 @@ protected Entity computeNext() {
com.google.datastore.v1.LookupResponse lookup(
final com.google.datastore.v1.LookupRequest requestPb) {
- Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_LOOKUP);
- try (Scope scope = traceUtil.getTracer().withSpan(span)) {
+ ReadOptions readOptions = requestPb.getReadOptions();
+ boolean isTransactional = readOptions.hasTransaction() || readOptions.hasNewTransaction();
+ String spanName = (isTransactional ? SPAN_NAME_TRANSACTION_LOOKUP : SPAN_NAME_LOOKUP);
+ com.google.cloud.datastore.telemetry.TraceUtil.Span span = otelTraceUtil.startSpan(spanName);
+
+ try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
return RetryHelper.runWithRetries(
- new Callable() {
- @Override
- public com.google.datastore.v1.LookupResponse call() throws DatastoreException {
- return datastoreRpc.lookup(requestPb);
- }
+ () -> {
+ com.google.datastore.v1.LookupResponse response = datastoreRpc.lookup(requestPb);
+ span.addEvent(
+ spanName + " complete.",
+ new ImmutableMap.Builder()
+ .put(ATTRIBUTES_KEY_RECEIVED, response.getFoundCount())
+ .put(ATTRIBUTES_KEY_MISSING, response.getMissingCount())
+ .put(ATTRIBUTES_KEY_DEFERRED, response.getDeferredCount())
+ .put(ATTRIBUTES_KEY_TRANSACTIONAL, isTransactional)
+ .put(
+ ATTRIBUTES_KEY_TRANSACTION_ID,
+ isTransactional ? readOptions.getTransaction().toStringUtf8() : "")
+ .build());
+ return response;
},
retrySettings,
requestPb.getReadOptions().getTransaction().isEmpty()
@@ -465,10 +613,10 @@ public com.google.datastore.v1.LookupResponse call() throws DatastoreException {
: TRANSACTION_OPERATION_EXCEPTION_HANDLER,
getOptions().getClock());
} catch (RetryHelperException e) {
- span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
+ span.end(e);
throw DatastoreException.translateAndThrow(e);
} finally {
- span.end(TraceUtil.END_SPAN_OPTIONS);
+ span.end();
}
}
@@ -492,8 +640,9 @@ public List reserveIds(Key... keys) {
com.google.datastore.v1.ReserveIdsResponse reserveIds(
final com.google.datastore.v1.ReserveIdsRequest requestPb) {
- Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_RESERVEIDS);
- try (Scope scope = traceUtil.getTracer().withSpan(span)) {
+ com.google.cloud.datastore.telemetry.TraceUtil.Span span =
+ otelTraceUtil.startSpan(SPAN_NAME_RESERVE_IDS);
+ try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
return RetryHelper.runWithRetries(
new Callable() {
@Override
@@ -505,10 +654,10 @@ public com.google.datastore.v1.ReserveIdsResponse call() throws DatastoreExcepti
EXCEPTION_HANDLER,
getOptions().getClock());
} catch (RetryHelperException e) {
- span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
+ span.end(e);
throw DatastoreException.translateAndThrow(e);
} finally {
- span.end(TraceUtil.END_SPAN_OPTIONS);
+ span.end();
}
}
@@ -602,20 +751,34 @@ private com.google.datastore.v1.CommitResponse commitMutation(
com.google.datastore.v1.CommitResponse commit(
final com.google.datastore.v1.CommitRequest requestPb) {
- Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_COMMIT);
- try (Scope scope = traceUtil.getTracer().withSpan(span)) {
- return RetryHelper.runWithRetries(
- () -> datastoreRpc.commit(requestPb),
- retrySettings,
- requestPb.getTransaction().isEmpty()
- ? EXCEPTION_HANDLER
- : TRANSACTION_OPERATION_EXCEPTION_HANDLER,
- getOptions().getClock());
+ final boolean isTransactional =
+ requestPb.hasTransaction() || requestPb.hasSingleUseTransaction();
+ final String spanName = isTransactional ? SPAN_NAME_TRANSACTION_COMMIT : SPAN_NAME_COMMIT;
+ com.google.cloud.datastore.telemetry.TraceUtil.Span span = otelTraceUtil.startSpan(spanName);
+ try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
+ CommitResponse response =
+ RetryHelper.runWithRetries(
+ () -> datastoreRpc.commit(requestPb),
+ retrySettings,
+ requestPb.getTransaction().isEmpty()
+ ? EXCEPTION_HANDLER
+ : TRANSACTION_OPERATION_EXCEPTION_HANDLER,
+ getOptions().getClock());
+ span.addEvent(
+ spanName + " complete.",
+ new ImmutableMap.Builder()
+ .put(ATTRIBUTES_KEY_DOCUMENT_COUNT, response.getMutationResultsCount())
+ .put(ATTRIBUTES_KEY_TRANSACTIONAL, isTransactional)
+ .put(
+ ATTRIBUTES_KEY_TRANSACTION_ID,
+ isTransactional ? requestPb.getTransaction().toStringUtf8() : "")
+ .build());
+ return response;
} catch (RetryHelperException e) {
- span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
+ span.end(e);
throw DatastoreException.translateAndThrow(e);
} finally {
- span.end(TraceUtil.END_SPAN_OPTIONS);
+ span.end();
}
}
@@ -626,24 +789,19 @@ ByteString requestTransactionId(
com.google.datastore.v1.BeginTransactionResponse beginTransaction(
final com.google.datastore.v1.BeginTransactionRequest requestPb) {
- Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_BEGINTRANSACTION);
- try (Scope scope = traceUtil.getTracer().withSpan(span)) {
+ com.google.cloud.datastore.telemetry.TraceUtil.Span span =
+ otelTraceUtil.startSpan(SPAN_NAME_BEGIN_TRANSACTION);
+ try (com.google.cloud.datastore.telemetry.TraceUtil.Scope scope = span.makeCurrent()) {
return RetryHelper.runWithRetries(
- new Callable() {
- @Override
- public com.google.datastore.v1.BeginTransactionResponse call()
- throws DatastoreException {
- return datastoreRpc.beginTransaction(requestPb);
- }
- },
+ () -> datastoreRpc.beginTransaction(requestPb),
retrySettings,
EXCEPTION_HANDLER,
getOptions().getClock());
} catch (RetryHelperException e) {
- span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
+ span.end(e);
throw DatastoreException.translateAndThrow(e);
} finally {
- span.end(TraceUtil.END_SPAN_OPTIONS);
+ span.end();
}
}
@@ -657,8 +815,9 @@ void rollbackTransaction(ByteString transaction) {
}
void rollback(final com.google.datastore.v1.RollbackRequest requestPb) {
- Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_ROLLBACK);
- try (Scope scope = traceUtil.getTracer().withSpan(span)) {
+ com.google.cloud.datastore.telemetry.TraceUtil.Span span =
+ otelTraceUtil.startSpan(SPAN_NAME_ROLLBACK);
+ try (Scope scope = span.makeCurrent()) {
RetryHelper.runWithRetries(
new Callable() {
@Override
@@ -670,11 +829,16 @@ public Void call() throws DatastoreException {
retrySettings,
EXCEPTION_HANDLER,
getOptions().getClock());
+ span.addEvent(
+ SPAN_NAME_ROLLBACK,
+ new ImmutableMap.Builder()
+ .put(ATTRIBUTES_KEY_TRANSACTION_ID, requestPb.getTransaction().toStringUtf8())
+ .build());
} catch (RetryHelperException e) {
- span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
+ span.end(e);
throw DatastoreException.translateAndThrow(e);
} finally {
- span.end(TraceUtil.END_SPAN_OPTIONS);
+ span.end();
}
}
}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOpenTelemetryOptions.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOpenTelemetryOptions.java
new file mode 100644
index 000000000..ac266562e
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOpenTelemetryOptions.java
@@ -0,0 +1,97 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.datastore;
+
+import io.opentelemetry.api.OpenTelemetry;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+public class DatastoreOpenTelemetryOptions {
+ private final boolean enabled;
+ private final @Nullable OpenTelemetry openTelemetry;
+
+ DatastoreOpenTelemetryOptions(Builder builder) {
+ this.enabled = builder.enabled;
+ this.openTelemetry = builder.openTelemetry;
+ }
+
+ public boolean isEnabled() {
+ return enabled;
+ }
+
+ @Nullable
+ public OpenTelemetry getOpenTelemetry() {
+ return openTelemetry;
+ }
+
+ @Nonnull
+ public DatastoreOpenTelemetryOptions.Builder toBuilder() {
+ return new DatastoreOpenTelemetryOptions.Builder(this);
+ }
+
+ @Nonnull
+ public static DatastoreOpenTelemetryOptions.Builder newBuilder() {
+ return new DatastoreOpenTelemetryOptions.Builder();
+ }
+
+ public static class Builder {
+
+ private boolean enabled;
+
+ @Nullable private OpenTelemetry openTelemetry;
+
+ private Builder() {
+ enabled = false;
+ openTelemetry = null;
+ }
+
+ private Builder(DatastoreOpenTelemetryOptions options) {
+ this.enabled = options.enabled;
+ this.openTelemetry = options.openTelemetry;
+ }
+
+ @Nonnull
+ public DatastoreOpenTelemetryOptions build() {
+ return new DatastoreOpenTelemetryOptions(this);
+ }
+
+ /**
+ * Sets whether tracing should be enabled.
+ *
+ * @param enabled Whether tracing should be enabled.
+ */
+ @Nonnull
+ public DatastoreOpenTelemetryOptions.Builder setTracingEnabled(boolean enabled) {
+ this.enabled = enabled;
+ return this;
+ }
+
+ /**
+ * Sets the {@link OpenTelemetry} to use with this Datastore instance. If telemetry collection
+ * is enabled, but an `OpenTelemetry` is not provided, the Datastore SDK will attempt to use the
+ * `GlobalOpenTelemetry`.
+ *
+ * @param openTelemetry The OpenTelemetry that should be used by this Datastore instance.
+ */
+ @Nonnull
+ public DatastoreOpenTelemetryOptions.Builder setOpenTelemetry(
+ @Nonnull OpenTelemetry openTelemetry) {
+ this.openTelemetry = openTelemetry;
+ return this;
+ }
+ }
+}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java
index 8437c3e22..1ea79298c 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java
@@ -18,19 +18,29 @@
import static com.google.cloud.datastore.Validator.validateNamespace;
+import com.google.api.core.BetaApi;
+import com.google.api.gax.grpc.ChannelPoolSettings;
+import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
+import com.google.api.gax.rpc.TransportChannelProvider;
import com.google.cloud.ServiceDefaults;
import com.google.cloud.ServiceOptions;
import com.google.cloud.ServiceRpc;
import com.google.cloud.TransportOptions;
import com.google.cloud.datastore.spi.DatastoreRpcFactory;
import com.google.cloud.datastore.spi.v1.DatastoreRpc;
+import com.google.cloud.datastore.spi.v1.GrpcDatastoreRpc;
import com.google.cloud.datastore.spi.v1.HttpDatastoreRpc;
+import com.google.cloud.datastore.v1.DatastoreSettings;
+import com.google.cloud.grpc.GrpcTransportOptions;
import com.google.cloud.http.HttpTransportOptions;
import com.google.common.base.MoreObjects;
import com.google.common.collect.ImmutableSet;
+import java.io.IOException;
import java.lang.reflect.Method;
import java.util.Objects;
import java.util.Set;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
public class DatastoreOptions extends ServiceOptions {
@@ -39,10 +49,20 @@ public class DatastoreOptions extends ServiceOptions SCOPES = ImmutableSet.of(DATASTORE_SCOPE);
private static final String DEFAULT_DATABASE_ID = "";
+ public static final String PROJECT_ID_ENV_VAR = "DATASTORE_PROJECT_ID";
+ public static final String LOCAL_HOST_ENV_VAR = "DATASTORE_EMULATOR_HOST";
+ public static final int INIT_CHANNEL_COUNT = 1;
+ public static final int MIN_CHANNEL_COUNT = 1;
+ public static final int MAX_CHANNEL_COUNT = 4;
+
+ private transient TransportChannelProvider channelProvider = null;
private final String namespace;
private final String databaseId;
+ private final transient @Nonnull DatastoreOpenTelemetryOptions openTelemetryOptions;
+ private final transient @Nonnull com.google.cloud.datastore.telemetry.TraceUtil traceUtil;
+
public static class DefaultDatastoreFactory implements DatastoreFactory {
private static final DatastoreFactory INSTANCE = new DefaultDatastoreFactory();
@@ -59,21 +79,47 @@ public static class DefaultDatastoreRpcFactory implements DatastoreRpcFactory {
@Override
public ServiceRpc create(DatastoreOptions options) {
- return new HttpDatastoreRpc(options);
+ try {
+ if (options.getTransportOptions() instanceof GrpcTransportOptions) {
+ return new GrpcDatastoreRpc(options);
+ } else {
+ return new HttpDatastoreRpc(options);
+ }
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
}
}
+ @Nonnull
+ com.google.cloud.datastore.telemetry.TraceUtil getTraceUtil() {
+ return traceUtil;
+ }
+
+ @BetaApi
+ @Nonnull
+ public DatastoreOpenTelemetryOptions getOpenTelemetryOptions() {
+ return openTelemetryOptions;
+ }
+
public static class Builder extends ServiceOptions.Builder {
private String namespace;
private String databaseId;
+ private TransportChannelProvider channelProvider = null;
+ private String host;
+ private TransportOptions transportOptions;
+
+ @Nullable private DatastoreOpenTelemetryOptions openTelemetryOptions = null;
private Builder() {}
private Builder(DatastoreOptions options) {
super(options);
- namespace = options.namespace;
- databaseId = options.databaseId;
+ this.namespace = options.namespace;
+ this.databaseId = options.databaseId;
+ this.openTelemetryOptions = options.openTelemetryOptions;
+ this.channelProvider = validateChannelProvider(options.channelProvider);
}
@Override
@@ -82,11 +128,47 @@ public Builder setTransportOptions(TransportOptions transportOptions) {
throw new IllegalArgumentException(
"Only http transport is allowed for " + API_SHORT_NAME + ".");
}
+ this.transportOptions = transportOptions;
return super.setTransportOptions(transportOptions);
}
+ /**
+ * Sets the transport to gRPC. Note this functionality is experimental and subject to change.
+ */
+ @BetaApi
+ public Builder setTransportOptions(GrpcTransportOptions transportOptions) {
+ this.transportOptions = transportOptions;
+ return super.setTransportOptions(transportOptions);
+ }
+
+ @Override
+ public Builder setHost(String host) {
+ this.host = host;
+ return super.setHost(host);
+ }
+
+ /**
+ * Sets the {@link TransportChannelProvider} to use with this Datastore client.
+ *
+ * This is only compatible with clients using a gRPC transport (see {@code
+ * DatastoreOptions#setTransportOptions(GrpcTransportOptions)} for more details).
+ *
+ *
This functionality is experimental and subject to change.
+ *
+ * @param channelProvider A InstantiatingGrpcChannelProvider object that defines the transport
+ * provider for this client.
+ */
+ @BetaApi
+ public Builder setChannelProvider(TransportChannelProvider channelProvider) {
+ this.channelProvider = validateChannelProvider(channelProvider);
+ return this;
+ }
+
@Override
public DatastoreOptions build() {
+ if (this.host == null && this.transportOptions instanceof GrpcTransportOptions) {
+ this.setHost(DatastoreSettings.getDefaultEndpoint());
+ }
return new DatastoreOptions(this);
}
@@ -100,33 +182,81 @@ public Builder setDatabaseId(String databaseId) {
this.databaseId = databaseId;
return this;
}
+
+ /**
+ * Sets the {@link DatastoreOpenTelemetryOptions} to be used for this Firestore instance.
+ *
+ * @param openTelemetryOptions The `DatastoreOpenTelemetryOptions` to use.
+ */
+ @BetaApi
+ @Nonnull
+ public Builder setOpenTelemetryOptions(
+ @Nonnull DatastoreOpenTelemetryOptions openTelemetryOptions) {
+ this.openTelemetryOptions = openTelemetryOptions;
+ return this;
+ }
+ }
+
+ private static TransportChannelProvider validateChannelProvider(
+ TransportChannelProvider channelProvider) {
+ if (channelProvider != null && !(channelProvider instanceof InstantiatingGrpcChannelProvider)) {
+ throw new IllegalArgumentException(
+ "Only GRPC channels are allowed for " + API_SHORT_NAME + ".");
+ }
+ return channelProvider;
}
private DatastoreOptions(Builder builder) {
super(DatastoreFactory.class, DatastoreRpcFactory.class, builder, new DatastoreDefaults());
+
+ this.openTelemetryOptions =
+ builder.openTelemetryOptions != null
+ ? builder.openTelemetryOptions
+ : DatastoreOpenTelemetryOptions.newBuilder().build();
+ this.traceUtil = com.google.cloud.datastore.telemetry.TraceUtil.getInstance(this);
+
namespace = MoreObjects.firstNonNull(builder.namespace, defaultNamespace());
databaseId = MoreObjects.firstNonNull(builder.databaseId, DEFAULT_DATABASE_ID);
+
+ if (getTransportOptions() instanceof HttpTransportOptions && builder.channelProvider != null) {
+ throw new IllegalArgumentException(
+ "Only gRPC transport allows setting of channel provider or credentials provider");
+ } else if (getTransportOptions() instanceof GrpcTransportOptions) {
+ // For grpc transport options, configure default gRPC Connection pool with minChannelCount = 1
+ // and maxChannelCount = 4
+ this.channelProvider =
+ builder.channelProvider != null
+ ? builder.channelProvider
+ : GrpcTransportOptions.setUpChannelProvider(
+ DatastoreSettings.defaultGrpcTransportProviderBuilder()
+ .setChannelPoolSettings(
+ ChannelPoolSettings.builder()
+ .setInitialChannelCount(INIT_CHANNEL_COUNT)
+ .setMinChannelCount(MIN_CHANNEL_COUNT)
+ .setMaxChannelCount(MAX_CHANNEL_COUNT)
+ .build()),
+ this);
+ }
+ }
+
+ public TransportChannelProvider getTransportChannelProvider() {
+ return channelProvider;
}
@Override
protected String getDefaultHost() {
- String host =
- System.getProperty(
- com.google.datastore.v1.client.DatastoreHelper.LOCAL_HOST_ENV_VAR,
- System.getenv(com.google.datastore.v1.client.DatastoreHelper.LOCAL_HOST_ENV_VAR));
+ String host = System.getProperty(LOCAL_HOST_ENV_VAR, System.getenv(LOCAL_HOST_ENV_VAR));
return host != null ? host : com.google.datastore.v1.client.DatastoreFactory.DEFAULT_HOST;
}
@Override
protected String getDefaultProject() {
- String projectId =
- System.getProperty(
- com.google.datastore.v1.client.DatastoreHelper.PROJECT_ID_ENV_VAR,
- System.getenv(com.google.datastore.v1.client.DatastoreHelper.PROJECT_ID_ENV_VAR));
+ String projectId = System.getProperty(PROJECT_ID_ENV_VAR, System.getenv(PROJECT_ID_ENV_VAR));
return projectId != null ? projectId : super.getDefaultProject();
}
private static class DatastoreDefaults implements ServiceDefaults {
+ private final TransportOptions TRANSPORT_OPTIONS = getDefaultTransportOptionsBuilder().build();
@Override
public DatastoreFactory getDefaultServiceFactory() {
@@ -140,7 +270,11 @@ public DatastoreRpcFactory getDefaultRpcFactory() {
@Override
public TransportOptions getDefaultTransportOptions() {
- return getDefaultHttpTransportOptions();
+ return TRANSPORT_OPTIONS;
+ }
+
+ public static HttpTransportOptions.Builder getDefaultTransportOptionsBuilder() {
+ return HttpTransportOptions.newBuilder();
}
}
@@ -148,6 +282,10 @@ public static HttpTransportOptions getDefaultHttpTransportOptions() {
return HttpTransportOptions.newBuilder().build();
}
+ public static GrpcTransportOptions getDefaultGrpcTransportOptions() {
+ return GrpcTransportOptions.newBuilder().build();
+ }
+
/** Returns the default namespace to be used by the datastore service. */
public String getNamespace() {
return namespace;
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java
index 8aef7f5c0..c3137a9a7 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java
@@ -17,11 +17,13 @@
package com.google.cloud.datastore;
import com.google.api.core.BetaApi;
+import com.google.api.core.InternalExtensionOnly;
import com.google.cloud.datastore.models.ExplainOptions;
import java.util.Iterator;
import java.util.List;
/** An interface to represent Google Cloud Datastore read operations. */
+@InternalExtensionOnly
public interface DatastoreReader {
/**
@@ -61,9 +63,7 @@ public interface DatastoreReader {
*
* @throws DatastoreException upon failure
*/
- default AggregationResults runAggregation(AggregationQuery query) {
- throw new UnsupportedOperationException("Not implemented.");
- }
+ AggregationResults runAggregation(AggregationQuery query);
/**
* Submits a {@link AggregationQuery} with a specified {@link
@@ -72,7 +72,5 @@ default AggregationResults runAggregation(AggregationQuery query) {
* @throws DatastoreException upon failure
*/
@BetaApi
- default AggregationResults runAggregation(AggregationQuery query, ExplainOptions explainOptions) {
- throw new UnsupportedOperationException("Not implemented.");
- }
+ AggregationResults runAggregation(AggregationQuery query, ExplainOptions explainOptions);
}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReaderWriter.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReaderWriter.java
index a51a5aa77..bc8700c70 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReaderWriter.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReaderWriter.java
@@ -16,5 +16,8 @@
package com.google.cloud.datastore;
+import com.google.api.core.InternalExtensionOnly;
+
/** An interface that combines both Google Cloud Datastore read and write operations. */
+@InternalExtensionOnly
public interface DatastoreReaderWriter extends DatastoreReader, DatastoreWriter {}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreUtils.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreUtils.java
new file mode 100644
index 000000000..e991fd51d
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreUtils.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.datastore;
+
+import com.google.api.core.InternalApi;
+import com.google.cloud.NoCredentials;
+import com.google.common.base.Strings;
+import java.net.InetAddress;
+import java.net.URL;
+
+@InternalApi
+public class DatastoreUtils {
+
+ public static boolean isEmulator(DatastoreOptions datastoreOptions) {
+ return isLocalHost(datastoreOptions.getHost())
+ || NoCredentials.getInstance().equals(datastoreOptions.getCredentials());
+ }
+
+ public static boolean isLocalHost(String host) {
+ if (Strings.isNullOrEmpty(host)) {
+ return false;
+ }
+ try {
+ String normalizedHost = "http://" + removeScheme(host);
+ InetAddress hostAddr = InetAddress.getByName(new URL(normalizedHost).getHost());
+ return hostAddr.isAnyLocalAddress() || hostAddr.isLoopbackAddress();
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ public static String removeScheme(String url) {
+ if (url != null) {
+ url = url.toLowerCase();
+ if (url.startsWith("https://")) {
+ return url.substring("https://".length());
+ } else if (url.startsWith("http://")) {
+ return url.substring("http://".length());
+ }
+ }
+ return url;
+ }
+}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreWriter.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreWriter.java
index 6c1d6fdbc..b414995e6 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreWriter.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreWriter.java
@@ -16,9 +16,11 @@
package com.google.cloud.datastore;
+import com.google.api.core.InternalExtensionOnly;
import java.util.List;
/** An interface to represent Google Cloud Datastore write operations. */
+@InternalExtensionOnly
public interface DatastoreWriter {
/**
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GrpcToDatastoreCodeTranslation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GrpcToDatastoreCodeTranslation.java
new file mode 100644
index 000000000..1d63fb19a
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GrpcToDatastoreCodeTranslation.java
@@ -0,0 +1,93 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.datastore;
+
+import com.google.api.gax.grpc.GrpcStatusCode;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.rpc.Code;
+import io.grpc.Status;
+import java.util.Map;
+import java.util.function.Function;
+
+final class GrpcToDatastoreCodeTranslation {
+ /** Mappings between gRPC status codes and their corresponding code numbers. */
+ private static final ImmutableList STATUS_CODE_MAPPINGS =
+ ImmutableList.of(
+ StatusCodeMapping.of(Code.OK.getNumber(), Status.Code.OK),
+ StatusCodeMapping.of(Code.DATA_LOSS.getNumber(), Status.Code.DATA_LOSS),
+ StatusCodeMapping.of(Code.INVALID_ARGUMENT.getNumber(), Status.Code.INVALID_ARGUMENT),
+ StatusCodeMapping.of(Code.OUT_OF_RANGE.getNumber(), Status.Code.OUT_OF_RANGE),
+ StatusCodeMapping.of(Code.UNAUTHENTICATED.getNumber(), Status.Code.UNAUTHENTICATED),
+ StatusCodeMapping.of(Code.PERMISSION_DENIED.getNumber(), Status.Code.PERMISSION_DENIED),
+ StatusCodeMapping.of(Code.NOT_FOUND.getNumber(), Status.Code.NOT_FOUND),
+ StatusCodeMapping.of(Code.ALREADY_EXISTS.getNumber(), Status.Code.ALREADY_EXISTS),
+ StatusCodeMapping.of(
+ Code.FAILED_PRECONDITION.getNumber(), Status.Code.FAILED_PRECONDITION),
+ StatusCodeMapping.of(Code.RESOURCE_EXHAUSTED.getNumber(), Status.Code.RESOURCE_EXHAUSTED),
+ StatusCodeMapping.of(Code.INTERNAL.getNumber(), Status.Code.INTERNAL),
+ StatusCodeMapping.of(Code.UNIMPLEMENTED.getNumber(), Status.Code.UNIMPLEMENTED),
+ StatusCodeMapping.of(Code.UNAVAILABLE.getNumber(), Status.Code.UNAVAILABLE),
+ StatusCodeMapping.of(Code.DEADLINE_EXCEEDED.getNumber(), Status.Code.DEADLINE_EXCEEDED),
+ StatusCodeMapping.of(Code.ABORTED.getNumber(), Status.Code.ABORTED),
+ StatusCodeMapping.of(Code.CANCELLED.getNumber(), Status.Code.CANCELLED),
+ StatusCodeMapping.of(Code.UNKNOWN.getNumber(), Status.Code.UNKNOWN));
+
+ /** Index our {@link StatusCodeMapping} for constant time lookup by {@link Status.Code} */
+ private static final Map GRPC_CODE_INDEX =
+ STATUS_CODE_MAPPINGS.stream()
+ .collect(
+ ImmutableMap.toImmutableMap(StatusCodeMapping::getGrpcCode, Function.identity()));
+
+ static int grpcCodeToDatastoreStatusCode(Status.Code code) {
+ StatusCodeMapping found = GRPC_CODE_INDEX.get(code);
+ // theoretically it's possible for gRPC to add a new code we haven't mapped here, if this
+ // happens fall through to our default of 0
+ if (found != null) {
+ return found.getDatastoreCode();
+ } else {
+ return 0;
+ }
+ }
+
+ /**
+ * Simple tuple class to bind together our corresponding http status code and {@link Status.Code}
+ * while providing easy access to the correct {@link GrpcStatusCode} where necessary.
+ */
+ private static final class StatusCodeMapping {
+
+ private final int datastoreCode;
+
+ private final Status.Code grpcCode;
+
+ private StatusCodeMapping(int datastoreCode, Status.Code grpcCode) {
+ this.datastoreCode = datastoreCode;
+ this.grpcCode = grpcCode;
+ }
+
+ public int getDatastoreCode() {
+ return datastoreCode;
+ }
+
+ public Status.Code getGrpcCode() {
+ return grpcCode;
+ }
+
+ static StatusCodeMapping of(int datastoreCode, Status.Code grpcCode) {
+ return new StatusCodeMapping(datastoreCode, grpcCode);
+ }
+ }
+}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/IncompleteKey.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/IncompleteKey.java
index db9973cb5..71e31b94d 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/IncompleteKey.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/IncompleteKey.java
@@ -91,9 +91,11 @@ public Key getParent() {
PathElement parent = ancestors.get(ancestors.size() - 1);
Key.Builder keyBuilder;
if (parent.hasName()) {
- keyBuilder = Key.newBuilder(getProjectId(), parent.getKind(), parent.getName());
+ keyBuilder =
+ Key.newBuilder(getProjectId(), parent.getKind(), parent.getName(), getDatabaseId());
} else {
- keyBuilder = Key.newBuilder(getProjectId(), parent.getKind(), parent.getId());
+ keyBuilder =
+ Key.newBuilder(getProjectId(), parent.getKind(), parent.getId(), getDatabaseId());
}
String namespace = getNamespace();
if (namespace != null) {
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Key.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Key.java
index 9e851d0cb..14fe264bb 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Key.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Key.java
@@ -145,7 +145,7 @@ public Object getNameOrId() {
/** Returns the key in an encoded form that can be used as part of a URL. */
public String toUrlSafe() {
try {
- return URLEncoder.encode(TextFormat.printToString(toPb()), UTF_8.name());
+ return URLEncoder.encode(TextFormat.printer().printToString(toPb()), UTF_8.name());
} catch (UnsupportedEncodingException e) {
throw new IllegalStateException("Unexpected encoding exception", e);
}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java
index 50433a6a9..ca5b240ad 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java
@@ -17,6 +17,7 @@
package com.google.cloud.datastore;
import com.google.api.core.BetaApi;
+import com.google.api.core.InternalExtensionOnly;
import com.google.cloud.datastore.models.ExplainMetrics;
import com.google.datastore.v1.QueryResultBatch;
import java.util.Iterator;
@@ -31,6 +32,7 @@
*
* @param the type of the results value.
*/
+@InternalExtensionOnly
public interface QueryResults extends Iterator {
/** Returns the actual class of the result's values. */
@@ -75,7 +77,5 @@ public interface QueryResults extends Iterator {
QueryResultBatch.MoreResultsType getMoreResults();
@BetaApi
- default Optional getExplainMetrics() {
- throw new UnsupportedOperationException("Not implemented.");
- }
+ Optional getExplainMetrics();
}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java
index c4a85caab..630ddd225 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java
@@ -16,13 +16,13 @@
package com.google.cloud.datastore;
import static com.google.cloud.BaseService.EXCEPTION_HANDLER;
-import static com.google.cloud.datastore.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY;
import com.google.api.core.InternalApi;
import com.google.api.gax.retrying.RetrySettings;
import com.google.cloud.RetryHelper;
import com.google.cloud.RetryHelper.RetryHelperException;
import com.google.cloud.datastore.spi.v1.DatastoreRpc;
+import com.google.cloud.datastore.telemetry.TraceUtil;
import com.google.datastore.v1.AllocateIdsRequest;
import com.google.datastore.v1.AllocateIdsResponse;
import com.google.datastore.v1.BeginTransactionRequest;
@@ -31,6 +31,7 @@
import com.google.datastore.v1.CommitResponse;
import com.google.datastore.v1.LookupRequest;
import com.google.datastore.v1.LookupResponse;
+import com.google.datastore.v1.ReadOptions;
import com.google.datastore.v1.ReserveIdsRequest;
import com.google.datastore.v1.ReserveIdsResponse;
import com.google.datastore.v1.RollbackRequest;
@@ -39,9 +40,6 @@
import com.google.datastore.v1.RunAggregationQueryResponse;
import com.google.datastore.v1.RunQueryRequest;
import com.google.datastore.v1.RunQueryResponse;
-import io.opencensus.common.Scope;
-import io.opencensus.trace.Span;
-import io.opencensus.trace.Status;
import java.util.concurrent.Callable;
/**
@@ -52,19 +50,19 @@
public class RetryAndTraceDatastoreRpcDecorator implements DatastoreRpc {
private final DatastoreRpc datastoreRpc;
- private final TraceUtil traceUtil;
+ private final com.google.cloud.datastore.telemetry.TraceUtil otelTraceUtil;
private final RetrySettings retrySettings;
private final DatastoreOptions datastoreOptions;
public RetryAndTraceDatastoreRpcDecorator(
DatastoreRpc datastoreRpc,
- TraceUtil traceUtil,
+ TraceUtil otelTraceUtil,
RetrySettings retrySettings,
DatastoreOptions datastoreOptions) {
this.datastoreRpc = datastoreRpc;
- this.traceUtil = traceUtil;
this.retrySettings = retrySettings;
this.datastoreOptions = datastoreOptions;
+ this.otelTraceUtil = otelTraceUtil;
}
@Override
@@ -105,20 +103,36 @@ public RunQueryResponse runQuery(RunQueryRequest request) {
@Override
public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request) {
- return invokeRpc(
- () -> datastoreRpc.runAggregationQuery(request), SPAN_NAME_RUN_AGGREGATION_QUERY);
+ ReadOptions readOptions = request.getReadOptions();
+ boolean isTransactional = readOptions.hasTransaction() || readOptions.hasNewTransaction();
+ String spanName =
+ (isTransactional
+ ? com.google.cloud.datastore.telemetry.TraceUtil
+ .SPAN_NAME_TRANSACTION_RUN_AGGREGATION_QUERY
+ : com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY);
+ return invokeRpc(() -> datastoreRpc.runAggregationQuery(request), spanName);
+ }
+
+ @Override
+ public void close() throws Exception {
+ datastoreRpc.close();
+ }
+
+ @Override
+ public boolean isClosed() {
+ return datastoreRpc.isClosed();
}
public O invokeRpc(Callable block, String startSpan) {
- Span span = traceUtil.startSpan(startSpan);
- try (Scope scope = traceUtil.getTracer().withSpan(span)) {
+ com.google.cloud.datastore.telemetry.TraceUtil.Span span = otelTraceUtil.startSpan(startSpan);
+ try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
return RetryHelper.runWithRetries(
block, this.retrySettings, EXCEPTION_HANDLER, this.datastoreOptions.getClock());
} catch (RetryHelperException e) {
- span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
+ span.end(e);
throw DatastoreException.translateAndThrow(e);
} finally {
- span.end(TraceUtil.END_SPAN_OPTIONS);
+ span.end();
}
}
}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java
index 30cd05759..5bde80ed6 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java
@@ -27,6 +27,7 @@
import com.google.api.core.ApiFunction;
import com.google.api.core.InternalApi;
+import com.google.api.core.InternalExtensionOnly;
import com.google.cloud.StringEnumType;
import com.google.cloud.StringEnumValue;
import com.google.cloud.Timestamp;
@@ -700,6 +701,7 @@ public String toString() {
*
* @param the type of result the query returns.
*/
+ @InternalExtensionOnly
public interface Builder {
/** Sets the namespace for the query. */
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java
deleted file mode 100644
index 57525d15d..000000000
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java
+++ /dev/null
@@ -1,76 +0,0 @@
-/*
- * Copyright 2020 Google LLC
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.google.cloud.datastore;
-
-import com.google.cloud.datastore.spi.v1.HttpDatastoreRpc;
-import io.opencensus.trace.EndSpanOptions;
-import io.opencensus.trace.Span;
-import io.opencensus.trace.Tracer;
-import io.opencensus.trace.Tracing;
-
-/**
- * Helper class for tracing utility. It is used for instrumenting {@link HttpDatastoreRpc} with
- * OpenCensus APIs.
- *
- * TraceUtil instances are created by the {@link TraceUtil#getInstance()} method.
- */
-public class TraceUtil {
- private final Tracer tracer = Tracing.getTracer();
- private static final TraceUtil traceUtil = new TraceUtil();
- static final String SPAN_NAME_ALLOCATEIDS = "CloudDatastoreOperation.allocateIds";
- static final String SPAN_NAME_TRANSACTION = "CloudDatastoreOperation.readWriteTransaction";
- static final String SPAN_NAME_BEGINTRANSACTION = "CloudDatastoreOperation.beginTransaction";
- static final String SPAN_NAME_COMMIT = "CloudDatastoreOperation.commit";
- static final String SPAN_NAME_LOOKUP = "CloudDatastoreOperation.lookup";
- static final String SPAN_NAME_RESERVEIDS = "CloudDatastoreOperation.reserveIds";
- static final String SPAN_NAME_ROLLBACK = "CloudDatastoreOperation.rollback";
- static final String SPAN_NAME_RUNQUERY = "CloudDatastoreOperation.runQuery";
- static final String SPAN_NAME_RUN_AGGREGATION_QUERY =
- "CloudDatastoreOperation.runAggregationQuery";
- static final EndSpanOptions END_SPAN_OPTIONS =
- EndSpanOptions.builder().setSampleToLocalSpanStore(true).build();
-
- /**
- * Starts a new span.
- *
- * @param spanName The name of the returned Span.
- * @return The newly created {@link Span}.
- */
- protected Span startSpan(String spanName) {
- return tracer.spanBuilder(spanName).startSpan();
- }
-
- /**
- * Return the global {@link Tracer}.
- *
- * @return The global {@link Tracer}.
- */
- public Tracer getTracer() {
- return tracer;
- }
-
- /**
- * Return TraceUtil Object.
- *
- * @return An instance of {@link TraceUtil}
- */
- public static TraceUtil getInstance() {
- return traceUtil;
- }
-
- private TraceUtil() {}
-}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Transaction.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Transaction.java
index 7b6a67a2d..697e7a6ff 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Transaction.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Transaction.java
@@ -17,6 +17,7 @@
package com.google.cloud.datastore;
import com.google.api.core.BetaApi;
+import com.google.api.core.InternalExtensionOnly;
import com.google.cloud.datastore.models.ExplainOptions;
import com.google.protobuf.ByteString;
import java.util.Iterator;
@@ -63,6 +64,7 @@
* This class too should not be treated as a thread safe class.
*/
@NotThreadSafe
+@InternalExtensionOnly
public interface Transaction extends DatastoreBatchWriter, DatastoreReaderWriter {
interface Response {
@@ -179,9 +181,7 @@ interface Response {
QueryResults run(Query query);
@BetaApi
- default QueryResults run(Query query, ExplainOptions explainOptions) {
- throw new UnsupportedOperationException("Not implemented.");
- }
+ QueryResults run(Query query, ExplainOptions explainOptions);
/**
* Datastore add operation. This method will also allocate id for any entity with an incomplete
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java
index f08a908ec..e730db81f 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java
@@ -20,6 +20,7 @@
import com.google.api.core.BetaApi;
import com.google.cloud.datastore.models.ExplainOptions;
+import com.google.cloud.datastore.telemetry.TraceUtil;
import com.google.common.collect.ImmutableList;
import com.google.datastore.v1.ReadOptions;
import com.google.datastore.v1.TransactionOptions;
@@ -28,6 +29,7 @@
import java.util.Iterator;
import java.util.List;
import java.util.Optional;
+import javax.annotation.Nonnull;
final class TransactionImpl extends BaseDatastoreBatchWriter implements Transaction {
@@ -37,6 +39,8 @@ final class TransactionImpl extends BaseDatastoreBatchWriter implements Transact
private final ReadOptionProtoPreparer readOptionProtoPreparer;
+ @Nonnull private final TraceUtil traceUtil;
+
static class ResponseImpl implements Transaction.Response {
private final com.google.datastore.v1.CommitResponse response;
@@ -78,6 +82,7 @@ public List getGeneratedKeys() {
transactionId = datastore.requestTransactionId(requestPb);
this.readOptionProtoPreparer = new ReadOptionProtoPreparer();
+ this.traceUtil = datastore.getOptions().getTraceUtil();
}
@Override
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ValueBuilder.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ValueBuilder.java
index 3c60ef409..315728147 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ValueBuilder.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ValueBuilder.java
@@ -16,6 +16,7 @@
package com.google.cloud.datastore;
+import com.google.api.core.InternalExtensionOnly;
import com.google.cloud.GcpLaunchStage;
/**
@@ -25,6 +26,7 @@
* @param the value type.
* @param the value type's associated builder.
*/
+@InternalExtensionOnly
public interface ValueBuilder, B extends ValueBuilder> {
ValueType getValueType();
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/DatastoreAdminClient.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/DatastoreAdminClient.java
index 2e58d6679..9ef7f1edf 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/DatastoreAdminClient.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/DatastoreAdminClient.java
@@ -1,5 +1,5 @@
/*
- * Copyright 2024 Google LLC
+ * Copyright 2025 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/DatastoreAdminSettings.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/DatastoreAdminSettings.java
index 4bae1f6f3..35f993cc8 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/DatastoreAdminSettings.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/DatastoreAdminSettings.java
@@ -1,5 +1,5 @@
/*
- * Copyright 2024 Google LLC
+ * Copyright 2025 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@@ -65,7 +65,9 @@
* The builder of this class is recursive, so contained classes are themselves builders. When
* build() is called, the tree of builders is called to create the complete settings object.
*
- *
For example, to set the total timeout of getIndex to 30 seconds:
+ *
For example, to set the
+ * [RetrySettings](https://cloud.google.com/java/docs/reference/gax/latest/com.google.api.gax.retrying.RetrySettings)
+ * of getIndex:
*
*
{@code
* // This snippet has been automatically generated and should be regarded as a code template only.
@@ -82,10 +84,47 @@
* .getIndexSettings()
* .getRetrySettings()
* .toBuilder()
- * .setTotalTimeout(Duration.ofSeconds(30))
+ * .setInitialRetryDelayDuration(Duration.ofSeconds(1))
+ * .setInitialRpcTimeoutDuration(Duration.ofSeconds(5))
+ * .setMaxAttempts(5)
+ * .setMaxRetryDelayDuration(Duration.ofSeconds(30))
+ * .setMaxRpcTimeoutDuration(Duration.ofSeconds(60))
+ * .setRetryDelayMultiplier(1.3)
+ * .setRpcTimeoutMultiplier(1.5)
+ * .setTotalTimeoutDuration(Duration.ofSeconds(300))
* .build());
* DatastoreAdminSettings datastoreAdminSettings = datastoreAdminSettingsBuilder.build();
* }
+ *
+ * Please refer to the [Client Side Retry
+ * Guide](https://github.com/googleapis/google-cloud-java/blob/main/docs/client_retries.md) for
+ * additional support in setting retries.
+ *
+ * To configure the RetrySettings of a Long Running Operation method, create an
+ * OperationTimedPollAlgorithm object and update the RPC's polling algorithm. For example, to
+ * configure the RetrySettings for exportEntities:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * DatastoreAdminSettings.Builder datastoreAdminSettingsBuilder =
+ * DatastoreAdminSettings.newBuilder();
+ * TimedRetryAlgorithm timedRetryAlgorithm =
+ * OperationalTimedPollAlgorithm.create(
+ * RetrySettings.newBuilder()
+ * .setInitialRetryDelayDuration(Duration.ofMillis(500))
+ * .setRetryDelayMultiplier(1.5)
+ * .setMaxRetryDelayDuration(Duration.ofMillis(5000))
+ * .setTotalTimeoutDuration(Duration.ofHours(24))
+ * .build());
+ * datastoreAdminSettingsBuilder
+ * .createClusterOperationSettings()
+ * .setPollingAlgorithm(timedRetryAlgorithm)
+ * .build();
+ * }
*/
@Generated("by gapic-generator-java")
public class DatastoreAdminSettings extends ClientSettings {
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/package-info.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/package-info.java
index bf70b678a..73f593170 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/package-info.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/package-info.java
@@ -1,5 +1,5 @@
/*
- * Copyright 2024 Google LLC
+ * Copyright 2025 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/DatastoreAdminStub.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/DatastoreAdminStub.java
index 90c3c45f9..e5a17cc7d 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/DatastoreAdminStub.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/DatastoreAdminStub.java
@@ -1,5 +1,5 @@
/*
- * Copyright 2024 Google LLC
+ * Copyright 2025 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/DatastoreAdminStubSettings.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/DatastoreAdminStubSettings.java
index 64ad035d8..77314c4bc 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/DatastoreAdminStubSettings.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/DatastoreAdminStubSettings.java
@@ -1,5 +1,5 @@
/*
- * Copyright 2024 Google LLC
+ * Copyright 2025 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@@ -21,6 +21,7 @@
import com.google.api.core.ApiFunction;
import com.google.api.core.ApiFuture;
import com.google.api.core.BetaApi;
+import com.google.api.core.ObsoleteApi;
import com.google.api.gax.core.GaxProperties;
import com.google.api.gax.core.GoogleCredentialsProvider;
import com.google.api.gax.core.InstantiatingExecutorProvider;
@@ -66,9 +67,9 @@
import com.google.longrunning.Operation;
import com.google.protobuf.Empty;
import java.io.IOException;
+import java.time.Duration;
import java.util.List;
import javax.annotation.Generated;
-import org.threeten.bp.Duration;
// AUTO-GENERATED DOCUMENTATION AND CLASS.
/**
@@ -85,7 +86,9 @@
* The builder of this class is recursive, so contained classes are themselves builders. When
* build() is called, the tree of builders is called to create the complete settings object.
*
- *
For example, to set the total timeout of getIndex to 30 seconds:
+ *
For example, to set the
+ * [RetrySettings](https://cloud.google.com/java/docs/reference/gax/latest/com.google.api.gax.retrying.RetrySettings)
+ * of getIndex:
*
*
{@code
* // This snippet has been automatically generated and should be regarded as a code template only.
@@ -102,10 +105,47 @@
* .getIndexSettings()
* .getRetrySettings()
* .toBuilder()
- * .setTotalTimeout(Duration.ofSeconds(30))
+ * .setInitialRetryDelayDuration(Duration.ofSeconds(1))
+ * .setInitialRpcTimeoutDuration(Duration.ofSeconds(5))
+ * .setMaxAttempts(5)
+ * .setMaxRetryDelayDuration(Duration.ofSeconds(30))
+ * .setMaxRpcTimeoutDuration(Duration.ofSeconds(60))
+ * .setRetryDelayMultiplier(1.3)
+ * .setRpcTimeoutMultiplier(1.5)
+ * .setTotalTimeoutDuration(Duration.ofSeconds(300))
* .build());
* DatastoreAdminStubSettings datastoreAdminSettings = datastoreAdminSettingsBuilder.build();
* }
+ *
+ * Please refer to the [Client Side Retry
+ * Guide](https://github.com/googleapis/google-cloud-java/blob/main/docs/client_retries.md) for
+ * additional support in setting retries.
+ *
+ * To configure the RetrySettings of a Long Running Operation method, create an
+ * OperationTimedPollAlgorithm object and update the RPC's polling algorithm. For example, to
+ * configure the RetrySettings for exportEntities:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * DatastoreAdminStubSettings.Builder datastoreAdminSettingsBuilder =
+ * DatastoreAdminStubSettings.newBuilder();
+ * TimedRetryAlgorithm timedRetryAlgorithm =
+ * OperationalTimedPollAlgorithm.create(
+ * RetrySettings.newBuilder()
+ * .setInitialRetryDelayDuration(Duration.ofMillis(500))
+ * .setRetryDelayMultiplier(1.5)
+ * .setMaxRetryDelayDuration(Duration.ofMillis(5000))
+ * .setTotalTimeoutDuration(Duration.ofHours(24))
+ * .build());
+ * datastoreAdminSettingsBuilder
+ * .createClusterOperationSettings()
+ * .setPollingAlgorithm(timedRetryAlgorithm)
+ * .build();
+ * }
*/
@Generated("by gapic-generator-java")
public class DatastoreAdminStubSettings extends StubSettings {
@@ -163,9 +203,7 @@ public String extractNextToken(ListIndexesResponse payload) {
@Override
public Iterable extractResources(ListIndexesResponse payload) {
- return payload.getIndexesList() == null
- ? ImmutableList.of()
- : payload.getIndexesList();
+ return payload.getIndexesList();
}
};
@@ -258,15 +296,6 @@ public DatastoreAdminStub createStub() throws IOException {
"Transport not supported: %s", getTransportChannelProvider().getTransportName()));
}
- /** Returns the endpoint set by the user or the the service's default endpoint. */
- @Override
- public String getEndpoint() {
- if (super.getEndpoint() != null) {
- return super.getEndpoint();
- }
- return getDefaultEndpoint();
- }
-
/** Returns the default service name. */
@Override
public String getServiceName() {
@@ -279,6 +308,7 @@ public static InstantiatingExecutorProvider.Builder defaultExecutorProviderBuild
}
/** Returns the default service endpoint. */
+ @ObsoleteApi("Use getEndpoint() instead")
public static String getDefaultEndpoint() {
return "datastore.googleapis.com:443";
}
@@ -419,21 +449,21 @@ public static class Builder extends StubSettings.Builder getIndexSettings() {
return listIndexesSettings;
}
- /** Returns the endpoint set by the user or the the service's default endpoint. */
- @Override
- public String getEndpoint() {
- if (super.getEndpoint() != null) {
- return super.getEndpoint();
- }
- return getDefaultEndpoint();
- }
-
@Override
public DatastoreAdminStubSettings build() throws IOException {
return new DatastoreAdminStubSettings(this);
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/GrpcDatastoreAdminCallableFactory.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/GrpcDatastoreAdminCallableFactory.java
index 9fa4e901b..fa1147044 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/GrpcDatastoreAdminCallableFactory.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/GrpcDatastoreAdminCallableFactory.java
@@ -1,5 +1,5 @@
/*
- * Copyright 2024 Google LLC
+ * Copyright 2025 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/GrpcDatastoreAdminStub.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/GrpcDatastoreAdminStub.java
index 0ab662814..1f6e840e0 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/GrpcDatastoreAdminStub.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/GrpcDatastoreAdminStub.java
@@ -1,5 +1,5 @@
/*
- * Copyright 2024 Google LLC
+ * Copyright 2025 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/HttpJsonDatastoreAdminCallableFactory.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/HttpJsonDatastoreAdminCallableFactory.java
index 0f992e4c2..e56561c69 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/HttpJsonDatastoreAdminCallableFactory.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/HttpJsonDatastoreAdminCallableFactory.java
@@ -1,5 +1,5 @@
/*
- * Copyright 2024 Google LLC
+ * Copyright 2025 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/HttpJsonDatastoreAdminStub.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/HttpJsonDatastoreAdminStub.java
index af095ad28..7c1a53ea3 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/HttpJsonDatastoreAdminStub.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/admin/v1/stub/HttpJsonDatastoreAdminStub.java
@@ -1,5 +1,5 @@
/*
- * Copyright 2024 Google LLC
+ * Copyright 2025 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/AggregationBuilder.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/AggregationBuilder.java
index ce23edcf0..632f44393 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/AggregationBuilder.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/AggregationBuilder.java
@@ -16,6 +16,8 @@
package com.google.cloud.datastore.aggregation;
+import com.google.api.core.InternalExtensionOnly;
+
/**
* An interface to represent the builders which build and customize {@link Aggregation} for {@link
* com.google.cloud.datastore.AggregationQuery}.
@@ -23,6 +25,7 @@
* Used by {@link
* com.google.cloud.datastore.AggregationQuery.Builder#addAggregation(AggregationBuilder)}.
*/
+@InternalExtensionOnly
public interface AggregationBuilder {
A build();
}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java
index 632b6633d..af2f23788 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java
@@ -23,7 +23,9 @@
/** Represents an {@link Aggregation} which returns count. */
public class CountAggregation extends Aggregation {
- /** @param alias Alias to used when running this aggregation. */
+ /**
+ * @param alias Alias to used when running this aggregation.
+ */
public CountAggregation(String alias) {
super(alias);
}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/ExecutionStats.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/ExecutionStats.java
index 52184a01a..6738d84e3 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/ExecutionStats.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/ExecutionStats.java
@@ -15,25 +15,27 @@
*/
package com.google.cloud.datastore.models;
+import static com.google.api.gax.util.TimeConversionUtils.toThreetenDuration;
+
import com.google.api.core.BetaApi;
import com.google.api.core.InternalApi;
+import com.google.api.core.ObsoleteApi;
import com.google.cloud.Structs;
import com.google.common.base.Objects;
import java.util.Map;
-import org.threeten.bp.Duration;
/** Model class for {@link com.google.datastore.v1.ExecutionStats} */
@BetaApi
public class ExecutionStats {
private final long resultsReturned;
- private final Duration executionDuration;
+ private final java.time.Duration executionDuration;
private final long readOperations;
private final Map debugStats;
@InternalApi
public ExecutionStats(com.google.datastore.v1.ExecutionStats proto) {
this.resultsReturned = proto.getResultsReturned();
- this.executionDuration = Duration.ofNanos(proto.getExecutionDuration().getNanos());
+ this.executionDuration = java.time.Duration.ofNanos(proto.getExecutionDuration().getNanos());
this.readOperations = proto.getReadOperations();
this.debugStats = Structs.asMap(proto.getDebugStats());
}
@@ -51,8 +53,14 @@ public Map getDebugStats() {
return debugStats;
}
+ /** This method is obsolete. Use {@link #getExecutionDurationJavaTime()} instead. */
+ @ObsoleteApi("Use getExecutionDurationJavaTime() instead")
+ public org.threeten.bp.Duration getExecutionDuration() {
+ return toThreetenDuration(getExecutionDurationJavaTime());
+ }
+
/** Returns the total time to execute the query in the backend. */
- public Duration getExecutionDuration() {
+ public java.time.Duration getExecutionDurationJavaTime() {
return executionDuration;
}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/DatastoreRpcFactory.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/DatastoreRpcFactory.java
index 0b7f9094b..acb85a61d 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/DatastoreRpcFactory.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/DatastoreRpcFactory.java
@@ -16,6 +16,7 @@
package com.google.cloud.datastore.spi;
+import com.google.api.core.InternalExtensionOnly;
import com.google.cloud.datastore.DatastoreOptions;
import com.google.cloud.spi.ServiceRpcFactory;
@@ -23,4 +24,5 @@
* An interface for Datastore RPC factory. Implementation will be loaded via {@link
* java.util.ServiceLoader}.
*/
+@InternalExtensionOnly
public interface DatastoreRpcFactory extends ServiceRpcFactory {}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java
index 33b8e11ea..f13e3873b 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java
@@ -16,8 +16,11 @@
package com.google.cloud.datastore.spi.v1;
+import com.google.api.core.InternalExtensionOnly;
+import com.google.api.gax.rpc.HeaderProvider;
import com.google.cloud.ServiceRpc;
import com.google.cloud.datastore.DatastoreException;
+import com.google.cloud.datastore.v1.DatastoreSettings;
import com.google.datastore.v1.AllocateIdsRequest;
import com.google.datastore.v1.AllocateIdsResponse;
import com.google.datastore.v1.BeginTransactionRequest;
@@ -36,7 +39,8 @@
import com.google.datastore.v1.RunQueryResponse;
/** Provides access to the remote Datastore service. */
-public interface DatastoreRpc extends ServiceRpc {
+@InternalExtensionOnly
+public interface DatastoreRpc extends ServiceRpc, AutoCloseable {
/**
* Sends an allocate IDs request.
@@ -93,7 +97,24 @@ BeginTransactionResponse beginTransaction(BeginTransactionRequest request)
*
* @throws DatastoreException upon failure
*/
- default RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request) {
- throw new UnsupportedOperationException("Not implemented.");
+ RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request);
+
+ @Override
+ void close() throws Exception;
+
+ /** Returns true if this background resource has been shut down. */
+ boolean isClosed();
+
+ // This class is needed solely to get access to protected method setInternalHeaderProvider()
+ class DatastoreSettingsBuilder extends DatastoreSettings.Builder {
+ DatastoreSettingsBuilder(DatastoreSettings settings) {
+ super(settings);
+ }
+
+ @Override
+ protected DatastoreSettings.Builder setInternalHeaderProvider(
+ HeaderProvider internalHeaderProvider) {
+ return super.setInternalHeaderProvider(internalHeaderProvider);
+ }
}
}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/GrpcDatastoreRpc.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/GrpcDatastoreRpc.java
new file mode 100644
index 000000000..ea9043bb9
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/GrpcDatastoreRpc.java
@@ -0,0 +1,202 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.datastore.spi.v1;
+
+import static com.google.cloud.datastore.DatastoreUtils.isEmulator;
+import static com.google.cloud.datastore.DatastoreUtils.removeScheme;
+import static com.google.cloud.datastore.spi.v1.RpcUtils.retrySettingSetter;
+import static java.util.concurrent.TimeUnit.SECONDS;
+
+import com.google.api.core.InternalApi;
+import com.google.api.gax.core.BackgroundResource;
+import com.google.api.gax.core.GaxProperties;
+import com.google.api.gax.grpc.ChannelPoolSettings;
+import com.google.api.gax.grpc.GrpcCallContext;
+import com.google.api.gax.grpc.GrpcTransportChannel;
+import com.google.api.gax.rpc.ClientContext;
+import com.google.api.gax.rpc.HeaderProvider;
+import com.google.api.gax.rpc.NoHeaderProvider;
+import com.google.api.gax.rpc.TransportChannel;
+import com.google.cloud.ServiceOptions;
+import com.google.cloud.datastore.DatastoreException;
+import com.google.cloud.datastore.DatastoreOptions;
+import com.google.cloud.datastore.v1.DatastoreSettings;
+import com.google.cloud.datastore.v1.stub.DatastoreStubSettings;
+import com.google.cloud.datastore.v1.stub.GrpcDatastoreStub;
+import com.google.cloud.grpc.GrpcTransportOptions;
+import com.google.common.base.Strings;
+import com.google.datastore.v1.AllocateIdsRequest;
+import com.google.datastore.v1.AllocateIdsResponse;
+import com.google.datastore.v1.BeginTransactionRequest;
+import com.google.datastore.v1.BeginTransactionResponse;
+import com.google.datastore.v1.CommitRequest;
+import com.google.datastore.v1.CommitResponse;
+import com.google.datastore.v1.LookupRequest;
+import com.google.datastore.v1.LookupResponse;
+import com.google.datastore.v1.ReserveIdsRequest;
+import com.google.datastore.v1.ReserveIdsResponse;
+import com.google.datastore.v1.RollbackRequest;
+import com.google.datastore.v1.RollbackResponse;
+import com.google.datastore.v1.RunAggregationQueryRequest;
+import com.google.datastore.v1.RunAggregationQueryResponse;
+import com.google.datastore.v1.RunQueryRequest;
+import com.google.datastore.v1.RunQueryResponse;
+import io.grpc.CallOptions;
+import io.grpc.ManagedChannel;
+import io.grpc.ManagedChannelBuilder;
+import java.io.IOException;
+import java.util.Collections;
+
+@InternalApi
+public class GrpcDatastoreRpc implements DatastoreRpc {
+
+ private final GrpcDatastoreStub datastoreStub;
+ private final ClientContext clientContext;
+ private boolean closed;
+
+ public GrpcDatastoreRpc(DatastoreOptions datastoreOptions) throws IOException {
+ try {
+ clientContext =
+ isEmulator(datastoreOptions)
+ ? getClientContextForEmulator(datastoreOptions)
+ : getClientContext(datastoreOptions);
+
+ /* For grpc transport options, configure default gRPC Connection pool with minChannelCount = 1 and maxChannelCount = 4 */
+ DatastoreStubSettings datastoreStubSettings =
+ DatastoreStubSettings.newBuilder(clientContext)
+ .applyToAllUnaryMethods(retrySettingSetter(datastoreOptions))
+ .setTransportChannelProvider(
+ DatastoreSettings.defaultGrpcTransportProviderBuilder()
+ .setChannelPoolSettings(
+ ChannelPoolSettings.builder()
+ .setInitialChannelCount(DatastoreOptions.INIT_CHANNEL_COUNT)
+ .setMinChannelCount(DatastoreOptions.MIN_CHANNEL_COUNT)
+ .setMaxChannelCount(DatastoreOptions.MAX_CHANNEL_COUNT)
+ .build())
+ .build())
+ .build();
+ datastoreStub = GrpcDatastoreStub.create(datastoreStubSettings);
+ } catch (IOException e) {
+ throw new IOException(e);
+ }
+ }
+
+ @Override
+ public void close() throws Exception {
+ if (!closed) {
+ datastoreStub.close();
+ for (BackgroundResource resource : clientContext.getBackgroundResources()) {
+ resource.close();
+ }
+ closed = true;
+ }
+ for (BackgroundResource resource : clientContext.getBackgroundResources()) {
+ resource.awaitTermination(1, SECONDS);
+ }
+ }
+
+ @Override
+ public AllocateIdsResponse allocateIds(AllocateIdsRequest request) {
+ return datastoreStub.allocateIdsCallable().call(request);
+ }
+
+ @Override
+ public BeginTransactionResponse beginTransaction(BeginTransactionRequest request)
+ throws DatastoreException {
+ return datastoreStub.beginTransactionCallable().call(request);
+ }
+
+ @Override
+ public CommitResponse commit(CommitRequest request) {
+ return datastoreStub.commitCallable().call(request);
+ }
+
+ @Override
+ public LookupResponse lookup(LookupRequest request) {
+ return datastoreStub.lookupCallable().call(request);
+ }
+
+ @Override
+ public ReserveIdsResponse reserveIds(ReserveIdsRequest request) {
+ return datastoreStub.reserveIdsCallable().call(request);
+ }
+
+ @Override
+ public RollbackResponse rollback(RollbackRequest request) {
+ return datastoreStub.rollbackCallable().call(request);
+ }
+
+ @Override
+ public RunQueryResponse runQuery(RunQueryRequest request) {
+ return datastoreStub.runQueryCallable().call(request);
+ }
+
+ @Override
+ public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request) {
+ return datastoreStub.runAggregationQueryCallable().call(request);
+ }
+
+ @Override
+ public boolean isClosed() {
+ return closed && datastoreStub.isShutdown();
+ }
+
+ private ClientContext getClientContextForEmulator(DatastoreOptions datastoreOptions)
+ throws IOException {
+ ManagedChannel managedChannel =
+ ManagedChannelBuilder.forTarget(removeScheme(datastoreOptions.getHost()))
+ .usePlaintext()
+ .build();
+ TransportChannel transportChannel = GrpcTransportChannel.create(managedChannel);
+ return ClientContext.newBuilder()
+ .setCredentials(null)
+ .setTransportChannel(transportChannel)
+ .setDefaultCallContext(GrpcCallContext.of(managedChannel, CallOptions.DEFAULT))
+ .setBackgroundResources(Collections.singletonList(transportChannel))
+ .build();
+ }
+
+ private ClientContext getClientContext(DatastoreOptions datastoreOptions) throws IOException {
+ HeaderProvider internalHeaderProvider =
+ DatastoreSettings.defaultApiClientHeaderProviderBuilder()
+ .setClientLibToken(
+ ServiceOptions.getGoogApiClientLibName(),
+ GaxProperties.getLibraryVersion(datastoreOptions.getClass()))
+ .setResourceToken(getResourceToken(datastoreOptions))
+ .build();
+
+ DatastoreSettingsBuilder settingsBuilder =
+ new DatastoreSettingsBuilder(DatastoreSettings.newBuilder().build());
+ settingsBuilder.setCredentialsProvider(
+ GrpcTransportOptions.setUpCredentialsProvider(datastoreOptions));
+ settingsBuilder.setTransportChannelProvider(datastoreOptions.getTransportChannelProvider());
+ settingsBuilder.setInternalHeaderProvider(internalHeaderProvider);
+ settingsBuilder.setHeaderProvider(
+ datastoreOptions.getMergedHeaderProvider(new NoHeaderProvider()));
+ return ClientContext.create(settingsBuilder.build());
+ }
+
+ private String getResourceToken(DatastoreOptions datastoreOptions) {
+ StringBuilder builder = new StringBuilder("project_id=");
+ builder.append(datastoreOptions.getProjectId());
+ if (!Strings.isNullOrEmpty(datastoreOptions.getDatabaseId())) {
+ builder.append("&database_id=");
+ builder.append(datastoreOptions.getDatabaseId());
+ }
+ return builder.toString();
+ }
+}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java
index fd3cdc658..ac39ad5ba 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java
@@ -19,9 +19,9 @@
import com.google.api.client.http.HttpRequest;
import com.google.api.client.http.HttpRequestInitializer;
import com.google.api.client.http.HttpTransport;
+import com.google.api.core.InternalApi;
import com.google.cloud.datastore.DatastoreException;
import com.google.cloud.datastore.DatastoreOptions;
-import com.google.cloud.datastore.TraceUtil;
import com.google.cloud.http.CensusHttpModule;
import com.google.cloud.http.HttpTransportOptions;
import com.google.datastore.v1.AllocateIdsRequest;
@@ -40,10 +40,12 @@
import com.google.datastore.v1.RunAggregationQueryResponse;
import com.google.datastore.v1.RunQueryRequest;
import com.google.datastore.v1.RunQueryResponse;
+import io.opencensus.trace.Tracing;
import java.io.IOException;
import java.net.InetAddress;
import java.net.URL;
+@InternalApi
public class HttpDatastoreRpc implements DatastoreRpc {
private final com.google.datastore.v1.client.Datastore client;
@@ -58,6 +60,7 @@ public HttpDatastoreRpc(DatastoreOptions options) {
.initializer(getHttpRequestInitializer(options, httpTransportOptions))
.transport(transport);
String normalizedHost = options.getHost() != null ? options.getHost().toLowerCase() : "";
+
if (isLocalHost(normalizedHost)) {
clientBuilder = clientBuilder.localHost(removeScheme(normalizedHost));
} else if (!removeScheme(com.google.datastore.v1.client.DatastoreFactory.DEFAULT_HOST)
@@ -80,8 +83,7 @@ public HttpDatastoreRpc(DatastoreOptions options) {
private HttpRequestInitializer getHttpRequestInitializer(
final DatastoreOptions options, HttpTransportOptions httpTransportOptions) {
// Open Census initialization
- CensusHttpModule censusHttpModule =
- new CensusHttpModule(TraceUtil.getInstance().getTracer(), true);
+ CensusHttpModule censusHttpModule = new CensusHttpModule(Tracing.getTracer(), true);
final HttpRequestInitializer censusHttpModuleHttpRequestInitializer =
censusHttpModule.getHttpRequestInitializer(
httpTransportOptions.getHttpRequestInitializer(options));
@@ -211,4 +213,14 @@ public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryReques
throw translate(ex);
}
}
+
+ @Override
+ public void close() throws Exception {
+ throw new UnsupportedOperationException("Not implemented.");
+ }
+
+ @Override
+ public boolean isClosed() {
+ throw new UnsupportedOperationException("Not implemented.");
+ }
}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/RpcUtils.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/RpcUtils.java
new file mode 100644
index 000000000..dee8d6920
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/RpcUtils.java
@@ -0,0 +1,34 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.datastore.spi.v1;
+
+import com.google.api.core.ApiFunction;
+import com.google.api.core.InternalApi;
+import com.google.api.gax.rpc.UnaryCallSettings;
+import com.google.cloud.datastore.DatastoreOptions;
+
+@InternalApi
+public class RpcUtils {
+ @InternalApi
+ static ApiFunction, Void> retrySettingSetter(
+ DatastoreOptions datastoreOptions) {
+ return builder -> {
+ builder.setRetrySettings(datastoreOptions.getRetrySettings());
+ return null;
+ };
+ }
+}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java
new file mode 100644
index 000000000..ebb630515
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java
@@ -0,0 +1,131 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.datastore.telemetry;
+
+import com.google.api.core.ApiFunction;
+import com.google.api.core.ApiFuture;
+import com.google.api.core.InternalApi;
+import com.google.cloud.datastore.telemetry.TraceUtil.Context;
+import io.grpc.ManagedChannelBuilder;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanBuilder;
+import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.api.trace.TracerProvider;
+import java.util.Map;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/**
+ * Tracing utility implementation, used to stub out tracing instrumentation when tracing is
+ * disabled.
+ */
+@InternalApi
+public class DisabledTraceUtil implements TraceUtil {
+ static class Span implements TraceUtil.Span {
+ @Override
+ public void end() {}
+
+ @Override
+ public void end(Throwable error) {}
+
+ @Override
+ public void endAtFuture(ApiFuture futureValue) {}
+
+ @Override
+ public TraceUtil.Span addEvent(String name) {
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span addEvent(String name, Map attributes) {
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span setAttribute(String key, int value) {
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span setAttribute(String key, String value) {
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span setAttribute(String key, boolean value) {
+ return this;
+ }
+
+ public io.opentelemetry.api.trace.Span getSpan() {
+ return null;
+ }
+
+ @Override
+ public Scope makeCurrent() {
+ return new Scope();
+ }
+ }
+
+ static class Context implements TraceUtil.Context {
+ @Override
+ public Scope makeCurrent() {
+ return new Scope();
+ }
+ }
+
+ static class Scope implements TraceUtil.Scope {
+ @Override
+ public void close() {}
+ }
+
+ @Nullable
+ @Override
+ public ApiFunction getChannelConfigurator() {
+ return null;
+ }
+
+ @Override
+ public Span startSpan(String spanName) {
+ return new Span();
+ }
+
+ @Override
+ public TraceUtil.Span startSpan(String spanName, TraceUtil.Span parentSpan) {
+ return new Span();
+ }
+
+ public SpanBuilder addSettingsAttributesToCurrentSpan(SpanBuilder spanBuilder) {
+ return getTracer().spanBuilder("TRACING_DISABLED_NO_OP");
+ }
+
+ @Nonnull
+ @Override
+ public TraceUtil.Span getCurrentSpan() {
+ return new Span();
+ }
+
+ @Nonnull
+ @Override
+ public TraceUtil.Context getCurrentContext() {
+ return new Context();
+ }
+
+ @Override
+ public Tracer getTracer() {
+ return TracerProvider.noop().get(LIBRARY_NAME);
+ }
+}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java
new file mode 100644
index 000000000..40fc7308e
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java
@@ -0,0 +1,323 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.datastore.telemetry;
+
+import com.google.api.core.ApiFunction;
+import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutureCallback;
+import com.google.api.core.ApiFutures;
+import com.google.api.core.InternalApi;
+import com.google.cloud.datastore.DatastoreOptions;
+import com.google.cloud.datastore.telemetry.TraceUtil.Context;
+import com.google.cloud.datastore.telemetry.TraceUtil.Scope;
+import com.google.cloud.datastore.telemetry.TraceUtil.Span;
+import com.google.common.base.Throwables;
+import io.grpc.ManagedChannelBuilder;
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.common.AttributesBuilder;
+import io.opentelemetry.api.trace.SpanBuilder;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.api.trace.Tracer;
+import java.util.Map;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/**
+ * Tracing utility implementation, used to stub out tracing instrumentation when tracing is enabled.
+ */
+@InternalApi
+public class EnabledTraceUtil implements TraceUtil {
+ private final Tracer tracer;
+ private final OpenTelemetry openTelemetry;
+ private final DatastoreOptions datastoreOptions;
+
+ EnabledTraceUtil(DatastoreOptions datastoreOptions) {
+ OpenTelemetry openTelemetry = datastoreOptions.getOpenTelemetryOptions().getOpenTelemetry();
+
+ // If tracing is enabled, but an OpenTelemetry instance is not provided, fall back
+ // to using GlobalOpenTelemetry.
+ if (openTelemetry == null) {
+ openTelemetry = GlobalOpenTelemetry.get();
+ }
+
+ this.datastoreOptions = datastoreOptions;
+ this.openTelemetry = openTelemetry;
+ this.tracer = openTelemetry.getTracer(LIBRARY_NAME);
+ }
+
+ public OpenTelemetry getOpenTelemetry() {
+ return openTelemetry;
+ }
+
+ @Override
+ @Nullable
+ public ApiFunction getChannelConfigurator() {
+ // TODO(jimit) Update this to return a gRPC Channel Configurator after gRPC upgrade.
+ return null;
+ }
+
+ static class Span implements TraceUtil.Span {
+ private final io.opentelemetry.api.trace.Span span;
+ private final String spanName;
+
+ public Span(io.opentelemetry.api.trace.Span span, String spanName) {
+ this.span = span;
+ this.spanName = spanName;
+ }
+
+ @Override
+ public io.opentelemetry.api.trace.Span getSpan() {
+ return this.span;
+ }
+
+ /** Ends this span. */
+ @Override
+ public void end() {
+ span.end();
+ }
+
+ /** Ends this span in an error. */
+ @Override
+ public void end(Throwable error) {
+ span.setStatus(StatusCode.ERROR, error.getMessage());
+ span.recordException(
+ error,
+ Attributes.builder()
+ .put("exception.message", error.getMessage())
+ .put("exception.type", error.getClass().getName())
+ .put("exception.stacktrace", Throwables.getStackTraceAsString(error))
+ .build());
+ span.end();
+ }
+
+ /**
+ * If an operation ends in the future, its relevant span should end _after_ the future has been
+ * completed. This method "appends" the span completion code at the completion of the given
+ * future. In order for telemetry info to be recorded, the future returned by this method should
+ * be completed.
+ */
+ @Override
+ public void endAtFuture(ApiFuture futureValue) {
+ io.opentelemetry.context.Context asyncContext = io.opentelemetry.context.Context.current();
+ ApiFutures.addCallback(
+ futureValue,
+ new ApiFutureCallback() {
+ @Override
+ public void onFailure(Throwable t) {
+ try (io.opentelemetry.context.Scope scope = asyncContext.makeCurrent()) {
+ span.addEvent(spanName + " failed.");
+ end(t);
+ }
+ }
+
+ @Override
+ public void onSuccess(T result) {
+ try (io.opentelemetry.context.Scope scope = asyncContext.makeCurrent()) {
+ span.addEvent(spanName + " succeeded.");
+ end();
+ }
+ }
+ });
+ }
+
+ /** Adds the given event to this span. */
+ @Override
+ public TraceUtil.Span addEvent(String name) {
+ span.addEvent(name);
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span addEvent(String name, Map attributes) {
+ AttributesBuilder attributesBuilder = Attributes.builder();
+ attributes.forEach(
+ (key, value) -> {
+ if (value instanceof Integer) {
+ attributesBuilder.put(key, (int) value);
+ } else if (value instanceof Long) {
+ attributesBuilder.put(key, (long) value);
+ } else if (value instanceof Double) {
+ attributesBuilder.put(key, (double) value);
+ } else if (value instanceof Float) {
+ attributesBuilder.put(key, (float) value);
+ } else if (value instanceof Boolean) {
+ attributesBuilder.put(key, (boolean) value);
+ } else if (value instanceof String) {
+ attributesBuilder.put(key, (String) value);
+ } else {
+ // OpenTelemetry APIs do not support any other type.
+ throw new IllegalArgumentException(
+ "Unknown attribute type:" + value.getClass().getSimpleName());
+ }
+ });
+ span.addEvent(name, attributesBuilder.build());
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span setAttribute(String key, int value) {
+ span.setAttribute(ATTRIBUTE_SERVICE_PREFIX + key, value);
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span setAttribute(String key, String value) {
+ span.setAttribute(ATTRIBUTE_SERVICE_PREFIX + key, value);
+ return this;
+ }
+
+ @Override
+ public TraceUtil.Span setAttribute(String key, boolean value) {
+ span.setAttribute(ATTRIBUTE_SERVICE_PREFIX + key, value);
+ return this;
+ }
+
+ @Override
+ public Scope makeCurrent() {
+ try (io.opentelemetry.context.Scope scope = span.makeCurrent()) {
+ return new Scope(scope);
+ }
+ }
+ }
+
+ static class Scope implements TraceUtil.Scope {
+ private final io.opentelemetry.context.Scope scope;
+
+ Scope(io.opentelemetry.context.Scope scope) {
+ this.scope = scope;
+ }
+
+ @Override
+ public void close() {
+ scope.close();
+ }
+ }
+
+ static class Context implements TraceUtil.Context {
+ private final io.opentelemetry.context.Context context;
+
+ Context(io.opentelemetry.context.Context context) {
+ this.context = context;
+ }
+
+ @Override
+ public Scope makeCurrent() {
+ try (io.opentelemetry.context.Scope scope = context.makeCurrent()) {
+ return new Scope(scope);
+ }
+ }
+ }
+
+ /** Applies the current Datastore instance settings as attributes to the current Span */
+ @Override
+ public SpanBuilder addSettingsAttributesToCurrentSpan(SpanBuilder spanBuilder) {
+ spanBuilder =
+ spanBuilder.setAllAttributes(
+ Attributes.builder()
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.databaseId",
+ datastoreOptions.getDatabaseId())
+ .put(ATTRIBUTE_SERVICE_PREFIX + "settings.host", datastoreOptions.getHost())
+ .build());
+
+ if (datastoreOptions.getCredentials() != null) {
+ spanBuilder =
+ spanBuilder.setAttribute(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.credentials.authenticationType",
+ datastoreOptions.getCredentials().getAuthenticationType());
+ }
+
+ if (datastoreOptions.getRetrySettings() != null) {
+ spanBuilder =
+ spanBuilder.setAllAttributes(
+ Attributes.builder()
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retrySettings.initialRetryDelay",
+ datastoreOptions.getRetrySettings().getInitialRetryDelay().toString())
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retrySettings.maxRetryDelay",
+ datastoreOptions.getRetrySettings().getMaxRetryDelay().toString())
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retrySettings.retryDelayMultiplier",
+ String.valueOf(datastoreOptions.getRetrySettings().getRetryDelayMultiplier()))
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retrySettings.maxAttempts",
+ String.valueOf(datastoreOptions.getRetrySettings().getMaxAttempts()))
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retrySettings.initialRpcTimeout",
+ datastoreOptions.getRetrySettings().getInitialRpcTimeout().toString())
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retrySettings.maxRpcTimeout",
+ datastoreOptions.getRetrySettings().getMaxRpcTimeout().toString())
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retrySettings.rpcTimeoutMultiplier",
+ String.valueOf(datastoreOptions.getRetrySettings().getRpcTimeoutMultiplier()))
+ .put(
+ ATTRIBUTE_SERVICE_PREFIX + "settings.retrySettings.totalTimeout",
+ datastoreOptions.getRetrySettings().getTotalTimeout().toString())
+ .build());
+ }
+
+ // Add the memory utilization of the client at the time this trace was collected.
+ long totalMemory = Runtime.getRuntime().totalMemory();
+ long freeMemory = Runtime.getRuntime().freeMemory();
+ double memoryUtilization = ((double) (totalMemory - freeMemory)) / totalMemory;
+ spanBuilder.setAttribute(
+ ATTRIBUTE_SERVICE_PREFIX + "memoryUtilization",
+ String.format("%.2f", memoryUtilization * 100) + "%");
+
+ return spanBuilder;
+ }
+
+ @Override
+ public Span startSpan(String spanName) {
+ SpanBuilder spanBuilder = tracer.spanBuilder(spanName).setSpanKind(SpanKind.PRODUCER);
+ io.opentelemetry.api.trace.Span span =
+ addSettingsAttributesToCurrentSpan(spanBuilder).startSpan();
+ return new Span(span, spanName);
+ }
+
+ @Override
+ public TraceUtil.Span startSpan(String spanName, TraceUtil.Span parentSpan) {
+ SpanBuilder spanBuilder =
+ tracer
+ .spanBuilder(spanName)
+ .setSpanKind(SpanKind.PRODUCER)
+ .setParent(io.opentelemetry.context.Context.current().with(parentSpan.getSpan()));
+ return new Span(addSettingsAttributesToCurrentSpan(spanBuilder).startSpan(), spanName);
+ }
+
+ @Nonnull
+ @Override
+ public TraceUtil.Span getCurrentSpan() {
+ return new Span(io.opentelemetry.api.trace.Span.current(), "");
+ }
+
+ @Nonnull
+ @Override
+ public TraceUtil.Context getCurrentContext() {
+ return new Context(io.opentelemetry.context.Context.current());
+ }
+
+ @Override
+ public Tracer getTracer() {
+ return this.tracer;
+ }
+}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java
new file mode 100644
index 000000000..fd616a733
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java
@@ -0,0 +1,167 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.datastore.telemetry;
+
+import com.google.api.core.ApiFunction;
+import com.google.api.core.ApiFuture;
+import com.google.api.core.InternalExtensionOnly;
+import com.google.cloud.datastore.DatastoreOptions;
+import io.grpc.ManagedChannelBuilder;
+import io.opentelemetry.api.trace.SpanBuilder;
+import io.opentelemetry.api.trace.Tracer;
+import java.util.Map;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/** Utility interface to manage OpenTelemetry tracing instrumentation based on the configuration. */
+@InternalExtensionOnly
+public interface TraceUtil {
+ static final String ATTRIBUTE_SERVICE_PREFIX = "gcp.datastore.";
+ static final String ENABLE_TRACING_ENV_VAR = "DATASTORE_ENABLE_TRACING";
+ static final String LIBRARY_NAME = "com.google.cloud.datastore";
+ static final String SPAN_NAME_LOOKUP = "Lookup";
+ static final String SPAN_NAME_ALLOCATE_IDS = "AllocateIds";
+ static final String SPAN_NAME_RESERVE_IDS = "ReserveIds";
+ static final String SPAN_NAME_COMMIT = "Commit";
+ static final String SPAN_NAME_RUN_QUERY = "RunQuery";
+ static final String SPAN_NAME_RUN_AGGREGATION_QUERY = "RunAggregationQuery";
+ static final String SPAN_NAME_TRANSACTION_RUN = "Transaction.Run";
+ static final String SPAN_NAME_BEGIN_TRANSACTION = "Transaction.Begin";
+ static final String SPAN_NAME_TRANSACTION_LOOKUP = "Transaction.Lookup";
+ static final String SPAN_NAME_TRANSACTION_COMMIT = "Transaction.Commit";
+ static final String SPAN_NAME_TRANSACTION_RUN_QUERY = "Transaction.RunQuery";
+ static final String SPAN_NAME_ROLLBACK = "Transaction.Rollback";
+ static final String SPAN_NAME_TRANSACTION_RUN_AGGREGATION_QUERY =
+ "Transaction.RunAggregationQuery";
+ static final String ATTRIBUTES_KEY_DOCUMENT_COUNT = "doc_count";
+ static final String ATTRIBUTES_KEY_TRANSACTIONAL = "transactional";
+ static final String ATTRIBUTES_KEY_TRANSACTION_ID = "transaction_id";
+ static final String ATTRIBUTES_KEY_READ_CONSISTENCY = "read_consistency";
+ static final String ATTRIBUTES_KEY_RECEIVED = "Received";
+ static final String ATTRIBUTES_KEY_MISSING = "Missing";
+ static final String ATTRIBUTES_KEY_DEFERRED = "Deferred";
+ static final String ATTRIBUTES_KEY_MORE_RESULTS = "mor_results";
+
+ /**
+ * Creates and returns an instance of the TraceUtil class.
+ *
+ * @param datastoreOptions The DatastoreOptions object that is requesting an instance of
+ * TraceUtil.
+ * @return An instance of the TraceUtil class.
+ */
+ static TraceUtil getInstance(@Nonnull DatastoreOptions datastoreOptions) {
+ boolean createEnabledInstance = datastoreOptions.getOpenTelemetryOptions().isEnabled();
+
+ // The environment variable can override options to enable/disable telemetry collection.
+ String enableTracingEnvVar = System.getenv(ENABLE_TRACING_ENV_VAR);
+ if (enableTracingEnvVar != null) {
+ if (enableTracingEnvVar.equalsIgnoreCase("true")
+ || enableTracingEnvVar.equalsIgnoreCase("on")) {
+ createEnabledInstance = true;
+ }
+ if (enableTracingEnvVar.equalsIgnoreCase("false")
+ || enableTracingEnvVar.equalsIgnoreCase("off")) {
+ createEnabledInstance = false;
+ }
+ }
+
+ if (createEnabledInstance) {
+ return new EnabledTraceUtil(datastoreOptions);
+ } else {
+ return new DisabledTraceUtil();
+ }
+ }
+
+ /** Returns a channel configurator for gRPC, or {@code null} if tracing is disabled. */
+ @Nullable
+ ApiFunction getChannelConfigurator();
+
+ /** Represents a trace span. */
+ interface Span {
+ /** Adds the given event to this span. */
+ Span addEvent(String name);
+
+ /** Adds the given event with the given attributes to this span. */
+ Span addEvent(String name, Map attributes);
+
+ /** Adds the given attribute to this span. */
+ Span setAttribute(String key, int value);
+
+ /** Adds the given attribute to this span. */
+ Span setAttribute(String key, String value);
+
+ /** Adds the given attribute to this span. */
+ Span setAttribute(String key, boolean value);
+
+ io.opentelemetry.api.trace.Span getSpan();
+
+ /** Marks this span as the current span. */
+ Scope makeCurrent();
+
+ /** Ends this span. */
+ void end();
+
+ /** Ends this span in an error. */
+ void end(Throwable error);
+
+ /**
+ * If an operation ends in the future, its relevant span should end _after_ the future has been
+ * completed. This method "appends" the span completion code at the completion of the given
+ * future. In order for telemetry info to be recorded, the future returned by this method should
+ * be completed.
+ */
+ void endAtFuture(ApiFuture futureValue);
+ }
+
+ /** Represents a trace context. */
+ interface Context {
+ /** Makes this context the current context. */
+ Scope makeCurrent();
+ }
+
+ /** Represents a trace scope. */
+ interface Scope extends AutoCloseable {
+ /** Closes the current scope. */
+ void close();
+ }
+
+ /** Starts a new span with the given name, sets it as the current span, and returns it. */
+ Span startSpan(String spanName);
+
+ /**
+ * Starts a new span with the given name and the span represented by the parentSpan as its parent,
+ * sets it as the current span and returns it.
+ */
+ Span startSpan(String spanName, Span parentSpan);
+
+ /**
+ * Adds common SpanAttributes to the current span, useful when hand-creating a new Span without
+ * using the TraceUtil.Span interface.
+ */
+ SpanBuilder addSettingsAttributesToCurrentSpan(SpanBuilder spanBuilder);
+
+ /** Returns the current span. */
+ @Nonnull
+ Span getCurrentSpan();
+
+ /** Returns the current Context. */
+ @Nonnull
+ Context getCurrentContext();
+
+ /** Returns the current OpenTelemetry Tracer when OpenTelemetry SDK is provided. */
+ Tracer getTracer();
+}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/LocalDatastoreHelper.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/LocalDatastoreHelper.java
index 2723325ee..9f5555d0f 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/LocalDatastoreHelper.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/LocalDatastoreHelper.java
@@ -1,5 +1,5 @@
/*
- * Copyright 2015 Google LLC
+ * Copyright 2024 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@@ -16,12 +16,15 @@
package com.google.cloud.datastore.testing;
+import static com.google.api.gax.util.TimeConversionUtils.toJavaTimeDuration;
import static com.google.common.base.MoreObjects.firstNonNull;
import com.google.api.core.InternalApi;
+import com.google.api.core.ObsoleteApi;
import com.google.cloud.NoCredentials;
import com.google.cloud.ServiceOptions;
import com.google.cloud.datastore.DatastoreOptions;
+import com.google.cloud.grpc.GrpcTransportOptions;
import com.google.cloud.testing.BaseEmulatorHelper;
import com.google.common.collect.ImmutableList;
import java.io.IOException;
@@ -38,7 +41,6 @@
import java.util.UUID;
import java.util.concurrent.TimeoutException;
import java.util.logging.Logger;
-import org.threeten.bp.Duration;
/**
* Utility to start and stop local Google Cloud Datastore emulators.
@@ -52,12 +54,13 @@ public class LocalDatastoreHelper extends BaseEmulatorHelper {
private final double consistency;
private final Path gcdPath;
private boolean storeOnDisk;
+ private boolean firestoreInDatastoreMode;
// Gcloud emulator settings
private static final String GCLOUD_CMD_TEXT = "gcloud beta emulators datastore start";
private static final String GCLOUD_CMD_PORT_FLAG = "--host-port=";
private static final String VERSION_PREFIX = "cloud-datastore-emulator ";
- private static final String MIN_VERSION = "2.0.2";
+ private static final String MIN_VERSION = "2.3.1";
// Downloadable emulator settings
private static final String BIN_NAME = "cloud-datastore-emulator/cloud_datastore_emulator";
@@ -73,6 +76,8 @@ public class LocalDatastoreHelper extends BaseEmulatorHelper {
private static final String PROJECT_FLAG = "--project=";
private static final double DEFAULT_CONSISTENCY = 0.9;
private static final String DEFAULT_PROJECT_ID = PROJECT_ID_PREFIX + UUID.randomUUID();
+ private static final String FIRESTORE_IN_DATASTORE_MODE_FLAG =
+ "--use-firestore-in-datastore-mode";
private static final Logger LOGGER = Logger.getLogger(LocalDatastoreHelper.class.getName());
@@ -101,6 +106,7 @@ public static class Builder {
private int port;
private Path dataDir;
private boolean storeOnDisk = true;
+ private boolean firestoreInDatastoreMode = false;
private String projectId;
private Builder() {}
@@ -109,6 +115,7 @@ private Builder(LocalDatastoreHelper helper) {
this.consistency = helper.consistency;
this.dataDir = helper.gcdPath;
this.storeOnDisk = helper.storeOnDisk;
+ this.firestoreInDatastoreMode = helper.firestoreInDatastoreMode;
}
public Builder setConsistency(double consistency) {
@@ -136,6 +143,11 @@ public Builder setStoreOnDisk(boolean storeOnDisk) {
return this;
}
+ public Builder setFirestoreInDatastoreMode(boolean firestoreInDatastoreMode) {
+ this.firestoreInDatastoreMode = firestoreInDatastoreMode;
+ return this;
+ }
+
/** Creates a {@code LocalDatastoreHelper} object. */
public LocalDatastoreHelper build() {
return new LocalDatastoreHelper(this);
@@ -151,26 +163,41 @@ private LocalDatastoreHelper(Builder builder) {
this.consistency = builder.consistency > 0 ? builder.consistency : DEFAULT_CONSISTENCY;
this.gcdPath = builder.dataDir;
this.storeOnDisk = builder.storeOnDisk;
+ this.firestoreInDatastoreMode = builder.firestoreInDatastoreMode;
String binName = BIN_NAME;
if (isWindows()) {
binName = BIN_NAME.replace("/", "\\");
}
List gcloudCommand = new ArrayList<>(Arrays.asList(GCLOUD_CMD_TEXT.split(" ")));
gcloudCommand.add(GCLOUD_CMD_PORT_FLAG + "localhost:" + getPort());
- gcloudCommand.add(CONSISTENCY_FLAG + builder.consistency);
gcloudCommand.add(PROJECT_FLAG + projectId);
+ if (builder.firestoreInDatastoreMode) {
+ gcloudCommand.add(FIRESTORE_IN_DATASTORE_MODE_FLAG);
+ } else {
+ // At most one of --consistency | --use-firestore-in-datastore-mode can be specified.
+ // --consistency will be ignored with --use-firestore-in-datastore-mode.
+ gcloudCommand.add(CONSISTENCY_FLAG + builder.consistency);
+ }
if (!builder.storeOnDisk) {
gcloudCommand.add("--no-store-on-disk");
}
+ if (builder.dataDir != null) {
+ gcloudCommand.add("--data-dir=" + getGcdPath());
+ }
GcloudEmulatorRunner gcloudRunner =
new GcloudEmulatorRunner(gcloudCommand, VERSION_PREFIX, MIN_VERSION);
List binCommand = new ArrayList<>(Arrays.asList(binName, "start"));
binCommand.add("--testing");
- binCommand.add(BIN_CMD_PORT_FLAG + getPort());
- binCommand.add(CONSISTENCY_FLAG + getConsistency());
- if (builder.dataDir != null) {
- gcloudCommand.add("--data-dir=" + getGcdPath());
+ if (builder.firestoreInDatastoreMode) {
+ // Downloadable emulator runner takes the flag in a different
+ // format: --firestore_in_datastore_mode
+ binCommand.add("--firestore_in_datastore_mode");
+ } else {
+ // At most one of --consistency | --firestore_in_datastore_mode can be specified.
+ // --consistency will be ignored with --firestore_in_datastore_mode.
+ binCommand.add(CONSISTENCY_FLAG + getConsistency());
}
+ binCommand.add(BIN_CMD_PORT_FLAG + getPort());
DownloadableEmulatorRunner downloadRunner =
new DownloadableEmulatorRunner(binCommand, EMULATOR_URL, MD5_CHECKSUM, ACCESS_TOKEN);
this.emulatorRunners = ImmutableList.of(gcloudRunner, downloadRunner);
@@ -215,6 +242,14 @@ public DatastoreOptions getOptions(String namespace) {
return optionsBuilder.setNamespace(namespace).build();
}
+ /**
+ * Returns a {@link DatastoreOptions} instance that sets the host to use the Datastore emulator on
+ * localhost. The transportOptions is set to {@code grpcTransportOptions}.
+ */
+ public DatastoreOptions getGrpcTransportOptions(GrpcTransportOptions grpcTransportOptions) {
+ return optionsBuilder.setTransportOptions(grpcTransportOptions).build();
+ }
+
public DatastoreOptions.Builder setNamespace(String namespace) {
return optionsBuilder.setNamespace(namespace);
}
@@ -234,6 +269,13 @@ public boolean isStoreOnDisk() {
return storeOnDisk;
}
+ /**
+ * Returns {@code true} use firestore-in-datastore-mode, otherwise {@code false} use native mode.
+ */
+ public boolean isFirestoreInDatastoreMode() {
+ return firestoreInDatastoreMode;
+ }
+
/**
* Creates a local Datastore helper with the specified settings for project ID and consistency.
*
@@ -307,6 +349,14 @@ public void reset() throws IOException {
sendPostRequest("/reset");
}
+ /** This method is obsolete. Use {@link #stopDuration(java.time.Duration)} instead */
+ @ObsoleteApi("Use stopDuration(java.time.Duration) instead")
+ @Override
+ public void stop(org.threeten.bp.Duration timeout)
+ throws IOException, InterruptedException, TimeoutException {
+ stopDuration(toJavaTimeDuration(timeout));
+ }
+
/**
* Stops the Datastore emulator.
*
@@ -319,15 +369,16 @@ public void reset() throws IOException {
* this value high to ensure proper shutdown, like 5 seconds or more.
*/
@Override
- public void stop(Duration timeout) throws IOException, InterruptedException, TimeoutException {
+ public void stopDuration(java.time.Duration timeout)
+ throws IOException, InterruptedException, TimeoutException {
sendPostRequest("/shutdown");
- waitForProcess(timeout);
+ waitForProcessDuration(timeout);
deleteRecursively(gcdPath);
}
/**
- * Stops the Datastore emulator. The same as {@link #stop(Duration)} but with timeout duration of
- * 20 seconds.
+ * Stops the Datastore emulator. The same as {@link #stopDuration(java.time.Duration)} but with
+ * timeout duration of 20 seconds.
*
* It is important to stop the emulator. Since the emulator runs in its own process, not
* stopping it might cause it to become orphan.
@@ -335,7 +386,7 @@ public void stop(Duration timeout) throws IOException, InterruptedException, Tim
*
It is not required to call {@link #reset()} before {@code stop()}.
*/
public void stop() throws IOException, InterruptedException, TimeoutException {
- stop(Duration.ofSeconds(20));
+ stopDuration(java.time.Duration.ofSeconds(20));
}
static void deleteRecursively(Path path) throws IOException {
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/DatastoreClient.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/DatastoreClient.java
new file mode 100644
index 000000000..e8bea055f
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/DatastoreClient.java
@@ -0,0 +1,1087 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.datastore.v1;
+
+import com.google.api.gax.core.BackgroundResource;
+import com.google.api.gax.rpc.UnaryCallable;
+import com.google.cloud.datastore.v1.stub.DatastoreStub;
+import com.google.cloud.datastore.v1.stub.DatastoreStubSettings;
+import com.google.datastore.v1.AllocateIdsRequest;
+import com.google.datastore.v1.AllocateIdsResponse;
+import com.google.datastore.v1.BeginTransactionRequest;
+import com.google.datastore.v1.BeginTransactionResponse;
+import com.google.datastore.v1.CommitRequest;
+import com.google.datastore.v1.CommitResponse;
+import com.google.datastore.v1.Key;
+import com.google.datastore.v1.LookupRequest;
+import com.google.datastore.v1.LookupResponse;
+import com.google.datastore.v1.Mutation;
+import com.google.datastore.v1.ReadOptions;
+import com.google.datastore.v1.ReserveIdsRequest;
+import com.google.datastore.v1.ReserveIdsResponse;
+import com.google.datastore.v1.RollbackRequest;
+import com.google.datastore.v1.RollbackResponse;
+import com.google.datastore.v1.RunAggregationQueryRequest;
+import com.google.datastore.v1.RunAggregationQueryResponse;
+import com.google.datastore.v1.RunQueryRequest;
+import com.google.datastore.v1.RunQueryResponse;
+import com.google.protobuf.ByteString;
+import java.io.IOException;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import javax.annotation.Generated;
+
+// AUTO-GENERATED DOCUMENTATION AND CLASS.
+/**
+ * Service Description: Each RPC normalizes the partition IDs of the keys in its input entities, and
+ * always returns entities with keys with normalized partition IDs. This applies to all keys and
+ * entities, including those in values, except keys with both an empty path and an empty or unset
+ * partition ID. Normalization of input keys sets the project ID (if not already set) to the project
+ * ID from the request.
+ *
+ *
This class provides the ability to make remote calls to the backing service through method
+ * calls that map to API methods. Sample code to get started:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * String projectId = "projectId-894832108";
+ * ReadOptions readOptions = ReadOptions.newBuilder().build();
+ * List keys = new ArrayList<>();
+ * LookupResponse response = datastoreClient.lookup(projectId, readOptions, keys);
+ * }
+ * }
+ *
+ * Note: close() needs to be called on the DatastoreClient object to clean up resources such as
+ * threads. In the example above, try-with-resources is used, which automatically calls close().
+ *
+ *
+ * Methods
+ *
+ * | Method |
+ * Description |
+ * Method Variants |
+ *
+ *
+ * Lookup |
+ * Looks up entities by key. |
+ *
+ * Request object method variants only take one parameter, a request object, which must be constructed before the call.
+ *
+ * "Flattened" method variants have converted the fields of the request object into function parameters to enable multiple ways to call the same method.
+ *
+ * Callable method variants take no parameters and return an immutable API callable object, which can be used to initiate calls to the service.
+ *
+ * |
+ *
+ *
+ * RunQuery |
+ * Queries for entities. |
+ *
+ * Request object method variants only take one parameter, a request object, which must be constructed before the call.
+ *
+ * Callable method variants take no parameters and return an immutable API callable object, which can be used to initiate calls to the service.
+ *
+ * runQueryCallable()
+ *
+ * |
+ *
+ *
+ * RunAggregationQuery |
+ * Runs an aggregation query. |
+ *
+ * Request object method variants only take one parameter, a request object, which must be constructed before the call.
+ *
+ * Callable method variants take no parameters and return an immutable API callable object, which can be used to initiate calls to the service.
+ *
+ * |
+ *
+ *
+ * BeginTransaction |
+ * Begins a new transaction. |
+ *
+ * Request object method variants only take one parameter, a request object, which must be constructed before the call.
+ *
+ * "Flattened" method variants have converted the fields of the request object into function parameters to enable multiple ways to call the same method.
+ *
+ * Callable method variants take no parameters and return an immutable API callable object, which can be used to initiate calls to the service.
+ *
+ * |
+ *
+ *
+ * Commit |
+ * Commits a transaction, optionally creating, deleting or modifying some entities. |
+ *
+ * Request object method variants only take one parameter, a request object, which must be constructed before the call.
+ *
+ * "Flattened" method variants have converted the fields of the request object into function parameters to enable multiple ways to call the same method.
+ *
+ * commit(String projectId, CommitRequest.Mode mode, List<Mutation> mutations)
+ * commit(String projectId, CommitRequest.Mode mode, ByteString transaction, List<Mutation> mutations)
+ *
+ * Callable method variants take no parameters and return an immutable API callable object, which can be used to initiate calls to the service.
+ *
+ * |
+ *
+ *
+ * Rollback |
+ * Rolls back a transaction. |
+ *
+ * Request object method variants only take one parameter, a request object, which must be constructed before the call.
+ *
+ * "Flattened" method variants have converted the fields of the request object into function parameters to enable multiple ways to call the same method.
+ *
+ * Callable method variants take no parameters and return an immutable API callable object, which can be used to initiate calls to the service.
+ *
+ * rollbackCallable()
+ *
+ * |
+ *
+ *
+ * AllocateIds |
+ * Allocates IDs for the given keys, which is useful for referencing an entity before it is inserted. |
+ *
+ * Request object method variants only take one parameter, a request object, which must be constructed before the call.
+ *
+ * "Flattened" method variants have converted the fields of the request object into function parameters to enable multiple ways to call the same method.
+ *
+ * Callable method variants take no parameters and return an immutable API callable object, which can be used to initiate calls to the service.
+ *
+ * |
+ *
+ *
+ * ReserveIds |
+ * Prevents the supplied keys' IDs from being auto-allocated by Cloud Datastore. |
+ *
+ * Request object method variants only take one parameter, a request object, which must be constructed before the call.
+ *
+ * "Flattened" method variants have converted the fields of the request object into function parameters to enable multiple ways to call the same method.
+ *
+ * Callable method variants take no parameters and return an immutable API callable object, which can be used to initiate calls to the service.
+ *
+ * reserveIdsCallable()
+ *
+ * |
+ *
+ *
+ *
+ * See the individual methods for example code.
+ *
+ *
Many parameters require resource names to be formatted in a particular way. To assist with
+ * these names, this class includes a format method for each type of name, and additionally a parse
+ * method to extract the individual identifiers contained within names that are returned.
+ *
+ *
This class can be customized by passing in a custom instance of DatastoreSettings to create().
+ * For example:
+ *
+ *
To customize credentials:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * DatastoreSettings datastoreSettings =
+ * DatastoreSettings.newBuilder()
+ * .setCredentialsProvider(FixedCredentialsProvider.create(myCredentials))
+ * .build();
+ * DatastoreClient datastoreClient = DatastoreClient.create(datastoreSettings);
+ * }
+ *
+ * To customize the endpoint:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * DatastoreSettings datastoreSettings =
+ * DatastoreSettings.newBuilder().setEndpoint(myEndpoint).build();
+ * DatastoreClient datastoreClient = DatastoreClient.create(datastoreSettings);
+ * }
+ *
+ * To use REST (HTTP1.1/JSON) transport (instead of gRPC) for sending and receiving requests over
+ * the wire:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * DatastoreSettings datastoreSettings = DatastoreSettings.newHttpJsonBuilder().build();
+ * DatastoreClient datastoreClient = DatastoreClient.create(datastoreSettings);
+ * }
+ *
+ * Please refer to the GitHub repository's samples for more quickstart code snippets.
+ */
+@Generated("by gapic-generator-java")
+public class DatastoreClient implements BackgroundResource {
+ private final DatastoreSettings settings;
+ private final DatastoreStub stub;
+
+ /** Constructs an instance of DatastoreClient with default settings. */
+ public static final DatastoreClient create() throws IOException {
+ return create(DatastoreSettings.newBuilder().build());
+ }
+
+ /**
+ * Constructs an instance of DatastoreClient, using the given settings. The channels are created
+ * based on the settings passed in, or defaults for any settings that are not set.
+ */
+ public static final DatastoreClient create(DatastoreSettings settings) throws IOException {
+ return new DatastoreClient(settings);
+ }
+
+ /**
+ * Constructs an instance of DatastoreClient, using the given stub for making calls. This is for
+ * advanced usage - prefer using create(DatastoreSettings).
+ */
+ public static final DatastoreClient create(DatastoreStub stub) {
+ return new DatastoreClient(stub);
+ }
+
+ /**
+ * Constructs an instance of DatastoreClient, using the given settings. This is protected so that
+ * it is easy to make a subclass, but otherwise, the static factory methods should be preferred.
+ */
+ protected DatastoreClient(DatastoreSettings settings) throws IOException {
+ this.settings = settings;
+ this.stub = ((DatastoreStubSettings) settings.getStubSettings()).createStub();
+ }
+
+ protected DatastoreClient(DatastoreStub stub) {
+ this.settings = null;
+ this.stub = stub;
+ }
+
+ public final DatastoreSettings getSettings() {
+ return settings;
+ }
+
+ public DatastoreStub getStub() {
+ return stub;
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Looks up entities by key.
+ *
+ *
Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * String projectId = "projectId-894832108";
+ * ReadOptions readOptions = ReadOptions.newBuilder().build();
+ * List keys = new ArrayList<>();
+ * LookupResponse response = datastoreClient.lookup(projectId, readOptions, keys);
+ * }
+ * }
+ *
+ * @param projectId Required. The ID of the project against which to make the request.
+ * @param readOptions The options for this lookup request.
+ * @param keys Required. Keys of entities to look up.
+ * @throws com.google.api.gax.rpc.ApiException if the remote call fails
+ */
+ public final LookupResponse lookup(String projectId, ReadOptions readOptions, List keys) {
+ LookupRequest request =
+ LookupRequest.newBuilder()
+ .setProjectId(projectId)
+ .setReadOptions(readOptions)
+ .addAllKeys(keys)
+ .build();
+ return lookup(request);
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Looks up entities by key.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * LookupRequest request =
+ * LookupRequest.newBuilder()
+ * .setProjectId("projectId-894832108")
+ * .setDatabaseId("databaseId1688905718")
+ * .setReadOptions(ReadOptions.newBuilder().build())
+ * .addAllKeys(new ArrayList())
+ * .setPropertyMask(PropertyMask.newBuilder().build())
+ * .build();
+ * LookupResponse response = datastoreClient.lookup(request);
+ * }
+ * }
+ *
+ * @param request The request object containing all of the parameters for the API call.
+ * @throws com.google.api.gax.rpc.ApiException if the remote call fails
+ */
+ public final LookupResponse lookup(LookupRequest request) {
+ return lookupCallable().call(request);
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Looks up entities by key.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * LookupRequest request =
+ * LookupRequest.newBuilder()
+ * .setProjectId("projectId-894832108")
+ * .setDatabaseId("databaseId1688905718")
+ * .setReadOptions(ReadOptions.newBuilder().build())
+ * .addAllKeys(new ArrayList())
+ * .setPropertyMask(PropertyMask.newBuilder().build())
+ * .build();
+ * ApiFuture future = datastoreClient.lookupCallable().futureCall(request);
+ * // Do something.
+ * LookupResponse response = future.get();
+ * }
+ * }
+ */
+ public final UnaryCallable lookupCallable() {
+ return stub.lookupCallable();
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Queries for entities.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * RunQueryRequest request =
+ * RunQueryRequest.newBuilder()
+ * .setProjectId("projectId-894832108")
+ * .setDatabaseId("databaseId1688905718")
+ * .setPartitionId(PartitionId.newBuilder().build())
+ * .setReadOptions(ReadOptions.newBuilder().build())
+ * .setPropertyMask(PropertyMask.newBuilder().build())
+ * .setExplainOptions(ExplainOptions.newBuilder().build())
+ * .build();
+ * RunQueryResponse response = datastoreClient.runQuery(request);
+ * }
+ * }
+ *
+ * @param request The request object containing all of the parameters for the API call.
+ * @throws com.google.api.gax.rpc.ApiException if the remote call fails
+ */
+ public final RunQueryResponse runQuery(RunQueryRequest request) {
+ return runQueryCallable().call(request);
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Queries for entities.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * RunQueryRequest request =
+ * RunQueryRequest.newBuilder()
+ * .setProjectId("projectId-894832108")
+ * .setDatabaseId("databaseId1688905718")
+ * .setPartitionId(PartitionId.newBuilder().build())
+ * .setReadOptions(ReadOptions.newBuilder().build())
+ * .setPropertyMask(PropertyMask.newBuilder().build())
+ * .setExplainOptions(ExplainOptions.newBuilder().build())
+ * .build();
+ * ApiFuture future = datastoreClient.runQueryCallable().futureCall(request);
+ * // Do something.
+ * RunQueryResponse response = future.get();
+ * }
+ * }
+ */
+ public final UnaryCallable runQueryCallable() {
+ return stub.runQueryCallable();
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Runs an aggregation query.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * RunAggregationQueryRequest request =
+ * RunAggregationQueryRequest.newBuilder()
+ * .setProjectId("projectId-894832108")
+ * .setDatabaseId("databaseId1688905718")
+ * .setPartitionId(PartitionId.newBuilder().build())
+ * .setReadOptions(ReadOptions.newBuilder().build())
+ * .setExplainOptions(ExplainOptions.newBuilder().build())
+ * .build();
+ * RunAggregationQueryResponse response = datastoreClient.runAggregationQuery(request);
+ * }
+ * }
+ *
+ * @param request The request object containing all of the parameters for the API call.
+ * @throws com.google.api.gax.rpc.ApiException if the remote call fails
+ */
+ public final RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request) {
+ return runAggregationQueryCallable().call(request);
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Runs an aggregation query.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * RunAggregationQueryRequest request =
+ * RunAggregationQueryRequest.newBuilder()
+ * .setProjectId("projectId-894832108")
+ * .setDatabaseId("databaseId1688905718")
+ * .setPartitionId(PartitionId.newBuilder().build())
+ * .setReadOptions(ReadOptions.newBuilder().build())
+ * .setExplainOptions(ExplainOptions.newBuilder().build())
+ * .build();
+ * ApiFuture future =
+ * datastoreClient.runAggregationQueryCallable().futureCall(request);
+ * // Do something.
+ * RunAggregationQueryResponse response = future.get();
+ * }
+ * }
+ */
+ public final UnaryCallable
+ runAggregationQueryCallable() {
+ return stub.runAggregationQueryCallable();
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Begins a new transaction.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * String projectId = "projectId-894832108";
+ * BeginTransactionResponse response = datastoreClient.beginTransaction(projectId);
+ * }
+ * }
+ *
+ * @param projectId Required. The ID of the project against which to make the request.
+ * @throws com.google.api.gax.rpc.ApiException if the remote call fails
+ */
+ public final BeginTransactionResponse beginTransaction(String projectId) {
+ BeginTransactionRequest request =
+ BeginTransactionRequest.newBuilder().setProjectId(projectId).build();
+ return beginTransaction(request);
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Begins a new transaction.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * BeginTransactionRequest request =
+ * BeginTransactionRequest.newBuilder()
+ * .setProjectId("projectId-894832108")
+ * .setDatabaseId("databaseId1688905718")
+ * .setTransactionOptions(TransactionOptions.newBuilder().build())
+ * .build();
+ * BeginTransactionResponse response = datastoreClient.beginTransaction(request);
+ * }
+ * }
+ *
+ * @param request The request object containing all of the parameters for the API call.
+ * @throws com.google.api.gax.rpc.ApiException if the remote call fails
+ */
+ public final BeginTransactionResponse beginTransaction(BeginTransactionRequest request) {
+ return beginTransactionCallable().call(request);
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Begins a new transaction.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * BeginTransactionRequest request =
+ * BeginTransactionRequest.newBuilder()
+ * .setProjectId("projectId-894832108")
+ * .setDatabaseId("databaseId1688905718")
+ * .setTransactionOptions(TransactionOptions.newBuilder().build())
+ * .build();
+ * ApiFuture future =
+ * datastoreClient.beginTransactionCallable().futureCall(request);
+ * // Do something.
+ * BeginTransactionResponse response = future.get();
+ * }
+ * }
+ */
+ public final UnaryCallable
+ beginTransactionCallable() {
+ return stub.beginTransactionCallable();
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Commits a transaction, optionally creating, deleting or modifying some entities.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * String projectId = "projectId-894832108";
+ * CommitRequest.Mode mode = CommitRequest.Mode.forNumber(0);
+ * List mutations = new ArrayList<>();
+ * CommitResponse response = datastoreClient.commit(projectId, mode, mutations);
+ * }
+ * }
+ *
+ * @param projectId Required. The ID of the project against which to make the request.
+ * @param mode The type of commit to perform. Defaults to `TRANSACTIONAL`.
+ * @param mutations The mutations to perform.
+ * When mode is `TRANSACTIONAL`, mutations affecting a single entity are applied in order.
+ * The following sequences of mutations affecting a single entity are not permitted in a
+ * single `Commit` request:
+ *
- `insert` followed by `insert` - `update` followed by `insert` - `upsert` followed by
+ * `insert` - `delete` followed by `update`
+ *
When mode is `NON_TRANSACTIONAL`, no two mutations may affect a single entity.
+ * @throws com.google.api.gax.rpc.ApiException if the remote call fails
+ */
+ public final CommitResponse commit(
+ String projectId, CommitRequest.Mode mode, List mutations) {
+ CommitRequest request =
+ CommitRequest.newBuilder()
+ .setProjectId(projectId)
+ .setMode(mode)
+ .addAllMutations(mutations)
+ .build();
+ return commit(request);
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Commits a transaction, optionally creating, deleting or modifying some entities.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * String projectId = "projectId-894832108";
+ * CommitRequest.Mode mode = CommitRequest.Mode.forNumber(0);
+ * ByteString transaction = ByteString.EMPTY;
+ * List mutations = new ArrayList<>();
+ * CommitResponse response = datastoreClient.commit(projectId, mode, transaction, mutations);
+ * }
+ * }
+ *
+ * @param projectId Required. The ID of the project against which to make the request.
+ * @param mode The type of commit to perform. Defaults to `TRANSACTIONAL`.
+ * @param transaction The identifier of the transaction associated with the commit. A transaction
+ * identifier is returned by a call to
+ * [Datastore.BeginTransaction][google.datastore.v1.Datastore.BeginTransaction].
+ * @param mutations The mutations to perform.
+ * When mode is `TRANSACTIONAL`, mutations affecting a single entity are applied in order.
+ * The following sequences of mutations affecting a single entity are not permitted in a
+ * single `Commit` request:
+ *
- `insert` followed by `insert` - `update` followed by `insert` - `upsert` followed by
+ * `insert` - `delete` followed by `update`
+ *
When mode is `NON_TRANSACTIONAL`, no two mutations may affect a single entity.
+ * @throws com.google.api.gax.rpc.ApiException if the remote call fails
+ */
+ public final CommitResponse commit(
+ String projectId, CommitRequest.Mode mode, ByteString transaction, List mutations) {
+ CommitRequest request =
+ CommitRequest.newBuilder()
+ .setProjectId(projectId)
+ .setMode(mode)
+ .setTransaction(transaction)
+ .addAllMutations(mutations)
+ .build();
+ return commit(request);
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Commits a transaction, optionally creating, deleting or modifying some entities.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * CommitRequest request =
+ * CommitRequest.newBuilder()
+ * .setProjectId("projectId-894832108")
+ * .setDatabaseId("databaseId1688905718")
+ * .addAllMutations(new ArrayList())
+ * .build();
+ * CommitResponse response = datastoreClient.commit(request);
+ * }
+ * }
+ *
+ * @param request The request object containing all of the parameters for the API call.
+ * @throws com.google.api.gax.rpc.ApiException if the remote call fails
+ */
+ public final CommitResponse commit(CommitRequest request) {
+ return commitCallable().call(request);
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Commits a transaction, optionally creating, deleting or modifying some entities.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * CommitRequest request =
+ * CommitRequest.newBuilder()
+ * .setProjectId("projectId-894832108")
+ * .setDatabaseId("databaseId1688905718")
+ * .addAllMutations(new ArrayList())
+ * .build();
+ * ApiFuture future = datastoreClient.commitCallable().futureCall(request);
+ * // Do something.
+ * CommitResponse response = future.get();
+ * }
+ * }
+ */
+ public final UnaryCallable commitCallable() {
+ return stub.commitCallable();
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Rolls back a transaction.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * String projectId = "projectId-894832108";
+ * ByteString transaction = ByteString.EMPTY;
+ * RollbackResponse response = datastoreClient.rollback(projectId, transaction);
+ * }
+ * }
+ *
+ * @param projectId Required. The ID of the project against which to make the request.
+ * @param transaction Required. The transaction identifier, returned by a call to
+ * [Datastore.BeginTransaction][google.datastore.v1.Datastore.BeginTransaction].
+ * @throws com.google.api.gax.rpc.ApiException if the remote call fails
+ */
+ public final RollbackResponse rollback(String projectId, ByteString transaction) {
+ RollbackRequest request =
+ RollbackRequest.newBuilder().setProjectId(projectId).setTransaction(transaction).build();
+ return rollback(request);
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Rolls back a transaction.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * RollbackRequest request =
+ * RollbackRequest.newBuilder()
+ * .setProjectId("projectId-894832108")
+ * .setDatabaseId("databaseId1688905718")
+ * .setTransaction(ByteString.EMPTY)
+ * .build();
+ * RollbackResponse response = datastoreClient.rollback(request);
+ * }
+ * }
+ *
+ * @param request The request object containing all of the parameters for the API call.
+ * @throws com.google.api.gax.rpc.ApiException if the remote call fails
+ */
+ public final RollbackResponse rollback(RollbackRequest request) {
+ return rollbackCallable().call(request);
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Rolls back a transaction.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * RollbackRequest request =
+ * RollbackRequest.newBuilder()
+ * .setProjectId("projectId-894832108")
+ * .setDatabaseId("databaseId1688905718")
+ * .setTransaction(ByteString.EMPTY)
+ * .build();
+ * ApiFuture future = datastoreClient.rollbackCallable().futureCall(request);
+ * // Do something.
+ * RollbackResponse response = future.get();
+ * }
+ * }
+ */
+ public final UnaryCallable rollbackCallable() {
+ return stub.rollbackCallable();
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Allocates IDs for the given keys, which is useful for referencing an entity before it is
+ * inserted.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * String projectId = "projectId-894832108";
+ * List keys = new ArrayList<>();
+ * AllocateIdsResponse response = datastoreClient.allocateIds(projectId, keys);
+ * }
+ * }
+ *
+ * @param projectId Required. The ID of the project against which to make the request.
+ * @param keys Required. A list of keys with incomplete key paths for which to allocate IDs. No
+ * key may be reserved/read-only.
+ * @throws com.google.api.gax.rpc.ApiException if the remote call fails
+ */
+ public final AllocateIdsResponse allocateIds(String projectId, List keys) {
+ AllocateIdsRequest request =
+ AllocateIdsRequest.newBuilder().setProjectId(projectId).addAllKeys(keys).build();
+ return allocateIds(request);
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Allocates IDs for the given keys, which is useful for referencing an entity before it is
+ * inserted.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * AllocateIdsRequest request =
+ * AllocateIdsRequest.newBuilder()
+ * .setProjectId("projectId-894832108")
+ * .setDatabaseId("databaseId1688905718")
+ * .addAllKeys(new ArrayList())
+ * .build();
+ * AllocateIdsResponse response = datastoreClient.allocateIds(request);
+ * }
+ * }
+ *
+ * @param request The request object containing all of the parameters for the API call.
+ * @throws com.google.api.gax.rpc.ApiException if the remote call fails
+ */
+ public final AllocateIdsResponse allocateIds(AllocateIdsRequest request) {
+ return allocateIdsCallable().call(request);
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Allocates IDs for the given keys, which is useful for referencing an entity before it is
+ * inserted.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * AllocateIdsRequest request =
+ * AllocateIdsRequest.newBuilder()
+ * .setProjectId("projectId-894832108")
+ * .setDatabaseId("databaseId1688905718")
+ * .addAllKeys(new ArrayList())
+ * .build();
+ * ApiFuture future =
+ * datastoreClient.allocateIdsCallable().futureCall(request);
+ * // Do something.
+ * AllocateIdsResponse response = future.get();
+ * }
+ * }
+ */
+ public final UnaryCallable allocateIdsCallable() {
+ return stub.allocateIdsCallable();
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Prevents the supplied keys' IDs from being auto-allocated by Cloud Datastore.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * String projectId = "projectId-894832108";
+ * List keys = new ArrayList<>();
+ * ReserveIdsResponse response = datastoreClient.reserveIds(projectId, keys);
+ * }
+ * }
+ *
+ * @param projectId Required. The ID of the project against which to make the request.
+ * @param keys Required. A list of keys with complete key paths whose numeric IDs should not be
+ * auto-allocated.
+ * @throws com.google.api.gax.rpc.ApiException if the remote call fails
+ */
+ public final ReserveIdsResponse reserveIds(String projectId, List keys) {
+ ReserveIdsRequest request =
+ ReserveIdsRequest.newBuilder().setProjectId(projectId).addAllKeys(keys).build();
+ return reserveIds(request);
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Prevents the supplied keys' IDs from being auto-allocated by Cloud Datastore.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * ReserveIdsRequest request =
+ * ReserveIdsRequest.newBuilder()
+ * .setProjectId("projectId-894832108")
+ * .setDatabaseId("databaseId1688905718")
+ * .addAllKeys(new ArrayList())
+ * .build();
+ * ReserveIdsResponse response = datastoreClient.reserveIds(request);
+ * }
+ * }
+ *
+ * @param request The request object containing all of the parameters for the API call.
+ * @throws com.google.api.gax.rpc.ApiException if the remote call fails
+ */
+ public final ReserveIdsResponse reserveIds(ReserveIdsRequest request) {
+ return reserveIdsCallable().call(request);
+ }
+
+ // AUTO-GENERATED DOCUMENTATION AND METHOD.
+ /**
+ * Prevents the supplied keys' IDs from being auto-allocated by Cloud Datastore.
+ *
+ * Sample code:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * ReserveIdsRequest request =
+ * ReserveIdsRequest.newBuilder()
+ * .setProjectId("projectId-894832108")
+ * .setDatabaseId("databaseId1688905718")
+ * .addAllKeys(new ArrayList())
+ * .build();
+ * ApiFuture future =
+ * datastoreClient.reserveIdsCallable().futureCall(request);
+ * // Do something.
+ * ReserveIdsResponse response = future.get();
+ * }
+ * }
+ */
+ public final UnaryCallable reserveIdsCallable() {
+ return stub.reserveIdsCallable();
+ }
+
+ @Override
+ public final void close() {
+ stub.close();
+ }
+
+ @Override
+ public void shutdown() {
+ stub.shutdown();
+ }
+
+ @Override
+ public boolean isShutdown() {
+ return stub.isShutdown();
+ }
+
+ @Override
+ public boolean isTerminated() {
+ return stub.isTerminated();
+ }
+
+ @Override
+ public void shutdownNow() {
+ stub.shutdownNow();
+ }
+
+ @Override
+ public boolean awaitTermination(long duration, TimeUnit unit) throws InterruptedException {
+ return stub.awaitTermination(duration, unit);
+ }
+}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/DatastoreSettings.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/DatastoreSettings.java
new file mode 100644
index 000000000..74054341e
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/DatastoreSettings.java
@@ -0,0 +1,304 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.datastore.v1;
+
+import com.google.api.core.ApiFunction;
+import com.google.api.core.BetaApi;
+import com.google.api.gax.core.GoogleCredentialsProvider;
+import com.google.api.gax.core.InstantiatingExecutorProvider;
+import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
+import com.google.api.gax.httpjson.InstantiatingHttpJsonChannelProvider;
+import com.google.api.gax.rpc.ApiClientHeaderProvider;
+import com.google.api.gax.rpc.ClientContext;
+import com.google.api.gax.rpc.ClientSettings;
+import com.google.api.gax.rpc.TransportChannelProvider;
+import com.google.api.gax.rpc.UnaryCallSettings;
+import com.google.cloud.datastore.v1.stub.DatastoreStubSettings;
+import com.google.datastore.v1.AllocateIdsRequest;
+import com.google.datastore.v1.AllocateIdsResponse;
+import com.google.datastore.v1.BeginTransactionRequest;
+import com.google.datastore.v1.BeginTransactionResponse;
+import com.google.datastore.v1.CommitRequest;
+import com.google.datastore.v1.CommitResponse;
+import com.google.datastore.v1.LookupRequest;
+import com.google.datastore.v1.LookupResponse;
+import com.google.datastore.v1.ReserveIdsRequest;
+import com.google.datastore.v1.ReserveIdsResponse;
+import com.google.datastore.v1.RollbackRequest;
+import com.google.datastore.v1.RollbackResponse;
+import com.google.datastore.v1.RunAggregationQueryRequest;
+import com.google.datastore.v1.RunAggregationQueryResponse;
+import com.google.datastore.v1.RunQueryRequest;
+import com.google.datastore.v1.RunQueryResponse;
+import java.io.IOException;
+import java.util.List;
+import javax.annotation.Generated;
+
+// AUTO-GENERATED DOCUMENTATION AND CLASS.
+/**
+ * Settings class to configure an instance of {@link DatastoreClient}.
+ *
+ * The default instance has everything set to sensible defaults:
+ *
+ *
+ * - The default service address (datastore.googleapis.com) and default port (443) are used.
+ *
- Credentials are acquired automatically through Application Default Credentials.
+ *
- Retries are configured for idempotent methods but not for non-idempotent methods.
+ *
+ *
+ * The builder of this class is recursive, so contained classes are themselves builders. When
+ * build() is called, the tree of builders is called to create the complete settings object.
+ *
+ *
For example, to set the
+ * [RetrySettings](https://cloud.google.com/java/docs/reference/gax/latest/com.google.api.gax.retrying.RetrySettings)
+ * of lookup:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * DatastoreSettings.Builder datastoreSettingsBuilder = DatastoreSettings.newBuilder();
+ * datastoreSettingsBuilder
+ * .lookupSettings()
+ * .setRetrySettings(
+ * datastoreSettingsBuilder
+ * .lookupSettings()
+ * .getRetrySettings()
+ * .toBuilder()
+ * .setInitialRetryDelayDuration(Duration.ofSeconds(1))
+ * .setInitialRpcTimeoutDuration(Duration.ofSeconds(5))
+ * .setMaxAttempts(5)
+ * .setMaxRetryDelayDuration(Duration.ofSeconds(30))
+ * .setMaxRpcTimeoutDuration(Duration.ofSeconds(60))
+ * .setRetryDelayMultiplier(1.3)
+ * .setRpcTimeoutMultiplier(1.5)
+ * .setTotalTimeoutDuration(Duration.ofSeconds(300))
+ * .build());
+ * DatastoreSettings datastoreSettings = datastoreSettingsBuilder.build();
+ * }
+ *
+ * Please refer to the [Client Side Retry
+ * Guide](https://github.com/googleapis/google-cloud-java/blob/main/docs/client_retries.md) for
+ * additional support in setting retries.
+ */
+@Generated("by gapic-generator-java")
+public class DatastoreSettings extends ClientSettings {
+
+ /** Returns the object with the settings used for calls to lookup. */
+ public UnaryCallSettings lookupSettings() {
+ return ((DatastoreStubSettings) getStubSettings()).lookupSettings();
+ }
+
+ /** Returns the object with the settings used for calls to runQuery. */
+ public UnaryCallSettings runQuerySettings() {
+ return ((DatastoreStubSettings) getStubSettings()).runQuerySettings();
+ }
+
+ /** Returns the object with the settings used for calls to runAggregationQuery. */
+ public UnaryCallSettings
+ runAggregationQuerySettings() {
+ return ((DatastoreStubSettings) getStubSettings()).runAggregationQuerySettings();
+ }
+
+ /** Returns the object with the settings used for calls to beginTransaction. */
+ public UnaryCallSettings
+ beginTransactionSettings() {
+ return ((DatastoreStubSettings) getStubSettings()).beginTransactionSettings();
+ }
+
+ /** Returns the object with the settings used for calls to commit. */
+ public UnaryCallSettings commitSettings() {
+ return ((DatastoreStubSettings) getStubSettings()).commitSettings();
+ }
+
+ /** Returns the object with the settings used for calls to rollback. */
+ public UnaryCallSettings rollbackSettings() {
+ return ((DatastoreStubSettings) getStubSettings()).rollbackSettings();
+ }
+
+ /** Returns the object with the settings used for calls to allocateIds. */
+ public UnaryCallSettings allocateIdsSettings() {
+ return ((DatastoreStubSettings) getStubSettings()).allocateIdsSettings();
+ }
+
+ /** Returns the object with the settings used for calls to reserveIds. */
+ public UnaryCallSettings reserveIdsSettings() {
+ return ((DatastoreStubSettings) getStubSettings()).reserveIdsSettings();
+ }
+
+ public static final DatastoreSettings create(DatastoreStubSettings stub) throws IOException {
+ return new DatastoreSettings.Builder(stub.toBuilder()).build();
+ }
+
+ /** Returns a builder for the default ExecutorProvider for this service. */
+ public static InstantiatingExecutorProvider.Builder defaultExecutorProviderBuilder() {
+ return DatastoreStubSettings.defaultExecutorProviderBuilder();
+ }
+
+ /** Returns the default service endpoint. */
+ public static String getDefaultEndpoint() {
+ return DatastoreStubSettings.getDefaultEndpoint();
+ }
+
+ /** Returns the default service scopes. */
+ public static List getDefaultServiceScopes() {
+ return DatastoreStubSettings.getDefaultServiceScopes();
+ }
+
+ /** Returns a builder for the default credentials for this service. */
+ public static GoogleCredentialsProvider.Builder defaultCredentialsProviderBuilder() {
+ return DatastoreStubSettings.defaultCredentialsProviderBuilder();
+ }
+
+ /** Returns a builder for the default gRPC ChannelProvider for this service. */
+ public static InstantiatingGrpcChannelProvider.Builder defaultGrpcTransportProviderBuilder() {
+ return DatastoreStubSettings.defaultGrpcTransportProviderBuilder();
+ }
+
+ /** Returns a builder for the default REST ChannelProvider for this service. */
+ @BetaApi
+ public static InstantiatingHttpJsonChannelProvider.Builder
+ defaultHttpJsonTransportProviderBuilder() {
+ return DatastoreStubSettings.defaultHttpJsonTransportProviderBuilder();
+ }
+
+ public static TransportChannelProvider defaultTransportChannelProvider() {
+ return DatastoreStubSettings.defaultTransportChannelProvider();
+ }
+
+ public static ApiClientHeaderProvider.Builder defaultApiClientHeaderProviderBuilder() {
+ return DatastoreStubSettings.defaultApiClientHeaderProviderBuilder();
+ }
+
+ /** Returns a new gRPC builder for this class. */
+ public static Builder newBuilder() {
+ return Builder.createDefault();
+ }
+
+ /** Returns a new REST builder for this class. */
+ public static Builder newHttpJsonBuilder() {
+ return Builder.createHttpJsonDefault();
+ }
+
+ /** Returns a new builder for this class. */
+ public static Builder newBuilder(ClientContext clientContext) {
+ return new Builder(clientContext);
+ }
+
+ /** Returns a builder containing all the values of this settings class. */
+ public Builder toBuilder() {
+ return new Builder(this);
+ }
+
+ protected DatastoreSettings(Builder settingsBuilder) throws IOException {
+ super(settingsBuilder);
+ }
+
+ /** Builder for DatastoreSettings. */
+ public static class Builder extends ClientSettings.Builder {
+
+ protected Builder() throws IOException {
+ this(((ClientContext) null));
+ }
+
+ protected Builder(ClientContext clientContext) {
+ super(DatastoreStubSettings.newBuilder(clientContext));
+ }
+
+ protected Builder(DatastoreSettings settings) {
+ super(settings.getStubSettings().toBuilder());
+ }
+
+ protected Builder(DatastoreStubSettings.Builder stubSettings) {
+ super(stubSettings);
+ }
+
+ private static Builder createDefault() {
+ return new Builder(DatastoreStubSettings.newBuilder());
+ }
+
+ private static Builder createHttpJsonDefault() {
+ return new Builder(DatastoreStubSettings.newHttpJsonBuilder());
+ }
+
+ public DatastoreStubSettings.Builder getStubSettingsBuilder() {
+ return ((DatastoreStubSettings.Builder) getStubSettings());
+ }
+
+ /**
+ * Applies the given settings updater function to all of the unary API methods in this service.
+ *
+ * Note: This method does not support applying settings to streaming methods.
+ */
+ public Builder applyToAllUnaryMethods(
+ ApiFunction, Void> settingsUpdater) {
+ super.applyToAllUnaryMethods(
+ getStubSettingsBuilder().unaryMethodSettingsBuilders(), settingsUpdater);
+ return this;
+ }
+
+ /** Returns the builder for the settings used for calls to lookup. */
+ public UnaryCallSettings.Builder lookupSettings() {
+ return getStubSettingsBuilder().lookupSettings();
+ }
+
+ /** Returns the builder for the settings used for calls to runQuery. */
+ public UnaryCallSettings.Builder runQuerySettings() {
+ return getStubSettingsBuilder().runQuerySettings();
+ }
+
+ /** Returns the builder for the settings used for calls to runAggregationQuery. */
+ public UnaryCallSettings.Builder
+ runAggregationQuerySettings() {
+ return getStubSettingsBuilder().runAggregationQuerySettings();
+ }
+
+ /** Returns the builder for the settings used for calls to beginTransaction. */
+ public UnaryCallSettings.Builder
+ beginTransactionSettings() {
+ return getStubSettingsBuilder().beginTransactionSettings();
+ }
+
+ /** Returns the builder for the settings used for calls to commit. */
+ public UnaryCallSettings.Builder commitSettings() {
+ return getStubSettingsBuilder().commitSettings();
+ }
+
+ /** Returns the builder for the settings used for calls to rollback. */
+ public UnaryCallSettings.Builder rollbackSettings() {
+ return getStubSettingsBuilder().rollbackSettings();
+ }
+
+ /** Returns the builder for the settings used for calls to allocateIds. */
+ public UnaryCallSettings.Builder
+ allocateIdsSettings() {
+ return getStubSettingsBuilder().allocateIdsSettings();
+ }
+
+ /** Returns the builder for the settings used for calls to reserveIds. */
+ public UnaryCallSettings.Builder reserveIdsSettings() {
+ return getStubSettingsBuilder().reserveIdsSettings();
+ }
+
+ @Override
+ public DatastoreSettings build() throws IOException {
+ return new DatastoreSettings(this);
+ }
+ }
+}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/gapic_metadata.json b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/gapic_metadata.json
new file mode 100644
index 000000000..02196d36e
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/gapic_metadata.json
@@ -0,0 +1,42 @@
+{
+ "schema": "1.0",
+ "comment": "This file maps proto services/RPCs to the corresponding library clients/methods",
+ "language": "java",
+ "protoPackage": "google.datastore.v1",
+ "libraryPackage": "com.google.cloud.datastore.v1",
+ "services": {
+ "Datastore": {
+ "clients": {
+ "grpc": {
+ "libraryClient": "DatastoreClient",
+ "rpcs": {
+ "AllocateIds": {
+ "methods": ["allocateIds", "allocateIds", "allocateIdsCallable"]
+ },
+ "BeginTransaction": {
+ "methods": ["beginTransaction", "beginTransaction", "beginTransactionCallable"]
+ },
+ "Commit": {
+ "methods": ["commit", "commit", "commit", "commitCallable"]
+ },
+ "Lookup": {
+ "methods": ["lookup", "lookup", "lookupCallable"]
+ },
+ "ReserveIds": {
+ "methods": ["reserveIds", "reserveIds", "reserveIdsCallable"]
+ },
+ "Rollback": {
+ "methods": ["rollback", "rollback", "rollbackCallable"]
+ },
+ "RunAggregationQuery": {
+ "methods": ["runAggregationQuery", "runAggregationQueryCallable"]
+ },
+ "RunQuery": {
+ "methods": ["runQuery", "runQueryCallable"]
+ }
+ }
+ }
+ }
+ }
+ }
+}
\ No newline at end of file
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/package-info.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/package-info.java
new file mode 100644
index 000000000..0484a7c04
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/package-info.java
@@ -0,0 +1,49 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * A client to Cloud Datastore API
+ *
+ * The interfaces provided are listed below, along with usage samples.
+ *
+ *
======================= DatastoreClient =======================
+ *
+ *
Service Description: Each RPC normalizes the partition IDs of the keys in its input entities,
+ * and always returns entities with keys with normalized partition IDs. This applies to all keys and
+ * entities, including those in values, except keys with both an empty path and an empty or unset
+ * partition ID. Normalization of input keys sets the project ID (if not already set) to the project
+ * ID from the request.
+ *
+ *
Sample for DatastoreClient:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * try (DatastoreClient datastoreClient = DatastoreClient.create()) {
+ * String projectId = "projectId-894832108";
+ * ReadOptions readOptions = ReadOptions.newBuilder().build();
+ * List keys = new ArrayList<>();
+ * LookupResponse response = datastoreClient.lookup(projectId, readOptions, keys);
+ * }
+ * }
+ */
+@Generated("by gapic-generator-java")
+package com.google.cloud.datastore.v1;
+
+import javax.annotation.Generated;
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/stub/DatastoreStub.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/stub/DatastoreStub.java
new file mode 100644
index 000000000..231289f18
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/stub/DatastoreStub.java
@@ -0,0 +1,84 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.datastore.v1.stub;
+
+import com.google.api.gax.core.BackgroundResource;
+import com.google.api.gax.rpc.UnaryCallable;
+import com.google.datastore.v1.AllocateIdsRequest;
+import com.google.datastore.v1.AllocateIdsResponse;
+import com.google.datastore.v1.BeginTransactionRequest;
+import com.google.datastore.v1.BeginTransactionResponse;
+import com.google.datastore.v1.CommitRequest;
+import com.google.datastore.v1.CommitResponse;
+import com.google.datastore.v1.LookupRequest;
+import com.google.datastore.v1.LookupResponse;
+import com.google.datastore.v1.ReserveIdsRequest;
+import com.google.datastore.v1.ReserveIdsResponse;
+import com.google.datastore.v1.RollbackRequest;
+import com.google.datastore.v1.RollbackResponse;
+import com.google.datastore.v1.RunAggregationQueryRequest;
+import com.google.datastore.v1.RunAggregationQueryResponse;
+import com.google.datastore.v1.RunQueryRequest;
+import com.google.datastore.v1.RunQueryResponse;
+import javax.annotation.Generated;
+
+// AUTO-GENERATED DOCUMENTATION AND CLASS.
+/**
+ * Base stub class for the Datastore service API.
+ *
+ * This class is for advanced usage and reflects the underlying API directly.
+ */
+@Generated("by gapic-generator-java")
+public abstract class DatastoreStub implements BackgroundResource {
+
+ public UnaryCallable lookupCallable() {
+ throw new UnsupportedOperationException("Not implemented: lookupCallable()");
+ }
+
+ public UnaryCallable runQueryCallable() {
+ throw new UnsupportedOperationException("Not implemented: runQueryCallable()");
+ }
+
+ public UnaryCallable
+ runAggregationQueryCallable() {
+ throw new UnsupportedOperationException("Not implemented: runAggregationQueryCallable()");
+ }
+
+ public UnaryCallable
+ beginTransactionCallable() {
+ throw new UnsupportedOperationException("Not implemented: beginTransactionCallable()");
+ }
+
+ public UnaryCallable commitCallable() {
+ throw new UnsupportedOperationException("Not implemented: commitCallable()");
+ }
+
+ public UnaryCallable rollbackCallable() {
+ throw new UnsupportedOperationException("Not implemented: rollbackCallable()");
+ }
+
+ public UnaryCallable allocateIdsCallable() {
+ throw new UnsupportedOperationException("Not implemented: allocateIdsCallable()");
+ }
+
+ public UnaryCallable reserveIdsCallable() {
+ throw new UnsupportedOperationException("Not implemented: reserveIdsCallable()");
+ }
+
+ @Override
+ public abstract void close();
+}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/stub/DatastoreStubSettings.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/stub/DatastoreStubSettings.java
new file mode 100644
index 000000000..a4554339d
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/v1/stub/DatastoreStubSettings.java
@@ -0,0 +1,535 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.datastore.v1.stub;
+
+import com.google.api.core.ApiFunction;
+import com.google.api.core.BetaApi;
+import com.google.api.core.ObsoleteApi;
+import com.google.api.gax.core.GaxProperties;
+import com.google.api.gax.core.GoogleCredentialsProvider;
+import com.google.api.gax.core.InstantiatingExecutorProvider;
+import com.google.api.gax.grpc.GaxGrpcProperties;
+import com.google.api.gax.grpc.GrpcTransportChannel;
+import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
+import com.google.api.gax.httpjson.GaxHttpJsonProperties;
+import com.google.api.gax.httpjson.HttpJsonTransportChannel;
+import com.google.api.gax.httpjson.InstantiatingHttpJsonChannelProvider;
+import com.google.api.gax.retrying.RetrySettings;
+import com.google.api.gax.rpc.ApiClientHeaderProvider;
+import com.google.api.gax.rpc.ClientContext;
+import com.google.api.gax.rpc.StatusCode;
+import com.google.api.gax.rpc.StubSettings;
+import com.google.api.gax.rpc.TransportChannelProvider;
+import com.google.api.gax.rpc.UnaryCallSettings;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Lists;
+import com.google.datastore.v1.AllocateIdsRequest;
+import com.google.datastore.v1.AllocateIdsResponse;
+import com.google.datastore.v1.BeginTransactionRequest;
+import com.google.datastore.v1.BeginTransactionResponse;
+import com.google.datastore.v1.CommitRequest;
+import com.google.datastore.v1.CommitResponse;
+import com.google.datastore.v1.LookupRequest;
+import com.google.datastore.v1.LookupResponse;
+import com.google.datastore.v1.ReserveIdsRequest;
+import com.google.datastore.v1.ReserveIdsResponse;
+import com.google.datastore.v1.RollbackRequest;
+import com.google.datastore.v1.RollbackResponse;
+import com.google.datastore.v1.RunAggregationQueryRequest;
+import com.google.datastore.v1.RunAggregationQueryResponse;
+import com.google.datastore.v1.RunQueryRequest;
+import com.google.datastore.v1.RunQueryResponse;
+import java.io.IOException;
+import java.time.Duration;
+import java.util.List;
+import javax.annotation.Generated;
+
+// AUTO-GENERATED DOCUMENTATION AND CLASS.
+/**
+ * Settings class to configure an instance of {@link DatastoreStub}.
+ *
+ * The default instance has everything set to sensible defaults:
+ *
+ *
+ * - The default service address (datastore.googleapis.com) and default port (443) are used.
+ *
- Credentials are acquired automatically through Application Default Credentials.
+ *
- Retries are configured for idempotent methods but not for non-idempotent methods.
+ *
+ *
+ * The builder of this class is recursive, so contained classes are themselves builders. When
+ * build() is called, the tree of builders is called to create the complete settings object.
+ *
+ *
For example, to set the
+ * [RetrySettings](https://cloud.google.com/java/docs/reference/gax/latest/com.google.api.gax.retrying.RetrySettings)
+ * of lookup:
+ *
+ *
{@code
+ * // This snippet has been automatically generated and should be regarded as a code template only.
+ * // It will require modifications to work:
+ * // - It may require correct/in-range values for request initialization.
+ * // - It may require specifying regional endpoints when creating the service client as shown in
+ * // https://cloud.google.com/java/docs/setup#configure_endpoints_for_the_client_library
+ * DatastoreStubSettings.Builder datastoreSettingsBuilder = DatastoreStubSettings.newBuilder();
+ * datastoreSettingsBuilder
+ * .lookupSettings()
+ * .setRetrySettings(
+ * datastoreSettingsBuilder
+ * .lookupSettings()
+ * .getRetrySettings()
+ * .toBuilder()
+ * .setInitialRetryDelayDuration(Duration.ofSeconds(1))
+ * .setInitialRpcTimeoutDuration(Duration.ofSeconds(5))
+ * .setMaxAttempts(5)
+ * .setMaxRetryDelayDuration(Duration.ofSeconds(30))
+ * .setMaxRpcTimeoutDuration(Duration.ofSeconds(60))
+ * .setRetryDelayMultiplier(1.3)
+ * .setRpcTimeoutMultiplier(1.5)
+ * .setTotalTimeoutDuration(Duration.ofSeconds(300))
+ * .build());
+ * DatastoreStubSettings datastoreSettings = datastoreSettingsBuilder.build();
+ * }
+ *
+ * Please refer to the [Client Side Retry
+ * Guide](https://github.com/googleapis/google-cloud-java/blob/main/docs/client_retries.md) for
+ * additional support in setting retries.
+ */
+@Generated("by gapic-generator-java")
+public class DatastoreStubSettings extends StubSettings {
+ /** The default scopes of the service. */
+ private static final ImmutableList DEFAULT_SERVICE_SCOPES =
+ ImmutableList.builder()
+ .add("https://www.googleapis.com/auth/cloud-platform")
+ .add("https://www.googleapis.com/auth/datastore")
+ .build();
+
+ private final UnaryCallSettings lookupSettings;
+ private final UnaryCallSettings runQuerySettings;
+ private final UnaryCallSettings
+ runAggregationQuerySettings;
+ private final UnaryCallSettings
+ beginTransactionSettings;
+ private final UnaryCallSettings commitSettings;
+ private final UnaryCallSettings rollbackSettings;
+ private final UnaryCallSettings allocateIdsSettings;
+ private final UnaryCallSettings reserveIdsSettings;
+
+ /** Returns the object with the settings used for calls to lookup. */
+ public UnaryCallSettings lookupSettings() {
+ return lookupSettings;
+ }
+
+ /** Returns the object with the settings used for calls to runQuery. */
+ public UnaryCallSettings runQuerySettings() {
+ return runQuerySettings;
+ }
+
+ /** Returns the object with the settings used for calls to runAggregationQuery. */
+ public UnaryCallSettings
+ runAggregationQuerySettings() {
+ return runAggregationQuerySettings;
+ }
+
+ /** Returns the object with the settings used for calls to beginTransaction. */
+ public UnaryCallSettings
+ beginTransactionSettings() {
+ return beginTransactionSettings;
+ }
+
+ /** Returns the object with the settings used for calls to commit. */
+ public UnaryCallSettings commitSettings() {
+ return commitSettings;
+ }
+
+ /** Returns the object with the settings used for calls to rollback. */
+ public UnaryCallSettings rollbackSettings() {
+ return rollbackSettings;
+ }
+
+ /** Returns the object with the settings used for calls to allocateIds. */
+ public UnaryCallSettings allocateIdsSettings() {
+ return allocateIdsSettings;
+ }
+
+ /** Returns the object with the settings used for calls to reserveIds. */
+ public UnaryCallSettings reserveIdsSettings() {
+ return reserveIdsSettings;
+ }
+
+ public DatastoreStub createStub() throws IOException {
+ if (getTransportChannelProvider()
+ .getTransportName()
+ .equals(GrpcTransportChannel.getGrpcTransportName())) {
+ return GrpcDatastoreStub.create(this);
+ }
+ if (getTransportChannelProvider()
+ .getTransportName()
+ .equals(HttpJsonTransportChannel.getHttpJsonTransportName())) {
+ return HttpJsonDatastoreStub.create(this);
+ }
+ throw new UnsupportedOperationException(
+ String.format(
+ "Transport not supported: %s", getTransportChannelProvider().getTransportName()));
+ }
+
+ /** Returns the default service name. */
+ @Override
+ public String getServiceName() {
+ return "datastore";
+ }
+
+ /** Returns a builder for the default ExecutorProvider for this service. */
+ public static InstantiatingExecutorProvider.Builder defaultExecutorProviderBuilder() {
+ return InstantiatingExecutorProvider.newBuilder();
+ }
+
+ /** Returns the default service endpoint. */
+ @ObsoleteApi("Use getEndpoint() instead")
+ public static String getDefaultEndpoint() {
+ return "datastore.googleapis.com:443";
+ }
+
+ /** Returns the default mTLS service endpoint. */
+ public static String getDefaultMtlsEndpoint() {
+ return "datastore.mtls.googleapis.com:443";
+ }
+
+ /** Returns the default service scopes. */
+ public static List getDefaultServiceScopes() {
+ return DEFAULT_SERVICE_SCOPES;
+ }
+
+ /** Returns a builder for the default credentials for this service. */
+ public static GoogleCredentialsProvider.Builder defaultCredentialsProviderBuilder() {
+ return GoogleCredentialsProvider.newBuilder()
+ .setScopesToApply(DEFAULT_SERVICE_SCOPES)
+ .setUseJwtAccessWithScope(true);
+ }
+
+ /** Returns a builder for the default gRPC ChannelProvider for this service. */
+ public static InstantiatingGrpcChannelProvider.Builder defaultGrpcTransportProviderBuilder() {
+ return InstantiatingGrpcChannelProvider.newBuilder()
+ .setMaxInboundMessageSize(Integer.MAX_VALUE);
+ }
+
+ /** Returns a builder for the default REST ChannelProvider for this service. */
+ @BetaApi
+ public static InstantiatingHttpJsonChannelProvider.Builder
+ defaultHttpJsonTransportProviderBuilder() {
+ return InstantiatingHttpJsonChannelProvider.newBuilder();
+ }
+
+ public static TransportChannelProvider defaultTransportChannelProvider() {
+ return defaultGrpcTransportProviderBuilder().build();
+ }
+
+ public static ApiClientHeaderProvider.Builder defaultGrpcApiClientHeaderProviderBuilder() {
+ return ApiClientHeaderProvider.newBuilder()
+ .setGeneratedLibToken("gapic", GaxProperties.getLibraryVersion(DatastoreStubSettings.class))
+ .setTransportToken(
+ GaxGrpcProperties.getGrpcTokenName(), GaxGrpcProperties.getGrpcVersion());
+ }
+
+ public static ApiClientHeaderProvider.Builder defaultHttpJsonApiClientHeaderProviderBuilder() {
+ return ApiClientHeaderProvider.newBuilder()
+ .setGeneratedLibToken("gapic", GaxProperties.getLibraryVersion(DatastoreStubSettings.class))
+ .setTransportToken(
+ GaxHttpJsonProperties.getHttpJsonTokenName(),
+ GaxHttpJsonProperties.getHttpJsonVersion());
+ }
+
+ public static ApiClientHeaderProvider.Builder defaultApiClientHeaderProviderBuilder() {
+ return DatastoreStubSettings.defaultGrpcApiClientHeaderProviderBuilder();
+ }
+
+ /** Returns a new gRPC builder for this class. */
+ public static Builder newBuilder() {
+ return Builder.createDefault();
+ }
+
+ /** Returns a new REST builder for this class. */
+ public static Builder newHttpJsonBuilder() {
+ return Builder.createHttpJsonDefault();
+ }
+
+ /** Returns a new builder for this class. */
+ public static Builder newBuilder(ClientContext clientContext) {
+ return new Builder(clientContext);
+ }
+
+ /** Returns a builder containing all the values of this settings class. */
+ public Builder toBuilder() {
+ return new Builder(this);
+ }
+
+ protected DatastoreStubSettings(Builder settingsBuilder) throws IOException {
+ super(settingsBuilder);
+
+ lookupSettings = settingsBuilder.lookupSettings().build();
+ runQuerySettings = settingsBuilder.runQuerySettings().build();
+ runAggregationQuerySettings = settingsBuilder.runAggregationQuerySettings().build();
+ beginTransactionSettings = settingsBuilder.beginTransactionSettings().build();
+ commitSettings = settingsBuilder.commitSettings().build();
+ rollbackSettings = settingsBuilder.rollbackSettings().build();
+ allocateIdsSettings = settingsBuilder.allocateIdsSettings().build();
+ reserveIdsSettings = settingsBuilder.reserveIdsSettings().build();
+ }
+
+ /** Builder for DatastoreStubSettings. */
+ public static class Builder extends StubSettings.Builder {
+ private final ImmutableList> unaryMethodSettingsBuilders;
+ private final UnaryCallSettings.Builder lookupSettings;
+ private final UnaryCallSettings.Builder runQuerySettings;
+ private final UnaryCallSettings.Builder
+ runAggregationQuerySettings;
+ private final UnaryCallSettings.Builder
+ beginTransactionSettings;
+ private final UnaryCallSettings.Builder commitSettings;
+ private final UnaryCallSettings.Builder rollbackSettings;
+ private final UnaryCallSettings.Builder
+ allocateIdsSettings;
+ private final UnaryCallSettings.Builder
+ reserveIdsSettings;
+ private static final ImmutableMap>
+ RETRYABLE_CODE_DEFINITIONS;
+
+ static {
+ ImmutableMap.Builder> definitions =
+ ImmutableMap.builder();
+ definitions.put(
+ "retry_policy_0_codes",
+ ImmutableSet.copyOf(
+ Lists.newArrayList(
+ StatusCode.Code.UNAVAILABLE, StatusCode.Code.DEADLINE_EXCEEDED)));
+ definitions.put(
+ "no_retry_1_codes", ImmutableSet.copyOf(Lists.newArrayList()));
+ RETRYABLE_CODE_DEFINITIONS = definitions.build();
+ }
+
+ private static final ImmutableMap RETRY_PARAM_DEFINITIONS;
+
+ static {
+ ImmutableMap.Builder definitions = ImmutableMap.builder();
+ RetrySettings settings = null;
+ settings =
+ RetrySettings.newBuilder()
+ .setInitialRetryDelayDuration(Duration.ofMillis(100L))
+ .setRetryDelayMultiplier(1.3)
+ .setMaxRetryDelayDuration(Duration.ofMillis(60000L))
+ .setInitialRpcTimeoutDuration(Duration.ofMillis(60000L))
+ .setRpcTimeoutMultiplier(1.0)
+ .setMaxRpcTimeoutDuration(Duration.ofMillis(60000L))
+ .setTotalTimeoutDuration(Duration.ofMillis(60000L))
+ .build();
+ definitions.put("retry_policy_0_params", settings);
+ settings =
+ RetrySettings.newBuilder()
+ .setInitialRpcTimeoutDuration(Duration.ofMillis(60000L))
+ .setRpcTimeoutMultiplier(1.0)
+ .setMaxRpcTimeoutDuration(Duration.ofMillis(60000L))
+ .setTotalTimeoutDuration(Duration.ofMillis(60000L))
+ .build();
+ definitions.put("no_retry_1_params", settings);
+ RETRY_PARAM_DEFINITIONS = definitions.build();
+ }
+
+ protected Builder() {
+ this(((ClientContext) null));
+ }
+
+ protected Builder(ClientContext clientContext) {
+ super(clientContext);
+
+ lookupSettings = UnaryCallSettings.newUnaryCallSettingsBuilder();
+ runQuerySettings = UnaryCallSettings.newUnaryCallSettingsBuilder();
+ runAggregationQuerySettings = UnaryCallSettings.newUnaryCallSettingsBuilder();
+ beginTransactionSettings = UnaryCallSettings.newUnaryCallSettingsBuilder();
+ commitSettings = UnaryCallSettings.newUnaryCallSettingsBuilder();
+ rollbackSettings = UnaryCallSettings.newUnaryCallSettingsBuilder();
+ allocateIdsSettings = UnaryCallSettings.newUnaryCallSettingsBuilder();
+ reserveIdsSettings = UnaryCallSettings.newUnaryCallSettingsBuilder();
+
+ unaryMethodSettingsBuilders =
+ ImmutableList.>of(
+ lookupSettings,
+ runQuerySettings,
+ runAggregationQuerySettings,
+ beginTransactionSettings,
+ commitSettings,
+ rollbackSettings,
+ allocateIdsSettings,
+ reserveIdsSettings);
+ initDefaults(this);
+ }
+
+ protected Builder(DatastoreStubSettings settings) {
+ super(settings);
+
+ lookupSettings = settings.lookupSettings.toBuilder();
+ runQuerySettings = settings.runQuerySettings.toBuilder();
+ runAggregationQuerySettings = settings.runAggregationQuerySettings.toBuilder();
+ beginTransactionSettings = settings.beginTransactionSettings.toBuilder();
+ commitSettings = settings.commitSettings.toBuilder();
+ rollbackSettings = settings.rollbackSettings.toBuilder();
+ allocateIdsSettings = settings.allocateIdsSettings.toBuilder();
+ reserveIdsSettings = settings.reserveIdsSettings.toBuilder();
+
+ unaryMethodSettingsBuilders =
+ ImmutableList.>of(
+ lookupSettings,
+ runQuerySettings,
+ runAggregationQuerySettings,
+ beginTransactionSettings,
+ commitSettings,
+ rollbackSettings,
+ allocateIdsSettings,
+ reserveIdsSettings);
+ }
+
+ private static Builder createDefault() {
+ Builder builder = new Builder(((ClientContext) null));
+
+ builder.setTransportChannelProvider(defaultTransportChannelProvider());
+ builder.setCredentialsProvider(defaultCredentialsProviderBuilder().build());
+ builder.setInternalHeaderProvider(defaultApiClientHeaderProviderBuilder().build());
+ builder.setMtlsEndpoint(getDefaultMtlsEndpoint());
+ builder.setSwitchToMtlsEndpointAllowed(true);
+
+ return initDefaults(builder);
+ }
+
+ private static Builder createHttpJsonDefault() {
+ Builder builder = new Builder(((ClientContext) null));
+
+ builder.setTransportChannelProvider(defaultHttpJsonTransportProviderBuilder().build());
+ builder.setCredentialsProvider(defaultCredentialsProviderBuilder().build());
+ builder.setInternalHeaderProvider(defaultHttpJsonApiClientHeaderProviderBuilder().build());
+ builder.setMtlsEndpoint(getDefaultMtlsEndpoint());
+ builder.setSwitchToMtlsEndpointAllowed(true);
+
+ return initDefaults(builder);
+ }
+
+ private static Builder initDefaults(Builder builder) {
+ builder
+ .lookupSettings()
+ .setRetryableCodes(RETRYABLE_CODE_DEFINITIONS.get("retry_policy_0_codes"))
+ .setRetrySettings(RETRY_PARAM_DEFINITIONS.get("retry_policy_0_params"));
+
+ builder
+ .runQuerySettings()
+ .setRetryableCodes(RETRYABLE_CODE_DEFINITIONS.get("retry_policy_0_codes"))
+ .setRetrySettings(RETRY_PARAM_DEFINITIONS.get("retry_policy_0_params"));
+
+ builder
+ .runAggregationQuerySettings()
+ .setRetryableCodes(RETRYABLE_CODE_DEFINITIONS.get("retry_policy_0_codes"))
+ .setRetrySettings(RETRY_PARAM_DEFINITIONS.get("retry_policy_0_params"));
+
+ builder
+ .beginTransactionSettings()
+ .setRetryableCodes(RETRYABLE_CODE_DEFINITIONS.get("no_retry_1_codes"))
+ .setRetrySettings(RETRY_PARAM_DEFINITIONS.get("no_retry_1_params"));
+
+ builder
+ .commitSettings()
+ .setRetryableCodes(RETRYABLE_CODE_DEFINITIONS.get("no_retry_1_codes"))
+ .setRetrySettings(RETRY_PARAM_DEFINITIONS.get("no_retry_1_params"));
+
+ builder
+ .rollbackSettings()
+ .setRetryableCodes(RETRYABLE_CODE_DEFINITIONS.get("no_retry_1_codes"))
+ .setRetrySettings(RETRY_PARAM_DEFINITIONS.get("no_retry_1_params"));
+
+ builder
+ .allocateIdsSettings()
+ .setRetryableCodes(RETRYABLE_CODE_DEFINITIONS.get("no_retry_1_codes"))
+ .setRetrySettings(RETRY_PARAM_DEFINITIONS.get("no_retry_1_params"));
+
+ builder
+ .reserveIdsSettings()
+ .setRetryableCodes(RETRYABLE_CODE_DEFINITIONS.get("retry_policy_0_codes"))
+ .setRetrySettings(RETRY_PARAM_DEFINITIONS.get("retry_policy_0_params"));
+
+ return builder;
+ }
+
+ /**
+ * Applies the given settings updater function to all of the unary API methods in this service.
+ *
+ * Note: This method does not support applying settings to streaming methods.
+ */
+ public Builder applyToAllUnaryMethods(
+ ApiFunction, Void> settingsUpdater) {
+ super.applyToAllUnaryMethods(unaryMethodSettingsBuilders, settingsUpdater);
+ return this;
+ }
+
+ public ImmutableList> unaryMethodSettingsBuilders() {
+ return unaryMethodSettingsBuilders;
+ }
+
+ /** Returns the builder for the settings used for calls to lookup. */
+ public UnaryCallSettings.Builder lookupSettings() {
+ return lookupSettings;
+ }
+
+ /** Returns the builder for the settings used for calls to runQuery. */
+ public UnaryCallSettings.Builder