diff --git a/aws_xray_sdk/ext/psycopg2/patch.py b/aws_xray_sdk/ext/psycopg2/patch.py index 0efb27fa..708bf45c 100644 --- a/aws_xray_sdk/ext/psycopg2/patch.py +++ b/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): diff --git a/tests/ext/psycopg2/test_psycopg2.py b/tests/ext/psycopg2/test_psycopg2.py index e9eb207d..4736b5c0 100644 --- a/tests/ext/psycopg2/test_psycopg2.py +++ b/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())