From 29cbe93ecdb987f10e90399d71dacfef6cb2497a Mon Sep 17 00:00:00 2001
From: surbhigarg92 <surbhigarg.92@gmail.com>
Date: Fri, 4 Oct 2024 14:20:16 +0530
Subject: [PATCH] chore: add directpath_enabled attribute

---
 .../cloud/spanner/BuiltInMetricsConstant.java |  1 +
 .../google/cloud/spanner/CompositeTracer.java |  2 +-
 .../cloud/spanner/spi/v1/GapicSpannerRpc.java | 30 +++++++++++--------
 .../spanner/spi/v1/HeaderInterceptor.java     |  3 ++
 4 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BuiltInMetricsConstant.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BuiltInMetricsConstant.java
index 4f8b091d550..f0596a1c20a 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BuiltInMetricsConstant.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BuiltInMetricsConstant.java
@@ -101,6 +101,7 @@ public class BuiltInMetricsConstant {
           DIRECT_PATH_ENABLED_KEY,
           DIRECT_PATH_USED_KEY);
 
+  public static boolean DIRECT_PATH_ENABLED;
   static Aggregation AGGREGATION_WITH_MILLIS_HISTOGRAM =
       Aggregation.explicitBucketHistogram(
           ImmutableList.of(
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/CompositeTracer.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/CompositeTracer.java
index 085a91fb88e..050116af8de 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/CompositeTracer.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/CompositeTracer.java
@@ -184,7 +184,7 @@ public void addAttributes(Map<String, String> attributes) {
     for (ApiTracer child : children) {
       if (child instanceof MetricsTracer) {
         MetricsTracer metricsTracer = (MetricsTracer) child;
-        attributes.forEach((key, value) -> metricsTracer.addAttributes(key, value));
+        metricsTracer.addAttributes(attributes);
       }
     }
   }
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java
index 2360b5d5173..02c6c71e7da 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java
@@ -29,6 +29,7 @@
 import com.google.api.gax.grpc.GrpcCallContext;
 import com.google.api.gax.grpc.GrpcCallSettings;
 import com.google.api.gax.grpc.GrpcStubCallableFactory;
+import com.google.api.gax.grpc.GrpcTransportChannel;
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.api.gax.retrying.ResultRetryAlgorithm;
@@ -63,6 +64,7 @@
 import com.google.cloud.grpc.GrpcTransportOptions;
 import com.google.cloud.spanner.AdminRequestsPerMinuteExceededException;
 import com.google.cloud.spanner.BackupId;
+import com.google.cloud.spanner.BuiltInMetricsConstant;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.Restore;
 import com.google.cloud.spanner.SpannerException;
@@ -396,18 +398,22 @@ public GapicSpannerRpc(final SpannerOptions options) {
       final String emulatorHost = System.getenv("SPANNER_EMULATOR_HOST");
 
       try {
-        this.spannerStub =
-            GrpcSpannerStub.create(
-                options
-                    .getSpannerStubSettings()
-                    .toBuilder()
-                    .setTransportChannelProvider(channelProvider)
-                    .setCredentialsProvider(credentialsProvider)
-                    .setStreamWatchdogProvider(watchdogProvider)
-                    .setTracerFactory(
-                        options.getApiTracerFactory(
-                            /* isAdminClient = */ false, isEmulatorEnabled(options, emulatorHost)))
-                    .build());
+        SpannerStubSettings spannerStubSettings =
+            options
+                .getSpannerStubSettings()
+                .toBuilder()
+                .setTransportChannelProvider(channelProvider)
+                .setCredentialsProvider(credentialsProvider)
+                .setStreamWatchdogProvider(watchdogProvider)
+                .setTracerFactory(
+                    options.getApiTracerFactory(
+                        /* isAdminClient = */ false, isEmulatorEnabled(options, emulatorHost)))
+                .build();
+        ClientContext clientContext = ClientContext.create(spannerStubSettings);
+        this.spannerStub = GrpcSpannerStub.create((clientContext));
+        BuiltInMetricsConstant.DIRECT_PATH_ENABLED =
+            ((GrpcTransportChannel) clientContext.getTransportChannel()).isDirectPath()
+                && isAttemptDirectPathXds;
         this.readRetrySettings =
             options.getSpannerStubSettings().streamingReadSettings().getRetrySettings();
         this.readRetryableCodes =
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java
index dd414bed397..1bf41ee2523 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java
@@ -228,6 +228,9 @@ private Map<String, String> getBuiltInMetricAttributes(String key, DatabaseName
           attributes.put(BuiltInMetricsConstant.DATABASE_KEY.getKey(), databaseName.getDatabase());
           attributes.put(
               BuiltInMetricsConstant.INSTANCE_ID_KEY.getKey(), databaseName.getInstance());
+          attributes.put(
+              BuiltInMetricsConstant.DIRECT_PATH_ENABLED_KEY.getKey(),
+              String.valueOf(BuiltInMetricsConstant.DIRECT_PATH_ENABLED));
           return attributes;
         });
   }