add support for psycopg3

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',

aws_xray_sdk/ext/psycopg/__init__.py

@@ -0,0 +1,4 @@
+from .patch import patch
+
+
+__all__ = ['patch']

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)

docs/index.rst

@@ -29,6 +29,7 @@ Currently supported web frameworks and libraries:
 * mysql-connector
 * pg8000
 * psycopg2
+* psycopg (psycopg3)
 * pymongo
 * pymysql
 * pynamodb

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())

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,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
 
@@ -130,6 +136,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}