Implement streaming callables

inlined · inlined · commit f1c7017e92e7 · 2025-02-04T07:56:24.000-08:00
diff --git a/samples/https_flask/functions/main.py b/samples/https_flask/functions/main.py
@@ -23,3 +23,13 @@ def world():
 @https_fn.on_request()
 def httpsflaskexample(request):
     return entrypoint(app, request)
+
+@https_fn.on_call()
+def callableexample(request: https_fn.Request):
+    return request.data
+
+@https_fn.on_call()
+def streamingcallable(request: https_fn.Request):
+    yield "Hello,"
+    yield "world!"
+    return request.data
diff --git a/src/firebase_functions/https_fn.py b/src/firebase_functions/https_fn.py
@@ -21,8 +21,10 @@
 import typing_extensions as _typing_extensions
 import enum as _enum
 import json as _json
+import inspect as _inspect
 import firebase_functions.private.util as _util
 import firebase_functions.core as _core
+import contextlib as _contextlib
 from functions_framework import logging as _logging
 
 from firebase_functions.options import HttpsOptions, _GLOBAL_OPTIONS
@@ -351,6 +353,12 @@ class CallableRequest(_typing.Generic[_core.T]):
 _C1 = _typing.Callable[[Request], Response]
 _C2 = _typing.Callable[[CallableRequest[_typing.Any]], _typing.Any]
 
+class _IterWithReturn:
+    def __init__(self, iterable):
+        self.iterable = iterable
+    
+    def __iter__(self):
+        self.value = yield from self.iterable
 
 def _on_call_handler(func: _C2, request: Request,
                      enforce_app_check: bool) -> Response:
@@ -401,7 +409,19 @@ def _on_call_handler(func: _C2, request: Request,
                     "Firebase-Instance-ID-Token"),
             )
         result = _core._with_init(func)(context)
-        return _jsonify(result=result)
+        if not _inspect.isgenerator(result):
+            return _jsonify(result=result)
+        
+        if request.headers.get("Accept") != "text/event-stream":
+            vals = _IterWithReturn(result)
+            for _ in vals:
+                next
+            return _jsonify(result=vals.value)
+
+        else:
+            return Response(_sse_encode_generator(result), content_type="text/plain")
+
+
     # Disable broad exceptions lint since we want to handle all exceptions here
     # and wrap as an HttpsError.
     # pylint: disable=broad-except
@@ -413,6 +433,19 @@ def _on_call_handler(func: _C2, request: Request,
         return _make_response(_jsonify(error=err._as_dict()), status)
 
 
+def _sse_encode_generator(gen: _typing.Generator):
+    iter = _IterWithReturn(gen)
+    try:
+        for chunk in iter:
+            yield f"data: %s\n\n" % _json.dumps(obj={"message": chunk})
+        yield f"data: %s\n\n" % _json.dumps(obj={"result": iter.value})
+    except Exception as err:
+        if not isinstance(err, HttpsError):
+            _logging.error("Unhandled error: %s", err)
+            err = HttpsError(FunctionsErrorCode.INTERNAL, "INTERNAL")
+        yield f"error: %s\n\n" % _json.dumps(obj={"error": err._as_dict()})
+    yield "END"
+
 @_util.copy_func_kwargs(HttpsOptions)
 def on_request(**kwargs) -> _typing.Callable[[_C1], _C1]:
     """
diff --git a/tests/test_https_fn.py b/tests/test_https_fn.py
@@ -3,8 +3,7 @@
 """
 
 import unittest
-from unittest.mock import Mock
-from flask import Flask, Request
+from flask import Flask, Request, jsonify as _jsonify
 from werkzeug.test import EnvironBuilder
 
 from firebase_functions import core, https_fn
@@ -25,7 +24,9 @@ def init():
             nonlocal hello
             hello = "world"
 
-        func = Mock(__name__="example_func")
+        @https_fn.on_request()
+        def func(_):
+            pass
 
         with app.test_request_context("/"):
             environ = EnvironBuilder(
@@ -37,9 +38,8 @@ def init():
                 },
             ).get_environ()
             request = Request(environ)
-            decorated_func = https_fn.on_request()(func)
 
-            decorated_func(request)
+            func(request)
 
         self.assertEqual(hello, "world")
 
@@ -53,7 +53,9 @@ def init():
             nonlocal hello
             hello = "world"
 
-        func = Mock(__name__="example_func")
+        @https_fn.on_call()
+        def func(_):
+            pass
 
         with app.test_request_context("/"):
             environ = EnvironBuilder(
@@ -65,8 +67,145 @@ def init():
                 },
             ).get_environ()
             request = Request(environ)
