Fix asyncpg tests to make it compatible with asyncpg>=0.29 (#1935)

* Fix asyncpg get_query to make it compatible with asyncpg>=0.29

* Keep testing with asyncpg<=0.28

* Upgrade flake8 in pre-commit

* Remove 0.28 test from PR framework list

---------

Co-authored-by: Colton Myers <colton@basepi.net>

Author: iurisilvio

.ci/.matrix_framework_full.yml

@@ -67,6 +67,7 @@ FRAMEWORK:
   - aiohttp-3.0
   - aiohttp-newest
   - aiopg-newest
+  - asyncpg-0.28
   - asyncpg-newest
   - tornado-newest
   - starlette-0.13

.pre-commit-config.yaml

@@ -11,7 +11,7 @@ repos:
         language_version: python3
         exclude: "(tests/utils/stacks/linenos.py|tests/utils/stacks/linenos2.py|tests/contrib/grpc/grpc_app/.*pb2.*.py)"
 -   repo: https://github.com/PyCQA/flake8
-    rev: 5.0.4
+    rev: 6.1.0
     hooks:
     -   id: flake8
         exclude: "(tests/utils/stacks/linenos.py|tests/utils/stacks/linenos2.py|tests/contrib/grpc/grpc_app/.*pb2.*.py)"

elasticapm/contrib/asgi.py

@@ -92,13 +92,15 @@ async def __call__(self, scope: "Scope", receive: "ASGIReceiveCallable", send: "
                 body = str(body_raw, errors="ignore")
 
                 # Dispatch to the ASGI callable
-                async def wrapped_receive():
+                async def new_wrapped_receive():
                     if messages:
                         return messages.pop(0)
 
                     # Once that's done we can just await any other messages.
                     return await receive()
 
+                wrapped_receive = new_wrapped_receive
+
             await set_context(lambda: self.get_data_from_request(scope, constants.TRANSACTION, body), "request")
 
         try:

elasticapm/instrumentation/packages/asyncio/asyncpg.py

@@ -55,14 +55,17 @@ class AsyncPGInstrumentation(AsyncAbstractInstrumentedModule):
         ("asyncpg.protocol.protocol", "Protocol.copy_out"),
     ]
 
-    def get_query(self, method, args):
+    def get_query(self, method, args, kwargs=None):
         if method in ["Protocol.query", "Protocol.copy_in", "Protocol.copy_out"]:
             return args[0]
-        else:
+        elif args:
             return args[0].query
+        else:
+            # asyncpg>=0.29 pass data as kwargs
+            return kwargs["state"].query
 
     async def call(self, module, method, wrapped, instance, args, kwargs):
-        query = self.get_query(method, args)
+        query = self.get_query(method, args, kwargs)
         name = extract_signature(query)
         sql_string = shorten(query, string_length=10000)
         context = {"db": {"type": "sql", "statement": sql_string}}

tests/requirements/reqs-asyncpg-0.28.txt

@@ -0,0 +1,2 @@
+asyncpg==0.28
+-r reqs-base.txt