From a83b482530b72bf49235568d2dfc252422de3f32 Mon Sep 17 00:00:00 2001 From: Nicholas Labich Date: Fri, 10 Jan 2020 15:34:50 -0500 Subject: [PATCH] Unwrap context provided to psycopg2.extensions.quote_ident --- aws_xray_sdk/ext/psycopg2/patch.py | 5 +++++ tests/ext/psycopg2/test_psycopg2.py | 15 +++++++++++++++ 2 files changed, 20 insertions(+) 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())