[DSM] Support schema expansion of maps with record values

DataDog · Oct 7, 2024 · 28d57b2 · 28d57b2
1 parent 85d4818
commit 28d57b2
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 7 deletions.
diff --git a/...nstrumentation/avro/src/main/java/datadog/trace/instrumentation/avro/SchemaExtractor.java b/...nstrumentation/avro/src/main/java/datadog/trace/instrumentation/avro/SchemaExtractor.java
@@ -50,11 +50,18 @@ public static boolean extractProperty(
             return false;
           };
         }
-
         break;
       case MAP:
         type = "object";
-        description = "Map type";
+        String keys = "string";
+        String values = getType(field.schema().getValueType().getType().getName());
+        if (values == "record") {
+          values = "#/components/schemas/" + field.schema().getValueType().getFullName();
+          if (!extractSchema(field.schema().getValueType(), builder, depth)) {
+            return false;
+          };
+        }
+        description = "Map type with " + keys + " keys and " + values + " values";
         break;
       case STRING:
         type = "string";

diff --git a/dd-java-agent/instrumentation/avro/src/test/groovy/AvroDatumReaderTest.groovy b/dd-java-agent/instrumentation/avro/src/test/groovy/AvroDatumReaderTest.groovy
@@ -23,8 +23,8 @@ class AvroDatumReaderTest extends AgentTestRunner {
 
     return true
   }
-  String schemaID = "1235632270399345373"
-  String openApiSchemaDef = "{\"components\":{\"schemas\":{\"TestRecord\":{\"properties\":{\"stringField\":{\"type\":\"string\"},\"intField\":{\"format\":\"int32\",\"type\":\"integer\"},\"longField\":{\"format\":\"int64\",\"type\":\"integer\"},\"floatField\":{\"format\":\"float\",\"type\":\"number\"},\"doubleField\":{\"format\":\"double\",\"type\":\"number\"},\"booleanField\":{\"type\":\"boolean\"},\"bytesField\":{\"format\":\"byte\",\"type\":\"string\"},\"nullField\":{\"type\":\"null\"},\"enumField\":{\"enum\":[\"A\",\"B\",\"C\"],\"type\":\"string\"},\"fixedField\":{\"type\":\"string\"},\"recordField\":{\"type\":\"#/components/schemas/NestedRecord\"},\"arrayField\":{\"items\":{\"type\":\"integer\"},\"type\":\"array\"},\"mapField\":{\"description\":\"Map type\",\"type\":\"object\"},\"arrayNestedField\":{\"items\":{\"type\":\"#/components/schemas/OtherNestedRecord\"},\"type\":\"array\"}},\"type\":\"object\"},\"NestedRecord\":{\"properties\":{\"nestedString\":{\"type\":\"string\"}},\"type\":\"object\"},\"OtherNestedRecord\":{\"properties\":{\"nestedString\":{\"type\":\"string\"}},\"type\":\"object\"}}},\"openapi\":\"3.0.0\"}"
+  String schemaID = "8924443781494069161"
+  String openApiSchemaDef = "{\"components\":{\"schemas\":{\"TestRecord\":{\"properties\":{\"stringField\":{\"type\":\"string\"},\"intField\":{\"format\":\"int32\",\"type\":\"integer\"},\"longField\":{\"format\":\"int64\",\"type\":\"integer\"},\"floatField\":{\"format\":\"float\",\"type\":\"number\"},\"doubleField\":{\"format\":\"double\",\"type\":\"number\"},\"booleanField\":{\"type\":\"boolean\"},\"bytesField\":{\"format\":\"byte\",\"type\":\"string\"},\"nullField\":{\"type\":\"null\"},\"enumField\":{\"enum\":[\"A\",\"B\",\"C\"],\"type\":\"string\"},\"fixedField\":{\"type\":\"string\"},\"recordField\":{\"type\":\"#/components/schemas/NestedRecord\"},\"arrayField\":{\"items\":{\"type\":\"integer\"},\"type\":\"array\"},\"mapField\":{\"description\":\"Map type with string keys and string values\",\"type\":\"object\"},\"arrayNestedField\":{\"items\":{\"type\":\"#/components/schemas/OtherNestedRecord\"},\"type\":\"array\"},\"mapNestedField\":{\"description\":\"Map type with string keys and #/components/schemas/ThirdTypeOfNestedRecord values\",\"type\":\"object\"}},\"type\":\"object\"},\"NestedRecord\":{\"properties\":{\"nestedString\":{\"type\":\"string\"}},\"type\":\"object\"},\"OtherNestedRecord\":{\"properties\":{\"nestedString\":{\"type\":\"string\"}},\"type\":\"object\"},\"ThirdTypeOfNestedRecord\":{\"properties\":{\"nestedString\":{\"type\":\"string\"}},\"type\":\"object\"}}},\"openapi\":\"3.0.0\"}"
   String schemaStr = '''
     {
       "type": "record",
@@ -43,7 +43,8 @@ class AvroDatumReaderTest extends AgentTestRunner {
         {"name": "recordField", "type": {"type": "record", "name": "NestedRecord", "fields": [{"name": "nestedString", "type": "string"}]}},
         {"name": "arrayField", "type": {"type": "array", "items": "int"}},
         {"name": "mapField", "type": {"type": "map", "values": "string"}},
-        {"name": "arrayNestedField", "type": { "type": "array", "items": {"type": "record", "name": "OtherNestedRecord", "fields": [{"name": "nestedString", "type": "string"}]}}}
+        {"name": "arrayNestedField", "type": { "type": "array", "items": {"type": "record", "name": "OtherNestedRecord", "fields": [{"name": "nestedString", "type": "string"}]}}},
+        {"name": "mapNestedField", "type": {"type": "map", "values": {"type": "record", "name": "ThirdTypeOfNestedRecord", "fields": [{"name": "nestedString", "type": "string"}]}}}
       ]
     }
     '''
@@ -91,12 +92,19 @@ class AvroDatumReaderTest extends AgentTestRunner {
     datum.put("mapField", map)
 
     // array of nested fields
-    GenericRecord nestedRecordA = new GenericData.Record(nestedSchemaDef)
+    GenericRecord nestedRecordA = new GenericData.Record(schemaDef.getField("arrayNestedField").schema().getElementType())
     nestedRecordA.put("nestedString", "a")
-    GenericRecord nestedRecordB = new GenericData.Record(nestedSchemaDef)
+    GenericRecord nestedRecordB = new GenericData.Record(schemaDef.getField("arrayNestedField").schema().getElementType())
     nestedRecordB.put("nestedString", "b")
     datum.put("arrayNestedField", Arrays.asList(nestedRecordA, nestedRecordB))
 
+    // map of nested fields
+    Map<String, GenericRecord> nestedMap = new HashMap<>()
+    GenericRecord nestedRecordC = new GenericData.Record(schemaDef.getField("mapNestedField").schema().getValueType())
+    nestedRecordC.put("nestedString", "a")
+    nestedMap.put("key1", nestedRecordC)
+    datum.put("mapNestedField", nestedMap)
+
     when:
     def bytes
     ByteArrayOutputStream out = new ByteArrayOutputStream()