feat(integrations): Add integration for aiomysql

.github/workflows/test-integrations-dbs.yml

@@ -56,6 +56,24 @@ jobs:
       SENTRY_PYTHON_TEST_POSTGRES_HOST: ${{ matrix.python-version == '3.6' && 'postgres' || 'localhost' }}
       SENTRY_PYTHON_TEST_POSTGRES_USER: postgres
       SENTRY_PYTHON_TEST_POSTGRES_PASSWORD: sentry
+    services:
+      mysql:
+        image: mysql:8.0
+        env:
+          MYSQL_ROOT_PASSWORD: root
+          MYSQL_DATABASE: test
+        options: >-
+          --health-cmd="mysqladmin ping"
+          --health-interval=10s
+          --health-timeout=5s
+          --health-retries=5
+        ports:
+          - 3306:3306
+    env:
+      SENTRY_PYTHON_TEST_MYSQL_HOST: ${{ matrix.python-version == '3.6' && 'mysql' || 'localhost' }}
+      SENTRY_PYTHON_TEST_MYSQL_USER: root
+      SENTRY_PYTHON_TEST_MYSQL_PASSWORD: root
+      SENTRY_PYTHON_TEST_MYSQL_DB: test
     # Use Docker container only for Python 3.6
     container: ${{ matrix.python-version == '3.6' && 'python:3.6' || null }}
     steps:
@@ -73,6 +91,10 @@ jobs:
       - name: Erase coverage
         run: |
           coverage erase
+      - name: Test aiomysql
+        run: |
+          set -x # print commands that are executed
+          ./scripts/runtox.sh "py${{ matrix.python-version }}-aiomysql"
       - name: Test asyncpg
         run: |
           set -x # print commands that are executed

scripts/populate_tox/config.py

@@ -47,6 +47,13 @@
         },
         "python": ">=3.7",
     },
+    "aiomysql": {
+        "package": "aiomysql",
+        "deps": {
+            "*": ["pytest-asyncio"],
+        },
+        "python": ">=3.7",
+    },
     "beam": {
         "package": "apache-beam",
         "python": ">=3.7",

scripts/split_tox_gh_actions/split_tox_gh_actions.py

@@ -46,6 +46,10 @@
     "asyncpg",
 }
 
+FRAMEWORKS_NEEDING_MYSQL = {
+    "aiomysql",
+}
+
 FRAMEWORKS_NEEDING_REDIS = {
     "celery",
 }
