Add botocore db.statement sanitization for dynamoDB

CHANGELOG.md

@@ -14,6 +14,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- `opentelemetry-instrumentation-botocore` Add optional db.statement query sanitization.
+  ([#1662](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1662))
 - Support `aio_pika` 9.x (([#1670](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1670])
 - `opentelemetry-instrumentation-redis` Add `sanitize_query` config option to allow query sanitization.  ([#1572](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1572))
 - `opentelemetry-instrumentation-elasticsearch` Add optional db.statement query sanitization. 

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/__init__.py

@@ -66,6 +66,8 @@ def request_hook(span, service_name, operation_name, api_params):
     def response_hook(span, service_name, operation_name, result):
         # response hook logic
 
+    dynamodb_sanitize_query (bool) - an optional dynamodb query sanitization flag, default is False
+
     # Instrument Botocore with hooks
     BotocoreInstrumentor().instrument(request_hook=request_hook, response_hook=response_hook)
 
@@ -76,6 +78,9 @@ def response_hook(span, service_name, operation_name, result):
     )
     ec2 = self.session.create_client("ec2", region_name="us-west-2")
     ec2.describe_instances()
+
+    # Instrument Botocore with dynamoDB sanitization enabled - default is False
+    BotocoreInstrumentor().instrument(dynamodb_sanitize_query=True)
 """
 
 import logging
@@ -117,6 +122,17 @@ def _patched_endpoint_prepare_request(wrapped, instance, args, kwargs):
     return wrapped(*args, **kwargs)
 
 
+class ExtensionsConfiguration:
+    def __init__(self, **kwargs):
+        dynamodb_sanitize_query = kwargs.get("dynamodb_sanitize_query", False)
+        self.dynamodb_configuration = {
+            "sanitize_query": dynamodb_sanitize_query
+        }
+
+    def get_dynamodb_configuration(self):
+        return self.dynamodb_configuration
+
+
 class BotocoreInstrumentor(BaseInstrumentor):
     """An instrumentor for Botocore.
 
@@ -139,6 +155,7 @@ def _instrument(self, **kwargs):
 
         self.request_hook = kwargs.get("request_hook")
         self.response_hook = kwargs.get("response_hook")
+        self.configuration = ExtensionsConfiguration(**kwargs)
 
         wrap_function_wrapper(
             "botocore.client",
@@ -161,7 +178,9 @@ def _patched_api_call(self, original_func, instance, args, kwargs):
         if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
             return original_func(*args, **kwargs)
 
-        call_context = _determine_call_context(instance, args)
+        call_context = _determine_call_context(
+            instance, args + (self.configuration,)
+        )
         if call_context is None:
             return original_func(*args, **kwargs)
 

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/dynamodb.py

@@ -62,14 +62,44 @@ def _conv_val_to_len(value) -> Optional[int]:
     return len(value) if value is not None else None
 
 
+_db_statement_operations = {
+    "BatchGetItem",
+    "BatchWriteItem",
+    "DeleteItem",
+    "GetItem",
+    "PutItem",
+    "Query",
+    "Scan",
+    "UpdateItem",
+}
+
+_sanitized_keys = (
+    "Item",
+    "Key",
+    "KeyConditionExpression",
+    "ExpressionAttributeValues",
+    "FilterExpression",
+)
+_sanitized_value = {"?": "?"}
+
+
+# pylint: disable=C0103
+def _conv_params_to_sanitized_str(params) -> str:
+    p = params.copy()
+    for key in p:
+        if key in _sanitized_keys:
+            p[key] = _sanitized_value
+
+    return str(p)
+
+
 ################################################################################
 # common request attributes
 ################################################################################
 
 _REQ_TABLE_NAME = ("TableName", _conv_val_to_single_attr_tuple)
 _REQ_REQITEMS_TABLE_NAMES = ("RequestItems", _conv_dict_to_key_tuple)
 
-
 _REQ_GLOBAL_SEC_INDEXES = ("GlobalSecondaryIndexes", _conv_list_to_json_list)
 _REQ_LOCAL_SEC_INDEXES = ("LocalSecondaryIndexes", _conv_list_to_json_list)
 
@@ -350,12 +380,25 @@ class _DynamoDbExtension(_AwsSdkExtension):
     def __init__(self, call_context: _AwsSdkCallContext):
         super().__init__(call_context)
         self._op = _OPERATION_MAPPING.get(call_context.operation)
+        self._configuration = (
+            call_context.configuration.get_dynamodb_configuration()
+        )
 
     def extract_attributes(self, attributes: _AttributeMapT):
         attributes[SpanAttributes.DB_SYSTEM] = DbSystemValues.DYNAMODB.value
         attributes[SpanAttributes.DB_OPERATION] = self._call_context.operation
         attributes[SpanAttributes.NET_PEER_NAME] = self._get_peer_name()
 
+        if self._call_context.operation in _db_statement_operations:
+            attributes[SpanAttributes.DB_STATEMENT] = str(
+                self._call_context.params
+            )
+
+            if self._configuration.get("sanitize_query"):
+                attributes[
+                    SpanAttributes.DB_STATEMENT
+                ] = _conv_params_to_sanitized_str(self._call_context.params)
+
         if self._op is None:
             return
 

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/types.py

@@ -24,8 +24,6 @@
 _BotoClientT = "botocore.client.BaseClient"
 _BotoResultT = Dict[str, Any]
 _BotoClientErrorT = "botocore.exceptions.ClientError"
-
-_OperationParamsT = Dict[str, Any]
 _AttributeMapT = Dict[str, AttributeValue]
 
 
@@ -44,6 +42,7 @@ class _AwsSdkCallContext:
         api_version: the API version of the called AWS service.
         span_name: the name used to create the span.
         span_kind: the kind used to create the span.
+        configuration: a class represents additional configuration for botocore extensions.
     """
 
     def __init__(self, client: _BotoClientT, args: Tuple[str, Dict[str, Any]]):
@@ -54,12 +53,19 @@ def __init__(self, client: _BotoClientT, args: Tuple[str, Dict[str, Any]]):
             _logger.warning("Could not get request params.")
             params = {}
 
+        try:
+            configurations = args[2]
+        except (IndexError, TypeError):
+            _logger.warning("Could not get request configurations.")
+            configurations = None
+
         boto_meta = client.meta
         service_model = boto_meta.service_model
 
         self.service = service_model.service_name.lower()  # type: str
         self.operation = operation  # type: str
         self.params = params  # type: Dict[str, Any]
+        self.configuration = configurations
 
         # 'operation' and 'service' are essential for instrumentation.
         # for all other attributes we extract them defensively. All of them should

instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_dynamodb.py

@@ -20,6 +20,8 @@
 
 from opentelemetry.instrumentation.botocore import BotocoreInstrumentor
 from opentelemetry.instrumentation.botocore.extensions.dynamodb import (
+    _conv_params_to_sanitized_str,
+    _db_statement_operations,
     _DynamoDbExtension,
 )
 from opentelemetry.semconv.trace import SpanAttributes
@@ -28,6 +30,60 @@
 
 # pylint: disable=too-many-public-methods
 
+default_table_name = "test_table"
+
+# General dynamoDB queries
+key_value_query = {"Key": {"username": "johndoe"}}
+
+secondary_index_query = {
+    "IndexName": "email-index",
+    "KeyConditionExpression": "email = :email",
+    "ExpressionAttributeValues": {":email": "[email protected]"},
+}
+
+scan_query = {
+    "FilterExpression": "age >= :age",
+    "ExpressionAttributeValues": {":age": 25},
+}
+
+projection_expression_query = {
+    "KeyConditionExpression": "username = :username",
+    "ExpressionAttributeValues": {":username": "johndoe"},
+    "ProjectionExpression": "username, email",
+}
+
+# Instrumentation tests queries
+delete_query = {
+    "TableName": default_table_name,
+    "Key": {"id": {"S": "1"}},
+    "ReturnConsumedCapacity": "TOTAL",
+    "ReturnItemCollectionMetrics": "SIZE",
+}
+
+get_query = {
+    "TableName": default_table_name,
+    "Key": {"id": {"S": "1"}},
+    "ConsistentRead": True,
+    "AttributesToGet": ["id"],
+    "ProjectionExpression": "1,2",
+    "ReturnConsumedCapacity": "TOTAL",
+}
+
+put_query = {
+    "TableName": default_table_name,
+    "Item": {"id": {"S": "1"}},
+    "ReturnConsumedCapacity": "TOTAL",
+    "ReturnItemCollectionMetrics": "SIZE",
+}
+
+update_query = {
+    "TableName": default_table_name,
+    "Key": {"id": {"S": "123"}},
+    "AttributeUpdates": {"id": {"Value": {"S": "456"}, "Action": "PUT"}},
+    "ReturnConsumedCapacity": "TOTAL",
+    "ReturnItemCollectionMetrics": "SIZE",
+}
+
 
 class TestDynamoDbExtension(TestBase):
     def setUp(self):
@@ -108,6 +164,8 @@ def assert_span(self, operation: str) -> Span:
             "dynamodb.us-west-2.amazonaws.com",
             span.attributes[SpanAttributes.NET_PEER_NAME],
         )
+        if operation in _db_statement_operations:
+            self.assertIn(SpanAttributes.DB_STATEMENT, span.attributes)
         return span
 
     def assert_table_names(self, span: Span, *table_names):
@@ -270,14 +328,11 @@ def test_create_table(self):
 
     @mock_dynamodb2
     def test_delete_item(self):
+        BotocoreInstrumentor().uninstrument()
+        BotocoreInstrumentor().instrument(dynamodb_sanitize_query=True)
         self._create_prepared_table()
 
-        self.client.delete_item(
-            TableName=self.default_table_name,
-            Key={"id": {"S": "1"}},
-            ReturnConsumedCapacity="TOTAL",
-            ReturnItemCollectionMetrics="SIZE",
-        )
+        self.client.delete_item(**delete_query)
 
         span = self.assert_span("DeleteItem")
         self.assert_table_names(span, self.default_table_name)
@@ -317,22 +372,21 @@ def test_describe_table(self):
 
     @mock_dynamodb2
     def test_get_item(self):
-        self._create_prepared_table()
-
-        self.client.get_item(
-            TableName=self.default_table_name,
-            Key={"id": {"S": "1"}},
-            ConsistentRead=True,
-            AttributesToGet=["id"],
-            ProjectionExpression="1,2",
-            ReturnConsumedCapacity="TOTAL",
-        )
+        # Reset instrumentation to use sanitized query
+        BotocoreInstrumentor().uninstrument()
+        BotocoreInstrumentor().instrument(dynamodb_sanitize_query=True)
 
+        self._create_prepared_table()
+        self.client.get_item(**get_query)
         span = self.assert_span("GetItem")
         self.assert_table_names(span, self.default_table_name)
         self.assert_consistent_read(span, True)
         self.assert_projection(span, "1,2")
         self.assert_consumed_capacity(span, self.default_table_name)
+        self.assertEqual(
+            span.attributes[SpanAttributes.DB_STATEMENT],
+            _conv_params_to_sanitized_str(get_query),
+        )
 
     @mock_dynamodb2
     def test_list_tables(self):
@@ -355,17 +409,16 @@ def test_list_tables(self):
     def test_put_item(self):
         table = "test_table"
         self._create_prepared_table(TableName=table)
-
-        self.client.put_item(
-            TableName=table,
-            Item={"id": {"S": "1"}, "idl": {"S": "2"}, "idg": {"S": "3"}},
-            ReturnConsumedCapacity="TOTAL",
-            ReturnItemCollectionMetrics="SIZE",
-        )
+        self.client.put_item(**put_query)
 
         span = self.assert_span("PutItem")
         self.assert_table_names(span, table)
         self.assert_consumed_capacity(span, table)
+        self.assertEqual(
+            span.attributes[SpanAttributes.DB_STATEMENT],
+            str(put_query),
+        )
+
         # moto does not seem to return these:
         # self.assert_item_coll_metrics(span)
 
@@ -448,17 +501,18 @@ def test_scan(self):
     def test_update_item(self):
         self._create_prepared_table()
 
-        self.client.update_item(
-            TableName=self.default_table_name,
-            Key={"id": {"S": "123"}},
-            AttributeUpdates={"id": {"Value": {"S": "456"}, "Action": "PUT"}},
-            ReturnConsumedCapacity="TOTAL",
-            ReturnItemCollectionMetrics="SIZE",
-        )
+        # Reset instrumentation to use sanitized query
+        BotocoreInstrumentor().uninstrument()
+        BotocoreInstrumentor().instrument(dynamodb_sanitize_query=True)
 
+        self.client.update_item(**update_query)
         span = self.assert_span("UpdateItem")
         self.assert_table_names(span, self.default_table_name)
         self.assert_consumed_capacity(span, self.default_table_name)
+        self.assertEqual(
+            span.attributes[SpanAttributes.DB_STATEMENT],
+            _conv_params_to_sanitized_str(update_query),
+        )
         # moto does not seem to return these:
         # self.assert_item_coll_metrics(span)
 
@@ -504,3 +558,34 @@ def test_update_table(self):
                 SpanAttributes.AWS_DYNAMODB_GLOBAL_SECONDARY_INDEX_UPDATES
             ],
         )
+
+    # pylint: disable=expression-not-assigned
+    def test_db_statement_sanitization(self):
+        sanitized_value = {"?": "?"}
+        self.assertEqual(
+            _conv_params_to_sanitized_str(key_value_query),
+            str({"Key": sanitized_value}),
+        ),
+        self.assertEqual(
+            _conv_params_to_sanitized_str(secondary_index_query),
+            str(
+                {
+                    "IndexName": "email-index",
+                    "KeyConditionExpression": sanitized_value,
+                    "ExpressionAttributeValues": sanitized_value,
+                }
+            ),
+        ),
+        self.assertEqual(
+            _conv_params_to_sanitized_str(scan_query),
+            str(
+                {
+                    "FilterExpression": sanitized_value,
+                    "ExpressionAttributeValues": sanitized_value,
+                }
+            ),
+        ),
+        self.assertNotEqual(
+            _conv_params_to_sanitized_str(projection_expression_query),
+            projection_expression_query,
+        )