Add support to intrument a single MySQL conn

docs/integrations/mysql.md

@@ -48,6 +48,7 @@ import mysql.connector
 
 logfire.configure()
 
+# To instrument the whole module:
 logfire.instrument_mysql()
 
 connection = mysql.connector.connect(
@@ -59,6 +60,9 @@ connection = mysql.connector.connect(
     use_pure=True,
 )
 
+# Or instrument just the connection:
+# connection = logfire.instrument_mysql(connection)
+
 with logfire.span('Create table and insert data'), connection.cursor() as cursor:
     cursor.execute(
         'CREATE TABLE IF NOT EXISTS test (id INT AUTO_INCREMENT PRIMARY KEY, num integer, data varchar(255));'

logfire-api/logfire_api/__init__.pyi

@@ -51,6 +51,7 @@ __all__ = [
     'instrument_sqlalchemy',
     'instrument_redis',
     'instrument_pymongo',
+    'instrument_mysql',
     'AutoTraceModule',
     'with_tags',
     'with_settings',
@@ -88,6 +89,7 @@ instrument_aiohttp_client = DEFAULT_LOGFIRE_INSTANCE.instrument_aiohttp_client
 instrument_sqlalchemy = DEFAULT_LOGFIRE_INSTANCE.instrument_sqlalchemy
 instrument_redis = DEFAULT_LOGFIRE_INSTANCE.instrument_redis
 instrument_pymongo = DEFAULT_LOGFIRE_INSTANCE.instrument_pymongo
+instrument_mysql = DEFAULT_LOGFIRE_INSTANCE.instrument_mysql
 shutdown = DEFAULT_LOGFIRE_INSTANCE.shutdown
 with_tags = DEFAULT_LOGFIRE_INSTANCE.with_tags
 with_settings = DEFAULT_LOGFIRE_INSTANCE.with_settings

logfire-api/logfire_api/_internal/integrations/mysql.pyi

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from opentelemetry.instrumentation.mysql import MySQLInstrumentor
+
+if TYPE_CHECKING:
+    from mysql.connector.abstracts import MySQLConnectionAbstract
+    from mysql.connector.pooling import PooledMySQLConnection
+    from typing_extensions import TypedDict, TypeVar, Unpack
+
+    MySQLConnection = TypeVar('MySQLConnection', PooledMySQLConnection, MySQLConnectionAbstract, None)
+
+    class MySQLInstrumentKwargs(TypedDict, total=False):
+        skip_dep_check: bool
+
+
+def instrument_mysql(
+    conn: MySQLConnection = None,
+    **kwargs: Unpack[MySQLInstrumentKwargs],
+) -> MySQLConnection:
+    """Instrument the `mysql` module or a specific MySQL connection so that spans are automatically created for each operation.
+
+    This function uses the OpenTelemetry MySQL Instrumentation library to instrument either the entire `mysql` module or a specific MySQL connection.
+
+    Args:
+        conn: The MySQL connection to instrument. If None, the entire `mysql` module is instrumented.
+        **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` methods.
+
+    Returns:
+        If a connection is provided, returns the instrumented connection. If no connection is provided, returns None.
+
+    See the `Logfire.instrument_mysql` method for details.
+    """

logfire-api/logfire_api/_internal/main.pyi

@@ -18,6 +18,7 @@ from .integrations.pymongo import PymongoInstrumentKwargs as PymongoInstrumentKw
 from .integrations.redis import RedisInstrumentKwargs as RedisInstrumentKwargs
 from .integrations.sqlalchemy import SQLAlchemyInstrumentKwargs as SQLAlchemyInstrumentKwargs
 from .integrations.starlette import StarletteInstrumentKwargs as StarletteInstrumentKwargs
+from .integrations.mysql import MySQLConnection as MySQLConnection, MySQLInstrumentKwargs as MySQLInstrumentKwargs
 from .json_encoder import logfire_json_dumps as logfire_json_dumps
 from .json_schema import JsonSchemaProperties as JsonSchemaProperties, attributes_json_schema as attributes_json_schema, attributes_json_schema_properties as attributes_json_schema_properties, create_json_schema as create_json_schema
 from .metrics import ProxyMeterProvider as ProxyMeterProvider
@@ -618,6 +619,22 @@ class Logfire:
         [OpenTelemetry pymongo Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/pymongo/pymongo.html)
             library, specifically `PymongoInstrumentor().instrument()`, to which it passes `**kwargs`.
         """
+    def instrument_mysql(self, conn: MySQLConnection, **kwargs: Unpack[MySQLInstrumentKwargs],
+    ) -> MySQLConnection:
+        """Instrument the `mysql` module or a specific MySQL connection so that spans are automatically created for each operation.
+
+        Uses the
+        [OpenTelemetry MySQL Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/mysql/mysql.html)
+        library.
+
+        Args:
+            conn: The `mysql` connection to instrument, or `None` to instrument all connections.
+            **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` methods.
+
+        Returns:
+            If a connection is provided, returns the instrumented connection. If no connection is provided, returns None.
+
+        """
     def instrument_redis(self, **kwargs: Unpack[RedisInstrumentKwargs]) -> None:
         """Instrument the `redis` module so that spans are automatically created for each operation.
 

logfire/_internal/integrations/mysql.py

@@ -5,15 +5,33 @@
 from opentelemetry.instrumentation.mysql import MySQLInstrumentor
 
 if TYPE_CHECKING:
-    from typing_extensions import TypedDict, Unpack
+    from mysql.connector.abstracts import MySQLConnectionAbstract
+    from mysql.connector.pooling import PooledMySQLConnection
+    from typing_extensions import TypedDict, TypeVar, Unpack
+
+    MySQLConnection = TypeVar('MySQLConnection', PooledMySQLConnection, MySQLConnectionAbstract, None)
 
     class MySQLInstrumentKwargs(TypedDict, total=False):
         skip_dep_check: bool
 
 
-def instrument_mysql(**kwargs: Unpack[MySQLInstrumentKwargs]) -> None:
-    """Instrument the `mysql` module so that spans are automatically created for each operation.
+def instrument_mysql(
+    conn: MySQLConnection = None,
+    **kwargs: Unpack[MySQLInstrumentKwargs],
+) -> MySQLConnection:
+    """Instrument the `mysql` module or a specific MySQL connection so that spans are automatically created for each operation.
+
+    This function uses the OpenTelemetry MySQL Instrumentation library to instrument either the entire `mysql` module or a specific MySQL connection.
+
+    Args:
+        conn: The MySQL connection to instrument. If None, the entire `mysql` module is instrumented.
+        **kwargs: Additional keyword arguments to pass to the OpenTelemetry `instrument` methods.
+
+    Returns:
+        If a connection is provided, returns the instrumented connection. If no connection is provided, returns None.
 
     See the `Logfire.instrument_mysql` method for details.
     """
-    MySQLInstrumentor().instrument(**kwargs)  # type: ignore[reportUnknownMemberType]
+    if conn is not None:
+        return MySQLInstrumentor().instrument_connection(conn)  # type: ignore[reportUnknownMemberType]
+    return MySQLInstrumentor().instrument(**kwargs)  # type: ignore[reportUnknownMemberType]