Add SNAPSHOT_REGEX_PLACEHOLDERS (#208)

E.g. depending on where snapshots run, the port may be 80 or 8080. Or the path in error.stack tags will be different.

Author: bwoebi

README.md

@@ -185,6 +185,8 @@ metrics incorrectly placed within the meta field, or other type errors related t
 that are not relevant to the test case. **Note that removing `span_id` is not permitted to allow span 
 ordering to be maintained.**
 
+- `SNAPSHOT_REGEX_PLACEHOLDERS` [`""`]: The regex expressions to replace by a placeholder. Expressed as a comma separated `key:value` list. Specifying `ba[rz]:placeholder` will change any occurrence of `bar` or `baz` to `{placeholder}`: `foobarbazqux` -> `foo{placeholder}{placeholder}qux`. This is in particular useful to strip path prefixes or other infrastructure dependent identifiers.
+
 - `DD_POOL_TRACE_CHECK_FAILURES` [`false`]: Set to `"true"` to pool Trace Check failures that occured within Test-Agent memory. These failures can be queried later using the `/test/trace_check/failures` endpoint. Can also be set using the `--pool-trace-check-failures=true` option.
 
 - `DD_DISABLE_ERROR_RESPONSES` [`false`]: Set to `"true"` to disable Test-Agent `<Response 400>` when a Trace Check fails, instead sending a valid `<Response 200>`. Recommended for use with the `DD_POOL_TRACE_CHECK_FAILURES` env variable. Can also be set using the `--disable-error-responses=true` option.

ddapm_test_agent/agent.py

@@ -9,6 +9,7 @@
 import logging
 import os
 import pprint
+import re
 import socket
 import sys
 from typing import Awaitable
@@ -91,6 +92,15 @@ def _parse_csv(s: str) -> List[str]:
     return [s.strip() for s in s.split(",") if s.strip() != ""]
 
 
+def _parse_map(s: str) -> Dict[str, str]:
+    """Return the values of a csv string.
+
+    >>> _parse_map("a:b,b:c,c:d")
+    {'a': 'b', 'b': 'c', 'c': 'd'}
+    """
+    return dict([s.strip().split(":", 2) for s in s.split(",") if s.strip()])
+
+
 def _session_token(request: Request) -> Optional[str]:
     token: Optional[str]
     if "X-Datadog-Test-Session-Token" in request.headers:
@@ -794,6 +804,15 @@ async def handle_snapshot(self, request: Request) -> web.Response:
         span_removes = list(default_span_removes | overrides)
         log.info("using removes %r", span_removes)
 
+        # Get the span attributes that are to be removed for this snapshot.
+        default_attribute_regex_replaces: Dict[str, str] = request.app["snapshot_regex_placeholders"]
+        regex_overrides = _parse_map(request.url.query.get("regex_placeholders", ""))
+        attribute_regex_replaces = dict(
+            (f"{{{key}}}", re.compile(regex))
+            for (key, regex) in (default_attribute_regex_replaces | regex_overrides).items()
+        )
+        log.info("using regex placeholders %r", span_removes)
+
         if "span_id" in span_removes:
             raise AssertionError("Cannot remove 'span_id' from spans")
 
@@ -841,7 +860,13 @@ async def handle_snapshot(self, request: Request) -> web.Response:
             elif received_traces:
                 # Create a new snapshot for the data received
                 with open(trace_snap_file, mode="w") as f:
-                    f.write(trace_snapshot.generate_snapshot(received_traces=received_traces, removed=span_removes))
+                    f.write(
+                        trace_snapshot.generate_snapshot(
+                            received_traces=received_traces,
+                            removed=span_removes,
+                            attribute_regex_replaces=attribute_regex_replaces,
+                        )
+                    )
                 log.info("wrote new trace snapshot to %r", os.path.abspath(trace_snap_file))
 
             # Get all stats buckets from the payloads since we don't care about the other fields (hostname, env, etc)
@@ -1115,6 +1140,7 @@ def make_app(
     pool_trace_check_failures: bool,
     disable_error_responses: bool,
     snapshot_removed_attrs: List[str],
+    snapshot_regex_placeholders: Dict[str, str],
 ) -> web.Application:
     agent = Agent()
     app = web.Application(
@@ -1190,6 +1216,7 @@ def make_app(
     app["pool_trace_check_failures"] = pool_trace_check_failures
     app["disable_error_responses"] = disable_error_responses
     app["snapshot_removed_attrs"] = snapshot_removed_attrs
+    app["snapshot_regex_placeholders"] = snapshot_regex_placeholders
     return app
 
 
@@ -1240,6 +1267,14 @@ def main(args: Optional[List[str]] = None) -> None:
             "with meta. or metrics."
         ),
     )
+    parser.add_argument(
+        "--snapshot-regex-placeholders",
+        type=_parse_map,
+        default=os.environ.get("SNAPSHOT_REGEX_PLACEHOLDERS", ""),
+        help=(
+            "Comma-separated list of placeholder:regex tuples where to remove the matching regexes with the placeholder."
+        ),
+    )
     parser.add_argument(
         "--enabled-checks",
         type=List[str],
@@ -1341,6 +1376,7 @@ def main(args: Optional[List[str]] = None) -> None:
         pool_trace_check_failures=parsed_args.pool_trace_check_failures,
         disable_error_responses=parsed_args.disable_error_responses,
         snapshot_removed_attrs=parsed_args.snapshot_removed_attrs,
+        snapshot_regex_placeholders=parsed_args.snapshot_regex_placeholders,
     )
 
     web.run_app(app, sock=apm_sock, port=parsed_args.port)

ddapm_test_agent/trace_snapshot.py

@@ -3,6 +3,7 @@
 import logging
 import operator
 import pprint
+from re import Pattern
 import textwrap
 from typing import Any
 from typing import Dict
@@ -459,40 +460,57 @@ def _ordered_span(s: Span) -> OrderedDictType[str, TopLevelSpanValue]:
     return d  # type: ignore
 
 
-def _snapshot_trace_str(trace: Trace, removed: Optional[List[str]] = None) -> str:
+def _walk_span_attributes_with_regex_replaces(
+    dictionary: Dict[str, Any], attribute_regex_replaces: Dict[str, Pattern[str]]
+) -> None:
+    for key, val in dictionary.items():
+        if isinstance(val, str):
+            for placeholder, pattern in attribute_regex_replaces.items():
+                dictionary[key] = pattern.sub(placeholder, dictionary[key])
+        elif isinstance(val, Dict):
+            _walk_span_attributes_with_regex_replaces(val, attribute_regex_replaces)
+        elif isinstance(val, List):
+            for v in val:
+                if isinstance(v, Dict):
+                    _walk_span_attributes_with_regex_replaces(v, attribute_regex_replaces)
+
+
+def _snapshot_trace_str(trace: Trace, removed: List[str], attribute_regex_replaces: Dict[str, Pattern[str]]) -> str:
     cmap = child_map(trace)
     stack: List[Tuple[int, Span]] = [(0, root_span(trace))]
     s = "[\n"
     while stack:
         prefix, span = stack.pop(0)
 
         # Remove any keys that are not needed for comparison
-        if removed:
-            for key in removed:
-                if key.startswith("meta."):
-                    span["meta"].pop(key[5:], None)
-                elif key.startswith("metrics."):
-                    span["metrics"].pop(key[8:], None)
-                elif key.startswith("span_links.attributes."):
-                    if "span_links" in span:
-                        for link in span["span_links"]:
-                            if "attributes" in link:
-                                link["attributes"].pop(key[22:], None)
-                elif key.startswith("span_links."):
-                    if "span_links" in span:
-                        for link in span["span_links"]:
-                            link.pop(key[11:], None)  # type: ignore
-                elif key.startswith("span_events.attributes."):
-                    if "span_events" in span:
-                        for event in span["span_events"]:
-                            if "attributes" in event:
-                                event["attributes"].pop(key[23:], None)
-                elif key.startswith("span_events."):
-                    if "span_events" in span:
-                        for event in span["span_events"]:
-                            event.pop(key[12:], None)  # type: ignore
-                else:
-                    span.pop(key, None)  # type: ignore
+        for key in removed:
+            if key.startswith("meta."):
+                span["meta"].pop(key[5:], None)
+            elif key.startswith("metrics."):
+                span["metrics"].pop(key[8:], None)
+            elif key.startswith("span_links.attributes."):
+                if "span_links" in span:
+                    for link in span["span_links"]:
+                        if "attributes" in link:
+                            link["attributes"].pop(key[22:], None)
+            elif key.startswith("span_links."):
+                if "span_links" in span:
+                    for link in span["span_links"]:
+                        link.pop(key[11:], None)  # type: ignore
+            elif key.startswith("span_events.attributes."):
+                if "span_events" in span:
+                    for event in span["span_events"]:
+                        if "attributes" in event:
+                            event["attributes"].pop(key[23:], None)
+            elif key.startswith("span_events."):
+                if "span_events" in span:
+                    for event in span["span_events"]:
+                        event.pop(key[12:], None)  # type: ignore
+            else:
+                span.pop(key, None)  # type: ignore
+
+        if attribute_regex_replaces:  # only walk if we actually have something to replace
+            _walk_span_attributes_with_regex_replaces(cast(Dict[str, Any], span), attribute_regex_replaces)
 
         for i, child in enumerate(reversed(cmap[span["span_id"]])):
             if i == 0:
@@ -507,15 +525,19 @@ def _snapshot_trace_str(trace: Trace, removed: Optional[List[str]] = None) -> st
     return s
 
 
-def _snapshot_json(traces: List[Trace], removed: Optional[List[str]] = None) -> str:
+def _snapshot_json(traces: List[Trace], removed: List[str], attribute_regex_replaces: Dict[str, Pattern[str]]) -> str:
     s = "["
     for t in traces:
-        s += _snapshot_trace_str(t, removed)
+        s += _snapshot_trace_str(t, removed, attribute_regex_replaces)
         if t != traces[-1]:
             s += ",\n"
     s += "]\n"
     return s
 
 
-def generate_snapshot(received_traces: List[Trace], removed: Optional[List[str]] = None) -> str:
-    return _snapshot_json(_normalize_traces(received_traces), removed)
+def generate_snapshot(
+    received_traces: List[Trace],
+    removed: Optional[List[str]] = None,
+    attribute_regex_replaces: Optional[Dict[str, Pattern[str]]] = None,
+) -> str:
+    return _snapshot_json(_normalize_traces(received_traces), removed or [], attribute_regex_replaces or {})

releasenotes/notes/SNAPSHOT_REGEX_PLACEHOLDERS-1d2c5e2711d1953b.yaml

@@ -0,0 +1,5 @@
+---
+features:
+  - |
+    Add a SNAPSHOT_REGEX_PLACEHOLDERS option to allow infrastructure independent
+    snapshots.

tests/conftest.py

@@ -90,6 +90,11 @@ def snapshot_removed_attrs() -> Generator[Set[str], None, None]:
     yield set()
 
 
+@pytest.fixture
+def snapshot_regex_placeholders() -> Generator[Dict[str, str], None, None]:
+    yield dict()
+
+
 @pytest.fixture
 async def agent_app(
     aiohttp_server,
@@ -104,6 +109,7 @@ async def agent_app(
     pool_trace_check_failures,
     disable_error_responses,
     snapshot_removed_attrs,
+    snapshot_regex_placeholders,
 ):
     app = await aiohttp_server(
         make_app(
@@ -118,6 +124,7 @@ async def agent_app(
             pool_trace_check_failures,
             disable_error_responses,
             snapshot_removed_attrs,
+            snapshot_regex_placeholders,
         )
     )
     yield app

tests/test_snapshot.py

@@ -477,6 +477,31 @@ async def test_removed_attributes_metrics(agent, tmp_path, snapshot_removed_attr
         assert "process_id" not in span[0]
 
 
+@pytest.mark.parametrize("snapshot_regex_placeholders", [{"addr": "localhost:8080", "path": "^/.*"}])
+async def test_with_regex_placeholders(agent, tmp_path, snapshot_removed_attrs, do_reference_v04_http_trace):
+    resp = await do_reference_v04_http_trace(token="test_case")
+    assert resp.status == 200
+
+    custom_dir = tmp_path / "custom"
+    custom_dir.mkdir()
+    custom_file_name = custom_dir / "custom_snapshot"
+    custom_file = custom_dir / "custom_snapshot.json"
+
+    resp = await agent.get(
+        "/test/session/snapshot", params={"test_session_token": "test_case", "file": str(custom_file_name)}
+    )
+    assert resp.status == 200, await resp.text()
+
+    assert os.path.exists(custom_file), custom_file
+    with open(custom_file, mode="r") as f:  # Check that the removed attributes are not present in the span
+        file_content = "".join(f.readlines())
+        assert file_content != ""
+        span = json.loads(file_content)
+        assert "http.request" == span[0][0]["name"]
+        assert "{path}" == span[0][0]["resource"]
+        assert "http://{addr}/users" == span[0][0]["meta"]["http.url"]
+
+
 ONE_SPAN_TRACE_NO_START = random_trace(1, remove_keys=["start"])
 TWO_SPAN_TRACE_NO_START = random_trace(2, remove_keys=["start"])
 FIVE_SPAN_TRACE_NO_START = random_trace(5, remove_keys=["start"])