add support for psycopg3

Author: Dilski

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

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/__init__.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)

tests/ext/psycopg/test_psycopg.py

@@ -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}