contrib.psycopg: handle extensions.adapt

this function does c-level checks of the type of the argument
so it must be a raw psycopg connection passed in. This adds
a hook to transparently downgrade.

Author: Stefan Tjarks

ddtrace/contrib/psycopg/patch.py

@@ -94,6 +94,28 @@ def _unroll_args(obj, scope=None):
     return func(obj, scope) if scope else func(obj)
 
 
+def _extensions_adapt(func, _, args, kwargs):
+    adapt = func(*args, **kwargs)
+    if hasattr(adapt, 'prepare'):
+        return AdapterWrapper(adapt)
+    return adapt
+
+
+class AdapterWrapper(wrapt.ObjectProxy):
+    def prepare(self, *args, **kwargs):
+        func = self.__wrapped__.prepare
+        if not args:
+            return func(*args, **kwargs)
+        conn = args[0]
+
+        # prepare performs a c-level check of the object type so
+        # we must be sure to pass in the actual db connection
+        if isinstance(conn, wrapt.ObjectProxy):
+            conn = conn.__wrapped__
+
+        return func(conn, *args[1:], **kwargs)
+
+
 # extension hooks
 _psycopg2_extensions = [
     (psycopg2.extensions.register_type,
@@ -105,4 +127,7 @@ def _unroll_args(obj, scope=None):
     (psycopg2._json.register_type,
      psycopg2._json, 'register_type',
      _extensions_register_type),
+    (psycopg2.extensions.adapt,
+     psycopg2.extensions, 'adapt',
+     _extensions_adapt),
 ]

tests/contrib/psycopg/test_psycopg.py

@@ -3,6 +3,8 @@
 
 # 3p
 import psycopg2
+from psycopg2 import _psycopg
+from psycopg2 import extensions
 from psycopg2 import extras
 from nose.tools import eq_
 
@@ -120,6 +122,23 @@ def test_manual_wrap_extension_types(self):
         #   TypeError: argument 2 must be a connection, cursor or None
         extras.register_default_json(conn)
 
+
+    def test_manual_wrap_extension_adapt(self):
+        conn, _ = self._get_conn_and_tracer()
+        # NOTE: this will crash if it doesn't work.
+        #   items = _ext.adapt([1, 2, 3])
+        #   items.prepare(conn)
+        #   TypeError: argument 2 must be a connection, cursor or None
+        items = extensions.adapt([1, 2, 3])
+        items.prepare(conn)
+
+        # NOTE: this will crash if it doesn't work.
+        #   binary = _ext.adapt(b'12345)
+        #   binary.prepare(conn)
+        #   TypeError: argument 2 must be a connection, cursor or None
+        binary = extensions.adapt(b'12345')
+        binary.prepare(conn)
+
     def test_connect_factory(self):
         tracer = get_dummy_tracer()