Merge pull request #67 from xhdnoah/compatible-mysql

feat: compatible with mysql clients

CHANGELOG.md

@@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file.
 
 The format is based on Keep a Changelog, and this project adheres to Semantic Versioning.
 
+## [1.2.8] - 2024-05-15
+
+- Compatible with both MySQLClient and PyMySQL
+
 ## [1.2.7] - 2024-04-28
 
 - Add MySQLClientInstrumentor used to automatic tracing the requests from mysqlClient, the `instrument()` method should be invoked before `connect()`.

peeweext/__init__.py

@@ -1 +1 @@
-__version__ = '1.2.7'
+__version__ = '1.2.8'

peeweext/otel.py

@@ -1,5 +1,27 @@
+from enum import Enum
 from opentelemetry import trace
-from opentelemetry.instrumentation.mysqlclient import MySQLClientInstrumentor
+
+mysql_connector = None
+
+
+class MySQLConnector(Enum):
+    mysqldb = 0
+    pymysql = 1
+
+
+try:
+    import MySQLdb
+    from opentelemetry.instrumentation.mysqlclient import MySQLClientInstrumentor
+
+    mysql_connector = MySQLConnector.mysqldb
+except ImportError:
+    try:
+        import pymysql
+        from opentelemetry.instrumentation.pymysql import PyMySQLInstrumentor
+
+        mysql_connector = MySQLConnector.pymysql
+    except ImportError:
+        pass
 
 
 def sync_once(func):
@@ -14,7 +36,11 @@ def wrapper(*args, **kwargs):
 
 @sync_once
 def otel_instrument(app=None):
-    if app is None or app.config.get_namespace("OTEL_").get("enable", False):
-        MySQLClientInstrumentor().instrument(
-            tracer_provider=trace.get_tracer_provider()
-        )
+    if app is not None and not app.config.get_namespace("OTEL_").get("enable", False):
+        return
+
+    tp = trace.get_tracer_provider()
+    if mysql_connector == MySQLConnector.mysqldb:
+        MySQLClientInstrumentor().instrument(tracer_provider=tp)
+    elif mysql_connector == MySQLConnector.pymysql:
+        PyMySQLInstrumentor().instrument(tracer_provider=tp)

requirements.txt

@@ -1,4 +1,5 @@
 peewee
 pendulum>=2.0.0
 blinker
-opentelemetry.instrumentation.mysqlclient
+opentelemetry-instrumentation-pymysql
+opentelemetry-instrumentation-mysqlclient