Add the logfire.instrument_mysql() (#341)

Co-authored-by: Alex Hall <[email protected]>

Author: aditkumar72

docs/integrations/index.md

@@ -25,6 +25,7 @@ Below you can see more details on how to use Logfire with some of the most popul
 | [Asyncpg](asyncpg.md)               | Databases               |
 | [Psycopg](psycopg.md)               | Databases               |
 | [PyMongo](pymongo.md)               | Databases               |
+| [MySQL](mysql.md)                   | Databases               |
 | [Redis](redis.md)                   | Databases               |
 | [Celery](celery.md)                 | Task Queue              |
 | [System Metrics](system_metrics.md) | System Metrics          |

docs/integrations/mysql.md

@@ -0,0 +1,86 @@
+# MySQL
+
+The [`logfire.instrument_mysql()`][logfire.Logfire.instrument_mysql] method can be used to instrument the [MySQL Connector/Python][mysql-connector] database driver with **Logfire**, creating a span for every query.
+
+## Installation
+
+Install `logfire` with the `mysql` extra:
+
+{{ install_logfire(extras=['mysql']) }}
+
+## Usage
+
+Let's setup a MySQL database using Docker and run a Python script that connects to the database using MySQL connector to
+demonstrate how to use **Logfire** with MySQL.
+
+### Setup a MySQL Database Using Docker
+
+First, we need to initialize a MySQL database. This can be easily done using Docker with the following command:
+
+```bash
+docker run --name mysql \
+    -e MYSQL_ROOT_PASSWORD=secret \
+    -e MYSQL_DATABASE=database \
+    -e MYSQL_USER=user \
+    -e MYSQL_PASSWORD=secret \
+    -p 3306:3306 -d mysql
+```
+
+This command accomplishes the following:
+
+- `--name mysql`: gives the container a name of "mysql".
+- `-e MYSQL_ROOT_PASSWORD=secret` sets the root password to "secret".
+- `-e MYSQL_DATABASE=database` creates a new database named "database".
+- `-e MYSQL_USER=user` creates a new user named "user".
+- `-e MYSQL_PASSWORD=secret` sets the password for the new user to "secret".
+- `-p 3306:3306` maps port 3306 inside Docker as port 3306 on the host machine.
+- `-d mysql` runs the container in the background and prints the container ID. The image is "mysql".
+
+### Run the Python script
+
+The following Python script connects to the MySQL database and executes some SQL queries:
+
+```py
+import logfire
+import mysql.connector
+
+logfire.configure()
+
+# To instrument the whole module:
+logfire.instrument_mysql()
+
+connection = mysql.connector.connect(
+    host="localhost",
+    user="user",
+    password="secret",
+    database="database",
+    port=3306,
+    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));'
+    )
+
+    # Insert some data
+    cursor.execute('INSERT INTO test (num, data) VALUES (%s, %s)', (100, 'abc'))
+    cursor.execute('INSERT INTO test (num, data) VALUES (%s, %s)', (200, 'def'))
+
+    # Query the data
+    cursor.execute('SELECT * FROM test')
+    results = cursor.fetchall()  # Fetch all rows
+    for row in results:
+        print(row)  # Print each row
+```
+
+[`logfire.instrument_mysql()`][logfire.Logfire.instrument_mysql] uses the
+**OpenTelemetry MySQL Instrumentation** package,
+which you can find more information about [here][opentelemetry-mysql].
+
+[opentelemetry-mysql]: https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/mysql/mysql.html
+[mysql]: https://www.mysql.com/
+[mysql-connector]: https://dev.mysql.com/doc/connector-python/en/

logfire-api/logfire_api/__init__.py

@@ -154,6 +154,7 @@ def shutdown(self, *args, **kwargs) -> None: ...
         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
 
         def no_auto_trace(x):

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/__init__.py

@@ -38,6 +38,7 @@
 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_trace_sample_rate = DEFAULT_LOGFIRE_INSTANCE.with_trace_sample_rate
@@ -112,6 +113,7 @@ def loguru_handler() -> dict[str, Any]:
     'instrument_sqlalchemy',
     'instrument_redis',
     'instrument_pymongo',
+    'instrument_mysql',
     'AutoTraceModule',
     'with_tags',
     'with_settings',

logfire/_internal/integrations/mysql.py

@@ -0,0 +1,37 @@
+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', bound=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.
+    """
+    if conn is not None:
+        return MySQLInstrumentor().instrument_connection(conn)  # type: ignore[reportUnknownMemberType]
+    return MySQLInstrumentor().instrument(**kwargs)  # type: ignore[reportUnknownMemberType]

logfire/_internal/main.py

@@ -68,6 +68,7 @@
     from .integrations.celery import CeleryInstrumentKwargs
     from .integrations.flask import FlaskInstrumentKwargs
     from .integrations.httpx import HTTPXInstrumentKwargs
+    from .integrations.mysql import MySQLConnection, MySQLInstrumentKwargs
     from .integrations.psycopg import PsycopgInstrumentKwargs
     from .integrations.pymongo import PymongoInstrumentKwargs
     from .integrations.redis import RedisInstrumentKwargs
@@ -1223,6 +1224,30 @@ def instrument_redis(self, **kwargs: Unpack[RedisInstrumentKwargs]) -> None:
         self._warn_if_not_initialized_for_instrumentation()
         return instrument_redis(**kwargs)
 
+    def instrument_mysql(
+        self,
+        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.
+
+        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.
+
+        """
+        from .integrations.mysql import instrument_mysql
+
+        self._warn_if_not_initialized_for_instrumentation()
+        return instrument_mysql(conn, **kwargs)
+
     def metric_counter(self, name: str, *, unit: str = '', description: str = '') -> Counter:
         """Create a counter metric.
 

mkdocs.yml

@@ -107,6 +107,7 @@ nav:
           - Asyncpg: integrations/asyncpg.md
           - Psycopg: integrations/psycopg.md
           - PyMongo: integrations/pymongo.md
+          - MySQL: integrations/mysql.md
           - Redis: integrations/redis.md
           - Celery: integrations/celery.md
           - System Metrics: integrations/system_metrics.md