added mechanism to set log group via env var

Author: William Armiros

aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/AWSXRayRecorderBuilder.java

@@ -18,6 +18,7 @@
 import com.amazonaws.xray.contexts.SegmentContextResolverChain;
 import com.amazonaws.xray.emitters.Emitter;
 import com.amazonaws.xray.entities.AWSLogReference;
+import com.amazonaws.xray.entities.StringValidator;
 import com.amazonaws.xray.listeners.SegmentListener;
 import com.amazonaws.xray.plugins.EC2Plugin;
 import com.amazonaws.xray.plugins.ECSPlugin;
@@ -50,6 +51,7 @@ public class AWSXRayRecorderBuilder {
         LogFactory.getLog(AWSXRayRecorderBuilder.class);
 
     private static final Map<String, Integer> ORIGIN_PRIORITY;
+    private static final String LOG_GROUP_KEY = "AWS_LOG_GROUP";
 
     static {
         HashMap<String, Integer> originPriority = new HashMap<>();
@@ -351,6 +353,16 @@ public AWSXRayRecorder build() {
             }
         });
 
+        String logGroupFromEnv = System.getenv(LOG_GROUP_KEY);
+        if (StringValidator.isNotNullOrBlank(logGroupFromEnv)) {
+            logger.info("Recording log group " + logGroupFromEnv + " from environment variable.");
+            AWSLogReference logReference = new AWSLogReference();
+            logReference.setLogGroup(logGroupFromEnv);
+            Set<AWSLogReference> logReferences = new HashSet<>();
+            logReferences.add(logReference);
+            client.addAllLogReferences(logReferences);
+        }
+
         return client;
     }
 }

aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/AWSXRayRecorderTest.java

@@ -24,6 +24,7 @@
 import com.amazonaws.xray.contexts.LambdaSegmentContextResolver;
 import com.amazonaws.xray.contexts.SegmentContextResolverChain;
 import com.amazonaws.xray.emitters.Emitter;
+import com.amazonaws.xray.entities.AWSLogReference;
 import com.amazonaws.xray.entities.Segment;
 import com.amazonaws.xray.entities.Subsegment;
 import com.amazonaws.xray.entities.TraceHeader;
@@ -51,6 +52,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.locks.ReentrantLock;
@@ -917,6 +919,18 @@ public void testBeginSegmentWithForcedSampling() {
 
         segment.setUser("user");
         assertThat(segment.getUser()).isEqualTo("user");  // Loose way to test that segment is real
+    }
+
+    @Test
+    public void testLogGroupFromEnvironment() {
+        environmentVariables.set("AWS_LOG_GROUP", "my-group");
+        AWSXRayRecorder recorder = AWSXRayRecorderBuilder.standard().build();
+        Segment segment = recorder.beginSegment("test");
+        AWSLogReference expected = new AWSLogReference();
+        expected.setLogGroup("my-group");
 
+        assertThat(segment.getAws()).containsKey("cloudwatch_logs");
+        Set<AWSLogReference> logReferences = (Set<AWSLogReference>) segment.getAws().get("cloudwatch_logs");
+        assertThat(logReferences).containsOnly(expected);
     }
 }