Merge pull request #1346 from newrelic/fix-datastore-spans

Fix Datastore Span Attributes

Author: hmstepanek

newrelic/core/database_node.py

@@ -288,7 +288,4 @@ def span_event(self, *args, **kwargs):
 
         self.agent_attributes["db.statement"] = sql
 
-        if self.target:
-            self.agent_attributes["db.collection"] = self.target
-
         return super(DatabaseNode, self).span_event(*args, **kwargs)

newrelic/core/datastore_node.py

@@ -136,12 +136,3 @@ def trace_node(self, stats, root, connections):
         return newrelic.core.trace_node.TraceNode(
             start_time=start_time, end_time=end_time, name=name, params=params, children=children, label=None
         )
-
-    def span_event(self, *args, **kwargs):
-        if self.operation:
-            self.agent_attributes["db.operation"] = self.operation
-
-        if self.target:
-            self.agent_attributes["db.collection"] = self.target
-
-        return super(DatastoreNode, self).span_event(*args, **kwargs)

newrelic/core/node_mixin.py

@@ -115,16 +115,30 @@ def span_event(self, *args, **kwargs):
         a_attrs = attrs[2]
 
         i_attrs["category"] = "datastore"
-        i_attrs["component"] = self.product
         i_attrs["span.kind"] = "client"
 
+        if self.product:
+            i_attrs["component"] = a_attrs["db.system"] = attribute.process_user_attribute("db.system", self.product)[1]
+        if self.operation:
+            a_attrs["db.operation"] = attribute.process_user_attribute("db.operation", self.operation)[1]
+        if self.target:
+            a_attrs["db.collection"] = attribute.process_user_attribute("db.collection", self.target)[1]
+
         if self.instance_hostname:
-            _, a_attrs["peer.hostname"] = attribute.process_user_attribute("peer.hostname", self.instance_hostname)
+            peer_hostname = attribute.process_user_attribute("peer.hostname", self.instance_hostname)[1]
         else:
-            a_attrs["peer.hostname"] = "Unknown"
+            peer_hostname = "Unknown"
+
+        a_attrs["peer.hostname"] = a_attrs["server.address"] = peer_hostname
+
+        peer_address = f"{peer_hostname}:{self.port_path_or_id or 'Unknown'}"
 
-        peer_address = f"{self.instance_hostname or 'Unknown'}:{self.port_path_or_id or 'Unknown'}"
+        a_attrs["peer.address"] = attribute.process_user_attribute("peer.address", peer_address)[1]
 
-        _, a_attrs["peer.address"] = attribute.process_user_attribute("peer.address", peer_address)
+        # Attempt to treat port_path_or_id as an integer, fallback to not including it
+        try:
+            a_attrs["server.port"] = int(self.port_path_or_id)
+        except Exception:
+            pass
 
         return attrs

newrelic/hooks/external_botocore.py

@@ -1055,7 +1055,6 @@ def _nr_dynamodb_datastore_trace_wrapper_(wrapped, instance, args, kwargs):
                 agent_attrs["cloud.resource_id"] = (
                     f"arn:{partition}:dynamodb:{region}:{account_id:012d}:table/{_target}"
                 )
-                agent_attrs["db.system"] = "DynamoDB"
 
         except Exception as e:
             _logger.debug("Failed to capture AWS DynamoDB info.", exc_info=True)

tests/agent_features/test_span_events.py

@@ -181,22 +181,48 @@ def test_database_db_statement_exclude():
 
 
 @pytest.mark.parametrize(
-    "trace_type,args,attrs",
+    "trace_type,args,i_attrs,a_attrs",
     (
         (
             DatastoreTrace,
-            ("db_product", "db_target", "db_operation"),
-            {"db.collection": "db_target", "db.operation": "db_operation"},
+            ("db_product", "db_target", "db_operation", "db_host", "1234", "db_name"),
+            {"component": "db_product", "span.kind": "client"},
+            {
+                "db.system": "db_product",
+                "db.operation": "db_operation",
+                "db.collection": "db_target",
+                "db.instance": "db_name",
+                "peer.address": "db_host:1234",
+                "peer.hostname": "db_host",
+                "server.address": "db_host",
+                "server.port": 1234,
+            },
         ),
         (
             DatabaseTrace,
             ("select 1 from db_table",),
-            {"db.collection": "db_table", "db.statement": "select ? from db_table"},
+            {"span.kind": "client"},
+            {"db.collection": "db_table", "db.statement": "select ? from db_table", "db.operation": "select"},
+        ),
+        (
+            DatabaseTrace,
+            ("select 1 from db_table", None, None, None, None, None, "db_host", "1234", "db_name"),
+            {"span.kind": "client"},
+            {
+                "db.statement": "select ? from db_table",
+                "db.operation": "select",
+                "db.collection": "db_table",
+                "db.instance": "db_name",
+                "peer.address": "db_host:1234",
+                "peer.hostname": "db_host",
+                "server.address": "db_host",
+                "server.port": 1234,
+            },
         ),
     ),
 )
-def test_datastore_database_trace_attrs(trace_type, args, attrs):
-    @validate_span_events(count=1, exact_agents=attrs)
+def test_datastore_database_trace_attrs(trace_type, args, i_attrs, a_attrs):
+    @validate_span_events(count=1, exact_intrinsics=i_attrs, exact_agents=a_attrs)
     @override_application_settings({"distributed_tracing.enabled": True, "span_events.enabled": True})
     @background_task(name="test_database_db_statement_exclude")
     def test():
@@ -293,8 +319,23 @@ def _test():
 @pytest.mark.parametrize(
     "kwarg_override,attribute_override",
     (
-        ({"host": "a" * 256}, {"peer.hostname": "a" * 255, "peer.address": "a" * 255}),
-        ({"port_path_or_id": "a" * 256, "host": "a"}, {"peer.hostname": "a", "peer.address": f"a:{'a' * 253}"}),
+        ({"product": "a" * 256}, {"component": "a" * 255, "db.system": "a" * 255}),
+        ({"target": "a" * 256}, {"db.collection": "a" * 255}),
+        ({"operation": "a" * 256}, {"db.operation": "a" * 255}),
+        ({"host": "a" * 256}, {"peer.hostname": "a" * 255, "peer.address": "a" * 255, "server.address": "a" * 255}),
+        (
+            {"port_path_or_id": "a" * 256, "host": "a"},
+            {"peer.hostname": "a", "peer.address": f"a:{'a' * 253}", "server.address": "a", "server.port": None},
+        ),
+        (
+            {"port_path_or_id": "1" * 256, "host": "a"},
+            {
+                "peer.hostname": "a",
+                "peer.address": f"a:{'1' * 253}",
+                "server.address": "a",
+                "server.port": int("1" * 256),
+            },
+        ),
         ({"database_name": "a" * 256}, {"db.instance": "a" * 255}),
     ),
 )
@@ -308,7 +349,16 @@ def test_datastore_span_limits(kwarg_override, attribute_override):
         "component": "library",
     }
 
-    exact_agents = {"db.instance": "db", "peer.hostname": "foo", "peer.address": "foo:1234"}
+    exact_agents = {
+        "db.collection": "table",
+        "db.instance": "db",
+        "db.operation": "operation",
+        "db.system": "library",
+        "peer.address": "foo:1234",
+        "peer.hostname": "foo",
+        "server.address": "foo",
+        "server.port": 1234,
+    }
 
     for k, v in attribute_override.items():
         if k in exact_agents: