Skip to content

Commit 5ea099f

Browse files
lqiu96lqiu96
committed
chore: Clean up code using transaction
1 parent 410dbbc commit 5ea099f

6 files changed

Lines changed: 55 additions & 68 deletions

File tree

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -176,12 +176,12 @@ static String extractStatusCode(Throwable throwable) {
176176
while (current != null) {
177177
if (current instanceof DatastoreException) {
178178
String reason = ((DatastoreException) current).getReason();
179-
if (reason != null && !reason.isEmpty()) {
179+
if (!Strings.isNullOrEmpty(reason)) {
180180
return reason;
181181
}
182182
}
183183
current = current.getCause();
184184
}
185-
return "UNKNOWN";
185+
return StatusCode.Code.UNKNOWN.toString();
186186
}
187187
}

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

Lines changed: 9 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,8 @@
3737
import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN_QUERY;
3838

3939
import com.google.api.core.BetaApi;
40-
import com.google.api.gax.grpc.GrpcStatusCode;
4140
import com.google.api.gax.retrying.RetrySettings;
41+
import com.google.api.gax.rpc.StatusCode;
4242
import com.google.cloud.BaseService;
4343
import com.google.cloud.ExceptionHandler;
4444
import com.google.cloud.RetryHelper;
@@ -241,24 +241,9 @@ private void recordAttempt(String status) {
241241
Map<String, String> attributes = new HashMap<>();
242242
attributes.put(TelemetryConstants.ATTRIBUTES_KEY_STATUS, status);
243243
attributes.put(
244-
TelemetryConstants.ATTRIBUTES_KEY_METHOD_NAME,
245-
TelemetryConstants.METHOD_COMMIT);
244+
TelemetryConstants.ATTRIBUTES_KEY_METHOD, TelemetryConstants.METHOD_TRANSACTION_COMMIT);
246245
metricsRecorder.recordTransactionAttemptCount(1, attributes);
247246
}
248-
249-
/**
250-
* Extracts the gRPC status code from the given exception. Falls back to "UNKNOWN" if the
251-
* status cannot be determined.
252-
*/
253-
private static String extractStatus(Exception ex) {
254-
if (ex instanceof DatastoreException) {
255-
String reason = ((DatastoreException) ex).getReason();
256-
if (reason != null && !reason.isEmpty()) {
257-
return reason;
258-
}
259-
}
260-
return "UNKNOWN";
261-
}
262247
}
263248

