diff --git a/dd-java-agent/src/main/java/com/datadoghq/trace/agent/integration/MongoHelper.java b/dd-java-agent/src/main/java/com/datadoghq/trace/agent/integration/MongoHelper.java index 454ccca5a4..1d9bcf21e0 100644 --- a/dd-java-agent/src/main/java/com/datadoghq/trace/agent/integration/MongoHelper.java +++ b/dd-java-agent/src/main/java/com/datadoghq/trace/agent/integration/MongoHelper.java @@ -9,6 +9,7 @@ import java.util.Arrays; import java.util.List; import java.util.Map; import lombok.extern.slf4j.Slf4j; +import org.bson.BsonArray; import org.bson.BsonDocument; import org.bson.BsonString; import org.bson.BsonValue; @@ -18,7 +19,8 @@ import org.jboss.byteman.rule.Rule; @Slf4j public class MongoHelper extends DDAgentTracingHelper { - private static final List WHILDCARD_FIELDS = Arrays.asList("ordered", "insert"); + private static final List WHILDCARD_FIELDS = + Arrays.asList("ordered", "insert", "count", "find"); private static final BsonValue HIDDEN_CAR = new BsonString("?"); public MongoHelper(final Rule rule) { @@ -48,27 +50,45 @@ public class MongoHelper extends DDAgentTracingHelper entry : origin.entrySet()) { if (WHILDCARD_FIELDS.contains(entry.getKey())) { normalized.put(entry.getKey(), entry.getValue()); - } else if (entry.getValue().isDocument()) { - final BsonDocument child = new BsonDocument(); - normalized.put(entry.getKey(), child); - norm(entry.getValue().asDocument(), child); } else { - normalized.put(entry.getKey(), HIDDEN_CAR); + final BsonValue child = norm(entry.getValue()); + normalized.put(entry.getKey(), child); } } + return normalized; + } + + private BsonValue norm(final BsonArray origin) { + final BsonArray normalized = new BsonArray(); + for (final BsonValue value : origin) { + final BsonValue child = norm(value); + normalized.add(child); + } + return normalized; + } + + private BsonValue norm(final BsonValue origin) { + + final BsonValue normalized; + if (origin.isDocument()) { + normalized = norm(origin.asDocument()); + } else if (origin.isArray()) { + normalized = norm(origin.asArray()); + } else { + normalized = HIDDEN_CAR; + } + return normalized; } }