Fix Psycopg3 Api Incompatibilities (#1197)

* Fix DBAPI2 wrappers to allow additional kwargs to executemany

* Pass extra psycopg3 kwargs through executemany

* Rename arguments in psycopg3 wrappers to match upstream

* Add tests for psycopg3 returning param

---------

Co-authored-by: Uma Annamalai <[email protected]>

Author: TimPansino

newrelic/hooks/database_dbapi2.py

@@ -51,12 +51,13 @@ def execute(self, sql, parameters=DEFAULT, *args, **kwargs):
             ):
                 return self.__wrapped__.execute(sql, **kwargs)
 
-    def executemany(self, sql, seq_of_parameters):
+    def executemany(self, sql, seq_of_parameters, *args, **kwargs):
         try:
             seq_of_parameters = list(seq_of_parameters)
             parameters = seq_of_parameters[0]
         except (TypeError, IndexError):
             parameters = DEFAULT
+
         if parameters is not DEFAULT:
             with DatabaseTrace(
                 sql=sql,
@@ -66,7 +67,7 @@ def executemany(self, sql, seq_of_parameters):
                 sql_parameters=parameters,
                 source=self.__wrapped__.executemany,
             ):
-                return self.__wrapped__.executemany(sql, seq_of_parameters)
+                return self.__wrapped__.executemany(sql, seq_of_parameters, *args, **kwargs)
         else:
             with DatabaseTrace(
                 sql=sql,
@@ -75,7 +76,7 @@ def executemany(self, sql, seq_of_parameters):
                 cursor_params=self._nr_cursor_params,
                 source=self.__wrapped__.executemany,
             ):
-                return self.__wrapped__.executemany(sql, seq_of_parameters)
+                return self.__wrapped__.executemany(sql, seq_of_parameters, *args, **kwargs)
 
     def callproc(self, procname, parameters=DEFAULT):
         with DatabaseTrace(

newrelic/hooks/database_dbapi2_async.py

@@ -49,12 +49,13 @@ async def execute(self, sql, parameters=DEFAULT, *args, **kwargs):
             ):
                 return await self.__wrapped__.execute(sql, **kwargs)
 
-    async def executemany(self, sql, seq_of_parameters):
+    async def executemany(self, sql, seq_of_parameters, *args, **kwargs):
         try:
             seq_of_parameters = list(seq_of_parameters)
             parameters = seq_of_parameters[0]
         except (TypeError, IndexError):
             parameters = DEFAULT
+
         if parameters is not DEFAULT:
             with DatabaseTrace(
                 sql=sql,
@@ -64,7 +65,7 @@ async def executemany(self, sql, seq_of_parameters):
                 sql_parameters=parameters,
                 source=self.__wrapped__.executemany,
             ):
-                return await self.__wrapped__.executemany(sql, seq_of_parameters)
+                return await self.__wrapped__.executemany(sql, seq_of_parameters, *args, **kwargs)
         else:
             with DatabaseTrace(
                 sql=sql,
@@ -73,7 +74,7 @@ async def executemany(self, sql, seq_of_parameters):
                 cursor_params=self._nr_cursor_params,
                 source=self.__wrapped__.executemany,
             ):
