Add DB connection attributes in spans (#2274)

---------

Co-authored-by: Ivana Kellyerova <[email protected]>

Author: antonpirker

sentry_sdk/consts.py

@@ -68,6 +68,12 @@ class SPANDATA:
     See: https://develop.sentry.dev/sdk/performance/span-data-conventions/
     """
 
+    DB_NAME = "db.name"
+    """
+    The name of the database being accessed. For commands that switch the database, this should be set to the target database (even if the command fails).
+    Example: myDatabase
+    """
+
     DB_OPERATION = "db.operation"
     """
     The name of the operation being executed, e.g. the MongoDB command name such as findAndModify, or the SQL keyword.
@@ -118,6 +124,31 @@ class SPANDATA:
     Example: 418
     """
 
+    SERVER_ADDRESS = "server.address"
+    """
+    Name of the database host.
+    Example: example.com
+    """
+
+    SERVER_PORT = "server.port"
+    """
+    Logical server port number
+    Example: 80; 8080; 443
+    """
+
+    SERVER_SOCKET_ADDRESS = "server.socket.address"
+    """
+    Physical server IP address or Unix socket address.
+    Example: 10.5.3.2
+    """
+
+    SERVER_SOCKET_PORT = "server.socket.port"
+    """
+    Physical server port.
+    Recommended: If different than server.port.
+    Example: 16456
+    """
+
 
 class OP:
     CACHE_GET_ITEM = "cache.get_item"

sentry_sdk/integrations/django/__init__.py

@@ -612,7 +612,7 @@ def execute(self, sql, params=None):
         with record_sql_queries(
             hub, self.cursor, sql, params, paramstyle="format", executemany=False
         ) as span:
-            _set_db_system_on_span(span, self.db.vendor)
+            _set_db_data(span, self.db.vendor, self.db.get_connection_params())
             return real_execute(self, sql, params)
 
     def executemany(self, sql, param_list):
@@ -624,7 +624,7 @@ def executemany(self, sql, param_list):
         with record_sql_queries(
             hub, self.cursor, sql, param_list, paramstyle="format", executemany=True
         ) as span:
-            _set_db_system_on_span(span, self.db.vendor)
+            _set_db_data(span, self.db.vendor, self.db.get_connection_params())
             return real_executemany(self, sql, param_list)
 
     def connect(self):
@@ -637,7 +637,7 @@ def connect(self):
             hub.add_breadcrumb(message="connect", category="query")
 
         with hub.start_span(op=OP.DB, description="connect") as span:
-            _set_db_system_on_span(span, self.vendor)
+            _set_db_data(span, self.vendor, self.get_connection_params())
             return real_connect(self)
 
     CursorWrapper.execute = execute
@@ -646,8 +646,22 @@ def connect(self):
     ignore_logger("django.db.backends")
 
 
-# https://github.com/django/django/blob/6a0dc2176f4ebf907e124d433411e52bba39a28e/django/db/backends/base/base.py#L29
-# Avaliable in Django 1.8+
-def _set_db_system_on_span(span, vendor):
-    # type: (Span, str) -> None
+def _set_db_data(span, vendor, connection_params):
+    # type: (Span, str, Dict[str, str]) -> None
     span.set_data(SPANDATA.DB_SYSTEM, vendor)
+
+    db_name = connection_params.get("dbname") or connection_params.get("database")
+    if db_name is not None:
+        span.set_data(SPANDATA.DB_NAME, db_name)
+
+    server_address = connection_params.get("host")
+    if server_address is not None:
+        span.set_data(SPANDATA.SERVER_ADDRESS, server_address)
+
+    server_port = connection_params.get("port")
+    if server_port is not None:
+        span.set_data(SPANDATA.SERVER_PORT, server_port)
+
+    server_socket_address = connection_params.get("unix_socket")
+    if server_socket_address is not None:
+        span.set_data(SPANDATA.SERVER_SOCKET_ADDRESS, server_socket_address)

sentry_sdk/integrations/pymongo.py

@@ -85,6 +85,27 @@ def _strip_pii(command):
     return command
 
 
+def _get_db_data(event):
+    # type: (Any) -> Dict[str, Any]
+    data = {}
+
+    data[SPANDATA.DB_SYSTEM] = "mongodb"
+
+    db_name = event.database_name
+    if db_name is not None:
+        data[SPANDATA.DB_NAME] = db_name
+
+    server_address = event.connection_id[0]
+    if server_address is not None:
+        data[SPANDATA.SERVER_ADDRESS] = server_address
+
+    server_port = event.connection_id[1]
+    if server_port is not None:
+        data[SPANDATA.SERVER_PORT] = server_port
+
+    return data
+
+
 class CommandTracer(monitoring.CommandListener):
     def __init__(self):
         # type: () -> None
@@ -121,10 +142,10 @@ def started(self, event):
                 pass
 
             data = {"operation_ids": {}}  # type: Dict[str, Any]
-
             data["operation_ids"]["operation"] = event.operation_id
             data["operation_ids"]["request"] = event.request_id
-            data[SPANDATA.DB_SYSTEM] = "mongodb"
+
+            data.update(_get_db_data(event))
 
             try:
                 lsid = command.pop("lsid")["id"]

sentry_sdk/integrations/sqlalchemy.py

@@ -67,9 +67,7 @@ def _before_cursor_execute(
     span = ctx_mgr.__enter__()
 
     if span is not None:
-        db_system = _get_db_system(conn.engine.name)
-        if db_system is not None:
-            span.set_data(SPANDATA.DB_SYSTEM, db_system)
+        _set_db_data(span, conn)
         context._sentry_sql_span = span
 
 
@@ -128,3 +126,22 @@ def _get_db_system(name):
         return "oracle"
 
     return None
+
+
+def _set_db_data(span, conn):
+    # type: (Span, Any) -> None
+    db_system = _get_db_system(conn.engine.name)
+    if db_system is not None:
+        span.set_data(SPANDATA.DB_SYSTEM, db_system)
+
+    db_name = conn.engine.url.database
+    if db_name is not None:
+        span.set_data(SPANDATA.DB_NAME, db_name)
+
+    server_address = conn.engine.url.host
+    if server_address is not None:
+        span.set_data(SPANDATA.SERVER_ADDRESS, server_address)
+
+    server_port = conn.engine.url.port
+    if server_port is not None:
+        span.set_data(SPANDATA.SERVER_PORT, server_port)

test-requirements.txt

@@ -13,4 +13,3 @@ asttokens
 responses
 pysocks
 ipdb
-mockupdb

tests/integrations/django/test_basic.py

@@ -1,9 +1,10 @@
 from __future__ import absolute_import
 
 import json
+import os
+import random
 import re
 import pytest
-import random
 from functools import partial
 
 from werkzeug.test import Client
@@ -584,9 +585,7 @@ def test_django_connect_trace(sentry_init, client, capture_events, render_span_t
 
 @pytest.mark.forked
 @pytest_mark_django_db_decorator(transaction=True)
-def test_django_connect_breadcrumbs(
-    sentry_init, client, capture_events, render_span_tree
-):
+def test_django_connect_breadcrumbs(sentry_init, capture_events):
     """
     Verify we record a breadcrumb when opening a new database.
     """
@@ -620,6 +619,43 @@ def test_django_connect_breadcrumbs(
     ]
 
 
+@pytest.mark.forked
+@pytest_mark_django_db_decorator(transaction=True)
+def test_db_connection_span_data(sentry_init, client, capture_events):
+    sentry_init(
+        integrations=[DjangoIntegration()],
+        send_default_pii=True,
+        traces_sample_rate=1.0,
+    )
+    from django.db import connections
+
+    if "postgres" not in connections:
+        pytest.skip("postgres tests disabled")
+
+    # trigger Django to open a new connection by marking the existing one as None.
+    connections["postgres"].connection = None
+
+    events = capture_events()
+
+    content, status, headers = client.get(reverse("postgres_select"))
+    assert status == "200 OK"
+
+    (event,) = events
+
+    for span in event["spans"]:
+        if span.get("op") == "db":
+            data = span.get("data")
+            assert data.get(SPANDATA.DB_SYSTEM) == "postgresql"
+            assert (
+                data.get(SPANDATA.DB_NAME)
+                == connections["postgres"].get_connection_params()["database"]
+            )
+            assert data.get(SPANDATA.SERVER_ADDRESS) == os.environ.get(
+                "SENTRY_PYTHON_TEST_POSTGRES_HOST", "localhost"
+            )
+            assert data.get(SPANDATA.SERVER_PORT) == 5432
+
+
 @pytest.mark.parametrize(
     "transaction_style,client_url,expected_transaction,expected_source,expected_response",
     [
@@ -1059,11 +1095,7 @@ def dummy(a, b):
 @pytest_mark_django_db_decorator()
 @pytest.mark.skipif(DJANGO_VERSION < (1, 9), reason="Requires Django >= 1.9")
 def test_cache_spans_disabled_middleware(
-    sentry_init,
-    client,
-    capture_events,
-    use_django_caching_with_middlewares,
-    settings,
+    sentry_init, client, capture_events, use_django_caching_with_middlewares
 ):
     sentry_init(
         integrations=[
@@ -1141,11 +1173,7 @@ def test_cache_spans_disabled_templatetag(
 @pytest_mark_django_db_decorator()
 @pytest.mark.skipif(DJANGO_VERSION < (1, 9), reason="Requires Django >= 1.9")
 def test_cache_spans_middleware(
-    sentry_init,
-    client,
-    capture_events,
-    use_django_caching_with_middlewares,
-    settings,
+    sentry_init, client, capture_events, use_django_caching_with_middlewares
 ):
     sentry_init(
         integrations=[

tests/integrations/pymongo/test_pymongo.py

@@ -57,6 +57,9 @@ def test_transactions(sentry_init, capture_events, mongo_server, with_pii):
     }
     for span in find, insert_success, insert_fail:
         assert span["data"][SPANDATA.DB_SYSTEM] == "mongodb"
+        assert span["data"][SPANDATA.DB_NAME] == "test_db"
+        assert span["data"][SPANDATA.SERVER_ADDRESS] == "localhost"
+        assert span["data"][SPANDATA.SERVER_PORT] == mongo_server.port
         for field, value in common_tags.items():
             assert span["tags"][field] == value
 

tests/integrations/sqlalchemy/test_sqlalchemy.py

@@ -122,6 +122,9 @@ class Address(Base):
 
     for span in event["spans"]:
         assert span["data"][SPANDATA.DB_SYSTEM] == "sqlite"
+        assert span["data"][SPANDATA.DB_NAME] == ":memory:"
+        assert SPANDATA.SERVER_ADDRESS not in span["data"]
+        assert SPANDATA.SERVER_PORT not in span["data"]
 
     assert (
         render_span_tree(event)