Merge pull request #198 from UnitedIncome/psycopg2-quote-ident-fix

Unwrap context provided to psycopg2.extensions.quote_ident

Author: chanchiem

aws_xray_sdk/ext/psycopg2/patch.py

@@ -18,6 +18,11 @@ def patch():
         'register_type',
         _xray_register_type_fix
     )
+    wrapt.wrap_function_wrapper(
+        'psycopg2.extensions',
+        'quote_ident',
+        _xray_register_type_fix
+    )
 
 
 def _xray_traced_connect(wrapped, instance, args, kwargs):

tests/ext/psycopg2/test_psycopg2.py

@@ -1,6 +1,7 @@
 import psycopg2
 import psycopg2.extras
 import psycopg2.pool
+import psycopg2.sql
 
 import pytest
 import testing.postgresql
@@ -158,3 +159,17 @@ def test_register_extensions():
                                 ' user=' + dsn['user'])
         assert psycopg2.extras.register_uuid(None, conn)
         assert psycopg2.extras.register_uuid(None, conn.cursor())
+
+
+def test_query_as_string():
+    with testing.postgresql.Postgresql() as postgresql:
+        url = postgresql.url()
+        dsn = postgresql.dsn()
+        conn = psycopg2.connect('dbname=' + dsn['database'] +
+                                ' password=mypassword' +
+                                ' host=' + dsn['host'] +
+                                ' port=' + str(dsn['port']) +
+                                ' user=' + dsn['user'])
+        test_sql = psycopg2.sql.Identifier('test')
+        assert test_sql.as_string(conn)
+        assert test_sql.as_string(conn.cursor())