chore(debugging): ensure correct log probe configuration handling (#5143

)

This change ensures that the RCM adapter for dynamic instrumentation can
parse the payload attributes correctly.

ddtrace/debugging/_probe/model.py

@@ -214,3 +214,8 @@ class LogFunctionProbe(Probe, FunctionLocationMixin, LogProbeMixin, ProbeConditi
 
 LineProbe = Union[LogLineProbe, MetricLineProbe]
 FunctionProbe = Union[LogFunctionProbe, MetricFunctionProbe]
+
+
+class ProbeType(object):
+    LOG_PROBE = "LOG_PROBE"
+    METRIC_PROBE = "METRIC_PROBE"

ddtrace/debugging/_probe/remoteconfig.py

@@ -1,4 +1,3 @@
-from itertools import chain
 import time
 from typing import Any
 from typing import Callable
@@ -22,6 +21,7 @@
 from ddtrace.debugging._probe.model import MetricFunctionProbe
 from ddtrace.debugging._probe.model import MetricLineProbe
 from ddtrace.debugging._probe.model import Probe
+from ddtrace.debugging._probe.model import ProbeType
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.remoteconfig.client import ConfigMetadata
 from ddtrace.internal.utils.cache import LFUCache
@@ -33,6 +33,11 @@
 _EXPRESSION_CACHE = LFUCache()
 
 
+def xlate_keys(d, mapping):
+    # type: (Dict[str, Any], Dict[str, str]) -> Dict[str, Any]
+    return {mapping.get(k, k): v for k, v in d.items()}
+
+
 def _invalid_expression(_):
     """Forces probes with invalid expression/conditions to never trigger.
 
@@ -74,7 +79,7 @@ def compile_or_invalid(expr):
 
 def _compile_segment(segment):
     if segment.get("str", ""):
-        return LiteralTemplateSegment(str=segment["str"])
+        return LiteralTemplateSegment(str_value=segment["str"])
     elif segment.get("json", None) is not None:
         return ExpressionTemplateSegment(expr=_compile_expression(segment))
 
@@ -116,19 +121,37 @@ def _create_probe_based_on_location(args, attribs, line_class, function_class):
     return ProbeType(**args)
 
 
-def probe(_id, _type, attribs):
-    # type: (str, str, Dict[str, Any]) -> Probe
+def probe_factory(attribs):
+    # type: (Dict[str, Any]) -> Probe
     """
     Create a new Probe instance.
     """
-    if _type == "logProbes":
+    try:
+        _type = attribs["type"]
+        _id = attribs["id"]
+    except KeyError as e:
+        raise ValueError("Invalid probe attributes: %s" % e)
+
+    if _type == ProbeType.LOG_PROBE:
         take_snapshot = attribs.get("captureSnapshot", False)
 
         args = dict(
             probe_id=_id,
             condition=_compile_expression(attribs.get("when")),
             tags=dict(_.split(":", 1) for _ in attribs.get("tags", [])),
-            limits=CaptureLimits(**attribs.get("capture", None)) if attribs.get("capture", None) else None,
+            limits=CaptureLimits(
+                **xlate_keys(
+                    attribs["capture"],
+                    {
+                        "maxReferenceDepth": "max_level",
+                        "maxCollectionSize": "max_size",
+                        "maxLength": "max_len",
+                        "maxFieldDepth": "max_fields",
+                    },
+                )
+            )
+            if "capture" in attribs
+            else None,
             rate=DEFAULT_SNAPSHOT_PROBE_RATE
             if take_snapshot
             else DEFAULT_PROBE_RATE,  # TODO: should we take rate limit out of Probe?
@@ -140,7 +163,7 @@ def probe(_id, _type, attribs):
 
         return _create_probe_based_on_location(args, attribs, LogLineProbe, LogFunctionProbe)
 
-    elif _type == "metricProbes":
+    elif _type == ProbeType.METRIC_PROBE:
         args = dict(
             probe_id=_id,
             condition=_compile_expression(attribs.get("when")),
@@ -157,25 +180,10 @@ def probe(_id, _type, attribs):
     raise ValueError("Unknown probe type: %s" % _type)
 
 
-_METRIC_PREFIX = "metricProbe_"
-_LOG_PREFIX = "logProbe_"
-
-
-def _make_probes(probes, _type):
-    return [probe(p["id"], _type, p) for p in probes]
-
-
 @_filter_by_env_and_version
 def get_probes(config_id, config):
     # type: (str, dict) -> Iterable[Probe]
-
-    if config_id.startswith(_METRIC_PREFIX):
-        return chain(_make_probes([config], "metricProbes"))
-
-    if config_id.startswith(_LOG_PREFIX):
-        return chain(_make_probes([config], "logProbes"))
-
-    raise ValueError("Unsupported config id: %s" % config_id)
+    return [probe_factory(config)]
 
 
 log = get_logger(__name__)

tests/debugging/probe/test_remoteconfig.py

@@ -4,10 +4,13 @@
 import pytest
 
 from ddtrace.debugging._config import config
+from ddtrace.debugging._probe.model import LogProbeMixin
 from ddtrace.debugging._probe.model import Probe
+from ddtrace.debugging._probe.model import ProbeType
 from ddtrace.debugging._probe.remoteconfig import ProbePollerEvent
 from ddtrace.debugging._probe.remoteconfig import ProbeRCAdapter
 from ddtrace.debugging._probe.remoteconfig import _filter_by_env_and_version
+from ddtrace.debugging._probe.remoteconfig import probe_factory
 from ddtrace.internal.remoteconfig.client import ConfigMetadata
 from tests.debugging.utils import create_snapshot_line_probe
 from tests.utils import override_global_config
@@ -200,6 +203,7 @@ def validate_events(expected):
             config_metadata("metricProbe_probe2"),
             {
                 "id": "probe2",
+                "type": ProbeType.METRIC_PROBE,
                 "tags": ["foo:bar"],
                 "where": {"sourceFile": "tests/submod/stuff.p", "lines": ["36"]},
                 "metricName": "test.counter",
@@ -217,6 +221,7 @@ def validate_events(expected):
             config_metadata("logProbe_probe3"),
             {
                 "id": "probe3",
+                "type": ProbeType.LOG_PROBE,
                 "tags": ["foo:bar"],
                 "where": {"sourceFile": "tests/submod/stuff.p", "lines": ["36"]},
                 "template": "hello {#foo}",
@@ -234,7 +239,7 @@ def validate_events(expected):
 
         # testing two things:
         #  1. after sleep 0.5 probe status should report 2 probes
-        #  2. bad config id raises ValueError
+        #  2. bad config raises ValueError
         with pytest.raises(ValueError):
             adapter(config_metadata("not-supported"), {})
 
@@ -255,3 +260,34 @@ def validate_events(expected):
 
     finally:
         config.diagnostics_interval = old_interval
+
+
+def test_log_probe_attributes_parsing():
+    probe = probe_factory(
+        {
+            "id": "3d338829-21c4-4a8a-8a1a-71fbce995efa",
+            "version": 0,
+            "type": ProbeType.LOG_PROBE,
+            "language": "python",
+            "active": True,
+            "where": {
+                "sourceFile": "foo.py",
+                "lines": ["57"],
+            },
+            "tags": ["env:staging", "version:v12417452-d2552757"],
+            "template": "{weekID} {idea}",
+            "segments": [
+                {"dsl": "weekID", "json": {"eq": [{"ref": "weekID"}, 1]}},
+                {"str": " "},
+                {"dsl": "idea", "json": {"ref": "idea"}},
+            ],
+            "captureSnapshot": False,
+            "capture": {"maxReferenceDepth": 42, "maxLength": 43},
+            "sampling": {"snapshotsPerSecond": 5000},
+        }
+    )
+
+    assert isinstance(probe, LogProbeMixin)
+
+    assert probe.limits.max_level == 42
+    assert probe.limits.max_len == 43