[psycopg2] fix for quote_ident typing bug

ddtrace/contrib/psycopg/patch.py

@@ -93,6 +93,17 @@ def _unroll_args(obj, scope=None):
 
     return func(obj, scope) if scope else func(obj)
 
+def _extensions_quote_ident(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)
 
 def _extensions_adapt(func, _, args, kwargs):
     adapt = func(*args, **kwargs)
@@ -136,3 +147,11 @@ def prepare(self, *args, **kwargs):
          psycopg2._json, 'register_type',
          _extensions_register_type),
     ]
+
+# `quote_ident` attribute is only available for psycopg >= 2.7
+if getattr(psycopg2, 'extensions', None) and getattr(psycopg2.extensions,
+                                                     'quote_ident', None):
+    _psycopg2_extensions += [(psycopg2.extensions.quote_ident,
+     psycopg2.extensions, 'quote_ident',
+     _extensions_quote_ident),
+    ]

tests/contrib/psycopg/test_psycopg.py

@@ -22,8 +22,6 @@
 
 PSYCOPG_VERSION = tuple(map(int, psycopg2.__version__.split()[0].split('.')))
 TEST_PORT = str(POSTGRES_CONFIG['port'])
-
-
 class PsycopgCore(object):
 
     # default service
@@ -126,7 +124,6 @@ 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.
@@ -143,6 +140,17 @@ def test_manual_wrap_extension_adapt(self):
         binary = extensions.adapt(b'12345')
         binary.prepare(conn)
 
+    @skipIf(PSYCOPG_VERSION < (2, 7), 'quote_ident not available in psycopg2<2.7')
+    def test_manual_wrap_extension_quote_ident(self):
+        from ddtrace import patch_all
+        patch_all()
+        from psycopg2.extensions import quote_ident
+
+        # NOTE: this will crash if it doesn't work.
+        #   TypeError: argument 2 must be a connection or a cursor
+        conn = psycopg2.connect(**POSTGRES_CONFIG)
+        quote_ident('foo', conn)
+
     def test_connect_factory(self):
         tracer = get_dummy_tracer()
 
@@ -214,9 +222,21 @@ def test_patch_unpatch(self):
         assert spans, spans
         eq_(len(spans), 1)
 
+
 def test_backwards_compatibilty_v3():
     tracer = get_dummy_tracer()
     factory = connection_factory(tracer, service="my-postgres-db")
     conn = psycopg2.connect(connection_factory=factory, **POSTGRES_CONFIG)
     conn.cursor().execute("select 'blah'")
 
+
+@skipIf(PSYCOPG_VERSION < (2, 7), 'quote_ident not available in psycopg2<2.7')
+def test_manual_wrap_extension_quote_ident_standalone():
+    from ddtrace import patch_all
+    patch_all()
+    from psycopg2.extensions import quote_ident
+
+    # NOTE: this will crash if it doesn't work.
+    #   TypeError: argument 2 must be a connection or a cursor
+    conn = psycopg2.connect(**POSTGRES_CONFIG)
+    quote_ident('foo', conn)