feat: add tracer_provider arg

shanbay · May 10, 2024 · e0671ca · e0671ca
1 parent e6cc55c
commit e0671ca
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 7 deletions.
diff --git a/peeweext/flask.py b/peeweext/flask.py
@@ -1,6 +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
 
 
@@ -21,7 +22,10 @@ def init_app(self, app):
             config.get('model', 'peeweext.model.Model'))
         conn_params = config.get('conn_params', {})
 
-        MySQLClientInstrumentor().instrument()
+        if app.config.get_namespace("OTEL_").get("enable", False):
+            MySQLClientInstrumentor().instrument(
+                tracer_provider=trace.get_tracer_provider()
+            )
         # initialize private connection pool
         self._database = db_url.connect(config['db_url'], **conn_params)
 

diff --git a/peeweext/sea.py b/peeweext/sea.py
@@ -5,22 +5,26 @@
 from playhouse import db_url
 from peewee import DoesNotExist, DataError
 import grpc
+from opentelemetry import trace
 from opentelemetry.instrumentation.mysqlclient import MySQLClientInstrumentor
 
 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', {})
-        MySQLClientInstrumentor().instrument()
-        self.database = db_url.connect(config['db_url'], **conn_params)
+            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)
         self._try_setup_celery()
 
     def _get_db(self):
@@ -45,6 +49,7 @@ 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(
@@ -75,7 +80,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))