From 522f78c49ffbd0b72e82f1a7d31f2ae2ad4222b7 Mon Sep 17 00:00:00 2001
From: Shivansh Anand Srivastava <shivansh.srivastava@traceable.ai>
Date: Wed, 20 Sep 2023 15:45:54 +0530
Subject: [PATCH] Expose resource string attribute from SpanAttributeUtils
 (#44)

* Expose resource string attribute from SpanAttributeUtils

* Remove javax nullable

* Remove javax nullable
---
 .../datamodel/shared/SpanAttributeUtils.java  | 21 ++++++++++++++
 .../shared/SpanAttributeUtilsTest.java        | 29 +++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/data-model/src/main/java/org/hypertrace/core/datamodel/shared/SpanAttributeUtils.java b/data-model/src/main/java/org/hypertrace/core/datamodel/shared/SpanAttributeUtils.java
index 0adcbd4..47cfcc9 100644
--- a/data-model/src/main/java/org/hypertrace/core/datamodel/shared/SpanAttributeUtils.java
+++ b/data-model/src/main/java/org/hypertrace/core/datamodel/shared/SpanAttributeUtils.java
@@ -19,6 +19,8 @@
 import org.hypertrace.core.datamodel.AttributeValue;
 import org.hypertrace.core.datamodel.Attributes;
 import org.hypertrace.core.datamodel.Event;
+import org.hypertrace.core.datamodel.Resource;
+import org.hypertrace.core.datamodel.StructuredTrace;
 
 /**
  * Span being a very generic data structure with different attributes, reading attributes from Span
@@ -82,6 +84,25 @@ public static String getStringAttribute(Event event, String attributeKey) {
     return value == null ? null : value.getValue();
   }
 
+  public static Optional<String> getResourceStringAttribute(
+      StructuredTrace trace, Event event, String attributeKey) {
+    if (event.getResourceIndex() < 0
+        || event.getResourceIndex() >= trace.getResourceList().size()) {
+      return Optional.empty();
+    }
+
+    return Optional.of(trace.getResourceList().get(event.getResourceIndex()))
+        .map(Resource::getAttributes)
+        .flatMap(attributes -> getAttributeString(attributes, attributeKey));
+  }
+
+  private static Optional<String> getAttributeString(@Nullable Attributes attributes, String key) {
+    return Optional.ofNullable(attributes)
+        .map(Attributes::getAttributeMap)
+        .map(attributeMap -> attributeMap.get(key))
+        .map(AttributeValue::getValue);
+  }
+
   public static Optional<String> getStringAttributeIgnoreKeyCase(Event event, String attributeKey) {
     if (event == null) {
       return Optional.empty();
diff --git a/data-model/src/test/java/org/hypertrace/core/datamodel/shared/SpanAttributeUtilsTest.java b/data-model/src/test/java/org/hypertrace/core/datamodel/shared/SpanAttributeUtilsTest.java
index 1c75504..7857850 100644
--- a/data-model/src/test/java/org/hypertrace/core/datamodel/shared/SpanAttributeUtilsTest.java
+++ b/data-model/src/test/java/org/hypertrace/core/datamodel/shared/SpanAttributeUtilsTest.java
@@ -2,14 +2,20 @@
 
 import static java.util.Collections.emptyList;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import org.hypertrace.core.datamodel.AttributeValue;
 import org.hypertrace.core.datamodel.Attributes;
 import org.hypertrace.core.datamodel.Event;
+import org.hypertrace.core.datamodel.Resource;
+import org.hypertrace.core.datamodel.StructuredTrace;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
@@ -63,4 +69,27 @@ void testGetStringListAttribute_multiValuedList() {
     final List<String> result = SpanAttributeUtils.getStringListAttribute(event, "attributeKey");
     assertEquals(List.of("value1", "value2", "3"), result);
   }
+
+  @Test
+  void testGetResourceStringAttribute() {
+    final Event event = mock(Event.class);
+    final StructuredTrace trace = mock(StructuredTrace.class);
+    Resource resource =
+        Resource.newBuilder()
+            .setAttributes(
+                Attributes.newBuilder()
+                    .setAttributeMap(
+                        Map.of("key1", AttributeValue.newBuilder().setValue("value1").build()))
+                    .build())
+            .build();
+    when(event.getResourceIndex()).thenReturn(0);
+    when(trace.getResourceList()).thenReturn(Collections.singletonList(resource));
+    Optional<String> valueAbsent =
+        SpanAttributeUtils.getResourceStringAttribute(trace, event, "key2");
+    Optional<String> valuePresent =
+        SpanAttributeUtils.getResourceStringAttribute(trace, event, "key1");
+    assertFalse(valueAbsent.isPresent());
+    assertTrue(valuePresent.isPresent());
+    assertEquals("value1", valuePresent.get());
+  }
 }