patch.py

# 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