From 60fe5c4714181d8f571b27ac63af720d242dd155 Mon Sep 17 00:00:00 2001
From: xuhaidong <haidong.xu@shanbay.com>
Date: Wed, 15 May 2024 16:46:51 +0800
Subject: [PATCH] feat: compatible with mysql clients

---
 CHANGELOG.md         |  4 ++++
 peeweext/__init__.py |  2 +-
 peeweext/otel.py     | 36 +++++++++++++++++++++++++++++++-----
 requirements.txt     |  3 ++-
 4 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 11e250b..e2124e9 100644
--- a/CHANGELOG.md
+++ b/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()`.
diff --git a/peeweext/__init__.py b/peeweext/__init__.py
index af391e6..6a3fa6e 100644
--- a/peeweext/__init__.py
+++ b/peeweext/__init__.py
@@ -1 +1 @@
-__version__ = '1.2.7'
+__version__ = '1.2.8'
diff --git a/peeweext/otel.py b/peeweext/otel.py
index 1cd6c5c..494ce23 100644
--- a/peeweext/otel.py
+++ b/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)
diff --git a/requirements.txt b/requirements.txt
index 45f5ae9..95d0a4e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,5 @@
 peewee
 pendulum>=2.0.0
 blinker
-opentelemetry.instrumentation.mysqlclient
\ No newline at end of file
+opentelemetry-instrumentation-pymysql
+opentelemetry-instrumentation-mysqlclient
\ No newline at end of file