From 4753e3f53e4b86e7db71c30bfb8b55eba2aee900 Mon Sep 17 00:00:00 2001 From: Dillon Fearns Date: Fri, 21 Jul 2023 14:42:14 +0100 Subject: [PATCH 1/3] add support for psycopg3 --- aws_xray_sdk/core/patcher.py | 2 + aws_xray_sdk/ext/psycopg/__init__.py | 4 + aws_xray_sdk/ext/psycopg/patch.py | 69 ++++++++++ docs/index.rst | 1 + tests/ext/psycopg/__init__.py | 0 tests/ext/psycopg/test_psycopg.py | 189 +++++++++++++++++++++++++++ tox.ini | 7 + 7 files changed, 272 insertions(+) create mode 100644 aws_xray_sdk/ext/psycopg/__init__.py create mode 100644 aws_xray_sdk/ext/psycopg/patch.py create mode 100644 tests/ext/psycopg/__init__.py create mode 100644 tests/ext/psycopg/test_psycopg.py diff --git a/aws_xray_sdk/core/patcher.py b/aws_xray_sdk/core/patcher.py index 8e4687e3..04c4c2ec 100644 --- a/aws_xray_sdk/core/patcher.py +++ b/aws_xray_sdk/core/patcher.py @@ -23,6 +23,7 @@ 'pymongo', 'pymysql', 'psycopg2', + 'psycopg', 'pg8000', 'sqlalchemy_core', 'httpx', @@ -38,6 +39,7 @@ 'pymongo', 'pymysql', 'psycopg2', + 'psycopg', 'pg8000', 'sqlalchemy_core', 'httpx', diff --git a/aws_xray_sdk/ext/psycopg/__init__.py b/aws_xray_sdk/ext/psycopg/__init__.py new file mode 100644 index 00000000..0e327081 --- /dev/null +++ b/aws_xray_sdk/ext/psycopg/__init__.py @@ -0,0 +1,4 @@ +from .patch import patch + + +__all__ = ['patch'] diff --git a/aws_xray_sdk/ext/psycopg/patch.py b/aws_xray_sdk/ext/psycopg/patch.py new file mode 100644 index 00000000..eb3a178b --- /dev/null +++ b/aws_xray_sdk/ext/psycopg/patch.py @@ -0,0 +1,69 @@ +import copy +import re +import wrapt +from operator import methodcaller + +from aws_xray_sdk.ext.dbapi2 import XRayTracedConn, XRayTracedCursor + + +def patch(): + wrapt.wrap_function_wrapper( + 'psycopg', + 'connect', + _xray_traced_connect + ) + wrapt.wrap_function_wrapper( + 'psycopg.extensions', + 'register_type', + _xray_register_type_fix + ) + wrapt.wrap_function_wrapper( + 'psycopg.extensions', + 'quote_ident', + _xray_register_type_fix + ) + + wrapt.wrap_function_wrapper( + 'psycopg.extras', + 'register_default_jsonb', + _xray_register_default_jsonb_fix + ) + + +def _xray_traced_connect(wrapped, instance, args, kwargs): + conn = wrapped(*args, **kwargs) + parameterized_dsn = {c[0]: c[-1] for c in map(methodcaller('split', '='), conn.dsn.split(' '))} + meta = { + 'database_type': 'PostgreSQL', + 'url': 'postgresql://{}@{}:{}/{}'.format( + parameterized_dsn.get('user', 'unknown'), + parameterized_dsn.get('host', 'unknown'), + parameterized_dsn.get('port', 'unknown'), + parameterized_dsn.get('dbname', 'unknown'), + ), + 'user': parameterized_dsn.get('user', 'unknown'), + 'database_version': str(conn.server_version), + 'driver_version': 'Psycopg 3' + } + + return XRayTracedConn(conn, meta) + + +def _xray_register_type_fix(wrapped, instance, args, kwargs): + """Send the actual connection or curser to register type.""" + our_args = list(copy.copy(args)) + if len(our_args) == 2 and isinstance(our_args[1], (XRayTracedConn, XRayTracedCursor)): + our_args[1] = our_args[1].__wrapped__ + + return wrapped(*our_args, **kwargs) + + +def _xray_register_default_jsonb_fix(wrapped, instance, args, kwargs): + our_kwargs = dict() + for key, value in kwargs.items(): + if key == "conn_or_curs" and isinstance(value, (XRayTracedConn, XRayTracedCursor)): + # unwrap the connection or cursor to be sent to register_default_jsonb + value = value.__wrapped__ + our_kwargs[key] = value + + return wrapped(*args, **our_kwargs) diff --git a/docs/index.rst b/docs/index.rst index 3d5dced1..6f829717 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -29,6 +29,7 @@ Currently supported web frameworks and libraries: * mysql-connector * pg8000 * psycopg2 +* psycopg (psycopg3) * pymongo * pymysql * pynamodb diff --git a/tests/ext/psycopg/__init__.py b/tests/ext/psycopg/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/ext/psycopg/test_psycopg.py b/tests/ext/psycopg/test_psycopg.py new file mode 100644 index 00000000..abe1245f --- /dev/null +++ b/tests/ext/psycopg/test_psycopg.py @@ -0,0 +1,189 @@ +import psycopg +import psycopg.extras +import psycopg.pool +import psycopg.sql + +import pytest +import testing.postgresql + +from aws_xray_sdk.core import patch +from aws_xray_sdk.core import xray_recorder +from aws_xray_sdk.core.context import Context + +patch(('psycopg',)) + + +@pytest.fixture(autouse=True) +def construct_ctx(): + """ + Clean up context storage on each test run and begin a segment + so that later subsegment can be attached. After each test run + it cleans up context storage again. + """ + xray_recorder.configure(service='test', sampling=False, context=Context()) + xray_recorder.clear_trace_entities() + xray_recorder.begin_segment('name') + yield + xray_recorder.clear_trace_entities() + + +def test_execute_dsn_kwargs(): + q = 'SELECT 1' + with testing.postgresql.Postgresql() as postgresql: + url = postgresql.url() + dsn = postgresql.dsn() + conn = psycopg.connect(dbname=dsn['database'], + user=dsn['user'], + password='', + host=dsn['host'], + port=dsn['port']) + cur = conn.cursor() + cur.execute(q) + + subsegment = xray_recorder.current_segment().subsegments[0] + assert subsegment.name == 'execute' + sql = subsegment.sql + assert sql['database_type'] == 'PostgreSQL' + assert sql['user'] == dsn['user'] + assert sql['url'] == url + assert sql['database_version'] + + +def test_execute_dsn_kwargs_alt_dbname(): + """ + Psycopg supports database to be passed as `database` or `dbname` + """ + q = 'SELECT 1' + + with testing.postgresql.Postgresql() as postgresql: + url = postgresql.url() + dsn = postgresql.dsn() + conn = psycopg.connect(database=dsn['database'], + user=dsn['user'], + password='', + host=dsn['host'], + port=dsn['port']) + cur = conn.cursor() + cur.execute(q) + + subsegment = xray_recorder.current_segment().subsegments[0] + assert subsegment.name == 'execute' + sql = subsegment.sql + assert sql['database_type'] == 'PostgreSQL' + assert sql['user'] == dsn['user'] + assert sql['url'] == url + assert sql['database_version'] + + +def test_execute_dsn_string(): + q = 'SELECT 1' + with testing.postgresql.Postgresql() as postgresql: + url = postgresql.url() + dsn = postgresql.dsn() + conn = psycopg.connect('dbname=' + dsn['database'] + + ' password=mypassword' + + ' host=' + dsn['host'] + + ' port=' + str(dsn['port']) + + ' user=' + dsn['user']) + cur = conn.cursor() + cur.execute(q) + + subsegment = xray_recorder.current_segment().subsegments[0] + assert subsegment.name == 'execute' + sql = subsegment.sql + assert sql['database_type'] == 'PostgreSQL' + assert sql['user'] == dsn['user'] + assert sql['url'] == url + assert sql['database_version'] + + +def test_execute_in_pool(): + q = 'SELECT 1' + with testing.postgresql.Postgresql() as postgresql: + url = postgresql.url() + dsn = postgresql.dsn() + pool = psycopg.pool.SimpleConnectionPool(1, 1, + dbname=dsn['database'], + user=dsn['user'], + password='', + host=dsn['host'], + port=dsn['port']) + cur = pool.getconn(key=dsn['user']).cursor() + cur.execute(q) + + subsegment = xray_recorder.current_segment().subsegments[0] + assert subsegment.name == 'execute' + sql = subsegment.sql + assert sql['database_type'] == 'PostgreSQL' + assert sql['user'] == dsn['user'] + assert sql['url'] == url + assert sql['database_version'] + + +def test_execute_bad_query(): + q = 'SELECT blarg' + with testing.postgresql.Postgresql() as postgresql: + url = postgresql.url() + dsn = postgresql.dsn() + conn = psycopg.connect(dbname=dsn['database'], + user=dsn['user'], + password='', + host=dsn['host'], + port=dsn['port']) + cur = conn.cursor() + try: + cur.execute(q) + except Exception: + pass + + subsegment = xray_recorder.current_segment().subsegments[0] + assert subsegment.name == 'execute' + sql = subsegment.sql + assert sql['database_type'] == 'PostgreSQL' + assert sql['user'] == dsn['user'] + assert sql['url'] == url + assert sql['database_version'] + + exception = subsegment.cause['exceptions'][0] + assert exception.type == 'UndefinedColumn' + + +def test_register_extensions(): + with testing.postgresql.Postgresql() as postgresql: + url = postgresql.url() + dsn = postgresql.dsn() + conn = psycopg.connect('dbname=' + dsn['database'] + + ' password=mypassword' + + ' host=' + dsn['host'] + + ' port=' + str(dsn['port']) + + ' user=' + dsn['user']) + assert psycopg.extras.register_uuid(None, conn) + assert psycopg.extras.register_uuid(None, conn.cursor()) + + +def test_query_as_string(): + with testing.postgresql.Postgresql() as postgresql: + url = postgresql.url() + dsn = postgresql.dsn() + conn = psycopg.connect('dbname=' + dsn['database'] + + ' password=mypassword' + + ' host=' + dsn['host'] + + ' port=' + str(dsn['port']) + + ' user=' + dsn['user']) + test_sql = psycopg.sql.Identifier('test') + assert test_sql.as_string(conn) + assert test_sql.as_string(conn.cursor()) + + +def test_register_default_jsonb(): + with testing.postgresql.Postgresql() as postgresql: + url = postgresql.url() + dsn = postgresql.dsn() + conn = psycopg.connect('dbname=' + dsn['database'] + + ' password=mypassword' + + ' host=' + dsn['host'] + + ' port=' + str(dsn['port']) + + ' user=' + dsn['user']) + + assert psycopg.extras.register_default_jsonb(conn_or_curs=conn, loads=lambda x: x) + assert psycopg.extras.register_default_jsonb(conn_or_curs=conn.cursor(), loads=lambda x: x) diff --git a/tox.ini b/tox.ini index fbe9c333..365a6e4c 100644 --- a/tox.ini +++ b/tox.ini @@ -30,6 +30,8 @@ envlist = py{37,38,39,310,311}-ext-psycopg2 + py{37,38,39,310,311}-ext-psycopg + py{37,38,39,310,311}-ext-pymysql py{37,38,39,310,311}-ext-pynamodb @@ -94,6 +96,9 @@ deps = ext-psycopg2: psycopg2 ext-psycopg2: testing.postgresql + ext-psycopg: psycopg + ext-psycopg: testing.postgresql + ext-pg8000: pg8000 <= 1.20.0 ext-pg8000: testing.postgresql @@ -130,6 +135,8 @@ commands = ext-pg8000: coverage run --append --source aws_xray_sdk -m pytest tests/ext/pg8000 {posargs} ext-psycopg2: coverage run --append --source aws_xray_sdk -m pytest tests/ext/psycopg2 {posargs} + + ext-psycopg: coverage run --append --source aws_xray_sdk -m pytest tests/ext/psycopg {posargs} ext-pymysql: coverage run --append --source aws_xray_sdk -m pytest tests/ext/pymysql {posargs} From 9a92a55c1704fb2cefa756551f64888d8388635f Mon Sep 17 00:00:00 2001 From: Chris Steinle Date: Sun, 20 Aug 2023 15:10:10 +0100 Subject: [PATCH 2/3] change patching and remove unrelated tests --- aws_xray_sdk/ext/psycopg/patch.py | 44 +++--------------- tests/ext/psycopg/test_psycopg.py | 75 +++++-------------------------- 2 files changed, 17 insertions(+), 102 deletions(-) diff --git a/aws_xray_sdk/ext/psycopg/patch.py b/aws_xray_sdk/ext/psycopg/patch.py index eb3a178b..abceb884 100644 --- a/aws_xray_sdk/ext/psycopg/patch.py +++ b/aws_xray_sdk/ext/psycopg/patch.py @@ -1,9 +1,7 @@ -import copy -import re import wrapt from operator import methodcaller -from aws_xray_sdk.ext.dbapi2 import XRayTracedConn, XRayTracedCursor +from aws_xray_sdk.ext.dbapi2 import XRayTracedConn def patch(): @@ -12,27 +10,17 @@ def patch(): 'connect', _xray_traced_connect ) - wrapt.wrap_function_wrapper( - 'psycopg.extensions', - 'register_type', - _xray_register_type_fix - ) - wrapt.wrap_function_wrapper( - 'psycopg.extensions', - 'quote_ident', - _xray_register_type_fix - ) wrapt.wrap_function_wrapper( - 'psycopg.extras', - 'register_default_jsonb', - _xray_register_default_jsonb_fix + 'psycopg_pool.pool', + 'ConnectionPool._connect', + _xray_traced_connect ) def _xray_traced_connect(wrapped, instance, args, kwargs): conn = wrapped(*args, **kwargs) - parameterized_dsn = {c[0]: c[-1] for c in map(methodcaller('split', '='), conn.dsn.split(' '))} + parameterized_dsn = {c[0]: c[-1] for c in map(methodcaller('split', '='), conn.info.dsn.split(' '))} meta = { 'database_type': 'PostgreSQL', 'url': 'postgresql://{}@{}:{}/{}'.format( @@ -42,28 +30,8 @@ def _xray_traced_connect(wrapped, instance, args, kwargs): parameterized_dsn.get('dbname', 'unknown'), ), 'user': parameterized_dsn.get('user', 'unknown'), - 'database_version': str(conn.server_version), + 'database_version': str(conn.info.server_version), 'driver_version': 'Psycopg 3' } return XRayTracedConn(conn, meta) - - -def _xray_register_type_fix(wrapped, instance, args, kwargs): - """Send the actual connection or curser to register type.""" - our_args = list(copy.copy(args)) - if len(our_args) == 2 and isinstance(our_args[1], (XRayTracedConn, XRayTracedCursor)): - our_args[1] = our_args[1].__wrapped__ - - return wrapped(*our_args, **kwargs) - - -def _xray_register_default_jsonb_fix(wrapped, instance, args, kwargs): - our_kwargs = dict() - for key, value in kwargs.items(): - if key == "conn_or_curs" and isinstance(value, (XRayTracedConn, XRayTracedCursor)): - # unwrap the connection or cursor to be sent to register_default_jsonb - value = value.__wrapped__ - our_kwargs[key] = value - - return wrapped(*args, **our_kwargs) diff --git a/tests/ext/psycopg/test_psycopg.py b/tests/ext/psycopg/test_psycopg.py index abe1245f..df55ee2d 100644 --- a/tests/ext/psycopg/test_psycopg.py +++ b/tests/ext/psycopg/test_psycopg.py @@ -1,7 +1,6 @@ import psycopg -import psycopg.extras -import psycopg.pool import psycopg.sql +import psycopg_pool import pytest import testing.postgresql @@ -49,32 +48,6 @@ def test_execute_dsn_kwargs(): assert sql['database_version'] -def test_execute_dsn_kwargs_alt_dbname(): - """ - Psycopg supports database to be passed as `database` or `dbname` - """ - q = 'SELECT 1' - - with testing.postgresql.Postgresql() as postgresql: - url = postgresql.url() - dsn = postgresql.dsn() - conn = psycopg.connect(database=dsn['database'], - user=dsn['user'], - password='', - host=dsn['host'], - port=dsn['port']) - cur = conn.cursor() - cur.execute(q) - - subsegment = xray_recorder.current_segment().subsegments[0] - assert subsegment.name == 'execute' - sql = subsegment.sql - assert sql['database_type'] == 'PostgreSQL' - assert sql['user'] == dsn['user'] - assert sql['url'] == url - assert sql['database_version'] - - def test_execute_dsn_string(): q = 'SELECT 1' with testing.postgresql.Postgresql() as postgresql: @@ -102,14 +75,16 @@ def test_execute_in_pool(): with testing.postgresql.Postgresql() as postgresql: url = postgresql.url() dsn = postgresql.dsn() - pool = psycopg.pool.SimpleConnectionPool(1, 1, - dbname=dsn['database'], - user=dsn['user'], - password='', - host=dsn['host'], - port=dsn['port']) - cur = pool.getconn(key=dsn['user']).cursor() - cur.execute(q) + pool = psycopg_pool.ConnectionPool('dbname=' + dsn['database'] + + ' password=mypassword' + + ' host=' + dsn['host'] + + ' port=' + str(dsn['port']) + + ' user=' + dsn['user'], + min_size=1, + max_size=1) + with pool.connection() as conn: + cur = conn.cursor() + cur.execute(q) subsegment = xray_recorder.current_segment().subsegments[0] assert subsegment.name == 'execute' @@ -147,20 +122,6 @@ def test_execute_bad_query(): exception = subsegment.cause['exceptions'][0] assert exception.type == 'UndefinedColumn' - -def test_register_extensions(): - with testing.postgresql.Postgresql() as postgresql: - url = postgresql.url() - dsn = postgresql.dsn() - conn = psycopg.connect('dbname=' + dsn['database'] + - ' password=mypassword' + - ' host=' + dsn['host'] + - ' port=' + str(dsn['port']) + - ' user=' + dsn['user']) - assert psycopg.extras.register_uuid(None, conn) - assert psycopg.extras.register_uuid(None, conn.cursor()) - - def test_query_as_string(): with testing.postgresql.Postgresql() as postgresql: url = postgresql.url() @@ -173,17 +134,3 @@ def test_query_as_string(): test_sql = psycopg.sql.Identifier('test') assert test_sql.as_string(conn) assert test_sql.as_string(conn.cursor()) - - -def test_register_default_jsonb(): - with testing.postgresql.Postgresql() as postgresql: - url = postgresql.url() - dsn = postgresql.dsn() - conn = psycopg.connect('dbname=' + dsn['database'] + - ' password=mypassword' + - ' host=' + dsn['host'] + - ' port=' + str(dsn['port']) + - ' user=' + dsn['user']) - - assert psycopg.extras.register_default_jsonb(conn_or_curs=conn, loads=lambda x: x) - assert psycopg.extras.register_default_jsonb(conn_or_curs=conn.cursor(), loads=lambda x: x) From 609269735790b16a7651050880f90232e38df79f Mon Sep 17 00:00:00 2001 From: Chris Steinle Date: Wed, 23 Aug 2023 09:23:20 +0100 Subject: [PATCH 3/3] add missing psycopg[pool] dependency --- tox.ini | 1 + 1 file changed, 1 insertion(+) diff --git a/tox.ini b/tox.ini index 365a6e4c..659ef762 100644 --- a/tox.ini +++ b/tox.ini @@ -97,6 +97,7 @@ deps = ext-psycopg2: testing.postgresql ext-psycopg: psycopg + ext-psycopg: psycopg[pool] ext-psycopg: testing.postgresql ext-pg8000: pg8000 <= 1.20.0