diff --git a/src/functions_framework/__init__.py b/src/functions_framework/__init__.py index c2a52d74..d4575b57 100644 --- a/src/functions_framework/__init__.py +++ b/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 diff --git a/tests/test_functions.py b/tests/test_functions.py index 501ea488..81860cae 100644 --- a/tests/test_functions.py +++ b/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" diff --git a/tests/test_functions/http_streaming/main.py b/tests/test_functions/http_streaming/main.py new file mode 100644 index 00000000..4b249697 --- /dev/null +++ b/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()))