Merge pull request #393 from stj/master

psycopg2: Patch all imports of register_type

Author: Emanuele Palazzetti

ddtrace/contrib/psycopg/patch.py

@@ -94,9 +94,40 @@ 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,
      psycopg2.extensions, 'register_type',
      _extensions_register_type),
+    (psycopg2._psycopg.register_type,
+     psycopg2._psycopg, 'register_type',
+     _extensions_register_type),
+    (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_
 
@@ -115,6 +117,28 @@ def test_manual_wrap_extension_types(self):
         #   TypeError: argument 2 must be a connection, cursor or None
         extras.register_uuid(conn_or_curs=conn)
 
+        # NOTE: this will crash if it doesn't work.
+        #   _ext.register_default_json(conn)
+        #   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()
 

tox.ini

@@ -89,6 +89,7 @@ deps =
     aiohttp20: aiohttp>=2.0,<2.1
     aiohttp21: aiohttp>=2.1,<2.2
     aiohttp22: aiohttp>=2.2,<2.3
+    aiohttp23: aiohttp>=2.3,<2.4
     tornado40: tornado>=4.0,<4.1
     tornado41: tornado>=4.1,<4.2
     tornado42: tornado>=4.2,<4.3
@@ -97,6 +98,7 @@ deps =
     futures: futures>=3.0,<3.1
     aiohttp_jinja012: aiohttp_jinja2>=0.12,<0.13
     aiohttp_jinja013: aiohttp_jinja2>=0.13,<0.14
+    aiohttp_jinja014: aiohttp_jinja2>=0.14,<0.15
     blinker: blinker
     boto: boto
     boto: moto<1.0
@@ -193,8 +195,11 @@ deps =
     requests212: requests-mock>=1.3
     requests213: requests>=2.13,<2.14
     requests213: requests-mock>=1.3
+    requests218: requests>=2.18,<2.18
+    requests218: requests-mock>=1.4
     sqlalchemy10: sqlalchemy>=1.0,<1.1
-    sqlalchemy11: sqlalchemy==1.1.0b3
+    sqlalchemy11: sqlalchemy>=1.1,<1.2
+    sqlalchemy12: sqlalchemy>=1.2,<1.3
     webtest: WebTest
 
 # pass along test env variables