psycopg/psycopg_async: fix handling of ServerCursor / AsyncServerCursor (#2489)

ServerCursor and AsyncServerCursor execute method has a different
interface than Cursor / AsyncCursor. So implement one that works with
both in our ProxyCursor.

Author: xrmx

elasticapm/instrumentation/packages/asyncio/psycopg_async.py

@@ -55,8 +55,8 @@ def _bake_sql(self, sql):
     def extract_signature(self, sql):
         return extract_signature(sql)
 
-    async def execute(self, query, params=None, *, prepare=None, binary=None, **kwargs):
-        return await self._trace_sql(self.__wrapped__.execute, query, params, prepare=prepare, binary=binary, **kwargs)
+    async def execute(self, query, params=None, **kwargs):
+        return await self._trace_sql(self.__wrapped__.execute, query, params, **kwargs)
 
     async def executemany(self, query, params_seq, **kwargs):
         return await self._trace_sql(self.__wrapped__.executemany, query, params_seq, **kwargs)

elasticapm/instrumentation/packages/psycopg.py

@@ -55,8 +55,8 @@ def _bake_sql(self, sql):
     def extract_signature(self, sql):
         return extract_signature(sql)
 
-    def execute(self, query, params=None, *, prepare=None, binary=None, **kwargs):
-        return self._trace_sql(self.__wrapped__.execute, query, params, prepare=prepare, binary=binary, **kwargs)
+    def execute(self, query, params=None, **kwargs):
+        return self._trace_sql(self.__wrapped__.execute, query, params, **kwargs)
 
     def executemany(self, query, params_seq, **kwargs):
         return self._trace_sql(self.__wrapped__.executemany, query, params_seq, **kwargs)

tests/instrumentation/asyncio_tests/psycopg_tests.py

@@ -119,3 +119,10 @@ async def test_executemany(instrument, postgres_connection, elasticapm_client):
     assert span["action"] == "query"
     assert span["sync"] == False
     assert span["name"] == "INSERT INTO test"
+
+
+async def test_server_cursor_execute(instrument, postgres_connection, elasticapm_client):
+    cursor = postgres_connection.cursor(name="async_server_cursor")
+    assert isinstance(cursor, psycopg.AsyncServerCursor)
+    record = await cursor.execute(query="SELECT 1", params=None, binary=None)
+    assert record

tests/instrumentation/psycopg_tests.py

@@ -279,3 +279,12 @@ def test_psycopg_connection(instrument, elasticapm_transaction, postgres_connect
     host = os.environ.get("POSTGRES_HOST", "localhost")
     assert span["name"] == f"psycopg.connect {host}:5432"
     assert span["action"] == "connect"
+
+
+@pytest.mark.integrationtest
+@pytest.mark.skipif(not has_postgres_configured, reason="PostgresSQL not configured")
+def test_server_cursor_execute(instrument, postgres_connection, elasticapm_client):
+    cursor = postgres_connection.cursor(name="server_cursor")
+    assert isinstance(cursor, psycopg.ServerCursor)
+    record = cursor.execute(query="SELECT 1", params=None, binary=True)
+    assert record