Add Agent capability to parse the SketchPayload and return the parsed version from the /api/beta/sketches backend proxy

Author: zacharycmontoya

utils/interfaces/_agent.py

@@ -155,6 +155,20 @@ def get_metrics(self):
             for point in content:
                 yield data, point
 
+    def get_sketches(self):
+        """Attempts to fetch the sketches the agent will submit to the backend."""
+
+        all_data = self.get_data(path_filters="/api/beta/sketches")
+
+        for data in all_data:
+            if "sketches" not in data["request"]["content"]:
+                raise ValueError("sketches property is missing in agent payload")
+
+            content = data["request"]["content"]["sketches"]
+
+            for point in content:
+                yield data, point
+
     def get_dsm_data(self):
         return self.get_data(path_filters="/api/v0.1/pipeline_stats")
 

utils/proxy/_decoders/protobuf_schemas.py

@@ -16,3 +16,4 @@
 
 TracePayload = _messages["datadog.trace.AgentPayload"]
 MetricPayload = _messages["datadog.agentpayload.MetricPayload"]
+SketchPayload = _messages["datadog.agentpayload.SketchPayload"]

utils/proxy/_deserializer.py

@@ -26,7 +26,7 @@
     ExportLogsServiceRequest,
     ExportLogsServiceResponse,
 )
-from _decoders.protobuf_schemas import MetricPayload, TracePayload
+from _decoders.protobuf_schemas import MetricPayload, TracePayload, SketchPayload
 
 
 logger = logging.getLogger(__name__)
@@ -177,6 +177,8 @@ def json_load():
             return result
         if path == "/api/v2/series":
             return MessageToDict(MetricPayload.FromString(content))
+        if path == "/api/beta/sketches":
+            return MessageToDict(SketchPayload.FromString(content))
 
     if content_type == "application/x-www-form-urlencoded" and content == b"[]" and path == "/v0.4/traces":
         return []