@@ -99,6 +103,7 @@
         "gcp",
     ],
     "DBs": [
+        "aiomysql",
         "asyncpg",
         "clickhouse_driver",
         "pymongo",
@@ -324,6 +329,7 @@ def render_template(group, frameworks, py_versions):
         "frameworks": frameworks,
         "needs_clickhouse": bool(set(frameworks) & FRAMEWORKS_NEEDING_CLICKHOUSE),
         "needs_docker": bool(set(frameworks) & FRAMEWORKS_NEEDING_DOCKER),
+        "needs_mysql": bool(set(frameworks) & FRAMEWORKS_NEEDING_MYSQL),
         "needs_postgres": bool(set(frameworks) & FRAMEWORKS_NEEDING_POSTGRES),
         "needs_redis": bool(set(frameworks) & FRAMEWORKS_NEEDING_REDIS),
         "needs_java": bool(set(frameworks) & FRAMEWORKS_NEEDING_JAVA),

scripts/split_tox_gh_actions/templates/test_group.jinja

@@ -38,6 +38,27 @@
       SENTRY_PYTHON_TEST_POSTGRES_USER: postgres
       SENTRY_PYTHON_TEST_POSTGRES_PASSWORD: sentry
 
+    {% endif %}
+    {% if needs_mysql %}
+    services:
+      mysql:
+        image: mysql:8.0
+        env:
+          MYSQL_ROOT_PASSWORD: root
+          MYSQL_DATABASE: test
+        options: >-
+          --health-cmd="mysqladmin ping"
+          --health-interval=10s
+          --health-timeout=5s
+          --health-retries=5
+        ports:
+          - 3306:3306
+    env:
+      SENTRY_PYTHON_TEST_MYSQL_HOST: {% raw %}${{ matrix.python-version == '3.6' && 'mysql' || 'localhost' }}{% endraw %}
+      SENTRY_PYTHON_TEST_MYSQL_USER: root
+      SENTRY_PYTHON_TEST_MYSQL_PASSWORD: root
+      SENTRY_PYTHON_TEST_MYSQL_DB: test
+
     {% endif %}
     # Use Docker container only for Python 3.6
     {% raw %}container: ${{ matrix.python-version == '3.6' && 'python:3.6' || null }}{% endraw %}

sentry_sdk/integrations/aiomysql.py

@@ -0,0 +1,205 @@
+# -*- coding: utf-8 -*-
+from __future__ import annotations
+
+from typing import Any, TypeVar, Callable, Awaitable
+
+import sentry_sdk
+from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.integrations import _check_minimum_version, Integration, DidNotEnable
+from sentry_sdk.tracing_utils import add_query_source, record_sql_queries
+from sentry_sdk.utils import (
+    capture_internal_exceptions,
+    parse_version,
+)
+
+try:
+    import aiomysql  # type: ignore[import-untyped]
+    from aiomysql.cursors import Cursor  # type: ignore[import-untyped]
+except ImportError:
+    raise DidNotEnable("aiomysql not installed.")
+
+
+class AioMySQLIntegration(Integration):
+    identifier = "aiomysql"
+    origin = f"auto.db.{identifier}"
+    _record_params = False
+
+    def __init__(self, *, record_params: bool = False):
+        AioMySQLIntegration._record_params = record_params
+
+    @staticmethod
+    def setup_once() -> None:
+        aiomysql_version = parse_version(aiomysql.__version__)
+        _check_minimum_version(AioMySQLIntegration, aiomysql_version)
+
+        Cursor.execute = _wrap_execute(Cursor.execute)
+        Cursor.executemany = _wrap_executemany(Cursor.executemany)
+        aiomysql.connect = _wrap_connect(aiomysql.connect)
+
+
+T = TypeVar("T")
+
+
+def _normalize_query(query: str | bytes | bytearray) -> str:
+    if isinstance(query, (bytes, bytearray)):
+        query = query.decode("utf-8", errors="replace")
+    return " ".join(query.split())
+
+
+def _wrap_execute(f: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
+    """Wrap Cursor.execute to capture SQL queries."""
+
+    async def _inner(*args: Any, **kwargs: Any) -> T:
+        if sentry_sdk.get_client().get_integration(AioMySQLIntegration) is None:
+            return await f(*args, **kwargs)
+
+        cursor = args[0]
+
+        # Skip if flagged by executemany (avoids double-recording)
+        if getattr(cursor, "_sentry_skip_next_execute", False):
+            cursor._sentry_skip_next_execute = False
+            return await f(*args, **kwargs)
+
+        query = args[1] if len(args) > 1 else kwargs.get("query", "")
+        query_str = _normalize_query(query)
+        params = args[2] if len(args) > 2 else kwargs.get("args")
+
+        conn = _get_connection(cursor)
+
+        integration = sentry_sdk.get_client().get_integration(
+            AioMySQLIntegration
+        )
+        params_list = params if integration and integration._record_params else None
+        param_style = "pyformat" if params_list else None
+
+        with record_sql_queries(
+            cursor=None,
+            query=query_str,
+            params_list=params_list,
+            paramstyle=param_style,
+            executemany=False,
+            span_origin=AioMySQLIntegration.origin,
+        ) as span:
+            if conn:
+                _set_db_data(span, conn)
+            res = await f(*args, **kwargs)
+
+        with capture_internal_exceptions():
+            add_query_source(span)
+
+        return res
+
+    return _inner
+
+
+def _wrap_executemany(
+    f: Callable[..., Awaitable[T]]
+) -> Callable[..., Awaitable[T]]:
+    """Wrap Cursor.executemany to capture SQL queries."""
+
+    async def _inner(*args: Any, **kwargs: Any) -> T:
+        if sentry_sdk.get_client().get_integration(AioMySQLIntegration) is None:
+            return await f(*args, **kwargs)
+
+        cursor = args[0]
+        query = args[1] if len(args) > 1 else kwargs.get("query", "")
+        query_str = _normalize_query(query)
+        seq_of_params = args[2] if len(args) > 2 else kwargs.get("args")
+
+        conn = _get_connection(cursor)
+
+        integration = sentry_sdk.get_client().get_integration(
+            AioMySQLIntegration
+        )
+        params_list = (
+            seq_of_params if integration and integration._record_params else None
+        )
+        param_style = "pyformat" if params_list else None
+
+        # Prevent double-recording: _do_execute_many calls self.execute internally
+        cursor._sentry_skip_next_execute = True
+        try:
+            with record_sql_queries(
+                cursor=None,
+                query=query_str,
+                params_list=params_list,
+                paramstyle=param_style,
+                executemany=True,
+                span_origin=AioMySQLIntegration.origin,
+            ) as span:
+                if conn:
+                    _set_db_data(span, conn)
+                res = await f(*args, **kwargs)
+
+            with capture_internal_exceptions():
+                add_query_source(span)
+
+            return res
+        finally:
+            cursor._sentry_skip_next_execute = False
+
+    return _inner
+
+
+def _get_connection(cursor: Any) -> Any:
+    """Get the underlying connection from a cursor."""
+    return getattr(cursor, "_connection", None)
+
+
+def _wrap_connect(
+    f: Callable[..., Awaitable[T]]
+) -> Callable[..., Awaitable[T]]:
+    """Wrap aiomysql.connect to capture connection spans."""
+
+    async def _inner(*args: Any, **kwargs: Any) -> T:
+        if sentry_sdk.get_client().get_integration(AioMySQLIntegration) is None:
+            return await f(*args, **kwargs)
+
+        host = kwargs.get("host", "localhost")
+        port = kwargs.get("port") or 3306
+        user = kwargs.get("user")
+        db = kwargs.get("db") or kwargs.get("database")
+
+        with sentry_sdk.start_span(
+            op=OP.DB,
+            name="connect",
+            origin=AioMySQLIntegration.origin,
+        ) as span:
+            span.set_data(SPANDATA.DB_SYSTEM, "mysql")
+            span.set_data(SPANDATA.SERVER_ADDRESS, host)
+            span.set_data(SPANDATA.SERVER_PORT, port)
+            span.set_data(SPANDATA.DB_NAME, db)
+            span.set_data(SPANDATA.DB_USER, user)
+
+            with capture_internal_exceptions():
+                sentry_sdk.add_breadcrumb(
+                    message="connect",
+                    category="query",
+                    data=span._data,
+                )
+            res = await f(*args, **kwargs)
+
+        return res
+
+    return _inner
+
+
+def _set_db_data(span: Any, conn: Any) -> None:
+    """Set database-related span data from connection object."""
+    span.set_data(SPANDATA.DB_SYSTEM, "mysql")
+
+    host = getattr(conn, "host", None)
+    if host:
+        span.set_data(SPANDATA.SERVER_ADDRESS, host)
+
+    port = getattr(conn, "port", None)
+    if port:
+        span.set_data(SPANDATA.SERVER_PORT, port)
+
+    database = getattr(conn, "db", None)
+    if database:
+        span.set_data(SPANDATA.DB_NAME, database)
+
+    user = getattr(conn, "user", None)
+    if user:
+        span.set_data(SPANDATA.DB_USER, user)

tests/integrations/aiomysql/__init__.py

@@ -0,0 +1,4 @@
+import pytest
+
+pytest.importorskip("aiomysql")
+pytest.importorskip("pytest_asyncio")