feat(tracer): [SVLS-5672] DynamoDB PutItem pointers (#10824)

PutItem is a bit tricky since it doesn't have a separate primary key
section. The primary key is mixed into the Item itself. So we need a way
for our customers to identify the primary key field names for a table.
We'll start by doing this with configuration, and we may add a way for
the user to opt in to us making a (cached) DescribeTable call on their
behalf.

## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

Author: apiarian-datadog

ddtrace/_trace/trace_handlers.py

@@ -9,6 +9,7 @@
 
 import wrapt
 
+from ddtrace import config
 from ddtrace._trace._span_pointer import _SpanPointerDescription
 from ddtrace._trace.utils import extract_DD_context_from_messages
 from ddtrace._trace.utils_botocore.span_pointers import extract_span_pointers_from_successful_botocore_response
@@ -487,6 +488,7 @@ def _on_botocore_patched_api_call_success(ctx, response):
     set_botocore_response_metadata_tags(span, response)
 
     for span_pointer_description in extract_span_pointers_from_successful_botocore_response(
+        dynamodb_primary_key_names_for_tables=config.botocore.dynamodb_primary_key_names_for_tables,
         endpoint_name=ctx.get_item("endpoint_name"),
         operation_name=ctx.get_item("operation"),
         request_parameters=ctx.get_item("params"),

ddtrace/_trace/utils_botocore/span_pointers.py

@@ -3,6 +3,7 @@
 from typing import Dict
 from typing import List
 from typing import NamedTuple
+from typing import Set
 
 from ddtrace._trace._span_pointer import _SpanPointerDescription
 from ddtrace._trace._span_pointer import _SpanPointerDirection
@@ -13,7 +14,19 @@
 log = get_logger(__name__)
 
 
+_DynamoDBTableName = str
+_DynamoDBItemFieldName = str
+_DynamoDBItemTypeTag = str
+
+_DynamoDBItemValue = Dict[_DynamoDBItemTypeTag, Any]
+_DynamoDBItem = Dict[_DynamoDBItemFieldName, _DynamoDBItemValue]
+
+_DynamoDBItemPrimaryKeyValue = Dict[_DynamoDBItemTypeTag, str]  # must be length 1
+_DynamoDBItemPrimaryKey = Dict[_DynamoDBItemFieldName, _DynamoDBItemPrimaryKeyValue]
+
+
 def extract_span_pointers_from_successful_botocore_response(
+    dynamodb_primary_key_names_for_tables: Dict[_DynamoDBTableName, Set[_DynamoDBItemFieldName]],
     endpoint_name: str,
     operation_name: str,
     request_parameters: Dict[str, Any],
@@ -22,9 +35,148 @@ def extract_span_pointers_from_successful_botocore_response(
     if endpoint_name == "s3":
         return _extract_span_pointers_for_s3_response(operation_name, request_parameters, response)
 
+    if endpoint_name == "dynamodb":
+        return _extract_span_pointers_for_dynamodb_response(
+            dynamodb_primary_key_names_for_tables, operation_name, request_parameters
+        )
+
+    return []
+
+
+def _extract_span_pointers_for_dynamodb_response(
+    dynamodb_primary_key_names_for_tables: Dict[_DynamoDBTableName, Set[_DynamoDBItemFieldName]],
+    operation_name: str,
+    request_parameters: Dict[str, Any],
+) -> List[_SpanPointerDescription]:
+    if operation_name == "PutItem":
+        return _extract_span_pointers_for_dynamodb_putitem_response(
+            dynamodb_primary_key_names_for_tables, request_parameters
+        )
+
     return []
 
 
+def _extract_span_pointers_for_dynamodb_putitem_response(
+    dynamodb_primary_key_names_for_tables: Dict[_DynamoDBTableName, Set[_DynamoDBItemFieldName]],
+    request_parameters: Dict[str, Any],
+) -> List[_SpanPointerDescription]:
+    try:
+        table_name = request_parameters["TableName"]
+        item = request_parameters["Item"]
+
+        return [
+            _aws_dynamodb_item_span_pointer_description(
+                pointer_direction=_SpanPointerDirection.DOWNSTREAM,
+                table_name=table_name,
+                primary_key=_aws_dynamodb_item_primary_key_from_item(
+                    dynamodb_primary_key_names_for_tables[table_name], item
+                ),
+            )
+        ]
+
+    except Exception as e:
+        log.warning(
+            "failed to generate DynamoDB.PutItem span pointer: %s",
+            str(e),
+        )
+        return []
+
+
+def _aws_dynamodb_item_primary_key_from_item(
+    primary_key_field_names: Set[_DynamoDBItemFieldName],
+    item: _DynamoDBItem,
+) -> _DynamoDBItemPrimaryKey:
+    if len(primary_key_field_names) not in (1, 2):
+        raise ValueError(f"unexpected number of primary key fields: {len(primary_key_field_names)}")
+
+    return {
+        primary_key_field_name: _aws_dynamodb_extract_and_verify_primary_key_field_value_item(
+            item, primary_key_field_name
+        )
+        for primary_key_field_name in primary_key_field_names
+    }
+
+
+def _aws_dynamodb_item_span_pointer_description(
+    pointer_direction: _SpanPointerDirection,
+    table_name: _DynamoDBTableName,
+    primary_key: _DynamoDBItemPrimaryKey,
+) -> _SpanPointerDescription:
+    return _SpanPointerDescription(
+        pointer_kind="aws.dynamodb.item",
+        pointer_direction=pointer_direction,
+        pointer_hash=_aws_dynamodb_item_span_pointer_hash(table_name, primary_key),
+        extra_attributes={},
+    )
+
+
+def _aws_dynamodb_extract_and_verify_primary_key_field_value_item(
+    item: _DynamoDBItem,
+    primary_key_field_name: _DynamoDBItemFieldName,
+) -> _DynamoDBItemPrimaryKeyValue:
+    if primary_key_field_name not in item:
+        raise ValueError(f"missing primary key field: {primary_key_field_name}")
+
+    value_object = item[primary_key_field_name]
+
+    if len(value_object) != 1:
+        raise ValueError(f"primary key field {primary_key_field_name} must have exactly one value: {len(value_object)}")
+
+    value_type, value_data = next(iter(value_object.items()))
+    if value_type not in ("S", "N", "B"):
+        raise ValueError(f"unexpected primary key field {primary_key_field_name} value type: {value_type}")
+
+    if not isinstance(value_data, str):
+        raise ValueError(f"unexpected primary key field {primary_key_field_name} value data type: {type(value_data)}")
+
+    return {value_type: value_data}
+
+
+def _aws_dynamodb_item_span_pointer_hash(table_name: _DynamoDBTableName, primary_key: _DynamoDBItemPrimaryKey) -> str:
+    if len(primary_key) == 1:
+        key, value_object = next(iter(primary_key.items()))
+        encoded_key_1 = key.encode("utf-8")
+        encoded_value_1 = _aws_dynamodb_item_encode_primary_key_value(value_object)
+        encoded_key_2 = b""
+        encoded_value_2 = b""
+
+    elif len(primary_key) == 2:
+        (key_1, value_object_1), (key_2, value_object_2) = sorted(
+            primary_key.items(), key=lambda x: x[0].encode("utf-8")
+        )
+        encoded_key_1 = key_1.encode("utf-8")
+        encoded_value_1 = _aws_dynamodb_item_encode_primary_key_value(value_object_1)
+        encoded_key_2 = key_2.encode("utf-8")
+        encoded_value_2 = _aws_dynamodb_item_encode_primary_key_value(value_object_2)
+
+    else:
+        raise ValueError(f"unexpected number of primary key fields: {len(primary_key)}")
+
+    return _standard_hashing_function(
+        table_name.encode("utf-8"),
+        encoded_key_1,
+        encoded_value_1,
+        encoded_key_2,
+        encoded_value_2,
+    )
+
+
+def _aws_dynamodb_item_encode_primary_key_value(value_object: _DynamoDBItemPrimaryKeyValue) -> bytes:
+    if len(value_object) != 1:
+        raise ValueError(f"primary key value object must have exactly one field: {len(value_object)}")
+
+    value_type, value = next(iter(value_object.items()))
+
+    if value_type == "S":
+        return value.encode("utf-8")
+
+    if value_type in ("N", "B"):
+        # these should already be here as ASCII strings
+        return value.encode("ascii")
+
+    raise ValueError(f"unknown primary key value type: {value_type}")
+
+
 def _extract_span_pointers_for_s3_response(
     operation_name: str,
     request_parameters: Dict[str, Any],

ddtrace/contrib/botocore/__init__.py

@@ -108,6 +108,37 @@
 
    Default: ``128``
 
+
+.. py:data:: ddtrace.config.botocore['dynamodb_primary_key_names_for_tables']
+
+    This enables DynamoDB API calls to be instrumented with span pointers. Many
+    DynamoDB API calls do not include the Item's Primary Key fields as separate
+    values, so they need to be provided to the tracer separately. This field
+    should be structured as a ``dict`` keyed by the table names as ``str``.
+    Each value should be the ``set`` of primary key field names (as ``str``)
+    for the associated table. The set may have exactly one or two elements,
+    depending on the Table's Primary Key schema.
+
+    In python this would look like::
+
+        ddtrace.config.botocore['dynamodb_primary_key_names_for_tables'] = {
+            'table_name': {'key1', 'key2'},
+            'other_table': {'other_key'},
+        }
+
+    Can also be enabled with the ``DD_BOTOCORE_DYNAMODB_TABLE_PRIMARY_KEYS``
+    environment variable which is parsed as a JSON object with strings for keys
+    and lists of strings for values.
+
+    This would look something like::
+
+        export DD_BOTOCORE_DYNAMODB_TABLE_PRIMARY_KEYS='{
+            "table_name": ["key1", "key2"],
+            "other_table": ["other_key"]
+        }'
+
+    Default: ``{}``
+
 """
 
 

ddtrace/contrib/internal/botocore/patch.py

@@ -2,7 +2,9 @@
 Trace queries to aws api done via botocore client
 """
 import collections
+import json
 import os
+from typing import Dict  # noqa:F401
 from typing import List  # noqa:F401
 from typing import Set  # noqa:F401
 from typing import Union  # noqa:F401
@@ -59,6 +61,32 @@
 log = get_logger(__name__)
 
 
+def _load_dynamodb_primary_key_names_for_tables() -> Dict[str, Set[str]]:
+    try:
+        encoded_table_primary_keys = os.getenv("DD_BOTOCORE_DYNAMODB_TABLE_PRIMARY_KEYS", "{}")
+        raw_table_primary_keys = json.loads(encoded_table_primary_keys)
+
+        table_primary_keys = {}
+        for table, primary_keys in raw_table_primary_keys.items():
+            if not isinstance(table, str):
+                raise ValueError(f"expected string table name: {table}")
+
+            if not isinstance(primary_keys, list):
+                raise ValueError(f"expected list of primary keys: {primary_keys}")
+
+            unique_primary_keys = set(primary_keys)
+            if not len(unique_primary_keys) == len(primary_keys):
+                raise ValueError(f"expected unique primary keys: {primary_keys}")
+
+            table_primary_keys[table] = unique_primary_keys
+
+        return table_primary_keys
+
+    except Exception as e:
+        log.warning("failed to load DD_BOTOCORE_DYNAMODB_TABLE_PRIMARY_KEYS: %s", e)
+        return {}
+
+
 # Botocore default settings
 config._add(
     "botocore",
@@ -73,6 +101,7 @@
         "instrument_internals": asbool(os.getenv("DD_BOTOCORE_INSTRUMENT_INTERNALS", default=False)),
         "propagation_enabled": asbool(os.getenv("DD_BOTOCORE_PROPAGATION_ENABLED", default=False)),
         "empty_poll_enabled": asbool(os.getenv("DD_BOTOCORE_EMPTY_POLL_ENABLED", default=True)),
+        "dynamodb_primary_key_names_for_tables": _load_dynamodb_primary_key_names_for_tables(),
     },
 )
 

releasenotes/notes/span-pointers-aws-dynamodb-putitem-f4be08418bc99093.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    botocore: Adds span pointers for successful DynamoDB ``PutItem`` spans. Table Primary Keys need to be provided with the ``ddtrace.config.botocore.dynamodb_primary_key_names_for_tables`` option or the ``DD_BOTOCORE_DYNAMODB_TABLE_PRIMARY_KEYS`` environment variable.