Misc updates after wrapt upgrade

Author: TimPansino

newrelic/common/object_wrapper.py

@@ -21,11 +21,16 @@
 
 import inspect
 
+from newrelic.packages.wrapt import (  # noqa: F401
+    BaseObjectProxy,
+    apply_patch,
+    resolve_path,
+    wrap_object,
+    wrap_object_attribute,
+)
 from newrelic.packages.wrapt import BoundFunctionWrapper as _BoundFunctionWrapper
 from newrelic.packages.wrapt import CallableObjectProxy as _CallableObjectProxy
 from newrelic.packages.wrapt import FunctionWrapper as _FunctionWrapper
-from newrelic.packages.wrapt import ObjectProxy as _ObjectProxy
-from newrelic.packages.wrapt import apply_patch, resolve_path, wrap_object, wrap_object_attribute  # noqa: F401
 
 # We previously had our own pure Python implementation of the generic
 # object wrapper but we now defer to using the wrapt module as its C
@@ -44,7 +49,7 @@
 # ObjectProxy or FunctionWrapper should be used going forward.
 
 
-class ObjectProxy(_ObjectProxy):
+class ObjectProxy(BaseObjectProxy):
     """
     This class provides method overrides for all object wrappers used by the
     agent. These methods allow attributes to be defined with the special prefix

newrelic/hooks/database_dbapi2.py

@@ -89,6 +89,9 @@ def callproc(self, procname, parameters=DEFAULT):
             else:
                 return self.__wrapped__.callproc(procname)
 
+    def __iter__(self):
+        return iter(self.__wrapped__)
+
 
 class ConnectionWrapper(ObjectProxy):
     __cursor_wrapper__ = CursorWrapper

tests/datastore_psycopg/test_cursor.py

@@ -98,7 +98,7 @@ async def _execute(connection, cursor, row_type, wrapper):
     # Consume inserted records to check that returning param functions
     records = []
     while True:
-        records.append(cursor.fetchone())
+        records.append(await maybe_await(cursor.fetchone()))
         if not cursor.nextset():
             break
     assert len(records) == len(params)
@@ -140,7 +140,7 @@ async def _exercise_db(connection, row_factory=None, use_cur_context=False, row_
     try:
         cursor = connection.cursor(**kwargs)
         if use_cur_context:
-            if hasattr(cursor, "__aenter__"):
+            if hasattr(cursor.__wrapped__, "__aenter__"):
                 async with cursor:
                     await _execute(connection, cursor, row_type, wrapper)
             else:

tests/datastore_psycopg/test_register.py

@@ -32,7 +32,7 @@ def test():
         psycopg.types.json.set_json_loads(loads=lambda x: x, context=connection)
         psycopg.types.json.set_json_loads(loads=lambda x: x, context=cursor)
 
-    if hasattr(connection, "__aenter__"):
+    if hasattr(connection.__wrapped__, "__aenter__"):
 
         async def coro():
             async with connection:
@@ -69,7 +69,7 @@ async def test():
 
         await maybe_await(cursor.execute(f"DROP TYPE if exists {type_name}"))
 
-    if hasattr(connection, "__aenter__"):
+    if hasattr(connection.__wrapped__, "__aenter__"):
 
         async def coro():
             async with connection:

tests/datastore_psycopg/test_rollback.py

@@ -57,7 +57,7 @@
 
 async def _exercise_db(connection):
     try:
-        if hasattr(connection, "__aenter__"):
+        if hasattr(connection.__wrapped__, "__aenter__"):
             async with connection:
                 raise RuntimeError("error")
         else: