fix typeerror by patch calls to register_default_jsonb

Author: srprash

aws_xray_sdk/ext/psycopg2/patch.py

@@ -23,11 +23,11 @@ def patch():
         _xray_register_type_fix
     )
 
-    # wrapt.wrap_function_wrapper(
-    #     'psycopg2.extras',
-    #     'register_default_jsonb',
-    #     _xray_register_default_jsonb_fix
-    # )
+    wrapt.wrap_function_wrapper(
+        'psycopg2.extras',
+        'register_default_jsonb',
+        _xray_register_default_jsonb_fix
+    )
 
 
 def _xray_traced_connect(wrapped, instance, args, kwargs):
@@ -58,12 +58,12 @@ def _xray_register_type_fix(wrapped, instance, args, kwargs):
     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)
+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/psycopg2/test_psycopg2.py

@@ -184,7 +184,6 @@ def test_register_default_jsonb():
                                 ' host=' + dsn['host'] +
                                 ' port=' + str(dsn['port']) +
                                 ' user=' + dsn['user'])
-        # test_sql = psycopg2.sql.Identifier('test')
 
-        assert psycopg2.extras.register_default_jsonb(None, conn)
-        assert psycopg2.extras.register_uuid(None, conn.cursor())
+        assert psycopg2.extras.register_default_jsonb(conn_or_curs=conn, loads=lambda x: x)
+        assert psycopg2.extras.register_default_jsonb(conn_or_curs=conn.cursor(), loads=lambda x: x)