diff --git a/changelog.d/stream-large-metadata-responses.fixed.md b/changelog.d/stream-large-metadata-responses.fixed.md new file mode 100644 index 000000000..1bd388935 --- /dev/null +++ b/changelog.d/stream-large-metadata-responses.fixed.md @@ -0,0 +1 @@ +Metadata endpoint streams responses above 20 MiB so uncompressed US metadata no longer exceeds Cloud Run's 32 MiB non-streamed response cap (fixes the /us/metadata 500 for clients that do not negotiate gzip). diff --git a/policyengine_api/fastapi_routes/metadata.py b/policyengine_api/fastapi_routes/metadata.py index 0984c7204..40dcbcf10 100644 --- a/policyengine_api/fastapi_routes/metadata.py +++ b/policyengine_api/fastapi_routes/metadata.py @@ -10,9 +10,16 @@ ensure_supported_country, ) from policyengine_api.fastapi_routes.dependencies import NativeRouteDependencies -from policyengine_api.fastapi_routes.responses import LegacyJSONResponse +from policyengine_api.fastapi_routes.responses import ( + LegacyJSONResponse, + render_legacy_json, +) from policyengine_api.json_types import JSONObject -from starlette.responses import Response +from policyengine_api.utils.streaming_json import ( + iter_body_chunks, + should_stream_body, +) +from starlette.responses import Response, StreamingResponse class MetadataSuccessPayload(TypedDict): @@ -48,6 +55,12 @@ def metadata(country_id: str) -> Response: "message": None, "result": metadata_reader.get_metadata(country_id), } - return LegacyJSONResponse(payload) + body = render_legacy_json(payload) + if should_stream_body(body): + return StreamingResponse( + iter_body_chunks(body), + media_type=LegacyJSONResponse.media_type, + ) + return Response(body, media_type=LegacyJSONResponse.media_type) return router diff --git a/policyengine_api/fastapi_routes/responses.py b/policyengine_api/fastapi_routes/responses.py index 32cd7d785..07c3e57a5 100644 --- a/policyengine_api/fastapi_routes/responses.py +++ b/policyengine_api/fastapi_routes/responses.py @@ -8,10 +8,16 @@ from starlette.responses import Response +def render_legacy_json(content: JSONValue) -> bytes: + """Serialize once with the same JSON encoder used by legacy metadata.""" + + return json.dumps(content).encode("utf-8") + + class LegacyJSONResponse(Response): """Serialize once with the same JSON encoder used by legacy metadata.""" media_type = "application/json" def render(self, content: JSONValue) -> bytes: - return json.dumps(content).encode("utf-8") + return render_legacy_json(content) diff --git a/policyengine_api/routes/metadata_routes.py b/policyengine_api/routes/metadata_routes.py index 8dd5465e4..cc1c310d0 100644 --- a/policyengine_api/routes/metadata_routes.py +++ b/policyengine_api/routes/metadata_routes.py @@ -2,6 +2,10 @@ from flask import Blueprint, Response from policyengine_api.utils.payload_validators import validate_country +from policyengine_api.utils.streaming_json import ( + iter_body_chunks, + should_stream_body, +) from policyengine_api.services.metadata_service import MetadataService metadata_bp = Blueprint("metadata", __name__) @@ -19,8 +23,18 @@ def get_metadata(country_id: str) -> Response: # Retrieve country metadata and add status and message to the response country_metadata = metadata_service.get_metadata(country_id) + body = json.dumps( + {"status": "ok", "message": None, "result": country_metadata} + ).encode("utf-8") + if should_stream_body(body): + return Response( + iter_body_chunks(body), + status=200, + mimetype="application/json", + direct_passthrough=True, + ) return Response( - json.dumps({"status": "ok", "message": None, "result": country_metadata}), + body, status=200, mimetype="application/json", ) diff --git a/policyengine_api/utils/streaming_json.py b/policyengine_api/utils/streaming_json.py new file mode 100644 index 000000000..53dc17216 --- /dev/null +++ b/policyengine_api/utils/streaming_json.py @@ -0,0 +1,27 @@ +"""Chunked delivery for JSON bodies too large to send with Content-Length. + +Cloud Run rejects HTTP/1 responses above 32 MiB unless they use +``Transfer-Encoding: chunked`` or another streaming mechanism +(https://docs.cloud.google.com/run/quotas). ``/us/metadata`` serializes to +~70 MB, so clients that do not negotiate gzip received an empty 500 from the +Google Frontend. Bodies above the threshold are therefore streamed; smaller +bodies keep their exact current framing, including Content-Length. +""" + +from __future__ import annotations + +from collections.abc import Iterator + +STREAMING_THRESHOLD_BYTES = 20 * 1024 * 1024 +STREAMING_CHUNK_BYTES = 1024 * 1024 + + +def should_stream_body(body: bytes) -> bool: + return len(body) >= STREAMING_THRESHOLD_BYTES + + +def iter_body_chunks( + body: bytes, chunk_size: int = STREAMING_CHUNK_BYTES +) -> Iterator[bytes]: + for start in range(0, len(body), chunk_size): + yield body[start : start + chunk_size] diff --git a/tests/unit/routes/test_metadata_routes_streaming.py b/tests/unit/routes/test_metadata_routes_streaming.py new file mode 100644 index 000000000..c9d8b5c9e --- /dev/null +++ b/tests/unit/routes/test_metadata_routes_streaming.py @@ -0,0 +1,81 @@ +"""Flask fallback metadata route: large bodies stream, small bodies keep Content-Length.""" + +import importlib +import json +import sys +from types import SimpleNamespace + +import pytest +from flask import Flask + +from policyengine_api.utils.streaming_json import STREAMING_THRESHOLD_BYTES + +SERVICE_MODULE = "policyengine_api.services.metadata_service" +ROUTE_MODULE = "policyengine_api.routes.metadata_routes" + + +class _MetadataService: + metadata_by_country: dict = {} + + def get_metadata(self, country_id: str): + return self.metadata_by_country[country_id] + + +def _load_metadata_blueprint_with_fake_service(): + """Import the real route module against a fake service module. + + The real ``MetadataService`` imports every country package at module + import, so the route module is loaded with the fake installed first and + the module cache is restored afterwards. + """ + + sentinel = object() + original_route_module = sys.modules.get(ROUTE_MODULE, sentinel) + original_service_module = sys.modules.get(SERVICE_MODULE, sentinel) + sys.modules.pop(ROUTE_MODULE, None) + sys.modules[SERVICE_MODULE] = SimpleNamespace(MetadataService=_MetadataService) + try: + return importlib.import_module(ROUTE_MODULE).metadata_bp + finally: + if original_route_module is sentinel: + sys.modules.pop(ROUTE_MODULE, None) + else: + sys.modules[ROUTE_MODULE] = original_route_module + if original_service_module is sentinel: + sys.modules.pop(SERVICE_MODULE, None) + else: + sys.modules[SERVICE_MODULE] = original_service_module + + +@pytest.fixture +def flask_client(): + app = Flask(__name__) + app.config["TESTING"] = True + app.register_blueprint(_load_metadata_blueprint_with_fake_service()) + return app.test_client() + + +def test_flask_metadata_streams_bodies_above_cloud_run_response_cap(flask_client): + large_value = "x" * (STREAMING_THRESHOLD_BYTES + 1024) + _MetadataService.metadata_by_country = {"us": {"large_value": large_value}} + + response = flask_client.get("/us/metadata") + + assert response.status_code == 200 + assert response.mimetype == "application/json" + assert response.headers.get("Content-Length") is None + assert json.loads(response.get_data()) == { + "status": "ok", + "message": None, + "result": {"large_value": large_value}, + } + + +def test_flask_metadata_below_streaming_threshold_keeps_content_length(flask_client): + _MetadataService.metadata_by_country = {"uk": {"small_value": "x" * 2_000}} + + response = flask_client.get("/uk/metadata") + + assert response.status_code == 200 + assert int(response.headers["Content-Length"]) == len(response.get_data()) + assert json.loads(response.get_data())["result"] == {"small_value": "x" * 2_000} diff --git a/tests/unit/test_stage6_native_metadata.py b/tests/unit/test_stage6_native_metadata.py index db7901c2a..0de2bf9f1 100644 --- a/tests/unit/test_stage6_native_metadata.py +++ b/tests/unit/test_stage6_native_metadata.py @@ -16,6 +16,7 @@ RouteImplementationSettings, ) from policyengine_api.request_context import REQUEST_ID_HEADER +from policyengine_api.utils.streaming_json import STREAMING_THRESHOLD_BYTES from policyengine_api.utils.payload_validators import validate_country @@ -275,3 +276,46 @@ def test_native_metadata_failure_is_500_without_exception_details(): assert response.headers["access-control-allow-origin"] == ( "https://app.policyengine.org" ) + + +def test_native_metadata_streams_bodies_above_cloud_run_response_cap(): + # Cloud Run drops HTTP/1 responses above 32 MiB unless they stream + # (https://docs.cloud.google.com/run/quotas); /us/metadata is ~70 MB raw. + large_value = "x" * (STREAMING_THRESHOLD_BYTES + 1024) + reader = _MetadataReader({"us": {"large_value": large_value}}) + + response = _native_client(reader).get( + "/us/metadata", + headers={"Accept-Encoding": "identity"}, + ) + + assert response.status_code == 200 + assert "content-length" not in response.headers + assert response.headers["content-type"] == "application/json" + assert response.json()["result"]["large_value"] == large_value + + +def test_native_metadata_streams_with_gzip_negotiation(): + large_value = "x" * (STREAMING_THRESHOLD_BYTES + 1024) + reader = _MetadataReader({"us": {"large_value": large_value}}) + + response = _native_client(reader).get( + "/us/metadata", + headers={"Accept-Encoding": "gzip"}, + ) + + assert response.status_code == 200 + assert response.headers["content-encoding"] == "gzip" + assert response.json()["result"]["large_value"] == large_value + + +def test_native_metadata_below_streaming_threshold_keeps_content_length(): + reader = _MetadataReader({"us": {"small_value": "x" * 2_000}}) + + response = _native_client(reader).get( + "/us/metadata", + headers={"Accept-Encoding": "identity"}, + ) + + assert response.status_code == 200 + assert int(response.headers["content-length"]) == len(response.content) diff --git a/tests/unit/test_streaming_json.py b/tests/unit/test_streaming_json.py new file mode 100644 index 000000000..2e3534e39 --- /dev/null +++ b/tests/unit/test_streaming_json.py @@ -0,0 +1,25 @@ +from policyengine_api.utils.streaming_json import ( + STREAMING_CHUNK_BYTES, + STREAMING_THRESHOLD_BYTES, + iter_body_chunks, + should_stream_body, +) + + +def test_should_stream_body_only_at_threshold(): + assert not should_stream_body(b"x" * (STREAMING_THRESHOLD_BYTES - 1)) + assert should_stream_body(b"x" * STREAMING_THRESHOLD_BYTES) + + +def test_iter_body_chunks_reassembles_exactly(): + body = bytes(range(256)) * 4 * 1024 # 1 MiB, non-uniform content + chunks = list(iter_body_chunks(body, 100_000)) + assert b"".join(chunks) == body + assert all(len(chunk) <= 100_000 for chunk in chunks) + + +def test_iter_body_chunks_default_chunk_size_covers_body(): + body = b"y" * (STREAMING_CHUNK_BYTES * 2 + 5) + chunks = list(iter_body_chunks(body)) + assert len(chunks) == 3 + assert b"".join(chunks) == body