mirror of
https://github.com/open-telemetry/opentelemetry-python-contrib.git
synced 2025-08-03 04:10:48 +08:00
40 lines
883 B
Python
40 lines
883 B
Python
# 3p
|
|
from ddtrace.vendor import wrapt
|
|
import pymysql
|
|
|
|
# project
|
|
from ddtrace import Pin
|
|
from ddtrace.contrib.dbapi import TracedConnection
|
|
from ...ext import net, db
|
|
|
|
CONN_ATTR_BY_TAG = {
|
|
net.TARGET_HOST: 'host',
|
|
net.TARGET_PORT: 'port',
|
|
db.USER: 'user',
|
|
db.NAME: 'db',
|
|
}
|
|
|
|
|
|
def patch():
|
|
wrapt.wrap_function_wrapper('pymysql', 'connect', _connect)
|
|
|
|
|
|
def unpatch():
|
|
if isinstance(pymysql.connect, wrapt.ObjectProxy):
|
|
pymysql.connect = pymysql.connect.__wrapped__
|
|
|
|
|
|
def _connect(func, instance, args, kwargs):
|
|
conn = func(*args, **kwargs)
|
|
return patch_conn(conn)
|
|
|
|
|
|
def patch_conn(conn):
|
|
tags = {t: getattr(conn, a, '') for t, a in CONN_ATTR_BY_TAG.items()}
|
|
pin = Pin(service='pymysql', app='pymysql', tags=tags)
|
|
|
|
# grab the metadata from the conn
|
|
wrapped = TracedConnection(conn, pin=pin)
|
|
pin.onto(wrapped)
|
|
return wrapped
|