feat(datastore): add database_id to metric attributes and fix IT tests

Author: lqiu96

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

@@ -254,7 +254,7 @@ public T call() throws DatastoreException {
     private void recordAttempt(String status) {
       Map<String, String> attributes =
           TelemetryUtils.buildMetricAttributes(
-              TelemetryConstants.METHOD_TRANSACTION_COMMIT, status);
+              TelemetryConstants.METHOD_TRANSACTION_COMMIT, status, datastore.getOptions().getDatabaseId());
       metricsRecorder.recordTransactionAttemptCount(1, attributes);
     }
   }
@@ -292,7 +292,7 @@ public <T> T runInTransaction(
     } finally {
       long latencyMs = stopwatch.elapsed(TimeUnit.MILLISECONDS);
       Map<String, String> attributes =
-          TelemetryUtils.buildMetricAttributes(TelemetryConstants.METHOD_TRANSACTION_RUN, status);
+          TelemetryUtils.buildMetricAttributes(TelemetryConstants.METHOD_TRANSACTION_RUN, status, getOptions().getDatabaseId());
       metricsRecorder.recordTransactionLatency(latencyMs, attributes);
       span.end();
     }
@@ -801,7 +801,7 @@ private <T> T runWithObservability(
 
     DatastoreOptions options = getOptions();
     Callable<T> attemptCallable =
-        TelemetryUtils.attemptMetricsCallable(callable, metricsRecorder, methodName);
+        TelemetryUtils.attemptMetricsCallable(callable, metricsRecorder, methodName, options.getDatabaseId());
     try (TraceUtil.Scope ignored = span.makeCurrent()) {
       return RetryHelper.runWithRetries(
           attemptCallable, retrySettings, exceptionHandler, options.getClock());
@@ -811,7 +811,7 @@ private <T> T runWithObservability(
       throw DatastoreException.translateAndThrow(e);
     } finally {
       TelemetryUtils.recordOperationMetrics(
-          metricsRecorder, operationStopwatch, methodName, operationStatus);
+          metricsRecorder, operationStopwatch, methodName, operationStatus, options.getDatabaseId());
       span.end();
     }
   }

java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java

@@ -205,7 +205,7 @@ <O> O invokeRpc(Callable<O> block, String startSpan, String methodName) {
     String operationStatus = StatusCode.Code.UNKNOWN.toString();
     try (TraceUtil.Scope ignored = span.makeCurrent()) {
       Callable<O> callable =
-          TelemetryUtils.attemptMetricsCallable(block, metricsRecorder, methodName);
+          TelemetryUtils.attemptMetricsCallable(block, metricsRecorder, methodName, datastoreOptions.getDatabaseId());
       O result =
           RetryHelper.runWithRetries(
               callable, this.retrySettings, EXCEPTION_HANDLER, this.datastoreOptions.getClock());
@@ -217,7 +217,7 @@ <O> O invokeRpc(Callable<O> block, String startSpan, String methodName) {
       throw DatastoreException.translateAndThrow(e);
     } finally {
       TelemetryUtils.recordOperationMetrics(
-          metricsRecorder, stopwatch, methodName, operationStatus);
+          metricsRecorder, stopwatch, methodName, operationStatus, datastoreOptions.getDatabaseId());
       span.end();
     }
   }

java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TelemetryUtils.java

@@ -42,11 +42,14 @@ private TelemetryUtils() {}
    * @param status The status of the operation or attempt.
    * @return The map of attributes.
    */
-  public static Map<String, String> buildMetricAttributes(String methodName, String status) {
+  public static Map<String, String> buildMetricAttributes(String methodName, String status, String databaseId) {
     Map<String, String> attributes = new HashMap<>();
     attributes.put(TelemetryConstants.ATTRIBUTES_KEY_METHOD, methodName);
     attributes.put(TelemetryConstants.ATTRIBUTES_KEY_STATUS, status);
     attributes.put(TelemetryConstants.ATTRIBUTES_KEY_SERVICE, TelemetryConstants.SERVICE_VALUE);
+    if (databaseId != null) {
+      attributes.put(TelemetryConstants.ATTRIBUTES_KEY_DATABASE_ID, databaseId);
+    }
     return attributes;
   }
 
@@ -65,9 +68,10 @@ public static void recordOperationMetrics(
       DatastoreMetricsRecorder metricsRecorder,
       Stopwatch operationStopwatch,
       String methodName,
-      String status) {
+      String status,
+      String databaseId) {
     if (methodName != null) {
-      Map<String, String> attributes = buildMetricAttributes(methodName, status);
+      Map<String, String> attributes = buildMetricAttributes(methodName, status, databaseId);
       metricsRecorder.recordOperationLatency(
           operationStopwatch.elapsed(TimeUnit.MILLISECONDS), attributes);
       metricsRecorder.recordOperationCount(1, attributes);
@@ -87,7 +91,7 @@ public static void recordOperationMetrics(
    * @return A wrapped callable that includes attempt-level metrics recording.
    */
   public static <T> Callable<T> attemptMetricsCallable(
-      Callable<T> callable, DatastoreMetricsRecorder metricsRecorder, String methodName) {
+      Callable<T> callable, DatastoreMetricsRecorder metricsRecorder, String methodName, String databaseId) {
     return () -> {
       Stopwatch stopwatch = Stopwatch.createStarted();
       String status = StatusCode.Code.UNKNOWN.toString();
@@ -99,7 +103,7 @@ public static <T> Callable<T> attemptMetricsCallable(
         status = DatastoreException.extractStatusCode(e);
         throw e;
       } finally {
-        Map<String, String> attributes = buildMetricAttributes(methodName, status);
+        Map<String, String> attributes = buildMetricAttributes(methodName, status, databaseId);
         metricsRecorder.recordAttemptLatency(stopwatch.elapsed(TimeUnit.MILLISECONDS), attributes);
         metricsRecorder.recordAttemptCount(1, attributes);
       }

java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ITDatastoreBuiltInAndCustomMetrics.java

@@ -182,9 +182,11 @@ public void bothBackendsActive_recorderIsComposite() {
    */
   @Test
   public void builtInMetricsExport_isDisabledByDefault() {
+    DatastoreOptions defaultOptions = DatastoreOptions.newBuilder()
+            .setProjectId(PROJECT_ID)
+            .build();
     assertThat(
-            datastore
-                .getOptions()
+            defaultOptions
                 .getOpenTelemetryOptions()
                 .isExportBuiltinMetricsToGoogleCloudMonitoring())
         .isFalse();