testing the call to register_default_jsonb

Author: srprash

aws_xray_sdk/ext/psycopg2/patch.py

@@ -23,14 +23,11 @@ def patch():
         _xray_register_type_fix
     )
 
-    # We should in general not patch internal methods (pyscopg2._json.register_type in this case).
-    # But there isn't a better workaround than this for the time being.
-    # Please see the GH issue for details: https://github.com/aws/aws-xray-sdk-python/issues/243
-    wrapt.wrap_function_wrapper(
-        'psycopg2._json',
-        'register_type',
-        _xray_register_type_fix
-    )
+    # wrapt.wrap_function_wrapper(
+    #     'psycopg2.extras',
+    #     'register_default_jsonb',
+    #     _xray_register_default_jsonb_fix
+    # )
 
 
 def _xray_traced_connect(wrapped, instance, args, kwargs):
@@ -59,3 +56,14 @@ def _xray_register_type_fix(wrapped, instance, args, kwargs):
         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/psycopg2/test_psycopg2.py

@@ -173,3 +173,18 @@ def test_query_as_string():
         test_sql = psycopg2.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 = psycopg2.connect('dbname=' + dsn['database'] +
+                                ' password=mypassword' +
+                                ' 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())