From 3a29206a7132cc7806df108c08db1536092e2c19 Mon Sep 17 00:00:00 2001
From: xuhaidong <haidong.xu@shanbay.com>
Date: Fri, 10 May 2024 14:49:23 +0800
Subject: [PATCH] refactor: otel instrument

---
 peeweext/flask.py |  8 ++------
 peeweext/otel.py  | 20 ++++++++++++++++++++
 peeweext/sea.py   | 20 +++++++-------------
 3 files changed, 29 insertions(+), 19 deletions(-)
 create mode 100644 peeweext/otel.py

diff --git a/peeweext/flask.py b/peeweext/flask.py
index 6eaa18c..a0580f3 100644
--- a/peeweext/flask.py
+++ b/peeweext/flask.py
@@ -1,8 +1,7 @@
 from werkzeug.local import LocalProxy
 from werkzeug.utils import import_string, cached_property
 from playhouse import db_url
-from opentelemetry import trace
-from opentelemetry.instrumentation.mysqlclient import MySQLClientInstrumentor
+from .otel import otel_instrument
 
 
 class UninitializedException(Exception):
@@ -22,10 +21,7 @@ def init_app(self, app):
             config.get('model', 'peeweext.model.Model'))
         conn_params = config.get('conn_params', {})
 
-        if app.config.get_namespace("OTEL_").get("enable", False):
-            MySQLClientInstrumentor().instrument(
-                tracer_provider=trace.get_tracer_provider()
-            )
+        otel_instrument(app)
         # initialize private connection pool
         self._database = db_url.connect(config['db_url'], **conn_params)
 
diff --git a/peeweext/otel.py b/peeweext/otel.py
new file mode 100644
index 0000000..1cd6c5c
--- /dev/null
+++ b/peeweext/otel.py
@@ -0,0 +1,20 @@
+from opentelemetry import trace
+from opentelemetry.instrumentation.mysqlclient import MySQLClientInstrumentor
+
+
+def sync_once(func):
+    def wrapper(*args, **kwargs):
+        if not wrapper._done:
+            wrapper._done = True
+            return func(*args, **kwargs)
+
+    wrapper._done = False
+    return wrapper
+
+
+@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()
+        )
diff --git a/peeweext/sea.py b/peeweext/sea.py
index bfe24f6..c337132 100644
--- a/peeweext/sea.py
+++ b/peeweext/sea.py
@@ -5,26 +5,21 @@
 from playhouse import db_url
 from peewee import DoesNotExist, DataError
 import grpc
-from opentelemetry import trace
-from opentelemetry.instrumentation.mysqlclient import MySQLClientInstrumentor
-
+from .otel import otel_instrument
 from .validation import ValidationError
 
 
 class Peeweext:
-    def __init__(self, ns="PW_"):
+    def __init__(self, ns='PW_'):
         self.ns = ns
 
     def init_app(self, app):
         config = app.config.get_namespace(self.ns)
         self.model_class = import_string(
-            config.get("model", "peeweext.model.Model"))
-        conn_params = config.get("conn_params", {})
-        if app.config.get_namespace("OTEL_").get("enable", False):
-            MySQLClientInstrumentor().instrument(
-                tracer_provider=trace.get_tracer_provider()
-            )
-        self.database = db_url.connect(config["db_url"], **conn_params)
+            config.get('model', 'peeweext.model.Model'))
+        conn_params = config.get('conn_params', {})
+        otel_instrument(app)
+        self.database = db_url.connect(config['db_url'], **conn_params)
         self._try_setup_celery()
 
     def _get_db(self):
@@ -49,7 +44,6 @@ def close_db(self):
     def _try_setup_celery(self):
         try:
             from celery.signals import task_prerun, task_postrun
-
             task_prerun.connect(
                 lambda *arg, **kw: self.connect_db(), weak=False)
             task_postrun.connect(
@@ -80,7 +74,7 @@ def __call__(self, servicer, request, context):
             return self.handler(servicer, request, context)
         except DoesNotExist:
             context.set_code(grpc.StatusCode.NOT_FOUND)
-            context.set_details("Record Not Found")
+            context.set_details('Record Not Found')
         except (ValidationError, DataError) as e:
             context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
             context.set_details(str(e))