Fix UnknownRemoteOperation for Database instrumentation.

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py

@@ -1,7 +1,8 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
+import re
 from logging import DEBUG, Logger, getLogger
-from typing import Optional
+from typing import Match, Optional
 from urllib.parse import ParseResult, urlparse
 
 from amazon.opentelemetry.distro._aws_attribute_keys import (
@@ -19,6 +20,7 @@
     UNKNOWN_REMOTE_SERVICE,
     UNKNOWN_SERVICE,
     extract_api_path_value,
+    get_dialect_keywords,
     get_egress_operation,
     get_ingress_operation,
     is_key_present,
@@ -38,6 +40,7 @@
 # Pertinent OTEL attribute keys
 _SERVICE_NAME: str = ResourceAttributes.SERVICE_NAME
 _DB_OPERATION: str = SpanAttributes.DB_OPERATION
+_DB_STATEMENT: str = SpanAttributes.DB_STATEMENT
 _DB_SYSTEM: str = SpanAttributes.DB_SYSTEM
 _FAAS_INVOKED_NAME: str = SpanAttributes.FAAS_INVOKED_NAME
 _FAAS_TRIGGER: str = SpanAttributes.FAAS_TRIGGER
@@ -189,9 +192,12 @@ def _set_remote_service_and_operation(span: ReadableSpan, attributes: BoundedAtt
     elif is_key_present(span, _RPC_SERVICE) or is_key_present(span, _RPC_METHOD):
         remote_service = _normalize_service_name(span, _get_remote_service(span, _RPC_SERVICE))
         remote_operation = _get_remote_operation(span, _RPC_METHOD)
-    elif is_key_present(span, _DB_SYSTEM) or is_key_present(span, _DB_OPERATION):
+    elif is_key_present(span, _DB_SYSTEM) or is_key_present(span, _DB_OPERATION) or is_key_present(span, _DB_STATEMENT):
         remote_service = _get_remote_service(span, _DB_SYSTEM)
-        remote_operation = _get_remote_operation(span, _DB_OPERATION)
+        if is_key_present(span, _DB_OPERATION):
+            remote_operation = _get_remote_operation(span, _DB_OPERATION)
+        else:
+            remote_operation = _get_db_statement_remote_operation(span, _DB_STATEMENT)
     elif is_key_present(span, _FAAS_INVOKED_NAME) or is_key_present(span, _FAAS_TRIGGER):
         remote_service = _get_remote_service(span, _FAAS_INVOKED_NAME)
         remote_operation = _get_remote_operation(span, _FAAS_TRIGGER)
@@ -232,6 +238,28 @@ def _get_remote_operation(span: ReadableSpan, remote_operation_key: str) -> str:
     return remote_operation
 
 
+# If no db.operation attribute provided in the span,
+# we use db.statement to compute a valid remote operation in a best-effort manner.
+# To do this, we take the first substring of the statement
+# and compare to a regex list of known SQL keywords.
+# The substring length is determined by the longest known SQL keywords.
+def _get_db_statement_remote_operation(span: ReadableSpan, statement_key: str) -> str:
+    remote_operation: str = span.attributes.get(statement_key)
+
+    if remote_operation is None:
+        return UNKNOWN_REMOTE_OPERATION
+
+    # Remove all whitespace and newline characters from the beginning of remote_operation
+    # and retrieve the first 27 characters
+    remote_operation = remote_operation.lstrip()[:27]
+    dialect_keywords = get_dialect_keywords()
+    pattern: str = r"^(?:" + "|".join(dialect_keywords) + r")\b"
+    match: Optional[Match[str]] = re.match(pattern, remote_operation.upper())
+    remote_operation = match.group(0) if match else UNKNOWN_REMOTE_OPERATION
+
+    return remote_operation
+
+
 def _normalize_service_name(span: ReadableSpan, service_name: str) -> str:
     """
     TODO: Determine if problems in Java instrumentation are relevant here. Do we need normalization? If so, probably we

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_span_processing_util.py

@@ -1,6 +1,9 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 """Utility module designed to support shared logic across AWS Span Processors."""
+import json
+import os
+
 from amazon.opentelemetry.distro._aws_attribute_keys import AWS_CONSUMER_PARENT_SPAN_KIND, AWS_LOCAL_OPERATION
 from opentelemetry.sdk.trace import InstrumentationScope, ReadableSpan
 from opentelemetry.semconv.trace import MessagingOperationValues, SpanAttributes
@@ -96,6 +99,15 @@ def is_local_root(span: ReadableSpan) -> bool:
     return span.parent is None or not span.parent.is_valid or span.parent.is_remote
 
 
+# Get valid keywords retrieved from db.statement if no db.operation value is identified
+def get_dialect_keywords():
+    current_dir = os.path.dirname(__file__)
+    file_path = os.path.join(current_dir, "configuration/dialect_keywords.json")
+    with open(file_path, "r", encoding="utf-8") as json_file:
+        keywords_data = json.load(json_file)
+    return keywords_data["keywords"]
+
+
 def _is_sqs_receive_message_consumer_span(span: ReadableSpan) -> bool:
     """To identify the SQS consumer spans produced by AWS SDK instrumentation"""
     messaging_operation: str = span.attributes.get(SpanAttributes.MESSAGING_OPERATION)