mirror of
https://github.com/open-telemetry/opentelemetry-python-contrib.git
synced 2025-07-31 22:23:12 +08:00
64 lines
1.6 KiB
Python
64 lines
1.6 KiB
Python
# 3p
|
|
import MySQLdb
|
|
|
|
from ddtrace.vendor.wrapt import wrap_function_wrapper as _w
|
|
|
|
# project
|
|
from ddtrace import Pin
|
|
from ddtrace.contrib.dbapi import TracedConnection
|
|
|
|
from ...ext import net, db
|
|
from ...utils.wrappers import unwrap as _u
|
|
|
|
KWPOS_BY_TAG = {
|
|
net.TARGET_HOST: ('host', 0),
|
|
db.USER: ('user', 1),
|
|
db.NAME: ('db', 3),
|
|
}
|
|
|
|
|
|
def patch():
|
|
# patch only once
|
|
if getattr(MySQLdb, '__datadog_patch', False):
|
|
return
|
|
setattr(MySQLdb, '__datadog_patch', True)
|
|
|
|
# `Connection` and `connect` are aliases for
|
|
# `Connect`; patch them too
|
|
_w('MySQLdb', 'Connect', _connect)
|
|
if hasattr(MySQLdb, 'Connection'):
|
|
_w('MySQLdb', 'Connection', _connect)
|
|
if hasattr(MySQLdb, 'connect'):
|
|
_w('MySQLdb', 'connect', _connect)
|
|
|
|
|
|
def unpatch():
|
|
if not getattr(MySQLdb, '__datadog_patch', False):
|
|
return
|
|
setattr(MySQLdb, '__datadog_patch', False)
|
|
|
|
# unpatch MySQLdb
|
|
_u(MySQLdb, 'Connect')
|
|
if hasattr(MySQLdb, 'Connection'):
|
|
_u(MySQLdb, 'Connection')
|
|
if hasattr(MySQLdb, 'connect'):
|
|
_u(MySQLdb, 'connect')
|
|
|
|
|
|
def _connect(func, instance, args, kwargs):
|
|
conn = func(*args, **kwargs)
|
|
return patch_conn(conn, *args, **kwargs)
|
|
|
|
|
|
def patch_conn(conn, *args, **kwargs):
|
|
tags = {t: kwargs[k] if k in kwargs else args[p]
|
|
for t, (k, p) in KWPOS_BY_TAG.items()
|
|
if k in kwargs or len(args) > p}
|
|
tags[net.TARGET_PORT] = conn.port
|
|
pin = Pin(service='mysql', app='mysql', tags=tags)
|
|
|
|
# grab the metadata from the conn
|
|
wrapped = TracedConnection(conn, pin=pin)
|
|
pin.onto(wrapped)
|
|
return wrapped
|