[uss_qualifier][draft stacked] Dss0030 concurrent dss queries with subscriptions

monitoring/monitorlib/fetch/__init__.py

@@ -8,6 +8,7 @@
 from enum import Enum
 from urllib.parse import urlparse
 
+import aiohttp
 import flask
 from loguru import logger
 import requests
@@ -145,6 +146,20 @@ def describe_response(resp: requests.Response) -> ResponseDescription:
     return ResponseDescription(**kwargs)
 
 
+def describe_aiohttp_response(
+    status: int, headers: Dict, resp_json: Dict, duration: datetime.timedelta
+) -> ResponseDescription:
+    kwargs = {
+        "code": status,
+        "headers": headers,
+        "elapsed_s": duration.total_seconds(),
+        "reported": StringBasedDateTime(datetime.datetime.utcnow()),
+        "json": resp_json,
+    }
+
+    return ResponseDescription(**kwargs)
+
+
 def describe_flask_response(resp: flask.Response, elapsed_s: float):
     headers = {k: v for k, v in resp.headers.items()}
     kwargs = {

monitoring/monitorlib/infrastructure.py

@@ -3,7 +3,7 @@
 import functools
 from typing import Dict, List, Optional
 import urllib.parse
-from aiohttp import ClientSession
+from aiohttp import ClientSession, ClientResponse
 
 import jwt
 import requests
@@ -189,19 +189,39 @@ def adjust_request_kwargs(self, url, method, kwargs):
             kwargs["timeout"] = self.timeout_seconds
         return kwargs
 
-    async def put(self, url, **kwargs):
+    async def put_with_headers(self, url, **kwargs):
+        """Issues a PUT and returns the status code, headers, and JSON body."""
         url = self._prefix_url + url
         if "auth" not in kwargs:
             kwargs = self.adjust_request_kwargs(url, "PUT", kwargs)
         async with self._client.put(url, **kwargs) as response:
-            return response.status, await response.json()
+            return (
+                response.status,
+                {k: v for k, v in response.headers.items()},
+                await response.json(),
+            )
 
-    async def get(self, url, **kwargs):
+    async def put(self, url, **kwargs):
+        """Issues a PUT and returns the status code and JSON body."""
+        (status, _, json) = await self.put_with_headers(url, **kwargs)
+        return status, json
+
+    async def get_with_headers(self, url, **kwargs):
+        """Issues a GET and returns the status code, headers, and JSON body."""
         url = self._prefix_url + url
         if "auth" not in kwargs:
             kwargs = self.adjust_request_kwargs(url, "GET", kwargs)
         async with self._client.get(url, **kwargs) as response:
-            return response.status, await response.json()
+            return (
+                response.status,
+                {k: v for k, v in response.headers.items()},
+                await response.json(),
+            )
+
+    async def get(self, url, **kwargs):
+        """Issues a GET and returns the status code and JSON body."""
+        (status, _, json) = await self.get_with_headers(url, **kwargs)
+        return status, json
 
     async def post(self, url, **kwargs):
         url = self._prefix_url + url
@@ -210,12 +230,22 @@ async def post(self, url, **kwargs):
         async with self._client.post(url, **kwargs) as response:
             return response.status, await response.json()
 
-    async def delete(self, url, **kwargs):
+    async def delete_with_headers(self, url, **kwargs):
+        """Issues a DELETE and returns the status code, headers, and JSON body."""
         url = self._prefix_url + url
         if "auth" not in kwargs:
             kwargs = self.adjust_request_kwargs(url, "DELETE", kwargs)
         async with self._client.delete(url, **kwargs) as response:
-            return response.status, await response.json()
+            return (
+                response.status,
+                {k: v for k, v in response.headers.items()},
+                await response.json(),
+            )
+
+    async def delete(self, url, **kwargs):
+        """Issues a DELETE and returns the status code and JSON body."""
+        (status, _, json) = await self.delete_with_headers(url, **kwargs)
+        return status, json
 
 
 def default_scopes(scopes: List[str]):