Stream metadata responses above Cloud Run's non-streamed size cap

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).

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

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)

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",
     )

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]

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}

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)

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