feat(integrations): Add integration for aiomysql

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

@@ -43,19 +43,33 @@ jobs:
         image: postgres
         env:
           POSTGRES_PASSWORD: sentry
-        # Set health checks to wait until postgres has started
         options: >-
           --health-cmd pg_isready
           --health-interval 10s
           --health-timeout 5s
           --health-retries 5
-        # Maps tcp port 5432 on service container to the host
         ports:
           - 5432:5432
+      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_POSTGRES_HOST: ${{ matrix.python-version == '3.6' && 'postgres' || 'localhost' }}
       SENTRY_PYTHON_TEST_POSTGRES_USER: postgres
       SENTRY_PYTHON_TEST_POSTGRES_PASSWORD: sentry
+      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 +87,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

.github/workflows/test-integrations-web-1.yml

@@ -43,13 +43,11 @@ jobs:
         image: postgres
         env:
           POSTGRES_PASSWORD: sentry
-        # Set health checks to wait until postgres has started
         options: >-
           --health-cmd pg_isready
           --health-interval 10s
           --health-timeout 5s
           --health-retries 5
-        # Maps tcp port 5432 on service container to the host
         ports:
           - 5432:5432
     env:

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

@@ -12,32 +12,54 @@
         # see https://github.com/actions/setup-python/issues/544#issuecomment-1332535877
         os: [ubuntu-22.04]
 
-    {% if needs_docker %}
+    {% if needs_postgres or needs_mysql or needs_docker %}
     services:
-      docker:
-        image: docker:dind  # Required for Docker network management
-        options: --privileged  # Required for Docker-in-Docker operations
-    {% endif %}
     {% if needs_postgres %}
-    services:
       postgres:
         image: postgres
         env:
           POSTGRES_PASSWORD: sentry
-        # Set health checks to wait until postgres has started
         options: >-
           --health-cmd pg_isready
           --health-interval 10s
           --health-timeout 5s
           --health-retries 5
-        # Maps tcp port 5432 on service container to the host
         ports:
           - 5432:5432
+    {% endif %}
+    {% if needs_mysql %}
+      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
+    {% endif %}
+    {% if needs_docker %}
+      docker:
+        image: docker:dind  # Required for Docker network management
+        options: --privileged  # Required for Docker-in-Docker operations
+    {% endif %}
+    {% endif %}
+    {% if needs_postgres or needs_mysql %}
     env:
+    {% if needs_postgres %}
       SENTRY_PYTHON_TEST_POSTGRES_HOST: {% raw %}${{ matrix.python-version == '3.6' && 'postgres' || 'localhost' }}{% endraw %}
       SENTRY_PYTHON_TEST_POSTGRES_USER: postgres
       SENTRY_PYTHON_TEST_POSTGRES_PASSWORD: sentry
-
+    {% endif %}
+    {% if needs_mysql %}
+      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 %}
     {% endif %}
     # Use Docker container only for Python 3.6
     {% raw %}container: ${{ matrix.python-version == '3.6' && 'python:3.6' || null }}{% endraw %}

sentry_sdk/integrations/__init__.py

@@ -66,6 +66,7 @@ def iter_default_integrations(
 
 _AUTO_ENABLING_INTEGRATIONS = [
     "sentry_sdk.integrations.aiohttp.AioHttpIntegration",
+    "sentry_sdk.integrations.aiomysql.AioMySQLIntegration",
     "sentry_sdk.integrations.anthropic.AnthropicIntegration",
     "sentry_sdk.integrations.ariadne.AriadneIntegration",
     "sentry_sdk.integrations.arq.ArqIntegration",
@@ -117,6 +118,7 @@ def iter_default_integrations(
 
 _MIN_VERSIONS = {
     "aiohttp": (3, 4),
+    "aiomysql": (0, 1, 1),
     "anthropic": (0, 16),
     "ariadne": (0, 20),
     "arq": (0, 23),

sentry_sdk/integrations/aiomysql.py

@@ -0,0 +1,198 @@
+# -*- 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.connection import Connection  # 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)
+
+        # Patch Connection._connect — this catches ALL connections:
+        #   - aiomysql.connect()
+        #   - aiomysql.create_pool() (pool.py does `from .connection import connect`
+        #     which ultimately calls Connection._connect)
+        #   - Reconnects
+        Connection._connect = _wrap_connect(Connection._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).
+        # Do NOT reset the flag here — it must stay True for the entire
+        # duration of executemany, which may call execute multiple times
+        # in a loop (non-INSERT fallback). Only _wrap_executemany's
+        # finally block should clear it.
+        if getattr(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 Connection._connect to capture connection spans."""
+
+    async def _inner(self: "Connection") -> T:
+        if sentry_sdk.get_client().get_integration(AioMySQLIntegration) is None:
+            return await f(self)
+
+        with sentry_sdk.start_span(
+            op=OP.DB,
+            name="connect",
+            origin=AioMySQLIntegration.origin,
+        ) as span:
+            _set_db_data(span, self)
+
+            with capture_internal_exceptions():
+                sentry_sdk.add_breadcrumb(
+                    message="connect",
+                    category="query",
+                    data=span._data,
+                )
+            res = await f(self)
+
+        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 is not None:
+        span.set_data(SPANDATA.SERVER_ADDRESS, host)
+
+    port = getattr(conn, "port", None)
+    if port is not None:
+        span.set_data(SPANDATA.SERVER_PORT, port)
+
+    database = getattr(conn, "db", None)
+    if database is not None:
+        span.set_data(SPANDATA.DB_NAME, database)
+
+    user = getattr(conn, "user", None)
+    if user is not None:
+        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")