Add Fastapi contrib

rollbar/__init__.py

@@ -84,6 +84,11 @@
 except ImportError:
     SanicRequest = None
 
+try:
+    from fastapi import Request as FastapiRequest
+except ImportError:
+    FastapiRequest = None
+
 try:
     from google.appengine.api.urlfetch import fetch as AppEngineFetch
 except ImportError:
@@ -1117,6 +1122,10 @@ def _build_request_data(request):
     if FalconRequest and isinstance(request, FalconRequest):
         return _build_falcon_request_data(request)
 
+    # fastapi
+    if FastapiRequest and isinstance(request, FastapiRequest):
+        return _build_fastapi_request_data(request)
+
     # Plain wsgi (should be last)
     if isinstance(request, dict) and 'wsgi.version' in request:
         return _build_wsgi_request_data(request)
@@ -1187,6 +1196,24 @@ def _build_django_request_data(request):
     return request_data
 
 
+def _build_fastapi_request_data(request):
+    """Fastapi relies on starlette which uses async functions to retrieve the request data and body.
+
+    So skipping for now since function callers are synchronous...
+    """
+    request_data = {
+        'url': str(request.url),
+        'method': request.method,
+        'GET': None,
+        'POST': None,
+        'user_ip': _asgi_extract_user_ip(request)
+    }
+
+    request_data['headers'] = request.headers
+
+    return request_data
+
+
 def _build_werkzeug_request_data(request):
     request_data = {
         'url': request.url,
@@ -1623,3 +1650,13 @@ def _wsgi_extract_user_ip(environ):
     if real_ip:
         return real_ip
     return environ['REMOTE_ADDR']
+
+
+def _asgi_extract_user_ip(request):
+    forwarded_for = request.headers.get('x-forwarded-for')
+    if forwarded_for:
+        return forwarded_for
+    real_ip = request.headers.get('x-real-ip')
+    if real_ip:
+        return real_ip
+    return request.client.host

rollbar/contrib/fastapi/__init__.py

@@ -0,0 +1,23 @@
+"""Integration with fastapi.
+
+See: https://fastapi.tiangolo.com/
+"""
+
+import rollbar
+
+from fastapi import Request
+
+
+def report_exception(request: Request):
+    rollbar.report_exc_info(request=request)
+
+
+def _hook(request, data):
+    data["framework"] = "fastapi"
+
+    if request:
+        endpoint = request.scope["endpoint"]
+        data["context"] = f"{endpoint.__module__}.{endpoint.__name__}"
+
+
+rollbar.BASE_DATA_HOOK = _hook

rollbar/examples/fastapi/app.py

@@ -0,0 +1,37 @@
+import rollbar
+import rollbar.contrib.fastapi
+import uvicorn
+
+from fastapi import FastAPI, Request, status
+from fastapi.responses import JSONResponse, Response
+
+app = FastAPI()
+
+
+rollbar.init("ACCESS_TOKEN", environment="development")
+
+
+@app.exception_handler(Exception)
+async def handle_unexpected_exceptions(request: Request, exc: Exception):
+    """This won't capture HTTPException."""
+    try:
+        raise exc
+    except Exception:
+        rollbar.contrib.fastapi.report_exception(request=request)
+
+    return Response(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR)
+
+
+@app.get("/")
+async def root(raise_exception: bool = False):
+    """Hello world api endpoint.
+
+    Use `?raise_exception=1` to raise an exception.
+    """
+    if raise_exception:
+        raise Exception("Testing exceptions")
+    return JSONResponse({"message": "Hello World"})
+
+
+if __name__ == "__main__":
+    uvicorn.run(app)

rollbar/test/fastapi_tests/test_fastapi.py

@@ -0,0 +1,188 @@
+"""
+Tests for fastapi instrumentation
+"""
+
+import json
+import os
+import sys
+
+try:
+    from unittest import mock
+except ImportError:
+    import mock
+
+import rollbar
+from rollbar.test import BaseTest
+
+# access token for https://rollbar.com/rollbar/pyrollbar
+TOKEN = "92c10f5616944b81a2e6f3c6493a0ec2"
+
+# Fastapi works on python +3.6
+ALLOWED_PYTHON_VERSION = sys.version_info[0] == 3 and sys.version_info[1] >= 6
+
+
+try:
+    import fastapi  # noqa: F401
+
+    FASTAPI_INSTALLED = True
+except ImportError:
+    FASTAPI_INSTALLED = False
+
+
+def create_app():
+    from fastapi import FastAPI
+
+    app = FastAPI()
+
+    @app.get("/")
+    def index():
+        return "Index page"
+
+    @app.get("/cause_error")
+    @app.post("/cause_error")
+    def cause_error():
+        raise Exception("Uh oh")
+
+    return app
+
+
+def init_rollbar(app):
+    import rollbar.contrib.fastapi
+    from fastapi import Request, status
+    from fastapi.responses import Response
+
+    rollbar.init(
+        TOKEN,
+        "fastapitest",
+        root=os.path.dirname(os.path.realpath(__file__)),
+        allow_logging_basic_config=True,
+        capture_email=True,
+        capture_username=True,
+    )
+
+    @app.exception_handler(Exception)
+    async def handle_unexpected_exceptions(request: Request, exc: Exception):
+        """This won't capture HTTPException."""
+        try:
+            raise exc
+        except Exception:
+            rollbar.contrib.fastapi.report_exception(request=request)
+
+        return Response(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR)
+
+
+if ALLOWED_PYTHON_VERSION and FASTAPI_INSTALLED:
+
+    from fastapi.testclient import TestClient
+
+    class FastapiTest(BaseTest):
+        def setUp(self):
+            super().setUp()
+            self.app = create_app()
+            init_rollbar(self.app)
+            self.client = TestClient(self.app, raise_server_exceptions=False)
+
+        def test_index(self):
+            resp = self.client.get("/")
+            self.assertEqual(resp.status_code, 200)
+
+        def assertStringEqual(self, left, right):
+            if sys.version_info[0] > 2:
+                if hasattr(left, "decode"):
+                    left = left.decode("ascii")
+                if hasattr(right, "decode"):
+                    right = right.decode("ascii")
+
+                return self.assertEqual(left, right)
+            else:
+                return self.assertEqual(left, right)
+
+        @mock.patch("rollbar.send_payload")
+        def test_uncaught(self, send_payload):
+            resp = self.client.get(
+                "/cause_error?foo=bar",
+                headers={"X-Real-Ip": "1.2.3.4", "User-Agent": "Fastapi Test"},
+            )
+            self.assertEqual(resp.status_code, 500)
+
+            self.assertEqual(send_payload.called, True)
+            payload = send_payload.call_args[0][0]
+            data = payload["data"]
+
+            self.assertIn("body", data)
+            self.assertEqual(data["body"]["trace"]["exception"]["class"], "Exception")
+            self.assertStringEqual(
+                data["body"]["trace"]["exception"]["message"], "Uh oh"
+            )
+
+            self.assertIn("request", data)
+            self.assertEqual(
+                data["request"]["url"], "http://testserver/cause_error?foo=bar"
+            )
+
+            self.assertEqual(data["request"]["user_ip"], "1.2.3.4")
+            self.assertEqual(data["request"]["method"], "GET")
+            self.assertEqual(data["request"]["headers"]["user-agent"], "Fastapi Test")
+
+        @mock.patch("rollbar.send_payload")
+        def test_uncaught_json_request(self, send_payload):
+            json_body = {"hello": "world"}
+            json_body_str = json.dumps(json_body)
+            resp = self.client.post(
+                "/cause_error",
+                data=json_body_str,
+                headers={
+                    "Content-Type": "application/json",
+                    "X-Forwarded-For": "5.6.7.8",
+                },
+            )
+
+            self.assertEqual(resp.status_code, 500)
+
+            self.assertEqual(send_payload.called, True)
+            payload = send_payload.call_args[0][0]
+            data = payload["data"]
+
+            self.assertIn("body", data)
+            self.assertEqual(data["body"]["trace"]["exception"]["class"], "Exception")
+            self.assertStringEqual(
+                data["body"]["trace"]["exception"]["message"], "Uh oh"
+            )
+
+            self.assertIn("request", data)
+            self.assertEqual(data["request"]["url"], "http://testserver/cause_error")
+            self.assertEqual(data["request"]["user_ip"], "5.6.7.8")
+            self.assertEqual(data["request"]["method"], "POST")
+
+        @mock.patch("rollbar.send_payload")
+        def test_uncaught_no_username_no_email(self, send_payload):
+            rollbar.SETTINGS["capture_email"] = False
+            rollbar.SETTINGS["capture_username"] = False
+
+            resp = self.client.get(
+                "/cause_error?foo=bar",
+                headers={"X-Real-Ip": "1.2.3.4", "User-Agent": "Fastapi Test"},
+            )
+            self.assertEqual(resp.status_code, 500)
+
+            self.assertEqual(send_payload.called, True)
+            payload = send_payload.call_args[0][0]
+            data = payload["data"]
+
+            self.assertIn("body", data)
+            self.assertEqual(data["body"]["trace"]["exception"]["class"], "Exception")
+            self.assertStringEqual(
+                data["body"]["trace"]["exception"]["message"], "Uh oh"
+            )
+
+            self.assertIn("request", data)
+            self.assertEqual(
+                data["request"]["url"], "http://testserver/cause_error?foo=bar"
+            )
+
+            self.assertEqual(data["request"]["user_ip"], "1.2.3.4")
+            self.assertEqual(data["request"]["method"], "GET")
+            self.assertEqual(data["request"]["headers"]["user-agent"], "Fastapi Test")
+
+            rollbar.SETTINGS["capture_email"] = True
+            rollbar.SETTINGS["capture_username"] = True

setup.py

@@ -1,7 +1,8 @@
-import re
 import os.path
+import re
 import sys
-from setuptools import setup, find_packages
+
+from setuptools import find_packages, setup
 
 HERE = os.path.abspath(os.path.dirname(__file__))
 
@@ -65,6 +66,7 @@
         "Framework :: Bottle",
         "Framework :: Django",
         "Framework :: Flask",
+        "Framework :: Fastapi",
         "Framework :: Pylons",
         "Framework :: Pyramid",
         "Framework :: Twisted",
@@ -81,5 +83,11 @@
         'requests>=0.12.1',
         'six>=1.9.0'
     ],
+    extra_requires={
+        "fastapi": [
+            "fastapi",
+            "uvicorn",
+        ],
+    },
     tests_require=tests_require,
     )