Allow to disable Transfer-Encoding: chunked (#1560)

Allow to disable Transfer-Encoding: chunked
2019-04-30 14:56:27 -07:00
parent 6a4a3f617f 7d6e60ab7d
commit ef6d78c580
4 changed files with 148 additions and 12 deletions
--- a/docs/sanic/response.md
+++ b/docs/sanic/response.md
@@ -60,6 +60,8 @@ async def index(request):
    return response.stream(streaming_fn, content_type='text/plain')
 ```

+See [Streaming](streaming.md) for more information.
+
 ## File Streaming
 For large files, a combination of File and Streaming above
 ```python
--- a/docs/sanic/streaming.md
+++ b/docs/sanic/streaming.md
@@ -117,3 +117,27 @@ async def index(request):

    return stream(stream_from_db)
 ```
+
+If a client supports HTTP/1.1, Sanic will use [chunked transfer encoding](https://en.wikipedia.org/wiki/Chunked_transfer_encoding); you can explicitly enable or disable it using `chunked` option of the `stream` function.
+
+## File Streaming
+
+Sanic provides `sanic.response.file_stream` function that is useful when you want to send a large file. It returns a `StreamingHTTPResponse` object and will use chunked transfer encoding by default; for this reason Sanic doesn't add `Content-Length` HTTP header in the response. If you want to use this header, you can disable chunked transfer encoding and add it manually:
+
+```python
+from aiofiles import os as async_os
+from sanic.response import file_stream
+
+@app.route("/")
+async def index(request):
+    file_path = "/srv/www/whatever.png"
+
+    file_stat = await async_os.stat(file_path)
+    headers = {"Content-Length": str(file_stat.st_size)}
+
+    return await file_stream(
+        file_path,
+        headers=headers,
+        chunked=False,
+    )
+```
--- a/sanic/response.py
+++ b/sanic/response.py
@@ -59,16 +59,23 @@ class StreamingHTTPResponse(BaseHTTPResponse):
        "status",
        "content_type",
        "headers",
+        "chunked",
        "_cookies",
    )

    def __init__(
-        self, streaming_fn, status=200, headers=None, content_type="text/plain"
+        self,
+        streaming_fn,
+        status=200,
+        headers=None,
+        content_type="text/plain",
+        chunked=True,
    ):
        self.content_type = content_type
        self.streaming_fn = streaming_fn
        self.status = status
        self.headers = CIMultiDict(headers or {})
+        self.chunked = chunked
        self._cookies = None

    async def write(self, data):
@@ -79,7 +86,10 @@ class StreamingHTTPResponse(BaseHTTPResponse):
        if type(data) != bytes:
            data = self._encode_body(data)

-        self.protocol.push_data(b"%x\r\n%b\r\n" % (len(data), data))
+        if self.chunked:
+            self.protocol.push_data(b"%x\r\n%b\r\n" % (len(data), data))
+        else:
+            self.protocol.push_data(data)
        await self.protocol.drain()

    async def stream(
@@ -88,6 +98,8 @@ class StreamingHTTPResponse(BaseHTTPResponse):
        """Streams headers, runs the `streaming_fn` callback that writes
        content to the response body, then finalizes the response body.
        """
+        if version != "1.1":
+            self.chunked = False
        headers = self.get_headers(
            version,
            keep_alive=keep_alive,
@@ -96,7 +108,8 @@ class StreamingHTTPResponse(BaseHTTPResponse):
        self.protocol.push_data(headers)
        await self.protocol.drain()
        await self.streaming_fn(self)
-        self.protocol.push_data(b"0\r\n\r\n")
+        if self.chunked:
+            self.protocol.push_data(b"0\r\n\r\n")
        # no need to await drain here after this write, because it is the
        # very last thing we write and nothing needs to wait for it.

@@ -109,8 +122,9 @@ class StreamingHTTPResponse(BaseHTTPResponse):
        if keep_alive and keep_alive_timeout is not None:
            timeout_header = b"Keep-Alive: %d\r\n" % keep_alive_timeout

-        self.headers["Transfer-Encoding"] = "chunked"
-        self.headers.pop("Content-Length", None)
+        if self.chunked and version == "1.1":
+            self.headers["Transfer-Encoding"] = "chunked"
+            self.headers.pop("Content-Length", None)
        self.headers["Content-Type"] = self.headers.get(
            "Content-Type", self.content_type
        )
@@ -327,6 +341,7 @@ async def file_stream(
    mime_type=None,
    headers=None,
    filename=None,
+    chunked=True,
    _range=None,
 ):
    """Return a streaming response object with file data.
@@ -336,6 +351,7 @@ async def file_stream(
    :param mime_type: Specific mime_type.
    :param headers: Custom Headers.
    :param filename: Override filename.
+    :param chunked: Enable or disable chunked transfer-encoding
    :param _range:
    """
    headers = headers or {}
@@ -383,6 +399,7 @@ async def file_stream(
        status=status,
        headers=headers,
        content_type=mime_type,
+        chunked=chunked,
    )


@@ -391,6 +408,7 @@ def stream(
    status=200,
    headers=None,
    content_type="text/plain; charset=utf-8",
+    chunked=True,
 ):
    """Accepts an coroutine `streaming_fn` which can be used to
    write chunks to a streaming response. Returns a `StreamingHTTPResponse`.
@@ -409,9 +427,14 @@ def stream(
        writes content to that response.
    :param mime_type: Specific mime_type.
    :param headers: Custom Headers.
+    :param chunked: Enable or disable chunked transfer-encoding
    """
    return StreamingHTTPResponse(
-        streaming_fn, headers=headers, content_type=content_type, status=status
+        streaming_fn,
+        headers=headers,
+        content_type=content_type,
+        status=status,
+        chunked=chunked,
    )


--- a/tests/test_response.py
+++ b/tests/test_response.py
@@ -192,22 +192,59 @@ def test_no_content(json_app):
 def streaming_app(app):
    @app.route("/")
    async def test(request):
-        return stream(sample_streaming_fn, content_type="text/csv")
+        return stream(
+            sample_streaming_fn,
+            headers={"Content-Length": "7"},
+            content_type="text/csv",
+        )

    return app


-def test_streaming_adds_correct_headers(streaming_app):
+@pytest.fixture
+def non_chunked_streaming_app(app):
+    @app.route("/")
+    async def test(request):
+        return stream(
+            sample_streaming_fn,
+            headers={"Content-Length": "7"},
+            content_type="text/csv",
+            chunked=False,
+        )
+
+    return app
+
+
+def test_chunked_streaming_adds_correct_headers(streaming_app):
    request, response = streaming_app.test_client.get("/")
    assert response.headers["Transfer-Encoding"] == "chunked"
    assert response.headers["Content-Type"] == "text/csv"
+    # Content-Length is not allowed by HTTP/1.1 specification
+    # when "Transfer-Encoding: chunked" is used
+    assert "Content-Length" not in response.headers


-def test_streaming_returns_correct_content(streaming_app):
+def test_chunked_streaming_returns_correct_content(streaming_app):
    request, response = streaming_app.test_client.get("/")
    assert response.text == "foo,bar"


+def test_non_chunked_streaming_adds_correct_headers(
+    non_chunked_streaming_app
+):
+    request, response = non_chunked_streaming_app.test_client.get("/")
+    assert "Transfer-Encoding" not in response.headers
+    assert response.headers["Content-Type"] == "text/csv"
+    assert response.headers["Content-Length"] == "7"
+
+
+def test_non_chunked_streaming_returns_correct_content(
+    non_chunked_streaming_app
+):
+    request, response = non_chunked_streaming_app.test_client.get("/")
+    assert response.text == "foo,bar"
+
+
@pytest.mark.parametrize("status", [200, 201, 400, 401])
 def test_stream_response_status_returns_correct_headers(status):
    response = StreamingHTTPResponse(sample_streaming_fn, status=status)
@@ -227,13 +264,27 @@ def test_stream_response_keep_alive_returns_correct_headers(
    assert b"Keep-Alive: %s\r\n" % str(keep_alive_timeout).encode() in headers


-def test_stream_response_includes_chunked_header():
+def test_stream_response_includes_chunked_header_http11():
    response = StreamingHTTPResponse(sample_streaming_fn)
-    headers = response.get_headers()
+    headers = response.get_headers(version="1.1")
    assert b"Transfer-Encoding: chunked\r\n" in headers


-def test_stream_response_writes_correct_content_to_transport(streaming_app):
+def test_stream_response_does_not_include_chunked_header_http10():
+    response = StreamingHTTPResponse(sample_streaming_fn)
+    headers = response.get_headers(version="1.0")
+    assert b"Transfer-Encoding: chunked\r\n" not in headers
+
+
+def test_stream_response_does_not_include_chunked_header_if_disabled():
+    response = StreamingHTTPResponse(sample_streaming_fn, chunked=False)
+    headers = response.get_headers(version="1.1")
+    assert b"Transfer-Encoding: chunked\r\n" not in headers
+
+
+def test_stream_response_writes_correct_content_to_transport_when_chunked(
+    streaming_app
+):
    response = StreamingHTTPResponse(sample_streaming_fn)
    response.protocol = MagicMock(HttpProtocol)
    response.protocol.transport = MagicMock(asyncio.Transport)
@@ -262,6 +313,42 @@ def test_stream_response_writes_correct_content_to_transport(streaming_app):
            b"0\r\n\r\n"
        )

+        assert len(response.protocol.transport.write.call_args_list) == 4
+
+        app.stop()
+
+    streaming_app.run(host=HOST, port=PORT)
+
+
+def test_stream_response_writes_correct_content_to_transport_when_not_chunked(
+    streaming_app,
+):
+    response = StreamingHTTPResponse(sample_streaming_fn)
+    response.protocol = MagicMock(HttpProtocol)
+    response.protocol.transport = MagicMock(asyncio.Transport)
+
+    async def mock_drain():
+        pass
+
+    def mock_push_data(data):
+        response.protocol.transport.write(data)
+
+    response.protocol.push_data = mock_push_data
+    response.protocol.drain = mock_drain
+
+    @streaming_app.listener("after_server_start")
+    async def run_stream(app, loop):
+        await response.stream(version="1.0")
+        assert response.protocol.transport.write.call_args_list[1][0][0] == (
+            b"foo,"
+        )
+
+        assert response.protocol.transport.write.call_args_list[2][0][0] == (
+            b"bar"
+        )
+
+        assert len(response.protocol.transport.write.call_args_list) == 3
+
        app.stop()

    streaming_app.run(host=HOST, port=PORT)