-                return await self.__wrapped__.executemany(sql, seq_of_parameters)
+                return await self.__wrapped__.executemany(sql, seq_of_parameters, *args, **kwargs)
 
     async def callproc(self, procname, parameters=DEFAULT):
         with DatabaseTrace(

newrelic/hooks/database_psycopg.py

@@ -87,49 +87,49 @@ def __enter__(self):
 
         return self
 
-    def execute(self, sql, parameters=DEFAULT, *args, **kwargs):
-        if hasattr(sql, "as_string"):
-            sql = sql.as_string(self)
+    def execute(self, query, params=DEFAULT, *args, **kwargs):
+        if hasattr(query, "as_string"):
+            query = query.as_string(self)
 
-        return super(CursorWrapper, self).execute(sql, parameters, *args, **kwargs)
+        return super(CursorWrapper, self).execute(query, params, *args, **kwargs)
 
-    def executemany(self, sql, seq_of_parameters):
-        if hasattr(sql, "as_string"):
-            sql = sql.as_string(self)
+    def executemany(self, query, params_seq, *args, **kwargs):
+        if hasattr(query, "as_string"):
+            query = query.as_string(self)
 
-        return super(CursorWrapper, self).executemany(sql, seq_of_parameters)
+        return super(CursorWrapper, self).executemany(query, params_seq, *args, **kwargs)
 
 
 class ConnectionSaveParamsWrapper(DBAPI2ConnectionWrapper):
 
     __cursor_wrapper__ = CursorWrapper
 
-    def execute(self, sql, parameters=DEFAULT, *args, **kwargs):
-        if hasattr(sql, "as_string"):
-            sql = sql.as_string(self)
+    def execute(self, query, params=DEFAULT, *args, **kwargs):
+        if hasattr(query, "as_string"):
+            query = query.as_string(self)
 
-        if parameters is not DEFAULT:
+        if params is not DEFAULT:
             with DatabaseTrace(
-                sql=sql,
+                sql=query,
                 dbapi2_module=self._nr_dbapi2_module,
                 connect_params=self._nr_connect_params,
                 cursor_params=None,
-                sql_parameters=parameters,
+                sql_parameters=params,
                 execute_params=(args, kwargs),
                 source=self.__wrapped__.execute,
             ):
-                cursor = self.__wrapped__.execute(sql, parameters, *args, **kwargs)
+                cursor = self.__wrapped__.execute(query, params, *args, **kwargs)
         else:
             with DatabaseTrace(
-                sql=sql,
+                sql=query,
                 dbapi2_module=self._nr_dbapi2_module,
                 connect_params=self._nr_connect_params,
                 cursor_params=None,
                 sql_parameters=None,
                 execute_params=(args, kwargs),
                 source=self.__wrapped__.execute,
             ):
-                cursor = self.__wrapped__.execute(sql, **kwargs)
+                cursor = self.__wrapped__.execute(query, **kwargs)
 
         return self.__cursor_wrapper__(cursor, self._nr_dbapi2_module, self._nr_connect_params, (args, kwargs))
 
@@ -233,49 +233,49 @@ async def __aenter__(self):
 
         return self
 
-    async def execute(self, sql, parameters=DEFAULT, *args, **kwargs):
-        if hasattr(sql, "as_string"):
-            sql = sql.as_string(self)
+    async def execute(self, query, params=DEFAULT, *args, **kwargs):
+        if hasattr(query, "as_string"):
+            query = query.as_string(self)
 
-        return await super(AsyncCursorWrapper, self).execute(sql, parameters, *args, **kwargs)
+        return await super(AsyncCursorWrapper, self).execute(query, params, *args, **kwargs)
 
-    async def executemany(self, sql, seq_of_parameters):
-        if hasattr(sql, "as_string"):
-            sql = sql.as_string(self)
+    async def executemany(self, query, params_seq, *args, **kwargs):
+        if hasattr(query, "as_string"):
+            query = query.as_string(self)
 
-        return await super(AsyncCursorWrapper, self).executemany(sql, seq_of_parameters)
+        return await super(AsyncCursorWrapper, self).executemany(query, params_seq, *args, **kwargs)
 
 
 class AsyncConnectionSaveParamsWrapper(DBAPI2AsyncConnectionWrapper):
 
     __cursor_wrapper__ = AsyncCursorWrapper
 
-    async def execute(self, sql, parameters=DEFAULT, *args, **kwargs):
-        if hasattr(sql, "as_string"):
-            sql = sql.as_string(self)
+    async def execute(self, query, params=DEFAULT, *args, **kwargs):
+        if hasattr(query, "as_string"):
+            query = query.as_string(self)
 
-        if parameters is not DEFAULT:
+        if params is not DEFAULT:
             with DatabaseTrace(
-                sql=sql,
+                sql=query,
                 dbapi2_module=self._nr_dbapi2_module,
                 connect_params=self._nr_connect_params,
                 cursor_params=None,
-                sql_parameters=parameters,
+                sql_parameters=params,
                 execute_params=(args, kwargs),
                 source=self.__wrapped__.execute,
             ):
-                cursor = await self.__wrapped__.execute(sql, parameters, *args, **kwargs)
+                cursor = await self.__wrapped__.execute(query, params, *args, **kwargs)
         else:
             with DatabaseTrace(
-                sql=sql,
+                sql=query,
                 dbapi2_module=self._nr_dbapi2_module,
                 connect_params=self._nr_connect_params,
                 cursor_params=None,
                 sql_parameters=None,
                 execute_params=(args, kwargs),
                 source=self.__wrapped__.execute,
             ):
-                cursor = await self.__wrapped__.execute(sql, **kwargs)
+                cursor = await self.__wrapped__.execute(query, **kwargs)
 
         return self.__cursor_wrapper__(cursor, self._nr_dbapi2_module, self._nr_connect_params, (args, kwargs))
 

tests/datastore_psycopg/test_cursor.py

@@ -99,9 +99,17 @@ async def _execute(connection, cursor, row_type, wrapper):
     sql = "create table %s (a integer, b real, c text)" % DB_SETTINGS["table_name"]
     await maybe_await(cursor.execute(wrapper(sql)))
 
-    sql = "insert into %s " % DB_SETTINGS["table_name"] + "values (%s, %s, %s)"
+    sql = "insert into %s " % DB_SETTINGS["table_name"] + "values (%s, %s, %s) returning a, b, c"
     params = [(1, 1.0, "1.0"), (2, 2.2, "2.2"), (3, 3.3, "3.3")]
-    await maybe_await(cursor.executemany(wrapper(sql), params))
+    await maybe_await(cursor.executemany(wrapper(sql), params, returning=True))
+
+    # Consume inserted records to check that returning param functions
+    records = []
+    while True:
+        records.append(cursor.fetchone())
+        if not cursor.nextset():
+            break
+    assert len(records) == len(params)
 
     sql = "select * from %s" % DB_SETTINGS["table_name"]
     await maybe_await(cursor.execute(wrapper(sql)))