add support for psycopg3

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)

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)