From 4c09a2c80a416cb4d33b570e583eac8d6028d1f2 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 25 Apr 2023 14:06:50 +0200 Subject: [PATCH] feat: Add `db.system` to sqlalchemy spans --- sentry_sdk/integrations/sqlalchemy.py | 25 +++++++++++++++++++ .../sqlalchemy/test_sqlalchemy.py | 4 +++ 2 files changed, 29 insertions(+) diff --git a/sentry_sdk/integrations/sqlalchemy.py b/sentry_sdk/integrations/sqlalchemy.py index 64e90aa187..2d6018d732 100644 --- a/sentry_sdk/integrations/sqlalchemy.py +++ b/sentry_sdk/integrations/sqlalchemy.py @@ -3,6 +3,7 @@ import re from sentry_sdk._types import TYPE_CHECKING +from sentry_sdk.consts import SPANDATA from sentry_sdk.hub import Hub from sentry_sdk.integrations import Integration, DidNotEnable from sentry_sdk.tracing_utils import record_sql_queries @@ -67,6 +68,9 @@ 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) context._sentry_sql_span = span @@ -102,3 +106,24 @@ def _handle_error(context, *args): if ctx_mgr is not None: execution_context._sentry_sql_span_manager = None ctx_mgr.__exit__(None, None, None) + + +# See: https://docs.sqlalchemy.org/en/20/dialects/index.html +def _get_db_system(name): + # type: (str) -> Optional[str] + if "sqlite" in name: + return "sqlite" + + if "postgres" in name: + return "postgresql" + + if "mariadb" in name: + return "mariadb" + + if "mysql" in name: + return "mysql" + + if "oracle" in name: + return "oracle" + + return None diff --git a/tests/integrations/sqlalchemy/test_sqlalchemy.py b/tests/integrations/sqlalchemy/test_sqlalchemy.py index d45ea36a19..ebd83f42fb 100644 --- a/tests/integrations/sqlalchemy/test_sqlalchemy.py +++ b/tests/integrations/sqlalchemy/test_sqlalchemy.py @@ -7,6 +7,7 @@ from sqlalchemy.orm import relationship, sessionmaker from sentry_sdk import capture_message, start_transaction, configure_scope +from sentry_sdk.consts import SPANDATA from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration from sentry_sdk.serializer import MAX_EVENT_BYTES from sentry_sdk.utils import json_dumps, MAX_STRING_LENGTH @@ -119,6 +120,9 @@ class Address(Base): (event,) = events + for span in event["spans"]: + assert span["data"][SPANDATA.DB_SYSTEM] == "sqlite" + assert ( render_span_tree(event) == """\