diff --git a/fastly_compute/requests/response.py b/fastly_compute/requests/response.py index 4963c5d..04deb91 100644 --- a/fastly_compute/requests/response.py +++ b/fastly_compute/requests/response.py @@ -6,7 +6,7 @@ from wit_world.imports import async_io, http_resp -from ..utils import read_response_body +from ..utils import create_body_reader from .exceptions import HTTPError @@ -85,7 +85,8 @@ def headers(self) -> dict[str, str]: def content(self) -> bytes: """Response body as bytes.""" if self._content is None: - self._content = self._read_body() + reader = create_body_reader(self._response_body) + self._content = reader.read() return self._content @property @@ -172,10 +173,6 @@ def encoding(self) -> str | None: """Response encoding.""" return self._parse_charset() - def _read_body(self) -> bytes: - """Read the complete response body from WIT.""" - return read_response_body(self._response_body) - def __bool__(self) -> bool: """Boolean evaluation returns ok status.""" return self.ok diff --git a/fastly_compute/utils.py b/fastly_compute/utils.py index ce29e92..acb6641 100644 --- a/fastly_compute/utils.py +++ b/fastly_compute/utils.py @@ -1,36 +1,82 @@ """Utility functions for fastly_compute package.""" +from io import BufferedReader, RawIOBase + from wit_world.imports import async_io, http_body -from fastly_compute.exceptions.types.error import Error -from fastly_compute.requests.exceptions import RequestException +class _RawBodyReader(RawIOBase): + """Raw I/O implementation that reads from a Fastly HTTP body. + + This class provides a streaming interface to read from Fastly HTTP bodies + without buffering the entire body in memory upfront. + """ + + def __init__(self, body_handle: async_io.Pollable, chunk_size: int = 4096): + """Initialize the reader. + + :param body_handle: Fastly HTTP body handle to read from + :param chunk_size: Size of chunks to read at a time (default: 4096) + """ + self._body = body_handle + self._chunk_size = chunk_size + self._closed = False + + def readable(self) -> bool: + """Return whether the stream is readable. + + :return: True if the stream is readable, False otherwise + """ + return not self._closed + + def readinto(self, b) -> int: + """Read up to len(b) bytes into writable buffer b. + + :param b: Writable buffer to read into + :return: Number of bytes read, or 0 on EOF + """ + if self._closed: + return 0 + + # Read a chunk from host into provided buffer + chunk_size = min(len(b), self._chunk_size) + chunk = http_body.read(self._body, chunk_size) + + if not chunk: + self._closed = True + return 0 + + # Copy chunk into the provided buffer + n = len(chunk) + b[:n] = chunk + return n + + +def create_body_reader( + body: async_io.Pollable, chunk_size: int = 4096 +) -> BufferedReader: + """Create a file-like reader for streaming a Fastly HTTP body. + + This function returns a BufferedReader that streams the body on-demand, + avoiding the need to buffer the entire body in memory. The returned reader + is compatible with the WSGI InputStream Protocol (PEP 3333). -def read_response_body( - response_body: async_io.Pollable, chunk_size: int = 4096 -) -> bytes: - """Read the complete response body from a WIT response body object. + The returned reader supports all standard file operations: read(), readline(), + readlines(), and iteration. Note that seeking is not supported as the body + is a forward-only stream. - Args: - response_body: WIT response body object to read from - chunk_size: Size of chunks to read at a time (default: 4096) + :param body: Fastly HTTP body handle to read from + :param chunk_size: Size of chunks to read at a time (default: 4096) + :return: A BufferedReader that streams the body content (WSGI compatible) - Returns: - Complete response body as bytes + Example:: - Raises: - RequestException: If there is a problem reading the response body. + def handle(request, body): + reader = create_body_reader(body) + # Read entire body + full_body = reader.read() + # Or read line by line + for line in reader: + process(line) """ - body_data: bytes = b"" - while True: - try: - chunk = http_body.read(response_body, chunk_size) - except Error as e: - raise RequestException.from_fastly_error(e, "http_body.read") from e - - if len(chunk) == 0: - break - else: - body_data += chunk - - return body_data + return BufferedReader(_RawBodyReader(body, chunk_size)) diff --git a/fastly_compute/wsgi.py b/fastly_compute/wsgi.py index 265aff2..cca949b 100644 --- a/fastly_compute/wsgi.py +++ b/fastly_compute/wsgi.py @@ -11,9 +11,9 @@ import sys import traceback from collections.abc import Callable -from io import BytesIO, Reader from typing import Any from urllib.parse import urlparse +from wsgiref.types import InputStream from wit_world.exports import HttpIncoming as WitHttpIncoming from wit_world.imports import async_io, http_body, http_req, http_resp @@ -25,11 +25,12 @@ from wit_world.imports.http_resp import send_downstream from fastly_compute.exceptions.types.error import CannotRead +from fastly_compute.utils import create_body_reader def serve_wsgi_request( req: http_req.Request, - body: Reader[bytes], + body: InputStream, app: Callable, handle_errors: bool = False, ) -> None: @@ -39,13 +40,11 @@ def serve_wsgi_request( specification, allowing any WSGI-compatible web framework to run on Fastly Compute. - Args: - req: Fastly HTTP request object from WIT bindings - body: Fastly HTTP body object from WIT bindings - app: WSGI application callable - handle_errors: If True, the wrapper will log exceptions and return 500; if not, - then it will be handled by the server (or will be handled by - the WSGI app/framework itself). + :param req: Fastly HTTP request object from WIT bindings + :param body: WSGI input stream containing the request body (PEP 3333 compliant) + :param app: WSGI application callable + :param handle_errors: If True, log exceptions and return 500; otherwise let + the server or WSGI app handle them """ response = http_resp.Response.new() response_body = http_body.new() @@ -208,17 +207,10 @@ def __call__(self): def handle(self, request: http_req.Request, body: async_io.Pollable) -> None: """Handle incoming HTTP requests by serving them through the WSGI app.""" - - def body_reader(body: async_io.Pollable) -> Reader[bytes]: - """Given a Fastly HTTP body object, return a file-like object - containing the body's content. - """ - return BytesIO(http_body.read(body, 2**32 - 1)) - with request: # Ensure dropping of request resource before trying to get another one. This dodges a crash. serve_wsgi_request( request, - body_reader(body), + create_body_reader(body), self.wsgi_app, handle_errors=self.handle_errors, ) @@ -241,7 +233,7 @@ def body_reader(body: async_io.Pollable) -> Reader[bytes]: with request: serve_wsgi_request( request, - body_reader(body), + create_body_reader(body), self.wsgi_app, handle_errors=self.handle_errors, )