diff --git a/aws_xray_sdk/core/patcher.py b/aws_xray_sdk/core/patcher.py index 3fe002a7..2541004f 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..abceb884 --- /dev/null +++ b/aws_xray_sdk/ext/psycopg/patch.py @@ -0,0 +1,37 @@ +import wrapt +from operator import methodcaller + +from aws_xray_sdk.ext.dbapi2 import XRayTracedConn + + +def patch(): + wrapt.wrap_function_wrapper( + 'psycopg', + 'connect', + _xray_traced_connect + ) + + wrapt.wrap_function_wrapper( + '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.info.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.info.server_version), + 'driver_version': 'Psycopg 3' + } + + return XRayTracedConn(conn, meta) 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..df55ee2d --- /dev/null +++ b/tests/ext/psycopg/test_psycopg.py @@ -0,0 +1,136 @@ +import psycopg +import psycopg.sql +import psycopg_pool + +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_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.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' + 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_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()) diff --git a/tox.ini b/tox.ini index 45421f99..37264f1c 100644 --- a/tox.ini +++ b/tox.ini @@ -30,6 +30,8 @@ envlist = py{37,38,39,310,311,312}-ext-psycopg2 + py{37,38,39,310,311}-ext-psycopg + py{37,38,39,310,311,312}-ext-pymysql py{37,38,39,310,311,312}-ext-pynamodb @@ -99,6 +101,10 @@ deps = ext-psycopg2: psycopg2 ext-psycopg2: testing.postgresql + ext-psycopg: psycopg + ext-psycopg: psycopg[pool] + ext-psycopg: testing.postgresql + ext-pg8000: pg8000 <= 1.20.0 ext-pg8000: testing.postgresql @@ -136,6 +142,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}