Psycopg types

ddtrace/contrib/dbapi/__init__.py

@@ -57,13 +57,6 @@ def __init__(self, conn):
         name = _get_vendor(conn)
         self._datadog_pin = Pin(service=name, app=name)
 
-    def execute(self, *args, **kwargs):
-        # this method only exists on some clients, so trigger an attribute
-        # error if it doesn't.
-        getattr(self.__wrapped__, 'execute')
-        # otherwise, keep going.
-        return self.cursor().execute(*args, **kwargs)
-
     def cursor(self, *args, **kwargs):
         cursor = self.__wrapped__.cursor(*args, **kwargs)
         pin = self._datadog_pin

ddtrace/contrib/psycopg/__init__.py

@@ -20,6 +20,6 @@
 with require_modules(required_modules) as missing_modules:
     if not missing_modules:
         from .connection import connection_factory
-        from .patch import patch
+        from .patch import patch, patch_conn
 
-        __all__ = ['connection_factory', 'patch']
+        __all__ = ['connection_factory', 'patch', 'patch_conn']

ddtrace/contrib/psycopg/patch.py

@@ -19,20 +19,18 @@ def patch():
     """ Patch monkey patches psycopg's connection function
         so that the connection's functions are traced.
     """
-    setattr(_connect, 'datadog_patched_func', psycopg2.connect)
-    wrapt.wrap_function_wrapper('psycopg2', 'connect', _connect)
+    wrapt.wrap_function_wrapper(psycopg2, 'connect', _connect)
+    _patch_extensions() # do this early just in case
 
-def unpatch():
-    """ Unpatch will undo any monkeypatching. """
-    connect = getattr(_connect, 'datadog_patched_func', None)
-    if connect is not None:
-        psycopg2.connect = connect
-
-def wrap(conn, service="postgres", tracer=None):
+def patch_conn(conn, service="postgres", tracer=None):
     """ Wrap will patch the instance so that it's queries
         are traced. Optionally set the service name of the
         connection.
     """
+    # ensure we've patched extensions (this is idempotent) in
+    # case we're only tracing some connections.
+    _patch_extensions()
+
     c = dbapi.TracedConnection(conn)
 
     # fetch tags from the dsn
@@ -45,15 +43,45 @@ def wrap(conn, service="postgres", tracer=None):
         "db.application" : dsn.get("application_name"),
     }
 
-    pin = Pin(
+    Pin(
         service=service,
         app="postgres",
         tracer=tracer,
-        tags=tags)
+        tags=tags).onto(c)
 
-    pin.onto(c)
     return c
 
+def _patch_extensions():
+    # we must patch extensions all the time (it's pretty harmless) so split
+    # from global patching of connections. must be idempotent.
+    for m, f, w in _extensions:
+        if not hasattr(m, f) or isinstance(getattr(m, f), wrapt.ObjectProxy):
+            continue
+        wrapt.wrap_function_wrapper(m, f, w)
+
+
+#
+# monkeypatch targets
+#
+
 def _connect(connect_func, _, args, kwargs):
-    db = connect_func(*args, **kwargs)
-    return wrap(db)
+    conn = connect_func(*args, **kwargs)
+    return patch_conn(conn)
+
+def _extensions_register_type(func, _, args, kwargs):
+    def _unroll_args(obj, scope=None):
+        return obj, scope
+    obj, scope = _unroll_args(*args, **kwargs)
+
+    # register_type performs a c-level check of the object
+    # type so we must be sure to pass in the actual db connection
+    if scope and isinstance(scope, wrapt.ObjectProxy):
+        scope = scope.__wrapped__
+
+    return func(obj, scope) if scope else func(obj)
+
+
+# extension hooks
+_extensions = [
+    (psycopg2.extensions, 'register_type', _extensions_register_type),
+]

ddtrace/pin.py

@@ -40,3 +40,9 @@ def onto(self, obj):
             return setattr(obj, '_datadog_pin', self)
         except AttributeError:
             log.warn("can't pin onto object", exc_info=True)
+
+    def __repr__(self):
+        return "Pin(service:%s,app:%s,name:%s)" % (
+            self.service,
+            self.app,
+            self.name)

tests/contrib/psycopg/test_psycopg.py

@@ -3,15 +3,17 @@
 
 # 3p
 import psycopg2
+from psycopg2 import extras
 from nose.tools import eq_
 
 # project
 from ddtrace import Tracer
 from ddtrace.contrib.psycopg import connection_factory
 
 # testing
-from ..config import POSTGRES_CONFIG
-from ...test_tracer import DummyWriter
+from tests.contrib.config import POSTGRES_CONFIG
+from tests.test_tracer import get_test_tracer
+from ddtrace.contrib.psycopg import patch_conn
 
 
 TEST_PORT = str(POSTGRES_CONFIG['port'])
@@ -71,21 +73,22 @@ def assert_conn_is_traced(tracer, db, service):
     eq_(span.span_type, "sql")
 
 def test_manual_wrap():
-    from ddtrace.contrib.psycopg.patch import wrap
-    db = psycopg2.connect(**POSTGRES_CONFIG)
-
-    writer = DummyWriter()
-    tracer = Tracer()
-    tracer.writer = writer
-    wrapped = wrap(db, service="foo", tracer=tracer)
+    conn = psycopg2.connect(**POSTGRES_CONFIG)
+    tracer = get_test_tracer()
+    wrapped = patch_conn(conn, service="foo", tracer=tracer)
     assert_conn_is_traced(tracer, wrapped, "foo")
 
-
+def test_manual_wrap_extension_types():
+    conn = psycopg2.connect(**POSTGRES_CONFIG)
+    tracer = get_test_tracer()
+    wrapped = patch_conn(conn, service="foo", tracer=tracer)
+    # NOTE: this will crash if it doesn't work.
+    #   _ext.register_type(_ext.UUID, conn_or_curs)
+    #   TypeError: argument 2 must be a connection, cursor or None
+    extras.register_uuid(conn_or_curs=wrapped)
 
 def test_connect_factory():
-    writer = DummyWriter()
-    tracer = Tracer()
-    tracer.writer = writer
+    tracer = get_test_tracer()
 
     services = ["db", "another"]
     for service in services:
@@ -94,7 +97,7 @@ def test_connect_factory():
         assert_conn_is_traced(tracer, db, service)
 
     # ensure we have the service types
-    services = writer.pop_services()
+    services = tracer.writer.pop_services()
     expected = {
         "db" : {"app":"postgres", "app_type":"db"},
         "another" : {"app":"postgres", "app_type":"db"},