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/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/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/spi/v1/HttpDatastoreRpc.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java
index fd3cdc658..d927a2d7f 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
@@ -21,7 +21,6 @@
import com.google.api.client.http.HttpTransport;
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,6 +39,7 @@
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;
@@ -80,8 +80,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));
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..06941c721
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java
@@ -0,0 +1,145 @@
+/*
+ * 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.SpanContext;
+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 SpanContext implements TraceUtil.SpanContext {
+ @Override
+ public io.opentelemetry.api.trace.SpanContext getSpanContext() {
+ return null;
+ }
+ }
+
+ 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.SpanContext parentSpanContext) {
+ 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();
+ }
+
+ @Nonnull
+ @Override
+ public TraceUtil.SpanContext getCurrentSpanContext() {
+ return new SpanContext();
+ }
+
+ @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..3b962754d
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java
@@ -0,0 +1,346 @@
+/*
+ * 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.SpanContext;
+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.Span;
+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 io.opentelemetry.context.Context;
+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 SpanContext implements TraceUtil.SpanContext {
+ private final io.opentelemetry.api.trace.SpanContext spanContext;
+
+ public SpanContext(io.opentelemetry.api.trace.SpanContext spanContext) {
+ this.spanContext = spanContext;
+ }
+
+ @Override
+ public io.opentelemetry.api.trace.SpanContext getSpanContext() {
+ return this.spanContext;
+ }
+ }
+
+ 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;
+ }
+
+ /** 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;
+ }
+
+ public io.opentelemetry.api.trace.Span getSpan() {
+ return this.span;
+ }
+
+ @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.SpanContext parentSpanContext) {
+ SpanBuilder spanBuilder =
+ tracer
+ .spanBuilder(spanName)
+ .setSpanKind(SpanKind.PRODUCER)
+ .setParent(
+ io.opentelemetry.context.Context.current()
+ .with(
+ io.opentelemetry.api.trace.Span.wrap(parentSpanContext.getSpanContext())));
+ io.opentelemetry.api.trace.Span span =
+ addSettingsAttributesToCurrentSpan(spanBuilder).startSpan();
+ return new Span(span, 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());
+ }
+
+ @Nonnull
+ @Override
+ public TraceUtil.SpanContext getCurrentSpanContext() {
+ return new SpanContext(io.opentelemetry.api.trace.Span.current().getSpanContext());
+ }
+
+ @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..dce53e952
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java
@@ -0,0 +1,168 @@
+/*
+ * 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.Span;
+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";
+ /**
+ * 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's context */
+ interface SpanContext {
+ io.opentelemetry.api.trace.SpanContext getSpanContext();
+ }
+
+ /** 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 parentSpanContext as its
+ * parents, sets it as the current span and returns it.
+ */
+ Span startSpan(String spanName, SpanContext parentSpanContext);
+
+ /**
+ * 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 SpanContext */
+ @Nonnull
+ SpanContext getCurrentSpanContext();
+
+ /** 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/RemoteDatastoreHelper.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/RemoteDatastoreHelper.java
index 596ce96d8..6167cedca 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/RemoteDatastoreHelper.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/RemoteDatastoreHelper.java
@@ -19,13 +19,16 @@
import com.google.api.core.InternalApi;
import com.google.api.gax.retrying.RetrySettings;
import com.google.cloud.datastore.Datastore;
+import com.google.cloud.datastore.DatastoreOpenTelemetryOptions;
import com.google.cloud.datastore.DatastoreOptions;
import com.google.cloud.datastore.Key;
import com.google.cloud.datastore.Query;
import com.google.cloud.datastore.QueryResults;
import com.google.cloud.datastore.StructuredQuery;
import com.google.cloud.http.HttpTransportOptions;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
import java.util.UUID;
+import javax.annotation.Nullable;
import org.threeten.bp.Duration;
/**
@@ -38,13 +41,13 @@
* RetrySettings#getTotalTimeout()} is {@code 120000} and {@link
* RetrySettings#getInitialRetryDelay()} is {@code 250}. {@link
* HttpTransportOptions#getConnectTimeout()} and {@link HttpTransportOptions#getReadTimeout()} are
- * both both set to {@code 60000}.
+ * both both set to {@code 60000}. If an OpenTelemetrySdk object is passed in, OpenTelemetry Trace
+ * collection will be enabled for the Client application.
*
* Internal testing use only
*/
@InternalApi
public class RemoteDatastoreHelper {
-
private final DatastoreOptions options;
private final Datastore datastore;
private final String namespace;
@@ -78,18 +81,30 @@ public static RemoteDatastoreHelper create() {
}
/** Creates a {@code RemoteStorageHelper} object. */
- public static RemoteDatastoreHelper create(String databaseId) {
+ public static RemoteDatastoreHelper create(
+ String databaseId, @Nullable OpenTelemetrySdk openTelemetrySdk) {
HttpTransportOptions transportOptions = DatastoreOptions.getDefaultHttpTransportOptions();
transportOptions =
transportOptions.toBuilder().setConnectTimeout(60000).setReadTimeout(60000).build();
- DatastoreOptions datastoreOption =
+ DatastoreOptions.Builder datastoreOptionBuilder =
DatastoreOptions.newBuilder()
.setDatabaseId(databaseId)
.setNamespace(UUID.randomUUID().toString())
.setRetrySettings(retrySettings())
- .setTransportOptions(transportOptions)
- .build();
- return new RemoteDatastoreHelper(datastoreOption);
+ .setTransportOptions(transportOptions);
+
+ if (openTelemetrySdk != null) {
+ datastoreOptionBuilder.setOpenTelemetryOptions(
+ DatastoreOpenTelemetryOptions.newBuilder()
+ .setOpenTelemetry(openTelemetrySdk)
+ .setTracingEnabled(true)
+ .build());
+ }
+ return new RemoteDatastoreHelper(datastoreOptionBuilder.build());
+ }
+
+ public static RemoteDatastoreHelper create(String databaseId) {
+ return create(databaseId, /*openTelemetrySdk=*/ null);
}
private static RetrySettings retrySettings() {
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java
index a545580e2..85703f739 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java
@@ -70,6 +70,20 @@ public void testHost() {
assertEquals("http://localhost:" + PORT, options.build().getHost());
}
+ @Test
+ public void testOpenTelemetryOptionsEnabled() {
+ options.setOpenTelemetryOptions(
+ DatastoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build());
+ assertTrue(options.build().getOpenTelemetryOptions().isEnabled());
+ }
+
+ @Test
+ public void testOpenTelemetryOptionsDisabled() {
+ options.setOpenTelemetryOptions(
+ DatastoreOpenTelemetryOptions.newBuilder().setTracingEnabled(false).build());
+ assertTrue(!options.build().getOpenTelemetryOptions().isEnabled());
+ }
+
@Test
public void testNamespace() {
assertTrue(options.build().getNamespace().isEmpty());
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java
deleted file mode 100644
index b86355afa..000000000
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java
+++ /dev/null
@@ -1,84 +0,0 @@
-/*
- * 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
- *
- * 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;
-
-import static com.google.cloud.datastore.TraceUtil.END_SPAN_OPTIONS;
-import static com.google.cloud.datastore.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY;
-import static com.google.common.truth.Truth.assertThat;
-import static com.google.rpc.Code.UNAVAILABLE;
-import static org.easymock.EasyMock.createNiceMock;
-import static org.easymock.EasyMock.createStrictMock;
-import static org.easymock.EasyMock.expect;
-import static org.easymock.EasyMock.replay;
-import static org.easymock.EasyMock.verify;
-
-import com.google.api.gax.retrying.RetrySettings;
-import com.google.cloud.datastore.spi.v1.DatastoreRpc;
-import com.google.datastore.v1.RunAggregationQueryRequest;
-import com.google.datastore.v1.RunAggregationQueryResponse;
-import io.opencensus.trace.Span;
-import io.opencensus.trace.Tracer;
-import org.junit.Before;
-import org.junit.Test;
-
-public class RetryAndTraceDatastoreRpcDecoratorTest {
-
- public static final int MAX_ATTEMPTS = 3;
- private DatastoreRpc mockDatastoreRpc;
- private TraceUtil mockTraceUtil;
- private DatastoreOptions datastoreOptions =
- DatastoreOptions.newBuilder().setProjectId("project-id").build();
- private RetrySettings retrySettings =
- RetrySettings.newBuilder().setMaxAttempts(MAX_ATTEMPTS).build();
-
- private RetryAndTraceDatastoreRpcDecorator datastoreRpcDecorator;
-
- @Before
- public void setUp() throws Exception {
- mockDatastoreRpc = createStrictMock(DatastoreRpc.class);
- mockTraceUtil = createStrictMock(TraceUtil.class);
- datastoreRpcDecorator =
- new RetryAndTraceDatastoreRpcDecorator(
- mockDatastoreRpc, mockTraceUtil, retrySettings, datastoreOptions);
- }
-
- @Test
- public void testRunAggregationQuery() {
- Span mockSpan = createStrictMock(Span.class);
- RunAggregationQueryRequest aggregationQueryRequest =
- RunAggregationQueryRequest.getDefaultInstance();
- RunAggregationQueryResponse aggregationQueryResponse =
- RunAggregationQueryResponse.getDefaultInstance();
-
- expect(mockDatastoreRpc.runAggregationQuery(aggregationQueryRequest))
- .andThrow(
- new DatastoreException(
- UNAVAILABLE.getNumber(), "API not accessible currently", UNAVAILABLE.name()))
- .times(2)
- .andReturn(aggregationQueryResponse);
- expect(mockTraceUtil.startSpan(SPAN_NAME_RUN_AGGREGATION_QUERY)).andReturn(mockSpan);
- expect(mockTraceUtil.getTracer()).andReturn(createNiceMock(Tracer.class));
- mockSpan.end(END_SPAN_OPTIONS);
-
- replay(mockDatastoreRpc, mockTraceUtil, mockSpan);
-
- RunAggregationQueryResponse actualAggregationQueryResponse =
- datastoreRpcDecorator.runAggregationQuery(aggregationQueryRequest);
-
- assertThat(actualAggregationQueryResponse).isSameInstanceAs(aggregationQueryResponse);
- verify(mockDatastoreRpc, mockTraceUtil, mockSpan);
- }
-}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java
new file mode 100644
index 000000000..bee54a1f0
--- /dev/null
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java
@@ -0,0 +1,1009 @@
+/*
+ * 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.it;
+
+import static com.google.cloud.datastore.aggregation.Aggregation.count;
+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_AGGREGATION_QUERY;
+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 static com.google.common.truth.Truth.assertThat;
+import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+import com.google.api.gax.rpc.NotFoundException;
+import com.google.cloud.datastore.AggregationQuery;
+import com.google.cloud.datastore.AggregationResult;
+import com.google.cloud.datastore.AggregationResults;
+import com.google.cloud.datastore.Datastore;
+import com.google.cloud.datastore.DatastoreOptions;
+import com.google.cloud.datastore.Entity;
+import com.google.cloud.datastore.IncompleteKey;
+import com.google.cloud.datastore.Key;
+import com.google.cloud.datastore.KeyFactory;
+import com.google.cloud.datastore.Query;
+import com.google.cloud.datastore.QueryResults;
+import com.google.cloud.datastore.ReadOption;
+import com.google.cloud.datastore.StructuredQuery;
+import com.google.cloud.datastore.StructuredQuery.PropertyFilter;
+import com.google.cloud.datastore.Transaction;
+import com.google.cloud.datastore.testing.RemoteDatastoreHelper;
+import com.google.cloud.opentelemetry.trace.TraceConfiguration;
+import com.google.cloud.opentelemetry.trace.TraceExporter;
+import com.google.cloud.trace.v1.TraceServiceClient;
+import com.google.common.base.Preconditions;
+import com.google.devtools.cloudtrace.v1.Trace;
+import com.google.devtools.cloudtrace.v1.TraceSpan;
+import com.google.testing.junit.testparameterinjector.TestParameter;
+import com.google.testing.junit.testparameterinjector.TestParameterInjector;
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.TraceFlags;
+import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
+import io.opentelemetry.sdk.common.CompletableResultCode;
+import io.opentelemetry.sdk.resources.Resource;
+import io.opentelemetry.sdk.trace.SdkTracerProvider;
+import io.opentelemetry.sdk.trace.export.BatchSpanProcessor;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.TreeMap;
+import java.util.concurrent.TimeUnit;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+// This End-to-End test verifies Client-side Tracing Functionality instrumented using the
+// OpenTelemetry API.
+// The test depends on the following external APIs/Services:
+// 1. Java OpenTelemetry SDK
+// 2. Cloud Trace Exporter
+// 3. TraceServiceClient from Cloud Trace API v1.
+//
+// Permissions required to run this test (https://cloud.google.com/trace/docs/iam#trace-roles):
+// 1. gcloud auth application-default login must be run with the test user.
+// 2. To write traces, test user must have one of roles/cloudtrace.[admin|agent|user] roles.
+// 3. To read traces, test user must have one of roles/cloudtrace.[admin|user] roles.
+//
+// Each test-case has the following workflow:
+// 1. OpenTelemetry SDK is initialized with Cloud Trace Exporter and 100% Trace Sampling
+// 2. On initialization, Datastore client is provided the OpenTelemetry SDK object from (1)
+// 3. A custom TraceID is generated and injected using a custom SpanContext
+// 4. Datastore operations are run inside a root TraceSpan created using the custom SpanContext from
+// (3).
+// 5. Traces are read-back using TraceServiceClient and verified against expected Call Stacks.
+@RunWith(TestParameterInjector.class)
+public class ITE2ETracingTest {
+
+ protected boolean isUsingGlobalOpenTelemetrySDK() {
+ return useGlobalOpenTelemetrySDK;
+ }
+
+ protected String datastoreNamedDatabase() {
+ return datastoreNamedDatabase;
+ }
+
+ // Helper class to track call-stacks in a trace
+ protected static class TraceContainer {
+
+ // Maps Span ID to TraceSpan
+ private final Map idSpanMap;
+
+ // Maps Parent Span ID to a list of Child SpanIDs, useful for top-down traversal
+ private final Map> parentChildIdMap;
+
+ // Tracks the Root Span ID
+ private long rootId;
+
+ public TraceContainer(String rootSpanName, Trace trace) {
+ idSpanMap = new TreeMap<>();
+ parentChildIdMap = new TreeMap<>();
+ for (TraceSpan span : trace.getSpansList()) {
+ long spanId = span.getSpanId();
+ idSpanMap.put(spanId, span);
+ if (rootSpanName.equals(span.getName())) {
+ rootId = span.getSpanId();
+ }
+
+ // Add self as a child of the parent span
+ if (!parentChildIdMap.containsKey(span.getParentSpanId())) {
+ parentChildIdMap.put(span.getParentSpanId(), new ArrayList<>());
+ }
+ parentChildIdMap.get(span.getParentSpanId()).add(spanId);
+ }
+ }
+
+ String spanName(long spanId) {
+ return idSpanMap.get(spanId).getName();
+ }
+
+ List childSpans(long spanId) {
+ return parentChildIdMap.get(spanId);
+ }
+
+ // This method only works for matching call stacks with traces which have children of distinct
+ // type at all levels. This is good enough as the intention is to validate if the e2e path is
+ // WAI - the intention is not to validate Cloud Trace's correctness w.r.t. durability of all
+ // kinds of traces.
+ boolean containsCallStack(String... callStack) throws RuntimeException {
+ List expectedCallStack = Arrays.asList(callStack);
+ if (expectedCallStack.isEmpty()) {
+ throw new RuntimeException("Input callStack is empty");
+ }
+ return dfsContainsCallStack(rootId, expectedCallStack);
+ }
+
+ // Depth-first check for call stack in the trace
+ private boolean dfsContainsCallStack(long spanId, List expectedCallStack) {
+ logger.info(
+ "span="
+ + spanName(spanId)
+ + ", expectedCallStack[0]="
+ + (expectedCallStack.isEmpty() ? "null" : expectedCallStack.get(0)));
+ if (expectedCallStack.isEmpty()) {
+ return false;
+ }
+ if (spanName(spanId).equals(expectedCallStack.get(0))) {
+ // Recursion termination
+ if (childSpans(spanId) == null) {
+ logger.info("No more children for " + spanName(spanId));
+ return expectedCallStack.size() <= 1;
+ } else {
+ // Examine the child spans
+ for (Long childSpan : childSpans(spanId)) {
+ int callStackListSize = expectedCallStack.size();
+ logger.info(
+ "childSpan="
+ + spanName(childSpan)
+ + ", expectedCallStackSize="
+ + callStackListSize);
+ if (dfsContainsCallStack(
+ childSpan,
+ expectedCallStack.subList(
+ /*fromIndexInclusive=*/ 1, /*toIndexExclusive*/ callStackListSize))) {
+ return true;
+ }
+ }
+ }
+ } else {
+ logger.info(spanName(spanId) + " didn't match " + expectedCallStack.get(0));
+ }
+ return false;
+ }
+ }
+
+ private static final Logger logger = Logger.getLogger(ITE2ETracingTest.class.getName());
+
+ private static final String RUN_AGGREGATION_QUERY_RPC_NAME = "RunAggregationQuery";
+
+ private static final String RUN_QUERY_RPC_NAME = "RunQuery";
+
+ private static final int NUM_TRACE_ID_BYTES = 32;
+
+ private static final int NUM_SPAN_ID_BYTES = 16;
+
+ private static final int GET_TRACE_RETRY_COUNT = 60;
+
+ private static final int GET_TRACE_RETRY_BACKOFF_MILLIS = 1000;
+
+ private static final int TRACE_FORCE_FLUSH_MILLIS = 5000;
+
+ private static final int TRACE_PROVIDER_SHUTDOWN_MILLIS = 1000;
+
+ private static Key KEY1;
+
+ private static Key KEY2;
+
+ private static Key KEY3;
+
+ private static Key KEY4;
+
+ // Random int generator for trace ID and span ID
+ private static Random random;
+
+ private static TraceExporter traceExporter;
+
+ // Required for reading back traces from Cloud Trace for validation
+ private static TraceServiceClient traceClient_v1;
+
+ // Custom SpanContext for each test, required for TraceID injection
+ private static SpanContext customSpanContext;
+
+ // Trace read back from Cloud Trace using traceClient_v1 for verification
+ private static Trace retrievedTrace;
+
+ private static String rootSpanName;
+ private static Tracer tracer;
+
+ // Required to set custom-root span
+ private static OpenTelemetrySdk openTelemetrySdk;
+
+ private static String projectId;
+
+ private static DatastoreOptions options;
+
+ private static Datastore datastore;
+
+ private static RemoteDatastoreHelper remoteDatastoreHelper;
+
+ @TestParameter boolean useGlobalOpenTelemetrySDK;
+
+ @TestParameter({
+ /*(default)*/
+ "",
+ "test-db"
+ })
+ String datastoreNamedDatabase;
+
+ @BeforeClass
+ public static void setup() throws IOException {
+ projectId = DatastoreOptions.getDefaultProjectId();
+ traceExporter =
+ TraceExporter.createWithConfiguration(
+ TraceConfiguration.builder().setProjectId(projectId).build());
+ traceClient_v1 = TraceServiceClient.create();
+ random = new Random();
+ }
+
+ @Before
+ public void before() throws Exception {
+ // Set up OTel SDK
+ Resource resource =
+ Resource.getDefault().merge(Resource.builder().put(SERVICE_NAME, "Sparky").build());
+
+ if (isUsingGlobalOpenTelemetrySDK()) {
+ openTelemetrySdk =
+ OpenTelemetrySdk.builder()
+ .setTracerProvider(
+ SdkTracerProvider.builder()
+ .setResource(resource)
+ .addSpanProcessor(BatchSpanProcessor.builder(traceExporter).build())
+ .setSampler(Sampler.alwaysOn())
+ .build())
+ .buildAndRegisterGlobal();
+ } else {
+ openTelemetrySdk =
+ OpenTelemetrySdk.builder()
+ .setTracerProvider(
+ SdkTracerProvider.builder()
+ .setResource(resource)
+ .addSpanProcessor(BatchSpanProcessor.builder(traceExporter).build())
+ .setSampler(Sampler.alwaysOn())
+ .build())
+ .build();
+ }
+
+ // Initialize the Datastore DB w/ the OTel SDK. Ideally we'd do this is the @BeforeAll method
+ // but because gRPC traces need to be deterministically force-flushed for every test
+ String namedDb = datastoreNamedDatabase();
+ logger.log(Level.INFO, "Integration test using named database " + namedDb);
+ remoteDatastoreHelper = RemoteDatastoreHelper.create(namedDb, openTelemetrySdk);
+ options = remoteDatastoreHelper.getOptions();
+ datastore = options.getService();
+
+ Preconditions.checkNotNull(
+ datastore,
+ "Error instantiating Datastore. Check that the service account credentials "
+ + "were properly set.");
+
+ String projectId = options.getProjectId();
+ String kind1 = "kind1";
+ KEY1 =
+ Key.newBuilder(projectId, kind1, "key1", options.getDatabaseId())
+ .setNamespace(options.getNamespace())
+ .build();
+ KEY2 =
+ Key.newBuilder(projectId, kind1, "key2", options.getDatabaseId())
+ .setNamespace(options.getNamespace())
+ .build();
+ KEY3 =
+ Key.newBuilder(projectId, kind1, "key3", options.getDatabaseId())
+ .setNamespace(options.getNamespace())
+ .build();
+ KEY4 =
+ Key.newBuilder(projectId, kind1, "key4", options.getDatabaseId())
+ .setNamespace(options.getNamespace())
+ .build();
+ // Set up the tracer for custom TraceID injection
+ rootSpanName =
+ String.format("%s%d", this.getClass().getSimpleName(), System.currentTimeMillis());
+ if (isUsingGlobalOpenTelemetrySDK()) {
+ tracer = GlobalOpenTelemetry.getTracer(rootSpanName);
+ } else {
+ tracer =
+ datastore
+ .getOptions()
+ .getOpenTelemetryOptions()
+ .getOpenTelemetry()
+ .getTracer(rootSpanName);
+ }
+
+ // Get up a new SpanContext (ergo TraceId) for each test
+ customSpanContext = getNewSpanContext();
+ assertNotNull(customSpanContext);
+ assertNull(retrievedTrace);
+ }
+
+ @After
+ public void after() throws Exception {
+ if (isUsingGlobalOpenTelemetrySDK()) {
+ GlobalOpenTelemetry.resetForTest();
+ }
+ remoteDatastoreHelper.deleteNamespace();
+ rootSpanName = null;
+ tracer = null;
+ retrievedTrace = null;
+ customSpanContext = null;
+ openTelemetrySdk = null;
+ }
+
+ @AfterClass
+ public static void teardown() throws Exception {
+ traceClient_v1.close();
+ }
+
+ // Generates a random hex string of length `numBytes`
+ private String generateRandomHexString(int numBytes) {
+ StringBuilder newTraceId = new StringBuilder();
+ while (newTraceId.length() < numBytes) {
+ newTraceId.append(Integer.toHexString(random.nextInt()));
+ }
+ return newTraceId.substring(0, numBytes);
+ }
+
+ protected String generateNewTraceId() {
+ return generateRandomHexString(NUM_TRACE_ID_BYTES);
+ }
+
+ // Generates a random 16-byte hex string
+ protected String generateNewSpanId() {
+ return generateRandomHexString(NUM_SPAN_ID_BYTES);
+ }
+
+ // Generates a new SpanContext w/ random traceId,spanId
+ protected SpanContext getNewSpanContext() {
+ String traceId = generateNewTraceId();
+ String spanId = generateNewSpanId();
+ logger.info("traceId=" + traceId + ", spanId=" + spanId);
+
+ return SpanContext.create(traceId, spanId, TraceFlags.getSampled(), TraceState.getDefault());
+ }
+
+ protected Span getNewRootSpanWithContext() {
+ // Execute the DB operation in the context of the custom root span.
+ return tracer
+ .spanBuilder(rootSpanName)
+ .setParent(Context.root().with(Span.wrap(customSpanContext)))
+ .startSpan();
+ }
+
+ protected void waitForTracesToComplete() throws Exception {
+ logger.info("Flushing traces...");
+ CompletableResultCode completableResultCode =
+ openTelemetrySdk.getSdkTracerProvider().forceFlush();
+ completableResultCode.join(TRACE_FORCE_FLUSH_MILLIS, TimeUnit.MILLISECONDS);
+ }
+
+ // Validates `retrievedTrace`. Cloud Trace indexes traces w/ eventual consistency, even when
+ // indexing traceId, therefore the test must retry a few times before the complete trace is
+ // available.
+ // For Transaction traces, there may be more spans than in the trace than specified in
+ // `callStack`. So `numExpectedSpans` is the expected total number of spans (and not just the
+ // spans in `callStack`)
+ protected void fetchAndValidateTrace(
+ String traceId, int numExpectedSpans, List> callStackList)
+ throws InterruptedException {
+ // Large enough count to accommodate eventually consistent Cloud Trace backend
+ int numRetries = GET_TRACE_RETRY_COUNT;
+ // Account for rootSpanName
+ numExpectedSpans++;
+
+ // Fetch traces
+ do {
+ try {
+ retrievedTrace = traceClient_v1.getTrace(projectId, traceId);
+ assertEquals(traceId, retrievedTrace.getTraceId());
+
+ logger.info(
+ "expectedSpanCount="
+ + numExpectedSpans
+ + ", retrievedSpanCount="
+ + retrievedTrace.getSpansCount());
+ } catch (NotFoundException notFound) {
+ logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms");
+ } catch (IndexOutOfBoundsException outOfBoundsException) {
+ logger.info("Call stack not found in trace. Retrying.");
+ }
+ if (retrievedTrace == null || numExpectedSpans != retrievedTrace.getSpansCount()) {
+ Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS);
+ }
+ } while (numRetries-- > 0
+ && (retrievedTrace == null || numExpectedSpans != retrievedTrace.getSpansCount()));
+
+ if (retrievedTrace == null || numExpectedSpans != retrievedTrace.getSpansCount()) {
+ throw new RuntimeException(
+ "Expected number of spans: "
+ + numExpectedSpans
+ + ", Actual number of spans: "
+ + (retrievedTrace != null
+ ? retrievedTrace.getSpansList().toString()
+ : "Trace NOT_FOUND"));
+ }
+
+ TraceContainer traceContainer = new TraceContainer(rootSpanName, retrievedTrace);
+
+ for (List callStack : callStackList) {
+ // Update all call stacks to be rooted at rootSpanName
+ ArrayList expectedCallStack = new ArrayList<>(callStack);
+
+ // numExpectedSpans should account for rootSpanName (not passed in callStackList)
+ expectedCallStack.add(0, rootSpanName);
+
+ // *May be* the full trace was returned
+ logger.info("Checking if TraceContainer contains the callStack");
+ String[] expectedCallList = new String[expectedCallStack.size()];
+ if (!traceContainer.containsCallStack(expectedCallStack.toArray(expectedCallList))) {
+ throw new RuntimeException(
+ "Expected spans: "
+ + Arrays.toString(expectedCallList)
+ + ", Actual spans: "
+ + (retrievedTrace != null
+ ? retrievedTrace.getSpansList().toString()
+ : "Trace NOT_FOUND"));
+ }
+ logger.severe("CallStack not found in TraceContainer.");
+ }
+ }
+
+ // Validates `retrievedTrace`. Cloud Trace indexes traces w/ eventual consistency, even when
+ // indexing traceId, therefore the test must retry a few times before the complete trace is
+ // available.
+ // For Non-Transaction traces, there is a 1:1 ratio of spans in `spanNames` and in the trace.
+ protected void fetchAndValidateTrace(String traceId, String... spanNames)
+ throws InterruptedException {
+ fetchAndValidateTrace(traceId, spanNames.length, Arrays.asList(Arrays.asList(spanNames)));
+ }
+
+ @Test
+ public void traceContainerTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ Entity entity = datastore.get(KEY1);
+ assertNull(entity);
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ Trace traceResp = null;
+ int expectedSpanCount = 2;
+
+ int numRetries = GET_TRACE_RETRY_COUNT;
+ do {
+ try {
+ traceResp = traceClient_v1.getTrace(projectId, customSpanContext.getTraceId());
+ if (traceResp.getSpansCount() == expectedSpanCount) {
+ logger.info("Success: Got " + expectedSpanCount + " spans.");
+ break;
+ }
+ } catch (NotFoundException notFoundException) {
+ Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS);
+ logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms");
+ }
+ logger.info(
+ "Trace Found. The trace did not contain "
+ + expectedSpanCount
+ + " spans. Going to retry.");
+ numRetries--;
+ } while (numRetries > 0);
+
+ // Make sure we got as many spans as we expected.
+ assertNotNull(traceResp);
+ assertEquals(expectedSpanCount, traceResp.getSpansCount());
+
+ TraceContainer traceCont = new TraceContainer(rootSpanName, traceResp);
+
+ // Contains exact path
+ assertTrue(traceCont.containsCallStack(rootSpanName, SPAN_NAME_LOOKUP));
+
+ // Top-level mismatch
+ assertFalse(traceCont.containsCallStack(SPAN_NAME_LOOKUP, RUN_QUERY_RPC_NAME));
+
+ // Leaf-level mismatch/missing
+ assertFalse(
+ traceCont.containsCallStack(
+ rootSpanName, SPAN_NAME_LOOKUP, RUN_AGGREGATION_QUERY_RPC_NAME));
+ }
+
+ @Test
+ public void lookupTraceTest() throws Exception {
+ // Make sure the test has a new SpanContext (and TraceId for injection)
+ assertNotNull(customSpanContext);
+
+ // Inject new trace ID
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ Entity entity = datastore.get(KEY1);
+ assertNull(entity);
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_LOOKUP);
+ }
+
+ @Test
+ public void allocateIdsTraceTest() throws Exception {
+ assertNotNull(customSpanContext);
+
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ String kind1 = "kind1";
+ KeyFactory keyFactory = datastore.newKeyFactory().setKind(kind1);
+ IncompleteKey pk1 = keyFactory.newKey();
+ Key key1 = datastore.allocateId(pk1);
+ assertEquals(key1.getProjectId(), pk1.getProjectId());
+ assertEquals(key1.getNamespace(), pk1.getNamespace());
+ assertEquals(key1.getAncestors(), pk1.getAncestors());
+ assertEquals(key1.getKind(), pk1.getKind());
+ assertTrue(key1.hasId());
+ assertFalse(key1.hasName());
+ assertEquals(Key.newBuilder(pk1, key1.getId()).build(), key1);
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_ALLOCATE_IDS);
+ }
+
+ @Test
+ public void reserveIdsTraceTest() throws Exception {
+ assertNotNull(customSpanContext);
+
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
+ Key key1 = keyFactory.newKey(10);
+ Key key2 = keyFactory.newKey("name");
+ List keyList = datastore.reserveIds(key1, key2);
+ assertEquals(2, keyList.size());
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_RESERVE_IDS);
+ }
+
+ @Test
+ public void commitTraceTest() throws Exception {
+ assertNotNull(customSpanContext);
+
+ Span rootSpan = getNewRootSpanWithContext();
+
+ Entity entity1 = Entity.newBuilder(KEY1).set("test_key", "test_value").build();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ Entity response = datastore.add(entity1);
+ assertEquals(entity1, response);
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_COMMIT);
+ }
+
+ @Test
+ public void putTraceTest() throws Exception {
+ assertNotNull(customSpanContext);
+
+ Span rootSpan = getNewRootSpanWithContext();
+
+ Entity entity1 = Entity.newBuilder(KEY1).set("test_key", "test_value").build();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ Entity response = datastore.put(entity1);
+ assertEquals(entity1, response);
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_COMMIT);
+ }
+
+ @Test
+ public void updateTraceTest() throws Exception {
+ assertNotNull(customSpanContext);
+
+ Entity entity1 = Entity.newBuilder(KEY1).set("test_field", "test_value1").build();
+ Entity entity2 = Entity.newBuilder(KEY2).set("test_field", "test_value2").build();
+ List entityList = new ArrayList<>();
+ entityList.add(entity1);
+ entityList.add(entity2);
+
+ List response = datastore.add(entity1, entity2);
+ assertEquals(entityList, response);
+
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ Entity entity1_update =
+ Entity.newBuilder(entity1).set("test_field", "new_test_value1").build();
+ Entity entity2_update =
+ Entity.newBuilder(entity2).set("test_field", "new_test_value1").build();
+ datastore.update(entity1_update, entity2_update);
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_COMMIT);
+ }
+
+ @Test
+ public void deleteTraceTest() throws Exception {
+ assertNotNull(customSpanContext);
+
+ Entity entity1 = Entity.newBuilder(KEY1).set("test_key", "test_value").build();
+ Entity response = datastore.put(entity1);
+ assertEquals(entity1, response);
+
+ Span rootSpan = getNewRootSpanWithContext();
+
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ datastore.delete(entity1.getKey());
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+ fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_COMMIT);
+ }
+
+ @Test
+ public void runQueryTraceTest() throws Exception {
+ Entity entity1 = Entity.newBuilder(KEY1).set("test_field", "test_value1").build();
+ Entity entity2 = Entity.newBuilder(KEY2).set("test_field", "test_value2").build();
+ List entityList = new ArrayList<>();
+ entityList.add(entity1);
+ entityList.add(entity2);
+
+ List response = datastore.add(entity1, entity2);
+ assertEquals(entityList, response);
+
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ PropertyFilter filter = PropertyFilter.eq("test_field", entity1.getValue("test_field"));
+ Query query =
+ Query.newEntityQueryBuilder().setKind(KEY1.getKind()).setFilter(filter).build();
+ QueryResults queryResults = datastore.run(query);
+ assertTrue(queryResults.hasNext());
+ assertEquals(entity1, queryResults.next());
+ assertFalse(queryResults.hasNext());
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_RUN_QUERY);
+ }
+
+ @Test
+ public void runAggregationQueryTraceTest() throws Exception {
+ Entity entity1 =
+ Entity.newBuilder(KEY1)
+ .set("pepper_name", "jalapeno")
+ .set("max_scoville_level", 10000)
+ .build();
+ Entity entity2 =
+ Entity.newBuilder(KEY2)
+ .set("pepper_name", "serrano")
+ .set("max_scoville_level", 25000)
+ .build();
+ Entity entity3 =
+ Entity.newBuilder(KEY3)
+ .set("pepper_name", "habanero")
+ .set("max_scoville_level", 350000)
+ .build();
+ Entity entity4 =
+ Entity.newBuilder(KEY4)
+ .set("pepper_name", "ghost")
+ .set("max_scoville_level", 1500000)
+ .build();
+
+ List entityList = new ArrayList<>();
+ entityList.add(entity1);
+ entityList.add(entity2);
+ entityList.add(entity3);
+ entityList.add(entity4);
+
+ List response = datastore.add(entity1, entity2, entity3, entity4);
+ assertEquals(entityList, response);
+
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ PropertyFilter mediumSpicyFilters = PropertyFilter.lt("max_scoville_level", 100000);
+ StructuredQuery mediumSpicyQuery =
+ Query.newEntityQueryBuilder()
+ .setKind(KEY1.getKind())
+ .setFilter(mediumSpicyFilters)
+ .build();
+ AggregationQuery countSpicyPeppers =
+ Query.newAggregationQueryBuilder()
+ .addAggregation(count().as("count"))
+ .over(mediumSpicyQuery)
+ .build();
+ AggregationResults results = datastore.runAggregation(countSpicyPeppers);
+ assertThat(results.size()).isEqualTo(1);
+ AggregationResult result = results.get(0);
+ assertThat(result.getLong("count")).isEqualTo(2L);
+ } finally {
+ rootSpan.end();
+ }
+
+ waitForTracesToComplete();
+
+ fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_RUN_AGGREGATION_QUERY);
+ }
+
+ @Test
+ public void newTransactionReadTest() throws Exception {
+ assertNotNull(customSpanContext);
+
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ Transaction transaction = datastore.newTransaction();
+ Entity entity = datastore.get(KEY1, ReadOption.transactionId(transaction.getTransactionId()));
+ transaction.commit();
+ assertNull(entity);
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ /*numExpectedSpans=*/ 3,
+ Arrays.asList(
+ Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION),
+ Collections.singletonList(SPAN_NAME_TRANSACTION_LOOKUP),
+ Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT)));
+ }
+
+ @Test
+ public void newTransactionQueryTest() throws Exception {
+ // Set up
+ Entity entity1 = Entity.newBuilder(KEY1).set("test_field", "test_value1").build();
+ Entity entity2 = Entity.newBuilder(KEY2).set("test_field", "test_value2").build();
+ List entityList = new ArrayList<>();
+ entityList.add(entity1);
+ entityList.add(entity2);
+
+ List response = datastore.add(entity1, entity2);
+ assertEquals(entityList, response);
+
+ assertNotNull(customSpanContext);
+
+ // Test
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ Transaction transaction = datastore.newTransaction();
+ PropertyFilter filter = PropertyFilter.eq("test_field", entity1.getValue("test_field"));
+ Query query =
+ Query.newEntityQueryBuilder().setKind(KEY1.getKind()).setFilter(filter).build();
+ QueryResults queryResults = transaction.run(query);
+ transaction.commit();
+ assertTrue(queryResults.hasNext());
+ assertEquals(entity1, queryResults.next());
+ assertFalse(queryResults.hasNext());
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ /*numExpectedSpans=*/ 3,
+ Arrays.asList(
+ Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION),
+ Collections.singletonList(SPAN_NAME_TRANSACTION_RUN_QUERY),
+ Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT)));
+ }
+
+ @Test
+ public void newTransactionReadWriteTraceTest() throws Exception {
+ // Set up
+ Entity entity1 = Entity.newBuilder(KEY1).set("pepper_type", "jalapeno").build();
+ Entity entity2 = Entity.newBuilder(KEY2).set("pepper_type", "habanero").build();
+ List entityList = new ArrayList<>();
+ entityList.add(entity1);
+ entityList.add(entity2);
+
+ List response = datastore.add(entity1, entity2);
+ assertEquals(entityList, response);
+
+ String simplified_spice_level = "not_spicy";
+ Entity entity1update =
+ Entity.newBuilder(entity1).set("spice_level", simplified_spice_level).build();
+
+ assertNotNull(customSpanContext);
+
+ // Test
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ Transaction transaction = datastore.newTransaction();
+ entity1 = transaction.get(KEY1);
+ switch (entity1.getString("pepper_type")) {
+ case "jalapeno":
+ simplified_spice_level = "mild";
+ break;
+
+ case "habanero":
+ simplified_spice_level = "hot";
+ break;
+ }
+ transaction.update(entity1update);
+ transaction.delete(KEY2);
+ transaction.commit();
+ assertFalse(transaction.isActive());
+ } finally {
+ rootSpan.end();
+ }
+
+ waitForTracesToComplete();
+
+ List list = datastore.fetch(KEY1, KEY2);
+ assertEquals(list.get(0), entity1update);
+ assertNull(list.get(1));
+
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ /*numExpectedSpans=*/ 3,
+ Arrays.asList(
+ Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION),
+ Collections.singletonList(SPAN_NAME_TRANSACTION_LOOKUP),
+ Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT)));
+ }
+
+ @Test
+ public void newTransactionRollbackTest() throws Exception {
+ // Set up
+ Entity entity1 = Entity.newBuilder(KEY1).set("pepper_type", "jalapeno").build();
+ Entity entity2 = Entity.newBuilder(KEY2).set("pepper_type", "habanero").build();
+ List entityList = new ArrayList<>();
+ entityList.add(entity1);
+ entityList.add(entity2);
+
+ List response = datastore.add(entity1, entity2);
+ assertEquals(entityList, response);
+
+ String simplified_spice_level = "not_spicy";
+ Entity entity1update =
+ Entity.newBuilder(entity1).set("spice_level", simplified_spice_level).build();
+
+ assertNotNull(customSpanContext);
+
+ // Test
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ Transaction transaction = datastore.newTransaction();
+ entity1 = transaction.get(KEY1);
+ switch (entity1.getString("pepper_type")) {
+ case "jalapeno":
+ simplified_spice_level = "mild";
+ break;
+
+ case "habanero":
+ simplified_spice_level = "hot";
+ break;
+ }
+ transaction.update(entity1update);
+ transaction.delete(KEY2);
+ transaction.rollback();
+ assertFalse(transaction.isActive());
+ } finally {
+ rootSpan.end();
+ }
+
+ waitForTracesToComplete();
+
+ List list = datastore.fetch(KEY1, KEY2);
+ assertEquals(list.get(0), entity1);
+ assertEquals(list.get(1), entity2);
+
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ /*numExpectedSpans=*/ 3,
+ Arrays.asList(
+ Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION),
+ Collections.singletonList(SPAN_NAME_TRANSACTION_LOOKUP),
+ Collections.singletonList(SPAN_NAME_ROLLBACK)));
+ }
+
+ @Test
+ public void runInTransactionQueryTest() throws Exception {
+ // Set up
+ Entity entity1 = Entity.newBuilder(KEY1).set("test_field", "test_value1").build();
+ Entity entity2 = Entity.newBuilder(KEY2).set("test_field", "test_value2").build();
+ List entityList = new ArrayList<>();
+ entityList.add(entity1);
+ entityList.add(entity2);
+
+ List response = datastore.add(entity1, entity2);
+ assertEquals(entityList, response);
+
+ assertNotNull(customSpanContext);
+
+ Span rootSpan = getNewRootSpanWithContext();
+ try (Scope ignored = rootSpan.makeCurrent()) {
+ PropertyFilter filter = PropertyFilter.eq("test_field", entity1.getValue("test_field"));
+ Query query =
+ Query.newEntityQueryBuilder().setKind(KEY1.getKind()).setFilter(filter).build();
+ Datastore.TransactionCallable callable =
+ transaction -> {
+ QueryResults queryResults = datastore.run(query);
+ assertTrue(queryResults.hasNext());
+ assertEquals(entity1, queryResults.next());
+ assertFalse(queryResults.hasNext());
+ return true;
+ };
+ datastore.runInTransaction(callable);
+ } finally {
+ rootSpan.end();
+ }
+ waitForTracesToComplete();
+
+ fetchAndValidateTrace(
+ customSpanContext.getTraceId(),
+ /*numExpectedSpans=*/ 4,
+ Arrays.asList(
+ Arrays.asList(SPAN_NAME_TRANSACTION_RUN, SPAN_NAME_BEGIN_TRANSACTION),
+ Arrays.asList(SPAN_NAME_TRANSACTION_RUN, SPAN_NAME_RUN_QUERY),
+ Arrays.asList(SPAN_NAME_TRANSACTION_RUN, SPAN_NAME_TRANSACTION_COMMIT)));
+ }
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITTracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITTracingTest.java
new file mode 100644
index 000000000..85ff4758b
--- /dev/null
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITTracingTest.java
@@ -0,0 +1,849 @@
+/*
+ * 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.it;
+
+import static com.google.cloud.datastore.aggregation.Aggregation.count;
+import static com.google.cloud.datastore.telemetry.TraceUtil.*;
+import static com.google.common.truth.Truth.assertThat;
+import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+import com.google.cloud.datastore.AggregationQuery;
+import com.google.cloud.datastore.AggregationResult;
+import com.google.cloud.datastore.AggregationResults;
+import com.google.cloud.datastore.Datastore;
+import com.google.cloud.datastore.DatastoreOpenTelemetryOptions;
+import com.google.cloud.datastore.DatastoreOptions;
+import com.google.cloud.datastore.Entity;
+import com.google.cloud.datastore.IncompleteKey;
+import com.google.cloud.datastore.Key;
+import com.google.cloud.datastore.KeyFactory;
+import com.google.cloud.datastore.Query;
+import com.google.cloud.datastore.QueryResults;
+import com.google.cloud.datastore.ReadOption;
+import com.google.cloud.datastore.StructuredQuery;
+import com.google.cloud.datastore.StructuredQuery.PropertyFilter;
+import com.google.cloud.datastore.Transaction;
+import com.google.cloud.datastore.testing.RemoteDatastoreHelper;
+import com.google.common.base.Preconditions;
+import com.google.testing.junit.testparameterinjector.TestParameter;
+import com.google.testing.junit.testparameterinjector.TestParameterInjector;
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
+import io.opentelemetry.sdk.OpenTelemetrySdkBuilder;
+import io.opentelemetry.sdk.common.CompletableResultCode;
+import io.opentelemetry.sdk.resources.Resource;
+import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter;
+import io.opentelemetry.sdk.trace.SdkTracerProvider;
+import io.opentelemetry.sdk.trace.SpanProcessor;
+import io.opentelemetry.sdk.trace.data.EventData;
+import io.opentelemetry.sdk.trace.data.SpanData;
+import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import javax.annotation.Nullable;
+import org.junit.After;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TestName;
+import org.junit.runner.RunWith;
+
+@RunWith(TestParameterInjector.class)
+public class ITTracingTest {
+ protected boolean isUsingGlobalOpenTelemetrySDK() {
+ return useGlobalOpenTelemetrySDK;
+ }
+
+ protected String datastoreNamedDatabase() {
+ return datastoreNamedDatabase;
+ }
+
+ private static final Logger logger =
+ Logger.getLogger(com.google.cloud.datastore.it.ITTracingTest.class.getName());
+
+ private static final int TRACE_FORCE_FLUSH_MILLIS = 1000;
+ private static final int TRACE_PROVIDER_SHUTDOWN_MILLIS = 1000;
+ private static final int IN_MEMORY_SPAN_EXPORTER_DELAY_MILLIS = 50;
+ private static final String SERVICE = "google.datastore.v1.Datastore/";
+
+ private static Key KEY1;
+
+ private static Key KEY2;
+
+ private static Key KEY3;
+
+ private static Key KEY4;
+
+ private static OpenTelemetrySdk openTelemetrySdk;
+
+ // We use an InMemorySpanExporter for testing which keeps all generated trace spans
+ // in memory so that we can check their correctness.
+ protected InMemorySpanExporter inMemorySpanExporter;
+ private static DatastoreOptions options;
+
+ protected Datastore datastore;
+ private static RemoteDatastoreHelper remoteDatastoreHelper;
+
+ @TestParameter boolean useGlobalOpenTelemetrySDK;
+
+ @TestParameter({
+ /*(default)*/
+ "",
+ "test-db"
+ })
+ String datastoreNamedDatabase;
+
+ Map spanNameToSpanId = new HashMap<>();
+ Map spanIdToParentSpanId = new HashMap<>();
+ Map spanNameToSpanData = new HashMap<>();
+
+ @Rule public TestName testName = new TestName();
+
+ @Before
+ public void before() {
+ inMemorySpanExporter = InMemorySpanExporter.create();
+
+ Resource resource =
+ Resource.getDefault().merge(Resource.builder().put(SERVICE_NAME, "Sparky").build());
+ SpanProcessor inMemorySpanProcessor = SimpleSpanProcessor.create(inMemorySpanExporter);
+ DatastoreOptions.Builder optionsBuilder = DatastoreOptions.newBuilder();
+ DatastoreOpenTelemetryOptions.Builder otelOptionsBuilder =
+ DatastoreOpenTelemetryOptions.newBuilder();
+ OpenTelemetrySdkBuilder openTelemetrySdkBuilder =
+ OpenTelemetrySdk.builder()
+ .setTracerProvider(
+ SdkTracerProvider.builder()
+ .setResource(resource)
+ .addSpanProcessor(inMemorySpanProcessor)
+ .setSampler(Sampler.alwaysOn())
+ .build());
+
+ if (isUsingGlobalOpenTelemetrySDK()) {
+ GlobalOpenTelemetry.resetForTest();
+ openTelemetrySdk = openTelemetrySdkBuilder.buildAndRegisterGlobal();
+ optionsBuilder.setOpenTelemetryOptions(otelOptionsBuilder.setTracingEnabled(true).build());
+ } else {
+ openTelemetrySdk = openTelemetrySdkBuilder.build();
+ optionsBuilder.setOpenTelemetryOptions(
+ otelOptionsBuilder.setTracingEnabled(true).setOpenTelemetry(openTelemetrySdk).build());
+ }
+
+ String namedDb = datastoreNamedDatabase();
+ logger.log(Level.INFO, "Integration test using named database " + namedDb);
+ remoteDatastoreHelper = RemoteDatastoreHelper.create(namedDb, openTelemetrySdk);
+ options = remoteDatastoreHelper.getOptions();
+ datastore = options.getService();
+
+ Preconditions.checkNotNull(
+ datastore,
+ "Error instantiating Datastore. Check that the service account credentials "
+ + "were properly set.");
+
+ String projectId = options.getProjectId();
+ String kind1 = "kind1";
+ KEY1 =
+ Key.newBuilder(projectId, kind1, "key1", options.getDatabaseId())
+ .setNamespace(options.getNamespace())
+ .build();
+ KEY2 =
+ Key.newBuilder(projectId, kind1, "key2", options.getDatabaseId())
+ .setNamespace(options.getNamespace())
+ .build();
+ KEY3 =
+ Key.newBuilder(projectId, kind1, "key3", options.getDatabaseId())
+ .setNamespace(options.getNamespace())
+ .build();
+ KEY4 =
+ Key.newBuilder(projectId, kind1, "key4", options.getDatabaseId())
+ .setNamespace(options.getNamespace())
+ .build();
+ cleanupTestSpanContext();
+ }
+
+ @After
+ public void after() throws Exception {
+ if (isUsingGlobalOpenTelemetrySDK()) {
+ GlobalOpenTelemetry.resetForTest();
+ }
+ remoteDatastoreHelper.deleteNamespace();
+ inMemorySpanExporter.reset();
+ CompletableResultCode completableResultCode =
+ openTelemetrySdk.getSdkTracerProvider().shutdown();
+ completableResultCode.join(TRACE_PROVIDER_SHUTDOWN_MILLIS, TimeUnit.MILLISECONDS);
+ openTelemetrySdk = null;
+ }
+
+ @AfterClass
+ public static void teardown() {}
+
+ void waitForTracesToComplete() throws Exception {
+ // The same way that querying the Cloud Trace backend may not give us the
+ // full trace on the first try, querying the in-memory traces may not result
+ // in the full trace immediately. Note that performing the `flush` is not
+ // enough. This doesn't pose an issue in practice, but can make tests flaky.
+ // Therefore, we're adding a delay to make sure we avoid any flakiness.
+ inMemorySpanExporter.flush().join(IN_MEMORY_SPAN_EXPORTER_DELAY_MILLIS, TimeUnit.MILLISECONDS);
+ TimeUnit.MILLISECONDS.sleep(IN_MEMORY_SPAN_EXPORTER_DELAY_MILLIS);
+
+ CompletableResultCode completableResultCode =
+ openTelemetrySdk.getSdkTracerProvider().forceFlush();
+ completableResultCode.join(TRACE_FORCE_FLUSH_MILLIS, TimeUnit.MILLISECONDS);
+ }
+
+ // Prepares all the spans in memory for inspection.
+ List prepareSpans() throws Exception {
+ waitForTracesToComplete();
+ List spans = inMemorySpanExporter.getFinishedSpanItems();
+ buildSpanMaps(spans);
+ printSpans();
+ return spans;
+ }
+
+ void buildSpanMaps(List spans) {
+ for (SpanData spanData : spans) {
+ spanNameToSpanData.put(spanData.getName(), spanData);
+ spanNameToSpanId.put(spanData.getName(), spanData.getSpanId());
+ spanIdToParentSpanId.put(spanData.getSpanId(), spanData.getParentSpanId());
+ }
+ }
+
+ // Returns the SpanData object for the span with the given name.
+ // Returns null if no span with the given name exists.
+ @Nullable
+ SpanData getSpanByName(String spanName) {
+ return spanNameToSpanData.get(spanName);
+ }
+
+ // Returns the SpanData object for the gRPC span with the given RPC name.
+ // Returns null if no such span exists.
+ @Nullable
+ SpanData getGrpcSpanByName(String rpcName) {
+ return getSpanByName(SERVICE + rpcName);
+ }
+
+ String grpcSpanName(String rpcName) {
+ return SERVICE + rpcName;
+ }
+
+ void assertSameTrace(SpanData... spans) {
+ if (spans.length > 1) {
+ String traceId = spans[0].getTraceId();
+ for (SpanData spanData : spans) {
+ assertEquals(traceId, spanData.getTraceId());
+ }
+ }
+ }
+
+ // Helper to see the spans in standard output while developing tests
+ void printSpans() {
+ for (SpanData spanData : spanNameToSpanData.values()) {
+ logger.log(
+ Level.FINE,
+ String.format(
+ "SPAN ID:%s, ParentID:%s, KIND:%s, TRACE ID:%s, NAME:%s, ATTRIBUTES:%s, EVENTS:%s\n",
+ spanData.getSpanId(),
+ spanData.getParentSpanId(),
+ spanData.getKind(),
+ spanData.getTraceId(),
+ spanData.getName(),
+ spanData.getAttributes().toString(),
+ spanData.getEvents().toString()));
+ }
+ }
+
+ // Asserts that the span hierarchy exists for the given span names. The hierarchy starts with the
+ // root span, followed
+ // by the child span, grandchild span, and so on. It also asserts that all the given spans belong
+ // to the same trace,
+ // and that datastore-generated spans contain the expected datastore attributes.
+ void assertSpanHierarchy(String... spanNamesHierarchy) {
+ List spanNames = Arrays.asList(spanNamesHierarchy);
+
+ for (int i = 0; i + 1 < spanNames.size(); ++i) {
+ String parentSpanName = spanNames.get(i);
+ String childSpanName = spanNames.get(i + 1);
+ SpanData parentSpan = getSpanByName(parentSpanName);
+ SpanData childSpan = getSpanByName(childSpanName);
+ assertNotNull(parentSpan);
+ assertNotNull(childSpan);
+ assertEquals(childSpan.getParentSpanId(), parentSpan.getSpanId());
+ assertSameTrace(childSpan, parentSpan);
+ // gRPC spans do not have datastore attributes.
+ if (!parentSpanName.startsWith(SERVICE)) {
+ assertHasExpectedAttributes(parentSpan);
+ }
+ if (!childSpanName.startsWith(SERVICE)) {
+ assertHasExpectedAttributes(childSpan);
+ }
+ }
+ }
+
+ void assertHasExpectedAttributes(SpanData spanData, String... additionalExpectedAttributes) {
+ // All datastore-generated spans have the settings attributes.
+ List expectedAttributes =
+ Arrays.asList(
+ "gcp.datastore.memoryUtilization",
+ "gcp.datastore.settings.host",
+ "gcp.datastore.settings.databaseId",
+ "gcp.datastore.settings.retrySettings.maxRpcTimeout",
+ "gcp.datastore.settings.retrySettings.retryDelayMultiplier",
+ "gcp.datastore.settings.retrySettings.initialRetryDelay",
+ "gcp.datastore.settings.credentials.authenticationType",
+ "gcp.datastore.settings.retrySettings.maxAttempts",
+ "gcp.datastore.settings.retrySettings.maxRetryDelay",
+ "gcp.datastore.settings.retrySettings.rpcTimeoutMultiplier",
+ "gcp.datastore.settings.retrySettings.totalTimeout",
+ "gcp.datastore.settings.retrySettings.initialRpcTimeout");
+
+ expectedAttributes.addAll(Arrays.asList(additionalExpectedAttributes));
+
+ Attributes spanAttributes = spanData.getAttributes();
+ for (String expectedAttribute : expectedAttributes) {
+ assertNotNull(spanAttributes.get(AttributeKey.stringKey(expectedAttribute)));
+ }
+ }
+
+ // Returns true if and only if the given span data contains an event with the given name and the
+ // given expected
+ // attributes.
+ boolean hasEvent(SpanData spanData, String eventName, @Nullable Attributes expectedAttributes) {
+ if (spanData == null) {
+ return false;
+ }
+
+ logger.log(
+ Level.INFO,
+ String.format(
+ "Checking if span named '%s' (ID='%s') contains an event named '%s'",
+ spanData.getName(), spanData.getSpanId(), eventName));
+
+ List events = spanData.getEvents();
+ for (EventData event : events) {
+ if (event.getName().equals(eventName)) {
+ if (expectedAttributes == null) {
+ return true;
+ }
+
+ // Make sure attributes also match.
+ Attributes eventAttributes = event.getAttributes();
+ return expectedAttributes.equals(eventAttributes);
+ }
+ }
+ return false;
+ }
+
+ void cleanupTestSpanContext() {
+ inMemorySpanExporter.reset();
+ spanNameToSpanId.clear();
+ spanIdToParentSpanId.clear();
+ spanNameToSpanData.clear();
+ }
+
+ // This is a POJO used for testing APIs that take a POJO.
+ public static class Pojo {
+ public int bar;
+
+ public Pojo() {
+ bar = 0;
+ }
+
+ public Pojo(int bar) {
+ this.bar = bar;
+ }
+
+ public int getBar() {
+ return bar;
+ }
+
+ public void setBar(int bar) {
+ this.bar = bar;
+ }
+ }
+
+ @Test
+ public void lookupTraceTest() throws Exception {
+ Entity entity = datastore.get(KEY1);
+ assertNull(entity);
+
+ waitForTracesToComplete();
+
+ List spans = prepareSpans();
+ assertEquals(1, spans.size());
+ assertSpanHierarchy(SPAN_NAME_LOOKUP);
+ SpanData span = getSpanByName(SPAN_NAME_LOOKUP);
+ assertTrue(
+ hasEvent(
+ span,
+ SPAN_NAME_LOOKUP + " complete.",
+ Attributes.builder()
+ .put("Received", 0)
+ .put("Missing", 1)
+ .put("Deferred", 0)
+ .put("transactional", false)
+ .put("transaction_id", "")
+ .build()));
+ }
+
+ @Test
+ public void allocateIdsTraceTest() throws Exception {
+ String kind1 = "kind1";
+ KeyFactory keyFactory = datastore.newKeyFactory().setKind(kind1);
+ IncompleteKey pk1 = keyFactory.newKey();
+ Key key1 = datastore.allocateId(pk1);
+
+ waitForTracesToComplete();
+
+ List spans = prepareSpans();
+ assertEquals(1, spans.size());
+ assertSpanHierarchy(SPAN_NAME_ALLOCATE_IDS);
+ }
+
+ @Test
+ public void reserveIdsTraceTest() throws Exception {
+ KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
+ Key key1 = keyFactory.newKey(10);
+ Key key2 = keyFactory.newKey("name");
+ List keyList = datastore.reserveIds(key1, key2);
+ assertEquals(2, keyList.size());
+
+ waitForTracesToComplete();
+
+ List spans = prepareSpans();
+ assertEquals(1, spans.size());
+ assertSpanHierarchy(SPAN_NAME_RESERVE_IDS);
+ }
+
+ @Test
+ public void commitTraceTest() throws Exception {
+ Entity entity1 = Entity.newBuilder(KEY1).set("test_key", "test_value").build();
+ Entity response = datastore.add(entity1);
+ assertEquals(entity1, response);
+
+ waitForTracesToComplete();
+
+ List spans = prepareSpans();
+ assertEquals(1, spans.size());
+ assertSpanHierarchy(SPAN_NAME_COMMIT);
+ }
+
+ @Test
+ public void putTraceTest() throws Exception {
+ Entity entity1 = Entity.newBuilder(KEY1).set("test_key", "test_value").build();
+ Entity response = datastore.put(entity1);
+ assertEquals(entity1, response);
+
+ waitForTracesToComplete();
+
+ List spans = prepareSpans();
+ assertEquals(1, spans.size());
+ assertSpanHierarchy(SPAN_NAME_COMMIT);
+ }
+
+ @Test
+ public void updateTraceTest() throws Exception {
+ Entity entity1 = Entity.newBuilder(KEY1).set("test_field", "test_value1").build();
+ Entity entity2 = Entity.newBuilder(KEY2).set("test_field", "test_value2").build();
+ List entityList = new ArrayList<>();
+ entityList.add(entity1);
+ entityList.add(entity2);
+
+ List response = datastore.add(entity1, entity2);
+ assertEquals(entityList, response);
+
+ List spans = prepareSpans();
+ assertEquals(1, spans.size());
+ assertSpanHierarchy(SPAN_NAME_COMMIT);
+
+ SpanData spanData = getSpanByName(SPAN_NAME_COMMIT);
+ assertTrue(
+ hasEvent(
+ spanData,
+ SPAN_NAME_COMMIT + " complete.",
+ Attributes.builder()
+ .put("doc_count", response.size())
+ .put("transactional", false)
+ .put("transaction_id", "")
+ .build()));
+
+ // Clean Up test span context to verify update spans
+ cleanupTestSpanContext();
+
+ Entity entity1_update = Entity.newBuilder(entity1).set("test_field", "new_test_value1").build();
+ Entity entity2_update = Entity.newBuilder(entity2).set("test_field", "new_test_value1").build();
+ datastore.update(entity1_update, entity2_update);
+
+ waitForTracesToComplete();
+
+ spans = prepareSpans();
+ assertEquals(1, spans.size());
+ assertSpanHierarchy(SPAN_NAME_COMMIT);
+ }
+
+ @Test
+ public void deleteTraceTest() throws Exception {
+ Entity entity1 = Entity.newBuilder(KEY1).set("test_key", "test_value").build();
+ Entity response = datastore.put(entity1);
+ assertEquals(entity1, response);
+
+ List spans = prepareSpans();
+ assertEquals(1, spans.size());
+ assertSpanHierarchy(SPAN_NAME_COMMIT);
+
+ SpanData spanData = getSpanByName(SPAN_NAME_COMMIT);
+ assertTrue(
+ hasEvent(
+ spanData,
+ SPAN_NAME_COMMIT + " complete.",
+ Attributes.builder()
+ .put("doc_count", 1)
+ .put("transactional", false)
+ .put("transaction_id", "")
+ .build()));
+
+ // Clean Up test span context to verify update spans
+ cleanupTestSpanContext();
+
+ datastore.delete(entity1.getKey());
+
+ waitForTracesToComplete();
+
+ spans = prepareSpans();
+ assertEquals(1, spans.size());
+ assertSpanHierarchy(SPAN_NAME_COMMIT);
+
+ spanData = getSpanByName(SPAN_NAME_COMMIT);
+ assertTrue(
+ hasEvent(
+ spanData,
+ SPAN_NAME_COMMIT + " complete.",
+ Attributes.builder()
+ .put("doc_count", 1)
+ .put("transactional", false)
+ .put("transaction_id", "")
+ .build()));
+ }
+
+ @Test
+ public void runQueryTraceTest() throws Exception {
+ Entity entity1 = Entity.newBuilder(KEY1).set("test_field", "test_value1").build();
+ Entity entity2 = Entity.newBuilder(KEY2).set("test_field", "test_value2").build();
+ List entityList = new ArrayList<>();
+ entityList.add(entity1);
+ entityList.add(entity2);
+
+ List response = datastore.add(entity1, entity2);
+ assertEquals(entityList, response);
+
+ // Clean Up test span context to verify RunQuery spans
+ cleanupTestSpanContext();
+
+ PropertyFilter filter = PropertyFilter.eq("test_field", entity1.getValue("test_field"));
+ Query query =
+ Query.newEntityQueryBuilder().setKind(KEY1.getKind()).setFilter(filter).build();
+ QueryResults queryResults = datastore.run(query);
+ assertTrue(queryResults.hasNext());
+ assertEquals(entity1, queryResults.next());
+ assertFalse(queryResults.hasNext());
+
+ waitForTracesToComplete();
+
+ List spans = prepareSpans();
+ assertEquals(1, spans.size());
+ assertSpanHierarchy(SPAN_NAME_RUN_QUERY);
+
+ SpanData span = getSpanByName(SPAN_NAME_RUN_QUERY);
+ assertTrue(
+ hasEvent(
+ span,
+ SPAN_NAME_RUN_QUERY + " complete.",
+ Attributes.builder()
+ .put("doc_count", 1)
+ .put("transactional", false)
+ .put("read_consistency", "READ_CONSISTENCY_UNSPECIFIED")
+ .put("more_results", "NO_MORE_RESULTS")
+ .put("transaction_id", "")
+ .build()));
+ }
+
+ @Test
+ public void runAggregationQueryTraceTest() throws Exception {
+ Entity entity1 =
+ Entity.newBuilder(KEY1)
+ .set("pepper_name", "jalapeno")
+ .set("max_scoville_level", 10000)
+ .build();
+ Entity entity2 =
+ Entity.newBuilder(KEY2)
+ .set("pepper_name", "serrano")
+ .set("max_scoville_level", 25000)
+ .build();
+ Entity entity3 =
+ Entity.newBuilder(KEY3)
+ .set("pepper_name", "habanero")
+ .set("max_scoville_level", 350000)
+ .build();
+ Entity entity4 =
+ Entity.newBuilder(KEY4)
+ .set("pepper_name", "ghost")
+ .set("max_scoville_level", 1500000)
+ .build();
+
+ List entityList = new ArrayList<>();
+ entityList.add(entity1);
+ entityList.add(entity2);
+ entityList.add(entity3);
+ entityList.add(entity4);
+
+ List response = datastore.add(entity1, entity2, entity3, entity4);
+ assertEquals(entityList, response);
+
+ // Clean Up test span context to verify RunAggregationQuery spans
+ cleanupTestSpanContext();
+
+ PropertyFilter mediumSpicyFilters = PropertyFilter.lt("max_scoville_level", 100000);
+ StructuredQuery mediumSpicyQuery =
+ Query.newEntityQueryBuilder().setKind(KEY1.getKind()).setFilter(mediumSpicyFilters).build();
+ AggregationQuery countSpicyPeppers =
+ Query.newAggregationQueryBuilder()
+ .addAggregation(count().as("count"))
+ .over(mediumSpicyQuery)
+ .build();
+ AggregationResults results = datastore.runAggregation(countSpicyPeppers);
+ assertThat(results.size()).isEqualTo(1);
+ AggregationResult result = results.get(0);
+ assertThat(result.getLong("count")).isEqualTo(2L);
+
+ waitForTracesToComplete();
+
+ List spans = prepareSpans();
+ assertEquals(1, spans.size());
+ assertSpanHierarchy(SPAN_NAME_RUN_AGGREGATION_QUERY);
+ }
+
+ @Test
+ public void newTransactionReadWriteTraceTest() throws Exception {
+ // Transaction.Begin
+ Transaction transaction = datastore.newTransaction();
+
+ // Transaction.Lookup
+ Entity entity = datastore.get(KEY1, ReadOption.transactionId(transaction.getTransactionId()));
+ assertNull(entity);
+
+ Entity updatedEntity = Entity.newBuilder(KEY1).set("test_field", "new_test_value1").build();
+ transaction.put(updatedEntity);
+
+ // Transaction.Commit
+ transaction.commit();
+
+ waitForTracesToComplete();
+
+ List spans = prepareSpans();
+ assertEquals(3, spans.size());
+
+ assertSpanHierarchy(SPAN_NAME_BEGIN_TRANSACTION);
+ assertSpanHierarchy(SPAN_NAME_TRANSACTION_LOOKUP);
+ SpanData span = getSpanByName(SPAN_NAME_TRANSACTION_LOOKUP);
+ assertTrue(
+ hasEvent(
+ span,
+ SPAN_NAME_TRANSACTION_LOOKUP + " complete.",
+ Attributes.builder()
+ .put("Deferred", 0)
+ .put("Missing", 1)
+ .put("Received", 0)
+ .put("transactional", true)
+ .put("transaction_id", transaction.getTransactionId().toStringUtf8())
+ .build()));
+
+ assertSpanHierarchy(SPAN_NAME_TRANSACTION_COMMIT);
+ span = getSpanByName(SPAN_NAME_TRANSACTION_COMMIT);
+ assertTrue(
+ hasEvent(
+ span,
+ SPAN_NAME_TRANSACTION_COMMIT + " complete.",
+ Attributes.builder()
+ .put("doc_count", 1)
+ .put("transactional", true)
+ .put("transaction_id", transaction.getTransactionId().toStringUtf8())
+ .build()));
+ }
+
+ @Test
+ public void newTransactionQueryTest() throws Exception {
+ Entity entity1 = Entity.newBuilder(KEY1).set("test_field", "test_value1").build();
+ Entity entity2 = Entity.newBuilder(KEY2).set("test_field", "test_value2").build();
+ List entityList = new ArrayList<>();
+ entityList.add(entity1);
+ entityList.add(entity2);
+
+ List response = datastore.add(entity1, entity2);
+ assertEquals(entityList, response);
+
+ // Clean Up test span context to verify Transaction RunQuery spans
+ cleanupTestSpanContext();
+
+ Transaction transaction = datastore.newTransaction();
+ PropertyFilter filter = PropertyFilter.eq("test_field", entity1.getValue("test_field"));
+ Query query =
+ Query.newEntityQueryBuilder().setKind(KEY1.getKind()).setFilter(filter).build();
+ QueryResults queryResults = transaction.run(query);
+ transaction.commit();
+ assertTrue(queryResults.hasNext());
+ assertEquals(entity1, queryResults.next());
+ assertFalse(queryResults.hasNext());
+
+ waitForTracesToComplete();
+
+ List spans = prepareSpans();
+ assertEquals(3, spans.size());
+
+ assertSpanHierarchy(SPAN_NAME_BEGIN_TRANSACTION);
+ assertSpanHierarchy(SPAN_NAME_TRANSACTION_RUN_QUERY);
+ assertSpanHierarchy(SPAN_NAME_TRANSACTION_COMMIT);
+ SpanData span = getSpanByName(SPAN_NAME_TRANSACTION_RUN_QUERY);
+ assertTrue(
+ hasEvent(
+ span,
+ SPAN_NAME_TRANSACTION_RUN_QUERY + " complete.",
+ Attributes.builder()
+ .put("doc_count", 1)
+ .put("transactional", true)
+ .put("read_consistency", "READ_CONSISTENCY_UNSPECIFIED")
+ .put("more_results", "NO_MORE_RESULTS")
+ .put("transaction_id", transaction.getTransactionId().toStringUtf8())
+ .build()));
+ }
+
+ @Test
+ public void newTransactionRollbackTest() throws Exception {
+ Entity entity1 = Entity.newBuilder(KEY1).set("pepper_type", "jalapeno").build();
+ Entity entity2 = Entity.newBuilder(KEY2).set("pepper_type", "habanero").build();
+ List entityList = new ArrayList<>();
+ entityList.add(entity1);
+ entityList.add(entity2);
+
+ List response = datastore.add(entity1, entity2);
+ assertEquals(entityList, response);
+
+ // Clean Up test span context to verify Transaction Rollback spans
+ cleanupTestSpanContext();
+
+ String simplified_spice_level = "not_spicy";
+ Entity entity1update =
+ Entity.newBuilder(entity1).set("spice_level", simplified_spice_level).build();
+ Transaction transaction = datastore.newTransaction();
+ entity1 = transaction.get(KEY1);
+ switch (entity1.getString("pepper_type")) {
+ case "jalapeno":
+ simplified_spice_level = "mild";
+ break;
+
+ case "habanero":
+ simplified_spice_level = "hot";
+ break;
+ }
+ transaction.update(entity1update);
+ transaction.delete(KEY2);
+ transaction.rollback();
+ assertFalse(transaction.isActive());
+
+ waitForTracesToComplete();
+
+ List spans = prepareSpans();
+ assertEquals(3, spans.size());
+
+ assertSpanHierarchy(SPAN_NAME_BEGIN_TRANSACTION);
+ assertSpanHierarchy(SPAN_NAME_TRANSACTION_LOOKUP);
+ SpanData span = getSpanByName(SPAN_NAME_TRANSACTION_LOOKUP);
+ assertTrue(
+ hasEvent(
+ span,
+ SPAN_NAME_TRANSACTION_LOOKUP + " complete.",
+ Attributes.builder()
+ .put("Deferred", 0)
+ .put("Missing", 0)
+ .put("Received", 1)
+ .put("transactional", true)
+ .put("transaction_id", transaction.getTransactionId().toStringUtf8())
+ .build()));
+
+ assertSpanHierarchy(SPAN_NAME_ROLLBACK);
+ span = getSpanByName(SPAN_NAME_ROLLBACK);
+ assertTrue(
+ hasEvent(
+ span,
+ SPAN_NAME_ROLLBACK,
+ Attributes.builder()
+ .put("transaction_id", transaction.getTransactionId().toStringUtf8())
+ .build()));
+ }
+
+ @Test
+ public void runInTransactionQueryTest() throws Exception {
+ // Set up
+ Entity entity1 = Entity.newBuilder(KEY1).set("test_field", "test_value1").build();
+ Entity entity2 = Entity.newBuilder(KEY2).set("test_field", "test_value2").build();
+ List entityList = new ArrayList<>();
+ entityList.add(entity1);
+ entityList.add(entity2);
+
+ List response = datastore.add(entity1, entity2);
+ assertEquals(entityList, response);
+
+ // Clean Up test span context to verify Transaction Rollback spans
+ cleanupTestSpanContext();
+
+ PropertyFilter filter = PropertyFilter.eq("test_field", entity1.getValue("test_field"));
+ Query query =
+ Query.newEntityQueryBuilder().setKind(KEY1.getKind()).setFilter(filter).build();
+ Datastore.TransactionCallable callable =
+ transaction -> {
+ QueryResults queryResults = datastore.run(query);
+ assertTrue(queryResults.hasNext());
+ assertEquals(entity1, queryResults.next());
+ assertFalse(queryResults.hasNext());
+ return true;
+ };
+ datastore.runInTransaction(callable);
+
+ waitForTracesToComplete();
+
+ List spans = prepareSpans();
+ assertEquals(4, spans.size());
+
+ // Since the runInTransaction method runs the TransactionCallable opaquely in a transaction
+ // there is no way for the API user to know the transaction ID, so we will not validate it here.
+ assertSpanHierarchy(SPAN_NAME_TRANSACTION_RUN, SPAN_NAME_BEGIN_TRANSACTION);
+ assertSpanHierarchy(SPAN_NAME_TRANSACTION_RUN, SPAN_NAME_RUN_QUERY);
+ assertSpanHierarchy(SPAN_NAME_TRANSACTION_RUN, SPAN_NAME_TRANSACTION_COMMIT);
+ }
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/DisabledTraceUtilTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/DisabledTraceUtilTest.java
new file mode 100644
index 000000000..89c91b3a7
--- /dev/null
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/DisabledTraceUtilTest.java
@@ -0,0 +1,54 @@
+/*
+ * 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 static com.google.common.truth.Truth.assertThat;
+
+import org.junit.Test;
+
+public class DisabledTraceUtilTest {
+ @Test
+ public void disabledTraceUtilDoesNotProvideChannelConfigurator() {
+ DisabledTraceUtil traceUtil = new DisabledTraceUtil();
+ assertThat(traceUtil.getChannelConfigurator()).isNull();
+ }
+
+ @Test
+ public void usesDisabledContext() {
+ DisabledTraceUtil traceUtil = new DisabledTraceUtil();
+ assertThat(traceUtil.getCurrentContext() instanceof DisabledTraceUtil.Context).isTrue();
+ }
+
+ @Test
+ public void usesDisabledSpan() {
+ DisabledTraceUtil traceUtil = new DisabledTraceUtil();
+ assertThat(traceUtil.getCurrentSpan() instanceof DisabledTraceUtil.Span).isTrue();
+ assertThat(traceUtil.startSpan("foo") instanceof DisabledTraceUtil.Span).isTrue();
+ assertThat(
+ traceUtil.startSpan("foo", traceUtil.getCurrentSpanContext())
+ instanceof DisabledTraceUtil.Span)
+ .isTrue();
+ }
+
+ @Test
+ public void usesDisabledScope() {
+ DisabledTraceUtil traceUtil = new DisabledTraceUtil();
+ assertThat(traceUtil.getCurrentContext().makeCurrent() instanceof DisabledTraceUtil.Scope)
+ .isTrue();
+ assertThat(traceUtil.getCurrentSpan().makeCurrent() instanceof DisabledTraceUtil.Scope)
+ .isTrue();
+ }
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/EnabledTraceUtilTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/EnabledTraceUtilTest.java
new file mode 100644
index 000000000..a3620bbc2
--- /dev/null
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/EnabledTraceUtilTest.java
@@ -0,0 +1,107 @@
+/*
+ * 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 static com.google.common.truth.Truth.assertThat;
+
+import com.google.cloud.NoCredentials;
+import com.google.cloud.datastore.DatastoreOpenTelemetryOptions;
+import com.google.cloud.datastore.DatastoreOptions;
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
+import org.junit.Before;
+import org.junit.Test;
+
+public class EnabledTraceUtilTest {
+ @Before
+ public void setUp() {
+ GlobalOpenTelemetry.resetForTest();
+ }
+
+ DatastoreOptions.Builder getBaseOptions() {
+ return DatastoreOptions.newBuilder()
+ .setProjectId("test-project")
+ .setCredentials(NoCredentials.getInstance());
+ }
+
+ DatastoreOptions getTracingEnabledOptions() {
+ return getBaseOptions()
+ .setOpenTelemetryOptions(
+ DatastoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build())
+ .build();
+ }
+
+ EnabledTraceUtil newEnabledTraceUtil() {
+ return new EnabledTraceUtil(getTracingEnabledOptions());
+ }
+
+ @Test
+ public void usesOpenTelemetryFromOptions() {
+ OpenTelemetrySdk myOpenTelemetrySdk = OpenTelemetrySdk.builder().build();
+ DatastoreOptions firestoreOptions =
+ getBaseOptions()
+ .setOpenTelemetryOptions(
+ DatastoreOpenTelemetryOptions.newBuilder()
+ .setTracingEnabled(true)
+ .setOpenTelemetry(myOpenTelemetrySdk)
+ .build())
+ .build();
+ EnabledTraceUtil traceUtil = new EnabledTraceUtil(firestoreOptions);
+ assertThat(traceUtil.getOpenTelemetry()).isEqualTo(myOpenTelemetrySdk);
+ }
+
+ @Test
+ public void usesGlobalOpenTelemetryIfOpenTelemetryInstanceNotProvided() {
+ OpenTelemetrySdk globalOpenTelemetrySdk = OpenTelemetrySdk.builder().buildAndRegisterGlobal();
+ DatastoreOptions firestoreOptions =
+ getBaseOptions()
+ .setOpenTelemetryOptions(
+ DatastoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build())
+ .build();
+ EnabledTraceUtil traceUtil = new EnabledTraceUtil(firestoreOptions);
+ assertThat(traceUtil.getOpenTelemetry()).isEqualTo(GlobalOpenTelemetry.get());
+ }
+
+ @Test
+ public void enabledTraceUtilProvidesChannelConfigurator() {
+ assertThat(newEnabledTraceUtil().getChannelConfigurator()).isNull();
+ }
+
+ @Test
+ public void usesEnabledContext() {
+ assertThat(newEnabledTraceUtil().getCurrentContext() instanceof EnabledTraceUtil.Context)
+ .isTrue();
+ }
+
+ @Test
+ public void usesEnabledSpan() {
+ EnabledTraceUtil traceUtil = newEnabledTraceUtil();
+ assertThat(traceUtil.getCurrentSpan() instanceof EnabledTraceUtil.Span).isTrue();
+ assertThat(traceUtil.startSpan("foo") != null).isTrue();
+ assertThat(
+ traceUtil.startSpan("foo", traceUtil.getCurrentSpanContext())
+ instanceof EnabledTraceUtil.Span)
+ .isTrue();
+ }
+
+ @Test
+ public void usesEnabledScope() {
+ EnabledTraceUtil traceUtil = newEnabledTraceUtil();
+ assertThat(traceUtil.getCurrentContext().makeCurrent() instanceof EnabledTraceUtil.Scope)
+ .isTrue();
+ assertThat(traceUtil.getCurrentSpan().makeCurrent() instanceof EnabledTraceUtil.Scope).isTrue();
+ }
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/TraceUtilTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/TraceUtilTest.java
new file mode 100644
index 000000000..f1cce8006
--- /dev/null
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/TraceUtilTest.java
@@ -0,0 +1,62 @@
+/*
+ * 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 static com.google.common.truth.Truth.assertThat;
+
+import com.google.cloud.NoCredentials;
+import com.google.cloud.datastore.DatastoreOpenTelemetryOptions;
+import com.google.cloud.datastore.DatastoreOptions;
+import org.junit.Test;
+
+public class TraceUtilTest {
+ @Test
+ public void defaultOptionsUseDisabledTraceUtil() {
+ TraceUtil traceUtil =
+ TraceUtil.getInstance(
+ DatastoreOptions.newBuilder()
+ .setProjectId("test-project")
+ .setCredentials(NoCredentials.getInstance())
+ .build());
+ assertThat(traceUtil instanceof DisabledTraceUtil).isTrue();
+ }
+
+ @Test
+ public void tracingDisabledOptionsUseDisabledTraceUtil() {
+ TraceUtil traceUtil =
+ TraceUtil.getInstance(
+ DatastoreOptions.newBuilder()
+ .setProjectId("test-project")
+ .setCredentials(NoCredentials.getInstance())
+ .setOpenTelemetryOptions(
+ DatastoreOpenTelemetryOptions.newBuilder().setTracingEnabled(false).build())
+ .build());
+ assertThat(traceUtil instanceof DisabledTraceUtil).isTrue();
+ }
+
+ @Test
+ public void tracingEnabledOptionsUseEnabledTraceUtil() {
+ TraceUtil traceUtil =
+ TraceUtil.getInstance(
+ DatastoreOptions.newBuilder()
+ .setProjectId("test-project")
+ .setCredentials(NoCredentials.getInstance())
+ .setOpenTelemetryOptions(
+ DatastoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build())
+ .build());
+ assertThat(traceUtil instanceof EnabledTraceUtil).isTrue();
+ }
+}
diff --git a/grpc-google-cloud-datastore-admin-v1/pom.xml b/grpc-google-cloud-datastore-admin-v1/pom.xml
index 69da06680..6527e96f3 100644
--- a/grpc-google-cloud-datastore-admin-v1/pom.xml
+++ b/grpc-google-cloud-datastore-admin-v1/pom.xml
@@ -4,13 +4,13 @@
4.0.0
com.google.api.grpc
grpc-google-cloud-datastore-admin-v1
- 2.20.3-SNAPSHOT
+ 2.21.2
grpc-google-cloud-datastore-admin-v1
GRPC library for google-cloud-datastore
com.google.cloud
google-cloud-datastore-parent
- 2.20.3-SNAPSHOT
+ 2.21.2
diff --git a/pom.xml b/pom.xml
index 407041f2b..fda01b132 100644
--- a/pom.xml
+++ b/pom.xml
@@ -4,7 +4,7 @@
com.google.cloud
google-cloud-datastore-parent
pom
- 2.20.3-SNAPSHOT
+ 2.21.2
Google Cloud Datastore Parent
https://github.com/googleapis/java-datastore
@@ -14,7 +14,7 @@
com.google.cloud
sdk-platform-java-config
- 3.32.0
+ 3.34.0
@@ -159,27 +159,27 @@
com.google.api.grpc
proto-google-cloud-datastore-admin-v1
- 2.20.3-SNAPSHOT
+ 2.21.2
com.google.api.grpc
grpc-google-cloud-datastore-admin-v1
- 2.20.3-SNAPSHOT
+ 2.21.2
com.google.cloud
google-cloud-datastore
- 2.20.3-SNAPSHOT
+ 2.21.2
com.google.api.grpc
proto-google-cloud-datastore-v1
- 0.111.3-SNAPSHOT
+ 0.112.2
com.google.cloud.datastore
datastore-v1-proto-client
- 2.20.3-SNAPSHOT
+ 2.21.2
com.google.api.grpc
diff --git a/proto-google-cloud-datastore-admin-v1/pom.xml b/proto-google-cloud-datastore-admin-v1/pom.xml
index f53996a6d..a8c91f4dd 100644
--- a/proto-google-cloud-datastore-admin-v1/pom.xml
+++ b/proto-google-cloud-datastore-admin-v1/pom.xml
@@ -4,13 +4,13 @@
4.0.0
com.google.api.grpc
proto-google-cloud-datastore-admin-v1
- 2.20.3-SNAPSHOT
+ 2.21.2
proto-google-cloud-datastore-admin-v1
Proto library for google-cloud-datastore
com.google.cloud
google-cloud-datastore-parent
- 2.20.3-SNAPSHOT
+ 2.21.2
diff --git a/proto-google-cloud-datastore-v1/pom.xml b/proto-google-cloud-datastore-v1/pom.xml
index 85f530c75..bf896574d 100644
--- a/proto-google-cloud-datastore-v1/pom.xml
+++ b/proto-google-cloud-datastore-v1/pom.xml
@@ -4,13 +4,13 @@
4.0.0
com.google.api.grpc
proto-google-cloud-datastore-v1
- 0.111.3-SNAPSHOT
+ 0.112.2
proto-google-cloud-datastore-v1
PROTO library for proto-google-cloud-datastore-v1
com.google.cloud
google-cloud-datastore-parent
- 2.20.3-SNAPSHOT
+ 2.21.2
diff --git a/samples/native-image-sample/README.md b/samples/native-image-sample/README.md
deleted file mode 100644
index 5f2cfbd27..000000000
--- a/samples/native-image-sample/README.md
+++ /dev/null
@@ -1,96 +0,0 @@
-# Datastore Sample Application with Native Image
-
-This application uses the [Google Cloud Datastore client library](https://cloud.google.com/datastore/docs/reference/libraries) and is compatible with Native Image compilation.
-
-This sample runs through some basic operations of creating/deleting entities, running queries, and running transaction code.
-
-## Setup Instructions
-
-You will need to follow these prerequisite steps in order to run the samples:
-
-1. If you have not already, [create a Google Cloud Platform Project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#creating_a_project).
-
-2. Install the [Google Cloud SDK](https://cloud.google.com/sdk/) which will allow you to run the sample with your project's credentials.
-
- Once installed, log in with Application Default Credentials using the following command:
-
- ```
- gcloud auth application-default login
- ```
-
- **Note:** Authenticating with Application Default Credentials is convenient to use during development, but we recommend [alternate methods of authentication](https://cloud.google.com/docs/authentication/production) during production use.
-
-3. Install the native image compiler.
-
- You can follow the [installation instructions](https://www.graalvm.org/docs/getting-started/#install-graalvm).
- After following the instructions, ensure that you install the native image extension installed by running:
-
- ```
- gu install native-image
- ```
-
- Once you finish following the instructions, verify that the default version of Java is set to the correct version by running `java -version` in a terminal.
-
- You will see something similar to the below output:
-
- ```
- $ java -version
-
- openjdk version "17.0.3" 2022-04-19
- OpenJDK Runtime Environment GraalVM CE 22.1.0 (build 17.0.3+7-jvmci-22.1-b06)
- OpenJDK 64-Bit Server VM GraalVM CE 22.1.0 (build 17.0.3+7-jvmci-22.1-b06, mixed mode, sharing)
- ```
-## Sample
-1. **(Optional)** If you wish to run the application against the [Datastore emulator](https://cloud.google.com/sdk/gcloud/reference/beta/emulators/datastore), ensure that you have the [Google Cloud SDK](https://cloud.google.com/sdk) installed.
-
- In a new terminal window, start the emulator via `gcloud`:
-
- ```
- gcloud beta emulators datastore start --host-port=localhost:9010
- ```
-
- Leave the emulator running in this terminal for now.
- In the next section, we will run the sample application against the Datastore emulator instance.
-
-2. Navigate to this directory and compile the application with the native image compiler.
-
- ```
- mvn package -P native -DskipTests
- ```
-
-3. **(Optional)** If you're using the emulator, export the `DATASTORE_EMULATOR_HOST` as an environment variable in your terminal.
-
- ```
- export DATASTORE_EMULATOR_HOST=localhost:9010
- ```
-
- The Datastore Client Libraries will detect this environment variable and automatically connect to the emulator instance if this variable is set.
-
-4. Run the application.
-
- ```
- ./target/native-image-sample
- ```
-
-5. The application will run through some basic Datastore operations and log some output statements.
-
- ```
- Successfully added entity.
- Reading entity: 1cf34cc1-2b8a-4945-9fc4-058f03dcd08e
- Successfully deleted entity: 1cf34cc1-2b8a-4945-9fc4-058f03dcd08e
- Run fake transaction code.
- Found entity:
- name=de4f36f4-3936-4252-98d3-e0d56d485254
- kind=test-kind
- namespace=nativeimage-test-namespace
- properties={description=StringValue{valueType=STRING, excludeFromIndexes=false, meaning=0, value=hello world}}
- Ran transaction callable.
- ```
-
-### Sample Integration test with Native Image Support
-
-In order to run the sample integration test as a native image, call the following command:
-
- ```
- mvn test -Pnative
- ```
diff --git a/samples/native-image-sample/pom.xml b/samples/native-image-sample/pom.xml
index 0f3c1919c..e69de29bb 100644
--- a/samples/native-image-sample/pom.xml
+++ b/samples/native-image-sample/pom.xml
@@ -1,141 +0,0 @@
-
-
- 4.0.0
- com.example.datastore
- native-image-sample
- Native Image Sample
- https://github.com/googleapis/java-datastore
-
-
-
- com.google.cloud.samples
- shared-configuration
- 1.2.0
-
-
-
-
- 1.8
- 1.8
- UTF-8
-
-
-
-
-
- com.google.cloud
- libraries-bom
- 26.41.0
- pom
- import
-
-
-
-
-
-
- com.google.cloud
- google-cloud-datastore
-
-
-
- junit
- junit
- 4.13.2
- test
-
-
- com.google.truth
- truth
- 1.4.3
- test
-
-
-
-
-
-
- org.apache.maven.plugins
- maven-jar-plugin
-
-
-
- com.example.datastore.NativeImageDatastoreSample
-
-
-
-
-
-
-
-
-
-
- native
-
-
-
- org.junit.vintage
- junit-vintage-engine
- 5.10.2
- test
-
-
- org.graalvm.buildtools
- junit-platform-native
- 0.10.2
- test
-
-
-
-
-
-
- org.apache.maven.plugins
- maven-surefire-plugin
-
- 3.2.5
-
-
- **/IT*
-
-
-
-
- org.graalvm.buildtools
- native-maven-plugin
- 0.10.2
- true
-
- com.example.datastore.NativeImageDatastoreSample
-
- --no-fallback
- --no-server
-
-
-
-
- build-native
-
- build
- test
-
- package
-
-
- test-native
-
- test
-
- test
-
-
-
-
-
-
-
-
\ No newline at end of file
diff --git a/samples/native-image-sample/src/main/java/com/example/datastore/NativeImageDatastoreSample.java b/samples/native-image-sample/src/main/java/com/example/datastore/NativeImageDatastoreSample.java
deleted file mode 100644
index 7ce5c900a..000000000
--- a/samples/native-image-sample/src/main/java/com/example/datastore/NativeImageDatastoreSample.java
+++ /dev/null
@@ -1,131 +0,0 @@
-/*
- * Copyright 2020-2021 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.example.datastore;
-
-import com.google.cloud.datastore.Datastore;
-import com.google.cloud.datastore.DatastoreOptions;
-import com.google.cloud.datastore.Entity;
-import com.google.cloud.datastore.Key;
-import com.google.cloud.datastore.Query;
-import com.google.cloud.datastore.QueryResults;
-import com.google.cloud.datastore.StructuredQuery;
-import com.google.cloud.datastore.Transaction;
-import java.time.Duration;
-import java.time.Instant;
-import java.util.UUID;
-
-/** Sample Datastore Application. */
-public class NativeImageDatastoreSample {
-
- /* Datastore namespace where entities will be created. */
- private static final String TEST_NAMESPACE = "nativeimage-test-namespace";
-
- /* Datastore kind used. */
- private static final String TEST_KIND = "test-kind";
-
- /** Entrypoint to the Datastore sample application. */
- public static void main(String[] args) {
- Instant startTime = Instant.now();
- Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
-
- String testId = UUID.randomUUID().toString();
-
- addEntity(datastore, testId);
- getEntity(datastore, testId);
- deleteEntity(datastore, testId);
-
- runTransaction(datastore);
-
- String id = UUID.randomUUID().toString();
- Key key = createKey(datastore, id);
- runTransactionCallable(datastore, key);
- Instant endTime = Instant.now();
- Duration duration = Duration.between(startTime, endTime);
- System.out.println("Duration: " + duration.toString());
- }
-
- static void addEntity(Datastore datastore, String id) {
- Key key = createKey(datastore, id);
- Entity entity = Entity.newBuilder(key).set("description", "hello world").build();
- datastore.add(entity);
- System.out.println("Successfully added entity.");
- }
-
- static void getEntity(Datastore datastore, String id) {
- Key key = createKey(datastore, id);
- Entity entity = datastore.get(key);
- System.out.println("Reading entity: " + entity.getKey().getName());
- }
-
- static void deleteEntity(Datastore datastore, String id) {
- Key key = createKey(datastore, id);
- datastore.delete(key);
-
- Entity entity = datastore.get(key);
- if (entity == null) {
- System.out.println("Successfully deleted entity: " + id);
- } else {
- throw new RuntimeException("Failed to delete entity: " + id);
- }
- }
-
- static void runTransactionCallable(Datastore datastore, Key entityKey) {
- datastore.runInTransaction(
- client -> {
- Entity entity = Entity.newBuilder(entityKey).set("description", "hello world").build();
- datastore.add(entity);
-
- StructuredQuery query =
- Query.newEntityQueryBuilder().setNamespace(TEST_NAMESPACE).setKind(TEST_KIND).build();
-
- QueryResults results = datastore.run(query);
- while (results.hasNext()) {
- Entity result = results.next();
- String name = result.getKey().getName();
- String kind = result.getKey().getKind();
- String namespace = result.getKey().getNamespace();
- System.out.println(
- "Found entity:"
- + "\n\t\tname="
- + name
- + "\n\t\tkind="
- + kind
- + "\n\t\tnamespace="
- + namespace
- + "\n\t\tproperties="
- + result.getProperties().toString());
- }
-
- datastore.delete(entityKey);
- return null;
- });
-
- System.out.println("Ran transaction callable.");
- }
-
- private static void runTransaction(Datastore datastore) {
- Transaction transaction = datastore.newTransaction();
- transaction.commit();
- transaction = datastore.newTransaction();
- transaction.rollback();
- System.out.println("Run fake transaction code.");
- }
-
- static Key createKey(Datastore datastore, String id) {
- return datastore.newKeyFactory().setNamespace(TEST_NAMESPACE).setKind(TEST_KIND).newKey(id);
- }
-}
diff --git a/samples/native-image-sample/src/test/java/com/example/datastore/ITNativeImageDatastoreSample.java b/samples/native-image-sample/src/test/java/com/example/datastore/ITNativeImageDatastoreSample.java
deleted file mode 100644
index 710f18367..000000000
--- a/samples/native-image-sample/src/test/java/com/example/datastore/ITNativeImageDatastoreSample.java
+++ /dev/null
@@ -1,75 +0,0 @@
-/*
- * 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
- *
- * 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.example.datastore;
-
-import static com.google.common.truth.Truth.assertThat;
-
-import com.google.cloud.datastore.Datastore;
-import com.google.cloud.datastore.DatastoreOptions;
-import com.google.cloud.datastore.Key;
-import java.io.ByteArrayOutputStream;
-import java.io.PrintStream;
-import java.util.UUID;
-import org.junit.Before;
-import org.junit.Test;
-
-/** Tests for {@link com.example.datastore.NativeImageDatastoreSample} */
-public class ITNativeImageDatastoreSample {
-
- private Datastore datastore;
- private ByteArrayOutputStream bout;
- private PrintStream out;
-
- @Before
- public void setUp() {
- datastore = DatastoreOptions.getDefaultInstance().getService();
- bout = new ByteArrayOutputStream();
- out = new PrintStream(bout);
- System.setOut(out);
- }
-
- @Test
- public void testAddAndGetEntity() {
- bout.reset();
- String testId = "test-id-" + UUID.randomUUID();
- NativeImageDatastoreSample.addEntity(datastore, testId);
- NativeImageDatastoreSample.getEntity(datastore, testId);
- assertThat(bout.toString()).contains("Reading entity: " + testId);
-
- NativeImageDatastoreSample.deleteEntity(datastore, testId);
- }
-
- @Test
- public void testRunTransactionalCallable() {
- bout.reset();
- String testId = "test-id-" + UUID.randomUUID();
- Key key = NativeImageDatastoreSample.createKey(datastore, testId);
- NativeImageDatastoreSample.runTransactionCallable(datastore, key);
- assertThat(bout.toString())
- .contains(
- "Found entity:"
- + "\n\t\tname="
- + testId
- + "\n\t\tkind=test-kind"
- + "\n\t\tnamespace=nativeimage-test-namespace"
- + "\n\t\tproperties={description=StringValue{valueType=STRING, excludeFromIndexes=false,"
- + " meaning=0, value=hello world}}\n"
- + "Ran transaction callable.");
-
- NativeImageDatastoreSample.deleteEntity(datastore, "test-id");
- }
-}
diff --git a/samples/pom.xml b/samples/pom.xml
index e81bec450..2e970d081 100644
--- a/samples/pom.xml
+++ b/samples/pom.xml
@@ -31,7 +31,6 @@
install-without-bom
snapshot
snippets
- native-image-sample
diff --git a/samples/snapshot/pom.xml b/samples/snapshot/pom.xml
index 93710c03b..5f2e900a8 100644
--- a/samples/snapshot/pom.xml
+++ b/samples/snapshot/pom.xml
@@ -28,7 +28,7 @@
com.google.cloud
google-cloud-datastore
- 2.20.3-SNAPSHOT
+ 2.21.2
diff --git a/samples/snippets/pom.xml b/samples/snippets/pom.xml
index f04cdc885..7e4de0be6 100644
--- a/samples/snippets/pom.xml
+++ b/samples/snippets/pom.xml
@@ -30,7 +30,7 @@
com.google.cloud
libraries-bom
- 26.41.0
+ 26.43.0
pom
import
diff --git a/versions.txt b/versions.txt
index 5d2ffbced..3cee21923 100644
--- a/versions.txt
+++ b/versions.txt
@@ -1,9 +1,9 @@
# Format:
# module:released-version:current-version
-google-cloud-datastore:2.20.2:2.20.3-SNAPSHOT
-google-cloud-datastore-bom:2.20.2:2.20.3-SNAPSHOT
-proto-google-cloud-datastore-v1:0.111.2:0.111.3-SNAPSHOT
-datastore-v1-proto-client:2.20.2:2.20.3-SNAPSHOT
-proto-google-cloud-datastore-admin-v1:2.20.2:2.20.3-SNAPSHOT
-grpc-google-cloud-datastore-admin-v1:2.20.2:2.20.3-SNAPSHOT
+google-cloud-datastore:2.21.2:2.21.2
+google-cloud-datastore-bom:2.21.2:2.21.2
+proto-google-cloud-datastore-v1:0.112.2:0.112.2
+datastore-v1-proto-client:2.21.2:2.21.2
+proto-google-cloud-datastore-admin-v1:2.21.2:2.21.2
+grpc-google-cloud-datastore-admin-v1:2.21.2:2.21.2