Skip to content

Commit

Permalink
refactor: otel instrument
Browse files Browse the repository at this point in the history
  • Loading branch information
xuhaidong committed May 10, 2024
1 parent e0671ca commit 3a29206
Show file tree
Hide file tree
Showing 3 changed files with 29 additions and 19 deletions.
8 changes: 2 additions & 6 deletions peeweext/flask.py
Original file line number Diff line number Diff line change
@@ -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):
Expand All @@ -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)

Expand Down
20 changes: 20 additions & 0 deletions peeweext/otel.py
Original file line number Diff line number Diff line change
@@ -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()
)
20 changes: 7 additions & 13 deletions peeweext/sea.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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(
Expand Down Expand Up @@ -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))
Expand Down

0 comments on commit 3a29206

Please sign in to comment.