From 1d7de3be77d50dd17c9c0ee5592aa7f7885e21b1 Mon Sep 17 00:00:00 2001
From: Jon Lund Steffensen <jon.steffensen@autodesk.com>
Date: Tue, 2 Apr 2024 17:17:46 -0700
Subject: [PATCH] Async handler support

---
 examples/todo/todo/handlers/todo_handlers.py |  2 +-
 flask_rebar/rebar.py                         |  2 +-
 setup.py                                     |  7 ++++++-
 tests/test_rebar.py                          | 18 ++++++++++++++++++
 4 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/examples/todo/todo/handlers/todo_handlers.py b/examples/todo/todo/handlers/todo_handlers.py
index 4298897..a31552a 100644
--- a/examples/todo/todo/handlers/todo_handlers.py
+++ b/examples/todo/todo/handlers/todo_handlers.py
@@ -1,4 +1,4 @@
-from enum import Enum
+import asyncio
 
 from todo.database import todo_id_sequence, todo_database
 from todo.app import rebar
diff --git a/flask_rebar/rebar.py b/flask_rebar/rebar.py
index dfaf46f..6abf0e9 100644
--- a/flask_rebar/rebar.py
+++ b/flask_rebar/rebar.py
@@ -176,7 +176,7 @@ def wrapped(*args: P.args, **kwargs: P.kwargs) -> Union[T, Response]:
         if headers_schema:
             g.validated_headers = get_header_params_or_400(schema=headers_schema)
 
-        rv: Any = f(*args, **kwargs)
+        rv: Any = current_app.ensure_sync(f)(*args, **kwargs)
 
         if not response_body_schema:
             return rv
diff --git a/setup.py b/setup.py
index 89c9aa4..1f78a87 100644
--- a/setup.py
+++ b/setup.py
@@ -20,6 +20,7 @@
     "sphinx_rtd_theme==1.2.2",
     "types-jsonschema==4.17.0.10",
     "types-setuptools==68.0.0.3",
+    "flask[async]>=2,<4",
 ]
 
 install_requires = [
@@ -42,7 +43,11 @@
         packages=find_packages(exclude=("test*", "examples")),
         package_data={"flask_rebar": ["py.typed"]},
         include_package_data=True,
-        extras_require={"dev": development, "enum": ["marshmallow-enum~=1.5"]},
+        extras_require={
+            "dev": development,
+            "enum": ["marshmallow-enum~=1.5"],
+            "async": ["flask[async]>=2,<4"],
+        },
         install_requires=install_requires,
         url="https://github.com/plangrid/flask-rebar",
         classifiers=[
diff --git a/tests/test_rebar.py b/tests/test_rebar.py
index 7c233e4..dacc440 100644
--- a/tests/test_rebar.py
+++ b/tests/test_rebar.py
@@ -7,6 +7,7 @@
     :copyright: Copyright 2018 PlanGrid, Inc., see AUTHORS.
     :license: MIT, see LICENSE for details.
 """
+import asyncio
 import json
 import unittest
 
@@ -864,3 +865,20 @@ def post_foo(*args, **kwargs):
         self.assertIsNotNone(swagger)  # really only care that it didn't barf
         swagger = SwaggerV3Generator().generate(registry)
         self.assertIsNotNone(swagger)
+
+    def test_async_handler(self):
+        rebar = Rebar()
+        registry = rebar.create_handler_registry()
+
+        @registry.handles(
+            rule="/async",
+            method="GET",
+            response_body_schema=DefaultResponseSchema(),
+        )
+        async def get_response():
+            await asyncio.sleep(0)
+            return DEFAULT_RESPONSE
+
+        app = create_rebar_app(rebar)
+        resp = app.test_client().get(path="/async")
+        self.assertEqual(resp.status_code, 200)