fix handling of responses with aiohttp4 (#788)

The not-yet-released aiohttp 4.0 changes a bit how HTTP exceptions work. They are no longer inheriting from Response.

Author: beniwohli

elasticapm/contrib/aiohttp/middleware.py

@@ -29,7 +29,7 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 import aiohttp
-from aiohttp.web import Response, middleware
+from aiohttp.web import HTTPException, Response, middleware
 
 import elasticapm
 from elasticapm.conf import constants
@@ -82,13 +82,14 @@ async def handle_request(request, handler):
  context={"request": get_data_from_request(request, elasticapm_client.config, constants.ERROR)}
  )
  elasticapm.set_transaction_result("HTTP 5xx", override=False)
- if isinstance(exc, Response):
+ elasticapm.set_context({"status_code": 500}, "response")
+ # some exceptions are response-like, e.g. have headers and status code. Let's try and capture them
+ if isinstance(exc, (Response, HTTPException)):
  elasticapm.set_context(
  lambda: get_data_from_response(exc, elasticapm_client.config, constants.ERROR), # noqa: F821
  "response",
  )
- else:
- elasticapm.set_context({"status_code": 500}, "response")
+
  raise
  finally:
  elasticapm_client.end_transaction()

elasticapm/contrib/aiohttp/utils.py

@@ -27,8 +27,9 @@
 # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+from typing import Union
 
-from aiohttp.web import Request, Response
+from aiohttp.web import HTTPException, Request, Response
 
 from elasticapm.conf import Config
 from elasticapm.utils import compat, get_url_dict
@@ -49,11 +50,11 @@ def get_data_from_request(request: Request, config: Config, event_type: str):
  return result
 
 
-def get_data_from_response(response: Response, config: Config, event_type: str):
+def get_data_from_response(response: Union[HTTPException, Response], config: Config, event_type: str):
  result = {}
-
- if isinstance(getattr(response, "status", None), compat.integer_types):
- result["status_code"] = response.status
+ status = getattr(response, "status", getattr(response, "status_code", None))
+ if isinstance(status, compat.integer_types):
+ result["status_code"] = status
  if config.capture_headers and getattr(response, "headers", None):
  headers = response.headers
  result["headers"] = {key: ";".join(headers.getall(key)) for key in compat.iterkeys(headers)}