From 55025773e659e5ab6f346932c9f00e89ac5e0b2c Mon Sep 17 00:00:00 2001
From: Nyenty Ayuk <ayuknyenty15@gmail.com>
Date: Thu, 30 Jun 2022 12:34:06 -0400
Subject: [PATCH] db.collection added to span events (#571)

* db.collection attributed added

Co-authored-by: Timothy Pansino <TimPansino@users.noreply.github.com>

* Fix db.collection gate

* Fix database_node gate

Co-authored-by: Timothy Pansino <TimPansino@users.noreply.github.com>
---
 newrelic/core/attribute.py               |  1 +
 newrelic/core/database_node.py           |  4 ++++
 newrelic/core/datastore_node.py          | 10 +++++++---
 tests/agent_features/test_span_events.py | 24 ++++++++++++++++++++++++
 4 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/newrelic/core/attribute.py b/newrelic/core/attribute.py
index 37f0afac9c..4c26739396 100644
--- a/newrelic/core/attribute.py
+++ b/newrelic/core/attribute.py
@@ -64,6 +64,7 @@
         'aws.lambda.arn',
         'aws.lambda.coldStart',
         'aws.lambda.eventSource.arn',
+        "db.collection",
         'db.instance',
         'db.operation',
         'db.statement',
diff --git a/newrelic/core/database_node.py b/newrelic/core/database_node.py
index b0f2d4aa40..3db38152f2 100644
--- a/newrelic/core/database_node.py
+++ b/newrelic/core/database_node.py
@@ -259,4 +259,8 @@ def span_event(self, *args, **kwargs):
                     'db.statement', sql, max_length=2000, ending='...')
 
         self.agent_attributes['db.statement'] = sql
+
+        if self.target:
+            self.agent_attributes["db.collection"] = self.target
+
         return super(DatabaseNode, self).span_event(*args, **kwargs)
diff --git a/newrelic/core/datastore_node.py b/newrelic/core/datastore_node.py
index 9135cc1e4d..290ec08464 100644
--- a/newrelic/core/datastore_node.py
+++ b/newrelic/core/datastore_node.py
@@ -21,9 +21,9 @@
 from newrelic.core.metric import TimeMetric
 
 _DatastoreNode = namedtuple('_DatastoreNode',
-        ['product', 'target', 'operation', 'children', 'start_time',
-        'end_time', 'duration', 'exclusive', 'host', 'port_path_or_id',
-        'database_name', 'guid', 'agent_attributes', 'user_attributes'])
+    ['product', 'target', 'operation', 'children', 'start_time',
+    'end_time', 'duration', 'exclusive', 'host', 'port_path_or_id',
+    'database_name', 'guid', 'agent_attributes', 'user_attributes'])
 
 
 class DatastoreNode(_DatastoreNode, DatastoreNodeMixin):
@@ -139,4 +139,8 @@ def trace_node(self, stats, root, connections):
     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)
diff --git a/tests/agent_features/test_span_events.py b/tests/agent_features/test_span_events.py
index 89581d5bb9..0024c1b8b8 100644
--- a/tests/agent_features/test_span_events.py
+++ b/tests/agent_features/test_span_events.py
@@ -203,6 +203,30 @@ def test_database_db_statement_exclude():
         pass
 
 
+@pytest.mark.parametrize('trace_type,args,attrs', (
+    (DatastoreTrace, ('db_product', 'db_target', 'db_operation'), {"db.collection": "db_target", "db.operation": "db_operation"}),
+    (DatabaseTrace, ("select 1 from db_table",), {"db.collection": "db_table", "db.statement": "select ? from db_table"}),
+))
+def test_datastore_database_trace_attrs(trace_type, args, attrs):
+    @validate_span_events(
+        count=1,
+        exact_agents=attrs,
+    )
+    @override_application_settings({
+        'distributed_tracing.enabled': True,
+        'span_events.enabled': True,
+    })
+    @background_task(name='test_database_db_statement_exclude')
+    def test():
+        transaction = current_transaction()
+        transaction._sampled = True
+
+        with trace_type(*args):
+            pass
+    
+    test()
+
+
 @pytest.mark.parametrize('exclude_url', (True, False))
 def test_external_spans(exclude_url):
     override_settings = {