fix: streaming requests cannot access request data (GoogleCloudPlatform#245)

Author: garethgeorge

src/functions_framework/__init__.py

@@ -294,10 +294,13 @@ def _configure_app(app, function, signature_type):
 def read_request(response):
     """
     Force the framework to read the entire request before responding, to avoid
-    connection errors when returning prematurely.
+    connection errors when returning prematurely. Skipped on streaming responses
+    as these may continue to operate on the request after they are returned.
     """
 
-    flask.request.get_data()
+    if not response.is_streamed:
+        flask.request.get_data()
+
     return response
 
 

tests/test_functions.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
+import io
 import json
 import pathlib
 import re
@@ -490,6 +490,18 @@ def test_function_returns_none():
     assert resp.status_code == 500
 
 
+def test_function_returns_stream():
+    source = TEST_FUNCTIONS_DIR / "http_streaming" / "main.py"
+    target = "function"
+
+    client = create_app(target, source).test_client()
+    resp = client.post("/", data="1\n2\n3\n4\n")
+
+    assert resp.status_code == 200
+    assert resp.is_streamed
+    assert resp.data.decode("utf-8") == "1.0\n3.0\n6.0\n10.0\n"
+
+
 def test_legacy_function_check_env(monkeypatch):
     source = TEST_FUNCTIONS_DIR / "http_check_env" / "main.py"
     target = "function"

tests/test_functions/http_streaming/main.py

@@ -0,0 +1,44 @@
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Function used in Worker tests of handling HTTP functions."""
+
+import flask
+
+from flask import Response, stream_with_context
+
+
+def function(request):
+    """Test HTTP function that reads a stream of integers and returns a stream
+    providing the sum of values read so far.
+
+    Args:
+      request: The HTTP request which triggered this function. Must contain a
+      stream of new line separated integers.
+
+    Returns:
+      Value and status code defined for the given mode.
+
+    Raises:
+      Exception: Thrown when requested in the incoming mode specification.
+    """
+    print("INVOKED THE STREAM FUNCTION!!!")
+
+    def generate():
+        sum_so_far = 0
+        for line in request.stream:
+            sum_so_far += float(line)
+            yield (str(sum_so_far) + "\n").encode("utf-8")
+
+    return Response(stream_with_context(generate()))