Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: support for transactional operations #1468

Merged
merged 3 commits into from
Jun 7, 2024
Merged

feat: support for transactional operations #1468

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e180a29
feat: support for transactional operations
jimit-j-shah May 30, 2024
96bf35c
fix: code review comments
jimit-j-shah Jun 5, 2024
946a805
fix: lint errors
jimit-j-shah Jun 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 33 additions & 16 deletions google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import com.google.cloud.ServiceOptions;
import com.google.cloud.datastore.execution.AggregationQueryExecutor;
import com.google.cloud.datastore.spi.v1.DatastoreRpc;
import com.google.cloud.datastore.telemetry.TraceUtil.Context;
import com.google.common.base.MoreObjects;
import com.google.common.base.Preconditions;
import com.google.common.collect.AbstractIterator;
Expand Down Expand Up @@ -52,6 +53,7 @@
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.Callable;
import javax.annotation.Nonnull;

final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datastore {

Expand Down Expand Up @@ -103,13 +105,18 @@ static class ReadWriteTransactionCallable<T> implements Callable<T> {
private final TransactionCallable<T> callable;
private volatile TransactionOptions options;
private volatile Transaction transaction;
@Nonnull private final Context transactionTraceContext;

ReadWriteTransactionCallable(
Datastore datastore, TransactionCallable<T> callable, TransactionOptions options) {
Datastore datastore,
TransactionCallable<T> callable,
TransactionOptions options,
@Nonnull Context parentTraceContext) {
this.datastore = datastore;
this.callable = callable;
this.options = options;
this.transaction = null;
this.transactionTraceContext = parentTraceContext;
jimit-j-shah marked this conversation as resolved.
Show resolved Hide resolved
}

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

@Override
public T call() throws DatastoreException {
transaction = datastore.newTransaction(options);
try {
try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored =
transactionTraceContext.makeCurrent()) {
transaction = datastore.newTransaction(options);
jimit-j-shah marked this conversation as resolved.
Show resolved Hide resolved
T value = callable.run(transaction);
transaction.commit();
return value;
Expand All @@ -154,36 +162,41 @@ public T call() throws DatastoreException {

@Override
public <T> T runInTransaction(final TransactionCallable<T> callable) {
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION);
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
com.google.cloud.datastore.telemetry.TraceUtil.Span span =
otelTraceUtil.startSpan(
com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN);
try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
return RetryHelper.runWithRetries(
new ReadWriteTransactionCallable<T>(this, callable, null),
new ReadWriteTransactionCallable<T>(this, callable, null, otelTraceUtil.currentContext()),
retrySettings,
TRANSACTION_EXCEPTION_HANDLER,
getOptions().getClock());
} catch (RetryHelperException e) {
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
span.end(e);
jimit-j-shah marked this conversation as resolved.
Show resolved Hide resolved
throw DatastoreException.translateAndThrow(e);
} finally {
span.end(TraceUtil.END_SPAN_OPTIONS);
span.end();
}
}

@Override
public <T> T runInTransaction(
final TransactionCallable<T> callable, TransactionOptions transactionOptions) {
Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION);
try (Scope scope = traceUtil.getTracer().withSpan(span)) {
com.google.cloud.datastore.telemetry.TraceUtil.Span span =
otelTraceUtil.startSpan(
com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN);
try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
return RetryHelper.runWithRetries(
new ReadWriteTransactionCallable<T>(this, callable, transactionOptions),
new ReadWriteTransactionCallable<T>(
this, callable, transactionOptions, otelTraceUtil.currentContext()),
jimit-j-shah marked this conversation as resolved.
Show resolved Hide resolved
retrySettings,
TRANSACTION_EXCEPTION_HANDLER,
getOptions().getClock());
} catch (RetryHelperException e) {
span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
span.end(e);
throw DatastoreException.translateAndThrow(e);
} finally {
span.end(TraceUtil.END_SPAN_OPTIONS);
span.end();
}
}

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

com.google.datastore.v1.CommitResponse commit(
final com.google.datastore.v1.CommitRequest requestPb) {
com.google.cloud.datastore.telemetry.TraceUtil.Span span =
otelTraceUtil.startSpan(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_COMMIT);
final boolean isTransactional =
requestPb.hasTransaction() || requestPb.hasSingleUseTransaction();
final String spanName =
isTransactional
? com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_COMMIT
: com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_COMMIT;
com.google.cloud.datastore.telemetry.TraceUtil.Span span = otelTraceUtil.startSpan(spanName);
span.setAttribute("isTransactional", requestPb.hasTransaction());
jimit-j-shah marked this conversation as resolved.
Show resolved Hide resolved

try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
return RetryHelper.runWithRetries(
() -> datastoreRpc.commit(requestPb),
Expand Down
10 changes: 9 additions & 1 deletion google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@

import com.google.api.core.BetaApi;
import com.google.cloud.datastore.models.ExplainOptions;
import com.google.cloud.datastore.telemetry.TraceUtil.Context;
import com.google.cloud.datastore.telemetry.TraceUtil.Scope;
import com.google.common.collect.ImmutableList;
import com.google.datastore.v1.ReadOptions;
import com.google.datastore.v1.TransactionOptions;
Expand All @@ -28,6 +30,7 @@
import java.util.Iterator;
import java.util.List;
import java.util.Optional;
import javax.annotation.Nonnull;

final class TransactionImpl extends BaseDatastoreBatchWriter implements Transaction {

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

private final ReadOptionProtoPreparer readOptionProtoPreparer;

@Nonnull private final Context transactionTraceContext;

static class ResponseImpl implements Transaction.Response {

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

transactionId = datastore.requestTransactionId(requestPb);
this.readOptionProtoPreparer = new ReadOptionProtoPreparer();
this.transactionTraceContext = datastore.getOptions().getTraceUtil().currentContext();
}

@Override
Expand All @@ -96,7 +102,9 @@ public Iterator<Entity> get(Key... keys) {
@Override
public List<Entity> fetch(Key... keys) {
validateActive();
return DatastoreHelper.fetch(this, keys);
try (Scope ignored = transactionTraceContext.makeCurrent()) {
return DatastoreHelper.fetch(this, keys);
}
}

@Override
Expand Down
2 changes: 2 additions & 0 deletions google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import com.google.api.core.InternalExtensionOnly;
import com.google.cloud.datastore.DatastoreOptions;
import io.grpc.ManagedChannelBuilder;
import io.opentelemetry.context.Context;
import java.util.Map;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
Expand All @@ -35,6 +36,7 @@ public interface TraceUtil {
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";
Expand Down
54 changes: 47 additions & 7 deletions google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,9 @@
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP;
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.common.truth.Truth.assertThat;
import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME;
import static org.junit.Assert.assertEquals;
Expand Down Expand Up @@ -749,6 +751,7 @@ public void transactionalLookupTest() throws Exception {
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();
Expand All @@ -757,18 +760,55 @@ public void transactionalLookupTest() throws Exception {

fetchAndValidateTrace(
customSpanContext.getTraceId(),
/*numExpectedSpans=*/ 2,
Collections.singletonList(Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION)));
/*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 runInTransactionQueryTest() 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<Entity> entityList = new ArrayList<>();
entityList.add(entity1);
entityList.add(entity2);

List<Entity> 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<Entity> query =
Query.newEntityQueryBuilder().setKind(KEY1.getKind()).setFilter(filter).build();
Datastore.TransactionCallable<Boolean> callable =
transaction -> {
QueryResults<Entity> 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=*/ 2,
Collections.singletonList(Collections.singletonList(SPAN_NAME_TRANSACTION_LOOKUP)));
/*numExpectedSpans=*/ 4,
Arrays.asList(
Collections.singletonList(SPAN_NAME_TRANSACTION_RUN),
Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION),
Collections.singletonList(SPAN_NAME_RUN_QUERY),
Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT)));
}

@Test
public void runInTransactionQueryTest() throws Exception {}

@Test
public void runInTransactionAggregationQueryTest() throws Exception {}

Expand Down
Loading