From 5c41d41c1ce36456a3778146b2d2f695d61531fa Mon Sep 17 00:00:00 2001
From: Andrea Lin <andrealin@google.com>
Date: Tue, 5 Jun 2018 12:47:56 -0700
Subject: [PATCH] Logging: set Trace in trace instead of label (#3301)

---
 .../cloud/logging/TraceLoggingEnhancer.java   |  4 +--
 .../cloud/logging/LoggingHandlerTest.java     | 32 +++++++++++++++++++
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/google-cloud-clients/google-cloud-logging/src/main/java/com/google/cloud/logging/TraceLoggingEnhancer.java b/google-cloud-clients/google-cloud-logging/src/main/java/com/google/cloud/logging/TraceLoggingEnhancer.java
index d659fefff8ef..c7b7308fa5e8 100644
--- a/google-cloud-clients/google-cloud-logging/src/main/java/com/google/cloud/logging/TraceLoggingEnhancer.java
+++ b/google-cloud-clients/google-cloud-logging/src/main/java/com/google/cloud/logging/TraceLoggingEnhancer.java
@@ -35,7 +35,7 @@ public TraceLoggingEnhancer(String prefix) {
   /**
    * Set the Trace ID associated with any logging done by the current thread.
    *
-   * @param id The traceID
+   * @param id The traceID, in the form projects/[PROJECT_ID]/traces/[TRACE_ID]
    */
   public static void setCurrentTraceId(String id) {
     if (id == null) {
@@ -58,7 +58,7 @@ public static String getCurrentTraceId() {
   public void enhanceLogEntry(com.google.cloud.logging.LogEntry.Builder builder) {
     String traceId = getCurrentTraceId();
     if (traceId != null) {
-      builder.addLabel(traceIdLabel, traceId);
+      builder.setTrace(traceId);
     }
   }
 }
diff --git a/google-cloud-clients/google-cloud-logging/src/test/java/com/google/cloud/logging/LoggingHandlerTest.java b/google-cloud-clients/google-cloud-logging/src/test/java/com/google/cloud/logging/LoggingHandlerTest.java
index 16985a4881df..638e8bf28468 100644
--- a/google-cloud-clients/google-cloud-logging/src/test/java/com/google/cloud/logging/LoggingHandlerTest.java
+++ b/google-cloud-clients/google-cloud-logging/src/test/java/com/google/cloud/logging/LoggingHandlerTest.java
@@ -131,6 +131,13 @@ public class LoggingHandlerTest {
       .addLabel("levelValue", String.valueOf(LoggingLevel.EMERGENCY.intValue()))
       .setTimestamp(123456789L)
       .build();
+  private static final LogEntry TRACE_ENTRY = LogEntry.newBuilder(StringPayload.of(MESSAGE))
+      .setSeverity(Severity.DEBUG)
+      .addLabel("levelName", "FINEST")
+      .addLabel("levelValue", String.valueOf(Level.FINEST.intValue()))
+      .setTrace("projects/projectId/traces/traceId")
+      .setTimestamp(123456789L)
+      .build();
   private static final String CONFIG_NAMESPACE = "com.google.cloud.logging.LoggingHandler";
   private static final ImmutableMap<String, String> CONFIG_MAP =
       ImmutableMap.<String, String>builder()
@@ -314,6 +321,31 @@ public void enhanceLogEntry(Builder builder) {
     handler.publish(newLogRecord(Level.FINEST, MESSAGE));
   }
 
+  @Test
+  public void testTraceEnhancedLogEntry() {
+    expect(options.getProjectId()).andReturn(PROJECT).anyTimes();
+    expect(options.getService()).andReturn(logging);
+    MonitoredResource resource = MonitoredResource.of("custom", ImmutableMap.<String, String>of());
+    logging.setFlushSeverity(Severity.ERROR);
+    expectLastCall().once();
+    logging.setWriteSynchronicity(Synchronicity.ASYNC);
+    expectLastCall().once();
+    logging.write(
+        ImmutableList.of(TRACE_ENTRY),
+        WriteOption.logName(LOG_NAME),
+        WriteOption.resource(resource),
+        WriteOption.labels(BASE_SEVERITY_MAP));
+    expectLastCall().once();
+    replay(options, logging);
+    LoggingEnhancer enhancer = new TraceLoggingEnhancer();
+    TraceLoggingEnhancer.setCurrentTraceId("projects/projectId/traces/traceId");
+    Handler handler =
+        new LoggingHandler(LOG_NAME, options, resource, Collections.singletonList(enhancer));
+    handler.setLevel(Level.ALL);
+    handler.setFormatter(new TestFormatter());
+    handler.publish(newLogRecord(Level.FINEST, MESSAGE));
+  }
+
   @Test
   public void testReportWriteError() {
     expect(options.getProjectId()).andReturn(PROJECT).anyTimes();