264249
@Override
@@ -274,12 +259,13 @@ public <T> T runInTransaction(
274259

275260
ReadWriteTransactionCallable<T> baseCallable =
276261
new ReadWriteTransactionCallable<>(this, callable, transactionOptions, metricsRecorder);
277-
Callable<T> transactionCallable =
278-
getOptions().getOpenTelemetryOptions().isTracingEnabled()
279-
? new TracedReadWriteTransactionCallable<>(baseCallable, span)
280-
: baseCallable;
281262

282-
String status = "OK";
263+
Callable<T> transactionCallable = baseCallable;
264+
if (getOptions().getOpenTelemetryOptions().isTracingEnabled()) {
265+
transactionCallable = new TracedReadWriteTransactionCallable<>(baseCallable, span);
266+
}
267+
268+
String status = StatusCode.Code.OK.toString();
283269
try (Scope ignored = span.makeCurrent()) {
284270
return RetryHelper.runWithRetries(
285271
transactionCallable,
@@ -295,14 +281,12 @@ public <T> T runInTransaction(
295281
Map<String, String> attributes = new HashMap<>();
296282
attributes.put(TelemetryConstants.ATTRIBUTES_KEY_STATUS, status);
297283
attributes.put(
298-
TelemetryConstants.ATTRIBUTES_KEY_METHOD_NAME,
299-
TelemetryConstants.METHOD_TRANSACTION_RUN);
284+
TelemetryConstants.ATTRIBUTES_KEY_METHOD, TelemetryConstants.METHOD_TRANSACTION_COMMIT);
300285
metricsRecorder.recordTransactionLatency(latencyMs, attributes);
301286
span.end();
302287
}
303288
}
304289

305-
306290
@Override
307291
public <T> QueryResults<T> run(Query<T> query) {
308292
return run(Optional.empty(), query, null);

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ public class TelemetryConstants {
3737
public static final String ATTRIBUTES_KEY_STATUS = "status";
3838

3939
/** Attribute key for the RPC method name (e.g. "Transaction.Run"). */
40-
public static final String ATTRIBUTES_KEY_METHOD_NAME = "method_name";
40+
public static final String ATTRIBUTES_KEY_METHOD = "method";
4141

4242
/* TODO(lawrenceqiu): For now, these are a duplicate of method names in TraceUtil. Those will use these eventually */
4343
// Format is not SnakeCase to keep backward compatibility with the existing values TraceUtil spans

‎java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreImplMetricsTest.java‎

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import com.google.cloud.ServiceOptions;
2323
import com.google.cloud.datastore.spi.DatastoreRpcFactory;
2424
import com.google.cloud.datastore.spi.v1.DatastoreRpc;
25+
import com.google.cloud.datastore.telemetry.TelemetryConstants;
2526
import com.google.datastore.v1.BeginTransactionRequest;
2627
import com.google.datastore.v1.BeginTransactionResponse;
2728
import com.google.datastore.v1.CommitRequest;
@@ -53,8 +54,7 @@
5354
public class DatastoreImplMetricsTest {
5455

5556
private static final String PROJECT_ID = "test-project";
56-
private static final String LATENCY_METRIC_NAME =
57-
"datastore.googleapis.com/transaction_latency";
57+
private static final String LATENCY_METRIC_NAME = "datastore.googleapis.com/transaction_latency";
5858
private static final String ATTEMPT_COUNT_METRIC_NAME =
5959
"datastore.googleapis.com/transaction_attempt_count";
6060

@@ -112,8 +112,8 @@ public void runInTransaction_recordsLatencyOnSuccess() {
112112
assertThat(point.getCount()).isEqualTo(1);
113113
assertThat(point.getSum()).isAtLeast(0.0);
114114
assertThat(point.getAttributes().get(AttributeKey.stringKey("status"))).isEqualTo("OK");
115-
assertThat(point.getAttributes().get(AttributeKey.stringKey("method_name")))
116-
.isEqualTo("Transaction.Run");
115+
assertThat(point.getAttributes().get(AttributeKey.stringKey("method")))
116+
.isEqualTo(TelemetryConstants.METHOD_TRANSACTION_COMMIT);
117117

118118
EasyMock.verify(rpcFactoryMock, rpcMock);
119119
}
@@ -141,8 +141,8 @@ public void runInTransaction_recordsPerAttemptCountOnSuccess() {
141141
.orElse(null);
142142
assertThat(point).isNotNull();
143143
assertThat(point.getValue()).isEqualTo(1);
144-
assertThat(point.getAttributes().get(AttributeKey.stringKey("method_name")))
145-
.isEqualTo("Commit");
144+
assertThat(point.getAttributes().get(AttributeKey.stringKey("method")))
145+
.isEqualTo(TelemetryConstants.METHOD_TRANSACTION_COMMIT);
146146

147147
EasyMock.verify(rpcFactoryMock, rpcMock);
148148
}
@@ -194,8 +194,7 @@ public Integer run(DatastoreReaderWriter transaction) {
194194
// Verify ABORTED attempt
195195
LongPointData abortedPoint =
196196
attemptMetric.get().getLongSumData().getPoints().stream()
197-
.filter(
198-
p -> "ABORTED".equals(p.getAttributes().get(AttributeKey.stringKey("status"))))
197+
.filter(p -> "ABORTED".equals(p.getAttributes().get(AttributeKey.stringKey("status"))))
199198
.findFirst()
200199
.orElse(null);
201200
assertThat(abortedPoint).isNotNull();
@@ -225,8 +224,7 @@ public Integer run(DatastoreReaderWriter transaction) {
225224
public void runInTransaction_recordsGrpcStatusCodeOnFailure() {
226225
// This test uses a separate set of nice mocks since the retry loop makes
227226
// multiple begin/rollback calls whose exact count depends on retry settings.
228-
DatastoreRpcFactory localRpcFactoryMock =
229-
EasyMock.createNiceMock(DatastoreRpcFactory.class);
227+
DatastoreRpcFactory localRpcFactoryMock = EasyMock.createNiceMock(DatastoreRpcFactory.class);
230228
DatastoreRpc localRpcMock = EasyMock.createNiceMock(DatastoreRpc.class);
231229

232230
InMemoryMetricReader localMetricReader = InMemoryMetricReader.create();
@@ -276,8 +274,7 @@ public void runInTransaction_recordsGrpcStatusCodeOnFailure() {
276274

277275
LongPointData abortedPoint =
278276
attemptMetric.get().getLongSumData().getPoints().stream()
279-
.filter(
280-
p -> "ABORTED".equals(p.getAttributes().get(AttributeKey.stringKey("status"))))
277+
.filter(p -> "ABORTED".equals(p.getAttributes().get(AttributeKey.stringKey("status"))))
281278
.findFirst()
282279
.orElse(null);
283280
assertThat(abortedPoint).isNotNull();

‎java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/MetricsRecorderTest.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,8 +40,7 @@ public class MetricsRecorderTest {
4040

4141
@Test
4242
public void defaultOptionsReturnNoOp() {
43-
DatastoreOpenTelemetryOptions options =
44-
DatastoreOpenTelemetryOptions.newBuilder().build();
43+
DatastoreOpenTelemetryOptions options = DatastoreOpenTelemetryOptions.newBuilder().build();
4544

4645
MetricsRecorder recorder = MetricsRecorder.getInstance(options);
4746

‎java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/OpenTelemetryMetricsRecorderTest.java‎

Lines changed: 33 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
import static com.google.common.truth.Truth.assertThat;
1919

20+
import com.google.api.gax.rpc.StatusCode;
2021
import io.opentelemetry.api.OpenTelemetry;
2122
import io.opentelemetry.api.common.AttributeKey;
2223
import io.opentelemetry.sdk.OpenTelemetrySdk;
@@ -52,21 +53,17 @@ public void setUp() {
5253
@Test
5354
public void recordTransactionLatency_recordsHistogramWithAttributes() {
5455
Map<String, String> attributes = new HashMap<>();
55-
attributes.put(TelemetryConstants.ATTRIBUTES_KEY_STATUS, "OK");
56+
attributes.put(TelemetryConstants.ATTRIBUTES_KEY_STATUS, StatusCode.Code.OK.toString());
5657
attributes.put(
57-
TelemetryConstants.ATTRIBUTES_KEY_METHOD_NAME,
58-
TelemetryConstants.METHOD_TRANSACTION_RUN);
58+
TelemetryConstants.ATTRIBUTES_KEY_METHOD, TelemetryConstants.METHOD_TRANSACTION_COMMIT);
5959

6060
recorder.recordTransactionLatency(150.0, attributes);
6161

6262
Collection<MetricData> metrics = metricReader.collectAllMetrics();
6363
MetricData latencyMetric =
6464
metrics.stream()
6565
.filter(
66-
m ->
67-
m.getName()
68-
.equals(
69-
TelemetryConstants.SERVICE_NAME + "/transaction_latency"))
66+
m -> m.getName().equals(TelemetryConstants.SERVICE_NAME + "/transaction_latency"))
7067
.findFirst()
7168
.orElse(null);
7269

@@ -80,18 +77,24 @@ public void recordTransactionLatency_recordsHistogramWithAttributes() {
8077
assertThat(point).isNotNull();
8178
assertThat(point.getSum()).isEqualTo(150.0);
8279
assertThat(point.getCount()).isEqualTo(1);
83-
assertThat(point.getAttributes().get(AttributeKey.stringKey("status"))).isEqualTo("OK");
84-
assertThat(point.getAttributes().get(AttributeKey.stringKey("method_name")))
85-
.isEqualTo("Transaction.Run");
80+
assertThat(
81+
point
82+
.getAttributes()
83+
.get(AttributeKey.stringKey(TelemetryConstants.ATTRIBUTES_KEY_STATUS)))
84+
.isEqualTo(StatusCode.Code.OK.toString());
85+
assertThat(
86+
point
87+
.getAttributes()
88+
.get(AttributeKey.stringKey(TelemetryConstants.ATTRIBUTES_KEY_METHOD)))
89+
.isEqualTo(TelemetryConstants.METHOD_TRANSACTION_COMMIT);
8690
}
8791

8892
@Test
8993
public void recordTransactionAttemptCount_recordsCounterWithAttributes() {
9094
Map<String, String> attributes = new HashMap<>();
91-
attributes.put(TelemetryConstants.ATTRIBUTES_KEY_STATUS, "ABORTED");
95+
attributes.put(TelemetryConstants.ATTRIBUTES_KEY_STATUS, StatusCode.Code.ABORTED.toString());
9296
attributes.put(
93-
TelemetryConstants.ATTRIBUTES_KEY_METHOD_NAME,
94-
TelemetryConstants.METHOD_COMMIT);
97+
TelemetryConstants.ATTRIBUTES_KEY_METHOD, TelemetryConstants.METHOD_TRANSACTION_COMMIT);
9598

9699
recorder.recordTransactionAttemptCount(1, attributes);
97100

@@ -101,8 +104,7 @@ public void recordTransactionAttemptCount_recordsCounterWithAttributes() {
101104
.filter(
102105
m ->
103106
m.getName()
104-
.equals(
105-
TelemetryConstants.SERVICE_NAME + "/transaction_attempt_count"))
107+
.equals(TelemetryConstants.SERVICE_NAME + "/transaction_attempt_count"))
106108
.findFirst()
107109
.orElse(null);
108110

@@ -114,24 +116,30 @@ public void recordTransactionAttemptCount_recordsCounterWithAttributes() {
114116
attemptMetric.getLongSumData().getPoints().stream().findFirst().orElse(null);
115117
assertThat(point).isNotNull();
116118
assertThat(point.getValue()).isEqualTo(1);
117-
assertThat(point.getAttributes().get(AttributeKey.stringKey("status"))).isEqualTo("ABORTED");
118-
assertThat(point.getAttributes().get(AttributeKey.stringKey("method_name")))
119-
.isEqualTo("Commit");
119+
assertThat(
120+
point
121+
.getAttributes()
122+
.get(AttributeKey.stringKey(TelemetryConstants.ATTRIBUTES_KEY_STATUS)))
123+
.isEqualTo(StatusCode.Code.ABORTED.toString());
124+
assertThat(
125+
point
126+
.getAttributes()
127+
.get(AttributeKey.stringKey(TelemetryConstants.ATTRIBUTES_KEY_METHOD)))
128+
.isEqualTo(TelemetryConstants.METHOD_TRANSACTION_COMMIT);
120129
}
121130

122131
@Test
123132
public void recordTransactionAttemptCount_multipleAttempts_accumulates() {
124133
Map<String, String> abortedAttributes = new HashMap<>();
125-
abortedAttributes.put(TelemetryConstants.ATTRIBUTES_KEY_STATUS, "ABORTED");
126134
abortedAttributes.put(
127-
TelemetryConstants.ATTRIBUTES_KEY_METHOD_NAME,
128-
TelemetryConstants.METHOD_COMMIT);
135+
TelemetryConstants.ATTRIBUTES_KEY_STATUS, StatusCode.Code.ABORTED.toString());
136+
abortedAttributes.put(
137+
TelemetryConstants.ATTRIBUTES_KEY_METHOD, TelemetryConstants.METHOD_TRANSACTION_COMMIT);
129138

130139
Map<String, String> okAttributes = new HashMap<>();
131-
okAttributes.put(TelemetryConstants.ATTRIBUTES_KEY_STATUS, "OK");
140+
okAttributes.put(TelemetryConstants.ATTRIBUTES_KEY_STATUS, StatusCode.Code.OK.toString());
132141
okAttributes.put(
133-
TelemetryConstants.ATTRIBUTES_KEY_METHOD_NAME,
134-
TelemetryConstants.METHOD_COMMIT);
142+
TelemetryConstants.ATTRIBUTES_KEY_METHOD, TelemetryConstants.METHOD_TRANSACTION_COMMIT);
135143

136144
// Simulate a retry scenario: first attempt ABORTED, second attempt OK
137145
recorder.recordTransactionAttemptCount(1, abortedAttributes);
@@ -143,8 +151,7 @@ public void recordTransactionAttemptCount_multipleAttempts_accumulates() {
143151
.filter(
144152
m ->
145153
m.getName()
146-
.equals(
147-
TelemetryConstants.SERVICE_NAME + "/transaction_attempt_count"))
154+
.equals(TelemetryConstants.SERVICE_NAME + "/transaction_attempt_count"))
148155
.findFirst()
149156
.orElse(null);
150157

0 commit comments

Comments
 (0)