Skip to content
This repository was archived by the owner on Feb 24, 2026. It is now read-only.

Commit 10341c0

Browse files
jimit-j-shahjimit-j-shah
committed
feat: support for transactional operations
- tested using newTransaction() and runInTransaction()
1 parent aa7a1dc commit 10341c0

4 files changed

Lines changed: 91 additions & 24 deletions

File tree

‎google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java‎

Lines changed: 33 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import com.google.cloud.ServiceOptions;
2626
import com.google.cloud.datastore.execution.AggregationQueryExecutor;
2727
import com.google.cloud.datastore.spi.v1.DatastoreRpc;
28+
import com.google.cloud.datastore.telemetry.TraceUtil.Context;
2829
import com.google.common.base.MoreObjects;
2930
import com.google.common.base.Preconditions;
3031
import com.google.common.collect.AbstractIterator;
@@ -52,6 +53,7 @@
5253
import java.util.Optional;
5354
import java.util.Set;
5455
import java.util.concurrent.Callable;
56+
import javax.annotation.Nonnull;
5557

5658
final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datastore {
5759

@@ -103,13 +105,18 @@ static class ReadWriteTransactionCallable<T> implements Callable<T> {
103105
private final TransactionCallable<T> callable;
104106
private volatile TransactionOptions options;
105107
private volatile Transaction transaction;
108+
@Nonnull private final Context transactionTraceContext;
106109

107110
ReadWriteTransactionCallable(
108-
Datastore datastore, TransactionCallable<T> callable, TransactionOptions options) {
111+
Datastore datastore,
112+
TransactionCallable<T> callable,
113+
TransactionOptions options,
114+
@Nonnull Context parentTraceContext) {
109115
this.datastore = datastore;
110116
this.callable = callable;
111117
this.options = options;
112118
this.transaction = null;
119+
this.transactionTraceContext = parentTraceContext;
113120
}
114121

115122
Datastore getDatastore() {
@@ -132,8 +139,9 @@ void setPrevTransactionId(ByteString transactionId) {
132139

133140
@Override
134141
public T call() throws DatastoreException {
135-
transaction = datastore.newTransaction(options);
136-
try {
142+
try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored =
143+
transactionTraceContext.makeCurrent()) {
144+
transaction = datastore.newTransaction(options);
137145
T value = callable.run(transaction);
138146
transaction.commit();
139147
return value;
@@ -154,36 +162,41 @@ public T call() throws DatastoreException {
154162

155163
@Override
156164
public <T> T runInTransaction(final TransactionCallable<T> callable) {
157-
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION);
158-
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
165+
com.google.cloud.datastore.telemetry.TraceUtil.Span span =
166+
otelTraceUtil.startSpan(
167+
com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN);
168+
try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
159169
return RetryHelper.runWithRetries(
160-
new ReadWriteTransactionCallable<T>(this, callable, null),
170+
new ReadWriteTransactionCallable<T>(this, callable, null, otelTraceUtil.currentContext()),
161171
retrySettings,
162172
TRANSACTION_EXCEPTION_HANDLER,
163173
getOptions().getClock());
164174
} catch (RetryHelperException e) {
165-
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
175+
span.end(e);
166176
throw DatastoreException.translateAndThrow(e);
167177
} finally {
168-
span.end(TraceUtil.END_SPAN_OPTIONS);
178+
span.end();
169179
}
170180
}
171181

172182
@Override
173183
public <T> T runInTransaction(
174184
final TransactionCallable<T> callable, TransactionOptions transactionOptions) {
175-
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION);
176-
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
185+
com.google.cloud.datastore.telemetry.TraceUtil.Span span =
186+
otelTraceUtil.startSpan(
187+
com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN);
188+
try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
177189
return RetryHelper.runWithRetries(
178-
new ReadWriteTransactionCallable<T>(this, callable, transactionOptions),
190+
new ReadWriteTransactionCallable<T>(
191+
this, callable, transactionOptions, otelTraceUtil.currentContext()),
179192
retrySettings,
180193
TRANSACTION_EXCEPTION_HANDLER,
181194
getOptions().getClock());
182195
} catch (RetryHelperException e) {
183-
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
196+
span.end(e);
184197
throw DatastoreException.translateAndThrow(e);
185198
} finally {
186-
span.end(TraceUtil.END_SPAN_OPTIONS);
199+
span.end();
187200
}
188201
}
189202

@@ -634,10 +647,14 @@ private com.google.datastore.v1.CommitResponse commitMutation(
634647

635648
com.google.datastore.v1.CommitResponse commit(
636649
final com.google.datastore.v1.CommitRequest requestPb) {
637-
com.google.cloud.datastore.telemetry.TraceUtil.Span span =
638-
otelTraceUtil.startSpan(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_COMMIT);
650+
final boolean isTransactional =
651+
requestPb.hasTransaction() || requestPb.hasSingleUseTransaction();
652+
final String spanName =
653+
isTransactional
654+
? com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_COMMIT
655+
: com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_COMMIT;
656+
com.google.cloud.datastore.telemetry.TraceUtil.Span span = otelTraceUtil.startSpan(spanName);
639657
span.setAttribute("isTransactional", requestPb.hasTransaction());
640-
641658
try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
642659
return RetryHelper.runWithRetries(
643660
() -> datastoreRpc.commit(requestPb),

‎google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
2020

2121
import com.google.api.core.BetaApi;
2222
import com.google.cloud.datastore.models.ExplainOptions;
23+
import com.google.cloud.datastore.telemetry.TraceUtil.Context;
24+
import com.google.cloud.datastore.telemetry.TraceUtil.Scope;
2325
import com.google.common.collect.ImmutableList;
2426
import com.google.datastore.v1.ReadOptions;
2527
import com.google.datastore.v1.TransactionOptions;
@@ -28,6 +30,7 @@
2830
import java.util.Iterator;
2931
import java.util.List;
3032
import java.util.Optional;
33+
import javax.annotation.Nonnull;
3134

3235
final class TransactionImpl extends BaseDatastoreBatchWriter implements Transaction {
3336

@@ -37,6 +40,8 @@ final class TransactionImpl extends BaseDatastoreBatchWriter implements Transact
3740

3841
private final ReadOptionProtoPreparer readOptionProtoPreparer;
3942

43+
@Nonnull private final Context transactionTraceContext;
44+
4045
static class ResponseImpl implements Transaction.Response {
4146

4247
private final com.google.datastore.v1.CommitResponse response;
@@ -78,6 +83,7 @@ public List<Key> getGeneratedKeys() {
7883

7984
transactionId = datastore.requestTransactionId(requestPb);
8085
this.readOptionProtoPreparer = new ReadOptionProtoPreparer();
86+
this.transactionTraceContext = datastore.getOptions().getTraceUtil().currentContext();
8187
}
8288

8389
@Override
@@ -96,7 +102,9 @@ public Iterator<Entity> get(Key... keys) {
96102
@Override
97103
public List<Entity> fetch(Key... keys) {
98104
validateActive();
99-
return DatastoreHelper.fetch(this, keys);
105+
try (Scope ignored = transactionTraceContext.makeCurrent()) {
106+
return DatastoreHelper.fetch(this, keys);
107+
}
100108
}
101109

102110
@Override

‎google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import com.google.api.core.InternalExtensionOnly;
2222
import com.google.cloud.datastore.DatastoreOptions;
2323
import io.grpc.ManagedChannelBuilder;
24+
import io.opentelemetry.context.Context;
2425
import java.util.Map;
2526
import javax.annotation.Nonnull;
2627
import javax.annotation.Nullable;
@@ -35,6 +36,7 @@ public interface TraceUtil {
3536
static final String SPAN_NAME_COMMIT = "Commit";
3637
static final String SPAN_NAME_RUN_QUERY = "RunQuery";
3738
static final String SPAN_NAME_RUN_AGGREGATION_QUERY = "RunAggregationQuery";
39+
static final String SPAN_NAME_TRANSACTION_RUN = "Transaction.run";
3840
static final String SPAN_NAME_BEGIN_TRANSACTION = "Transaction.Begin";
3941
static final String SPAN_NAME_TRANSACTION_LOOKUP = "Transaction.Lookup";
4042
static final String SPAN_NAME_TRANSACTION_COMMIT = "Transaction.Commit";

‎google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java‎

Lines changed: 47 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,9 @@
2222
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP;
2323
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY;
2424
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY;
25+
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_COMMIT;
2526
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_LOOKUP;
27+
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN;
2628
import static com.google.common.truth.Truth.assertThat;
2729
import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME;
2830
import static org.junit.Assert.assertEquals;
@@ -749,6 +751,7 @@ public void transactionalLookupTest() throws Exception {
749751
try (Scope ignored = rootSpan.makeCurrent()) {
750752
Transaction transaction = datastore.newTransaction();
751753
Entity entity = datastore.get(KEY1, ReadOption.transactionId(transaction.getTransactionId()));
754+
transaction.commit();
752755
assertNull(entity);
753756
} finally {
754757
rootSpan.end();
@@ -757,18 +760,55 @@ public void transactionalLookupTest() throws Exception {
757760

758761
fetchAndValidateTrace(
759762
customSpanContext.getTraceId(),
760-
/*numExpectedSpans=*/ 2,
761-
Collections.singletonList(Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION)));
763+
/*numExpectedSpans=*/ 3,
764+
Arrays.asList(
765+
Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION),
766+
Collections.singletonList(SPAN_NAME_TRANSACTION_LOOKUP),
767+
Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT)));
768+
}
769+
770+
@Test
771+
public void runInTransactionQueryTest() throws Exception {
772+
Entity entity1 = Entity.newBuilder(KEY1).set("test_field", "test_value1").build();
773+
Entity entity2 = Entity.newBuilder(KEY2).set("test_field", "test_value2").build();
774+
List<Entity> entityList = new ArrayList<>();
775+
entityList.add(entity1);
776+
entityList.add(entity2);
777+
778+
List<Entity> response = datastore.add(entity1, entity2);
779+
assertEquals(entityList, response);
780+
781+
assertNotNull(customSpanContext);
782+
783+
Span rootSpan = getNewRootSpanWithContext();
784+
try (Scope ignored = rootSpan.makeCurrent()) {
785+
PropertyFilter filter = PropertyFilter.eq("test_field", entity1.getValue("test_field"));
786+
Query<Entity> query =
787+
Query.newEntityQueryBuilder().setKind(KEY1.getKind()).setFilter(filter).build();
788+
Datastore.TransactionCallable<Boolean> callable =
789+
transaction -> {
790+
QueryResults<Entity> queryResults = datastore.run(query);
791+
assertTrue(queryResults.hasNext());
792+
assertEquals(entity1, queryResults.next());
793+
assertFalse(queryResults.hasNext());
794+
return true;
795+
};
796+
datastore.runInTransaction(callable);
797+
} finally {
798+
rootSpan.end();
799+
}
800+
waitForTracesToComplete();
762801

763802
fetchAndValidateTrace(
764803
customSpanContext.getTraceId(),
765-
/*numExpectedSpans=*/ 2,
766-
Collections.singletonList(Collections.singletonList(SPAN_NAME_TRANSACTION_LOOKUP)));
804+
/*numExpectedSpans=*/ 4,
805+
Arrays.asList(
806+
Collections.singletonList(SPAN_NAME_TRANSACTION_RUN),
807+
Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION),
808+
Collections.singletonList(SPAN_NAME_RUN_QUERY),
809+
Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT)));
767810
}
768811

769-
@Test
770-
public void runInTransactionQueryTest() throws Exception {}
771-
772812
@Test
773813
public void runInTransactionAggregationQueryTest() throws Exception {}
774814

0 commit comments

Comments
 (0)