Reduce memory consumption when building index by streaming response (#335)

Author: de-code

data_hub_api/docmaps/v2/api_router.py

@@ -1,6 +1,7 @@
 import logging
 
 from fastapi import APIRouter, HTTPException
+from fastapi.responses import StreamingResponse
 
 from data_hub_api.docmaps.v2.provider import DocmapsProvider
 
@@ -13,9 +14,12 @@ def create_docmaps_router(
 ) -> APIRouter:
     router = APIRouter()
 
-    @router.get("/v2/index")
-    def get_enhanced_preprints_docmaps_index():
-        return docmaps_provider.get_docmaps_index()
+    @router.get("/v2/index", response_class=StreamingResponse)
+    def stream_docmaps_index() -> StreamingResponse:
+        return StreamingResponse(
+            docmaps_provider.iter_docmaps_index_json_stream(),
+            media_type="application/json"
+        )
 
     @router.get("/v2/by-publisher/elife/get-by-manuscript-id")
     def get_enhanced_preprints_docmaps_by_manuscript_id_by_publisher_elife(manuscript_id: str):

data_hub_api/docmaps/v2/provider.py

@@ -1,3 +1,4 @@
+import json
 import logging
 from pathlib import Path
 from time import monotonic
@@ -70,6 +71,13 @@ def iter_docmaps_by_manuscript_id(
     def get_docmaps_by_manuscript_id(self, manuscript_id: str) -> Sequence[Docmap]:
         return list(self.iter_docmaps_by_manuscript_id(manuscript_id))
 
-    def get_docmaps_index(self) -> dict:
-        article_docmaps_list = list(self.iter_docmaps_by_manuscript_id())
-        return {'docmaps': article_docmaps_list}
+    def iter_docmaps_index_json_stream(self):
+        yield '{"docmaps":['
+        first = True
+        for docmap in self.iter_docmaps_by_manuscript_id():
+            if not first:
+                yield ','
+            else:
+                first = False
+            yield json.dumps(docmap)
+        yield ']}'

tests/unit_tests/docmaps/v2/api_router_test.py

@@ -1,3 +1,4 @@
+import json
 from unittest.mock import MagicMock
 import pytest
 
@@ -28,7 +29,9 @@ def test_should_return_json_with_docmaps_list_from_enhanced_preprint_provider(
         docmaps_provider_mock: MagicMock
     ):
         docmaps_index = [{'docmaps': [{'id': 'docmap_1'}, {'id': 'docmap_2'}]}]
-        docmaps_provider_mock.get_docmaps_index.return_value = docmaps_index
+        docmaps_provider_mock.iter_docmaps_index_json_stream.return_value = ([
+            json.dumps(docmaps_index)
+        ])
         client = create_test_client(docmaps_provider_mock)
         response = client.get('/v2/index')
         assert response.json() == docmaps_index

tests/unit_tests/docmaps/v2/provider_test.py

@@ -1,3 +1,4 @@
+import json
 from unittest.mock import patch, MagicMock
 from typing import Iterable, cast
 
@@ -19,6 +20,12 @@ def _iter_dict_from_bq_query_mock() -> Iterable[MagicMock]:
         yield mock
 
 
+def _get_docmaps_index_dict(provider: DocmapsProvider) -> dict:
+    return json.loads(
+        ''.join(provider.iter_docmaps_index_json_stream())
+    )
+
+
 class TestEnhancedPreprintsDocmapsProvider:
     def test_should_create_index_with_non_empty_docmaps(
         self,
@@ -27,7 +34,7 @@ def test_should_create_index_with_non_empty_docmaps(
         iter_dict_from_bq_query_mock.return_value = iter([
             DOCMAPS_QUERY_RESULT_ITEM_1
         ])
-        docmaps_index = DocmapsProvider().get_docmaps_index()
+        docmaps_index = _get_docmaps_index_dict(DocmapsProvider())
         assert docmaps_index['docmaps'] == [
             get_docmap_item_for_query_result_item(cast(ApiInput, DOCMAPS_QUERY_RESULT_ITEM_1))
         ]
@@ -42,8 +49,8 @@ def test_should_cache_docmaps_query_results(
         docmaps_provider = DocmapsProvider(
             query_results_cache=InMemorySingleObjectCache(max_age_in_seconds=10)
         )
-        docmaps_provider.get_docmaps_index()
-        docmaps_index = docmaps_provider.get_docmaps_index()
+        docmaps_index = _get_docmaps_index_dict(docmaps_provider)
+        docmaps_index = _get_docmaps_index_dict(docmaps_provider)
         assert iter_dict_from_bq_query_mock.call_count == 1
         assert docmaps_index['docmaps'] == [
             get_docmap_item_for_query_result_item(cast(ApiInput, DOCMAPS_QUERY_RESULT_ITEM_1))

tests/unit_tests/main_test.py

@@ -1,3 +1,4 @@
+import json
 from unittest.mock import patch, MagicMock
 from typing import Iterable, Sequence
 
@@ -47,7 +48,9 @@ def test_should_return_json_with_docmaps_list_from_enhanced_preprint_provider(
         enhanced_preprints_docmaps_provider_mock: MagicMock
     ):
         docmaps_index = [{'docmaps': [{'id': 'docmap_1'}, {'id': 'docmap_2'}]}]
-        enhanced_preprints_docmaps_provider_mock.get_docmaps_index.return_value = docmaps_index
+        enhanced_preprints_docmaps_provider_mock.iter_docmaps_index_json_stream.return_value = iter(
+            [json.dumps(docmaps_index)]
+        )
         client = TestClient(create_app())
         response = client.get('/enhanced-preprints/docmaps/v2/index')
         assert response.json() == docmaps_index