Repository: DataDog/dd-trace-py

Issue:
[psycopg2] The quote_ident function requires psycopg2.extensions.connection.
This is very related to https://github.com/DataDog/dd-trace-py/issues/383.

To reproduce:
```
from ddtrace import patch_all
patch_all()

import psycopg2
from psycopg2.extensions import quote_ident

conn = psycopg2.connect(dbname="test")
quote_ident('foo', conn) # Fails with TypeError: argument 2 must be a connection or a cursor
```


[psycopg2] The prepare functions require psycopg2.extensions.connection.
Here is an example:

```
from ddtrace import patch_all

patch_all()

import psycopg2
from psycopg2.extensions import adapt

conn = psycopg2.connect(dbname="test")

items = adapt([1, 2, 3])
items.prepare(conn) # Fails here

binary = adapt(b'12345')
binary.prepare(conn) # and here.
```

```
$ python example.py
Traceback (most recent call last):
  File "example.py", line 11, in <module>
    items.prepare(conn)
TypeError: argument 1 must be psycopg2.extensions.connection, not TracedConnection
```

This could probably be solved in a similar manner to #96 .


Discussion:

I too am having this issue. 

Resolving patch:
diff --git a/ddtrace/contrib/psycopg/patch.py b/ddtrace/contrib/psycopg/patch.py
--- a/ddtrace/contrib/psycopg/patch.py
+++ b/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),
+    ]


Test patch:
diff --git a/tests/contrib/psycopg/test_psycopg.py b/tests/contrib/psycopg/test_psycopg.py
--- a/tests/contrib/psycopg/test_psycopg.py
+++ b/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)