-            decorated_func = https_fn.on_call()(func)
-
-            decorated_func(request)
+            func(request)
 
         self.assertEqual("world", hello)
+
+    def test_callable_encoding(self):
+        app = Flask(__name__)
+
+        @https_fn.on_call()
+        def add(req: https_fn.CallableRequest[int]):
+            return req.data + 1
+        
+        with app.test_request_context("/"):
+            environ = EnvironBuilder(
+                method="POST",
+                json={
+                    "data": 1
+                }
+            ).get_environ()
+            request = Request(environ)
+
+            response = add(request)
+            self.assertEqual(response.status_code, 200)
+            self.assertEqual(response.get_json(), { "result": 2 })
+    
+    def test_callable_errors(self):
+        app = Flask(__name__)
+
+        @https_fn.on_call()
+        def throw_generic_error(req):
+            raise Exception("Invalid type")
+
+        @https_fn.on_call()
+        def throw_access_denied(req):
+            raise https_fn.HttpsError(https_fn.FunctionsErrorCode.PERMISSION_DENIED, "Permission is denied")
+        
+        with app.test_request_context("/"):
+            environ = EnvironBuilder(
+                method="POST",
+                json={
+                    "data": None
+                }
+            ).get_environ()
+            request = Request(environ)
+
+            response = throw_generic_error(request)
+            self.assertEqual(response.status_code, 500)
+            self.assertEqual(response.get_json(), { "error": { "message": "INTERNAL", "status": "INTERNAL" } })
+
+            response = throw_access_denied(request)
+            self.assertEqual(response.status_code, 403)
+            self.assertEqual(response.get_json(), { "error": { "message": "Permission is denied", "status": "PERMISSION_DENIED" }})
+
+    def test_yielding_without_streaming(self):
+        app = Flask(__name__)
+
+        @https_fn.on_call()
+        def yielder(req: https_fn.CallableRequest[int]):
+            yield from range(req.data)
+            return "OK"
+
+        @https_fn.on_call()
+        def yield_thrower(req: https_fn.CallableRequest[int]):
+            yield from range(req.data)
+            raise https_fn.HttpsError(https_fn.FunctionsErrorCode.PERMISSION_DENIED, "Can't read anymore")
+
+        with app.test_request_context("/"):
+            environ = EnvironBuilder(
+                method="POST",
+                json={
+                    "data": 5
+                }
+            ).get_environ()
+
+            request = Request(environ)
+            response = yielder(request)
+
+            self.assertEqual(response.status_code, 200)
+            self.assertEqual(response.get_json(), { "result": "OK" })
+
+        with app.test_request_context("/"):
+            environ = EnvironBuilder(
+                method="POST",
+                json={
+                    "data": 3
+                }
+            ).get_environ()
+
+            request = Request(environ)
+            response = yield_thrower(request)
+
+            self.assertEqual(response.status_code, 403)
+            self.assertEqual(response.get_json(), { "error": { "message": "Can't read anymore", "status": "PERMISSION_DENIED" }})
+
+
+    def test_yielding_with_streaming(self):
+        app = Flask(__name__)
+
+        @https_fn.on_call()
+        def yielder(req: https_fn.CallableRequest[int]):
+            yield from range(req.data)
+            return "OK"
+
+        @https_fn.on_call()
+        def yield_thrower(req: https_fn.CallableRequest[int]):
+            yield from range(req.data)
+            raise https_fn.HttpsError(https_fn.FunctionsErrorCode.INTERNAL, "Throwing")
+
+        with app.test_request_context("/"):
+            environ = EnvironBuilder(
+                method="POST",
+                json={
+                    "data": 2
+                },
+                headers={
+                    "accept": "text/event-stream"
+                }
+            ).get_environ()
+
+            request = Request(environ)
+            response = yielder(request)
+
+            self.assertEqual(response.status_code, 200)
+            chunks = list(response.response)
+            self.assertEqual(chunks, ['data: {"message": 0}\n\n', 'data: {"message": 1}\n\n', 'data: {"result": "OK"}\n\n', "END"])
+
+        with app.test_request_context("/"):
+            environ = EnvironBuilder(
+                method="POST",
+                json={
+                    "data": 2
+                },
+                headers={
+                    "accept": "text/event-stream"
+                }
+            ).get_environ()
+
+            request = Request(environ)
+            response = yield_thrower(request)
+
+            self.assertEqual(response.status_code, 200)
+            chunks = list(response.response)
+            self.assertEqual(chunks, ['data: {"message": 0}\n\n', 'data: {"message": 1}\n\n', 'error: {"error": {"status": "INTERNAL", "message": "Throwing"}}\n\n', "END"])
