protect against empty or null tag values (#852)

slackhq · Apr 10, 2024 · bb2b55a · bb2b55a
1 parent 74da443
commit bb2b55a
Show file tree

Hide file tree

Showing 5 changed files with 36 additions and 3 deletions.
diff --git a/astra/src/main/java/com/slack/astra/bulkIngestApi/opensearch/BulkApiRequestParser.java b/astra/src/main/java/com/slack/astra/bulkIngestApi/opensearch/BulkApiRequestParser.java
@@ -139,7 +139,11 @@ public static Trace.Span fromIngestDocument(
       if (!tagsContainServiceName && kv.getKey().equals(SERVICE_NAME_KEY)) {
         tagsContainServiceName = true;
       }
-      spanBuilder.addAllTags(SpanFormatter.convertKVtoProto(kv.getKey(), kv.getValue(), schema));
+      List<Trace.KeyValue> tags =
+          SpanFormatter.convertKVtoProto(kv.getKey(), kv.getValue(), schema);
+      if (tags != null) {
+        spanBuilder.addAllTags(tags);
+      }
     }
     if (!tagsContainServiceName) {
       spanBuilder.addTags(

diff --git a/astra/src/main/java/com/slack/astra/writer/SpanFormatter.java b/astra/src/main/java/com/slack/astra/writer/SpanFormatter.java
@@ -160,6 +160,10 @@ public static Trace.KeyValue makeTraceKV(String key, Object value, Schema.Schema
 
   public static List<Trace.KeyValue> convertKVtoProto(
       String key, Object value, Schema.IngestSchema schema) {
+    if (value == null || value.toString().isEmpty()) {
+      return null;
+    }
+
     if (schema.containsFields(key)) {
       List<Trace.KeyValue> tags = new ArrayList<>();
       Schema.SchemaField schemaFieldDef = schema.getFieldsMap().get(key);
@@ -185,7 +189,7 @@ public static List<Trace.KeyValue> convertKVtoProto(
     }
   }
 
-  public static Trace.KeyValue convertKVtoProto(String key, Object value) {
+  private static Trace.KeyValue convertKVtoProto(String key, Object value) {
     Trace.KeyValue.Builder tagBuilder = Trace.KeyValue.newBuilder();
     tagBuilder.setKey(key);
     if (value instanceof String) {

diff --git a/astra/src/test/java/com/slack/astra/schema/SpanFormatterWithSchemaTest.java b/astra/src/test/java/com/slack/astra/schema/SpanFormatterWithSchemaTest.java
@@ -127,6 +127,28 @@ public static void initializeSchema() {
     schema = Schema.IngestSchema.newBuilder().putAllFields(fields).build();
   }
 
+  @Test
+  public void parseIndexRequestWithNullValues() throws Exception {
+    final File schemaFile =
+        new File(getClass().getClassLoader().getResource("schema/test_schema.yaml").getFile());
+    Schema.IngestSchema schema = SchemaUtil.parseSchema(schemaFile.toPath());
+
+    byte[] rawRequest = getIndexRequestBytes("index_with_null_values");
+    List<IndexRequest> indexRequests = BulkApiRequestParser.parseBulkRequest(rawRequest);
+    AssertionsForClassTypes.assertThat(indexRequests.size()).isEqualTo(1);
+
+    Trace.Span doc =
+        BulkApiRequestParser.fromIngestDocument(
+            convertRequestToDocument(indexRequests.get(0)), schema);
+    assertThat(doc.getId().toStringUtf8()).isEqualTo("1");
+    assertThat(doc.getParentId().toStringUtf8()).isEqualTo("1");
+
+    List<String> docTags =
+        doc.getTagsList().stream().map(Trace.KeyValue::getKey).collect(Collectors.toList());
+    assertThat(docTags.size()).isEqualTo(2); // username,service_name
+    assertThat(docTags.contains("bucket")).isFalse();
+  }
+
   @Test
   public void parseIndexRequestToTraceProtoTest() throws Exception {
     final File schemaFile =
@@ -184,6 +206,7 @@ public void parseIndexRequestToTraceProtoTest() throws Exception {
     assertThat(doc2Tags.contains("service_name")).isEqualTo(true);
     assertThat(doc2Tags.contains("ip")).isEqualTo(true);
     assertThat(doc2Tags.contains("username")).isEqualTo(true);
+    assertThat(doc2Tags.contains("bucket")).isFalse();
     assertThat(doc2.getParentId().toStringUtf8()).isEqualTo("1");
     assertThat(doc2.getTraceId().toStringUtf8()).isEqualTo("2");
     assertThat(doc2.getName()).isEqualTo("check");

diff --git a/astra/src/test/resources/opensearchRequest/bulk/index_all_schema_fields.ndjson b/astra/src/test/resources/opensearchRequest/bulk/index_all_schema_fields.ndjson
@@ -1,6 +1,6 @@
 { "index" : { "_index" : "test", "_id" : "1" } }
 { "host" : "host1", "message" : "foo bar", "ip" : "192.168.1.1", "my_date" : "2014-09-01T12:00:00Z", "success" : true, "cost": 4.0, "amount" : 1.1, "amount_half_float" : 1.2, "value" : 42, "count" : 3, "count_scaled_long" : 80, "count_short" : 10, "bucket" : "20" }
 { "index" : { "_index" : "test", "_id" : "2" } }
-{ "parent_id" : "1", "trace_id" : "2", "name": "check", "ip" :  "::afff:4567:890a", "duration" : 20000, "username": "me" }
+{ "parent_id" : "1", "trace_id" : "2", "name": "check", "ip" :  "::afff:4567:890a", "duration" : 20000, "username": "me", "bucket": null }
 
 
diff --git a/astra/src/test/resources/opensearchRequest/bulk/index_with_null_values.ndjson b/astra/src/test/resources/opensearchRequest/bulk/index_with_null_values.ndjson
@@ -0,0 +1,2 @@
+{ "index" : { "_index" : "test", "_id" : "1" } }
+{ "parent_id" : "1", "username": "me", "bucket": null }