Backend dependency granularity for nosql databases (elastic#1639)

* use keyspace as db instance for cassandra

* service target fields for pymongo

* update changelog

Author: beniwohli

CHANGELOG.asciidoc

@@ -33,8 +33,9 @@ endif::[]
 
 // Unreleased changes go here
 // When the next release happens, nest these changes under the "Python Agent version 6.x" heading
-//[float]
-//===== Features
+[float]
+===== Features
+ * Add backend granularity data to SQL backends as well as Cassandra and pymongo {pull}1585[#1585], {pull}1639[#1639]
 
 [float]
 ===== Bug fixes

elasticapm/instrumentation/packages/cassandra.py

@@ -27,6 +27,7 @@
 #  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 #  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 #  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+from typing import Optional
 
 from elasticapm.instrumentation.packages.base import AbstractInstrumentedModule
 from elasticapm.instrumentation.packages.dbapi2 import extract_signature
@@ -49,6 +50,9 @@ def call(self, module, method, wrapped, instance, args, kwargs):
             else:
                 host = instance.endpoints_resolved[0].address
                 port = instance.endpoints_resolved[0].port
+            keyspace: Optional[str] = args[0] if args else kwargs.get("keyspace")
+            if keyspace:
+                context["db"] = {"instance": keyspace}
         else:
             hosts = list(instance.hosts)
             if hasattr(hosts[0], "endpoint"):
@@ -58,6 +62,9 @@ def call(self, module, method, wrapped, instance, args, kwargs):
                 # < cassandra-driver 3.18
                 host = hosts[0].address
                 port = instance.cluster.port
+            db_context = {}
+            if instance.keyspace:
+                db_context["instance"] = instance.keyspace
             span_action = "query"
             query = args[0] if args else kwargs.get("query")
             if hasattr(query, "query_string"):
@@ -70,7 +77,9 @@ def call(self, module, method, wrapped, instance, args, kwargs):
                 query_str = None
             if query_str:
                 name = extract_signature(query_str)
-                context["db"] = {"type": "sql", "statement": query_str}
+                db_context.update({"type": "sql", "statement": query_str})
+            if db_context:
+                context["db"] = db_context
         context["destination"] = {
             "address": host,
             "port": port,

elasticapm/instrumentation/packages/pymongo.py

@@ -83,13 +83,16 @@ def call(self, module, method, wrapped, instance, args, kwargs):
             "address": host,
             "port": port,
         }
+        context = {"destination": destination_info}
+        if instance.database.name:
+            context["db"] = {"instance": instance.database.name}
         with capture_span(
             signature,
             span_type="db",
             span_subtype="mongodb",
             span_action="query",
             leaf=True,
-            extra={"destination": destination_info},
+            extra=context,
         ):
             return wrapped(*args, **kwargs)
 
@@ -121,6 +124,9 @@ class PyMongoCursorInstrumentation(AbstractInstrumentedModule):
     def call(self, module, method, wrapped, instance, args, kwargs):
         collection = instance.collection
         signature = ".".join([collection.full_name, "cursor.refresh"])
+        context = {"destination": {}}
+        if instance.collection.database.name:
+            context["db"] = {"instance": instance.collection.database.name}
         with capture_span(
             signature,
             span_type="db",

tests/docker-compose.yml

@@ -41,6 +41,8 @@ services:
 
   mongodb36:
     image: mongo:3.6
+    ports:
+      - "27017:27017"
     volumes:
       - pymongodata36:/data/db
 

tests/instrumentation/cassandra_tests.py

@@ -36,6 +36,7 @@
 import socket
 
 from cassandra.cluster import Cluster
+from cassandra.protocol import ConfigurationException
 from cassandra.query import SimpleStatement
 
 from elasticapm.conf.constants import TRANSACTION
@@ -57,6 +58,11 @@ def cassandra_cluster():
 @pytest.fixture()
 def cassandra_session(cassandra_cluster):
     session = cassandra_cluster.connect()
+    try:
+        # make sure the keyspace doesn't exist
+        session.execute("DROP KEYSPACE testkeyspace;")
+    except ConfigurationException:
+        pass
     session.execute(
         """
         CREATE KEYSPACE testkeyspace
@@ -89,6 +95,38 @@ def test_cassandra_connect(instrument, elasticapm_client, cassandra_cluster):
     }
 
 
+def test_cassandra_connect_keyspace(instrument, elasticapm_client, cassandra_cluster):
+    session = cassandra_cluster.connect()
+    try:
+        session.execute(
+            """
+            CREATE KEYSPACE testkeyspace
+            WITH REPLICATION = { 'class' : 'SimpleStrategy' ,'replication_factor' : 1 }
+        """
+        )
+        elasticapm_client.begin_transaction("transaction.test")
+        sess = cassandra_cluster.connect("testkeyspace")
+        elasticapm_client.end_transaction("test")
+    finally:
+        session.execute("DROP KEYSPACE testkeyspace;")
+
+    transactions = elasticapm_client.events[TRANSACTION]
+    span = elasticapm_client.spans_for_transaction(transactions[0])[0]
+
+    assert span["type"] == "db"
+    assert span["subtype"] == "cassandra"
+    assert span["action"] == "connect"
+    assert span["duration"] > 0
+    assert span["name"] == "Cluster.connect"
+    assert span["context"]["destination"] == {
+        "address": socket.gethostbyname(os.environ.get("CASSANDRA_HOST", "localhost")),
+        "port": 9042,
+        "service": {"name": "", "resource": "cassandra/testkeyspace", "type": ""},
+    }
+    assert span["context"]["service"]["target"]["type"] == "cassandra"
+    assert span["context"]["service"]["target"]["name"] == "testkeyspace"
+
+
 def test_select_query_string(instrument, cassandra_session, elasticapm_client):
     elasticapm_client.begin_transaction("transaction.test")
     cassandra_session.execute("SELECT name from users")
@@ -99,12 +137,14 @@ def test_select_query_string(instrument, cassandra_session, elasticapm_client):
     assert span["subtype"] == "cassandra"
     assert span["action"] == "query"
     assert span["name"] == "SELECT FROM users"
-    assert span["context"]["db"] == {"statement": "SELECT name from users", "type": "sql"}
+    assert span["context"]["db"] == {"statement": "SELECT name from users", "type": "sql", "instance": "testkeyspace"}
     assert span["context"]["destination"] == {
         "address": socket.gethostbyname(os.environ.get("CASSANDRA_HOST", "localhost")),
         "port": 9042,
-        "service": {"name": "", "resource": "cassandra", "type": ""},
+        "service": {"name": "", "resource": "cassandra/testkeyspace", "type": ""},
     }
+    assert span["context"]["service"]["target"]["type"] == "cassandra"
+    assert span["context"]["service"]["target"]["name"] == "testkeyspace"
 
 
 def test_select_simple_statement(instrument, cassandra_session, elasticapm_client):
@@ -118,7 +158,7 @@ def test_select_simple_statement(instrument, cassandra_session, elasticapm_clien
     assert span["subtype"] == "cassandra"
     assert span["action"] == "query"
     assert span["name"] == "SELECT FROM users"
-    assert span["context"]["db"] == {"statement": "SELECT name from users", "type": "sql"}
+    assert span["context"]["db"] == {"statement": "SELECT name from users", "type": "sql", "instance": "testkeyspace"}
 
 
 def test_select_prepared_statement(instrument, cassandra_session, elasticapm_client):
@@ -132,7 +172,7 @@ def test_select_prepared_statement(instrument, cassandra_session, elasticapm_cli
     assert span["subtype"] == "cassandra"
     assert span["action"] == "query"
     assert span["name"] == "SELECT FROM users"
-    assert span["context"]["db"] == {"statement": "SELECT name from users", "type": "sql"}
+    assert span["context"]["db"] == {"statement": "SELECT name from users", "type": "sql", "instance": "testkeyspace"}
 
 
 def test_signature_create_keyspace():

tests/instrumentation/pymongo_tests.py

@@ -121,8 +121,10 @@ def test_collection_count_documents(instrument, elasticapm_client, mongo_databas
     assert span["context"]["destination"] == {
         "address": os.environ.get("MONGODB_HOST", "localhost"),
         "port": int(os.environ.get("MONGODB_PORT", 27017)),
-        "service": {"name": "", "resource": "mongodb", "type": ""},
+        "service": {"name": "", "resource": "mongodb/elasticapm_test", "type": ""},
     }
+    assert span["context"]["service"]["target"]["type"] == "mongodb"
+    assert span["context"]["service"]["target"]["name"] == "elasticapm_test"
 
 
 @pytest.mark.skipif(pymongo.version_tuple < (3, 7), reason="New in 3.7")
@@ -143,8 +145,10 @@ def test_collection_estimated_document_count(instrument, elasticapm_client, mong
     assert span["context"]["destination"] == {
         "address": os.environ.get("MONGODB_HOST", "localhost"),
         "port": int(os.environ.get("MONGODB_PORT", 27017)),
-        "service": {"name": "", "resource": "mongodb", "type": ""},
+        "service": {"name": "", "resource": "mongodb/elasticapm_test", "type": ""},
     }
+    assert span["context"]["service"]["target"]["type"] == "mongodb"
+    assert span["context"]["service"]["target"]["name"] == "elasticapm_test"
 
 
 @pytest.mark.integrationtest
@@ -254,7 +258,7 @@ def test_collection_find(instrument, elasticapm_client, mongo_database):
         assert span["context"]["destination"] == {
             "address": os.environ.get("MONGODB_HOST", "localhost"),
             "port": int(os.environ.get("MONGODB_PORT", 27017)),
-            "service": {"name": "", "resource": "mongodb", "type": ""},
+            "service": {"name": "", "resource/elasticapm_test": "mongodb", "type": ""},
         }
 
 
@@ -276,7 +280,7 @@ def test_collection_find_one(instrument, elasticapm_client, mongo_database):
     assert span["context"]["destination"] == {
         "address": os.environ.get("MONGODB_HOST", "localhost"),
         "port": int(os.environ.get("MONGODB_PORT", 27017)),
-        "service": {"name": "", "resource": "mongodb", "type": ""},
+        "service": {"name": "", "resource": "mongodb/elasticapm_test", "type": ""},
     }