Add manifest cache, PDF early limit check, and parallel S3 fetching

Author: femalves

config.py

@@ -21,6 +21,8 @@
 OPEN_SEARCH_INDEX = 'scan-explorer'
 OPEN_SEARCH_AGG_BUCKET_LIMIT = 10000
 
+REDIS_URL = 'redis://redis-backend:6379/1'
+
 ADS_SEARCH_SERVICE_URL = 'https://api.adsabs.harvard.edu/v1/search/query'
 ADS_SEARCH_SERVICE_TOKEN = '<CHANGE ME>'
 

requirements.txt

@@ -14,3 +14,4 @@ alembic==1.8.0
 img2pdf==0.4.4
 appmap>=1.1.0.dev0
 boto3==1.34.75
+redis==4.6.0

scan_explorer_service/tests/test_perf.py

@@ -0,0 +1,340 @@
+import unittest
+import json
+import time
+from flask import url_for
+from unittest.mock import patch, MagicMock
+from scan_explorer_service.tests.base import TestCaseDatabase
+from scan_explorer_service.models import Article, Base, Collection, Page
+
+
+class TestManifestCache(TestCaseDatabase):
+
+    def create_app(self):
+        from scan_explorer_service.app import create_app
+        return create_app(**{
+            'SQLALCHEMY_DATABASE_URI': self.postgresql_url,
+            'SQLALCHEMY_ECHO': False,
+            'TESTING': True,
+            'PROPAGATE_EXCEPTIONS': True,
+            'TRAP_BAD_REQUEST_ERRORS': True,
+            'PRESERVE_CONTEXT_ON_EXCEPTION': False
+        })
+
+    def setUp(self):
+        Base.metadata.drop_all(bind=self.app.db.engine)
+        Base.metadata.create_all(bind=self.app.db.engine)
+
+        import scan_explorer_service.views.manifest as m
+        m._redis_client = None
+
+        self.collection = Collection(type='type', journal='cacheJ', volume='0099')
+        self.app.db.session.add(self.collection)
+        self.app.db.session.commit()
+        self.app.db.session.refresh(self.collection)
+
+        self.article = Article(bibcode='2099CacheTest..001A',
+                               collection_id=self.collection.id)
+        self.app.db.session.add(self.article)
+        self.app.db.session.commit()
+        self.app.db.session.refresh(self.article)
+
+        self.page1 = Page(name='cp1', collection_id=self.collection.id,
+                          volume_running_page_num=1)
+        self.page1.width = 100
+        self.page1.height = 100
+        self.page1.label = '1'
+        self.page2 = Page(name='cp2', collection_id=self.collection.id,
+                          volume_running_page_num=2)
+        self.page2.width = 100
+        self.page2.height = 100
+        self.page2.label = '2'
+        self.app.db.session.add_all([self.page1, self.page2])
+        self.app.db.session.commit()
+
+        self.article.pages.append(self.page1)
+        self.article.pages.append(self.page2)
+        self.app.db.session.commit()
+
+    def tearDown(self):
+        import scan_explorer_service.views.manifest as m
+        m._redis_client = None
+        self.app.db.session.remove()
+        self.app.db.drop_all()
+
+    def _mock_redis(self):
+        mock_r = MagicMock()
+        store = {}
+
+        def mock_get(key):
+            entry = store.get(key)
+            if entry is None:
+                return None
+            val, exp = entry
+            if exp and time.monotonic() > exp:
+                del store[key]
+                return None
+            return val
+
+        def mock_setex(key, ttl, val):
+            store[key] = (val, time.monotonic() + ttl)
+
+        def mock_delete(key):
+            store.pop(key, None)
+
+        mock_r.get = mock_get
+        mock_r.setex = mock_setex
+        mock_r.delete = mock_delete
+        mock_r.ping.return_value = True
+
+        import scan_explorer_service.views.manifest as m
+        m._redis_client = mock_r
+        return mock_r, store
+
+    def test_cache_hit_returns_cached_json(self):
+        mock_r, store = self._mock_redis()
+        from scan_explorer_service.views.manifest import MANIFEST_CACHE_PREFIX
+        store[MANIFEST_CACHE_PREFIX + self.article.id] = ('{"@type":"sc:Manifest","cached":true}', time.monotonic() + 3600)
+
+        url = url_for("manifest.get_manifest", id=self.article.id)
+        r = self.client.get(url)
+        self.assertStatus(r, 200)
+        data = json.loads(r.data)
+        self.assertTrue(data.get('cached'))
+
+    def test_cache_hit_returns_correct_content_type(self):
+        mock_r, store = self._mock_redis()
+        from scan_explorer_service.views.manifest import MANIFEST_CACHE_PREFIX
+        store[MANIFEST_CACHE_PREFIX + self.collection.id] = ('{"@type":"sc:Manifest"}', time.monotonic() + 3600)
+
+        url = url_for("manifest.get_manifest", id=self.collection.id)
+        r = self.client.get(url)
+        self.assertStatus(r, 200)
+        self.assertIn('application/json', r.content_type)
+
+    def test_cache_miss_calls_setex(self):
+        mock_r, store = self._mock_redis()
+        original_setex = mock_r.setex
+        setex_calls = []
+
+        def tracking_setex(key, ttl, val):
+            setex_calls.append(key)
+            return original_setex(key, ttl, val)
+
+        mock_r.setex = tracking_setex
+
+        from scan_explorer_service.views.manifest import _cache_set, MANIFEST_CACHE_PREFIX
+        _cache_set(self.article.id, '{"@type":"sc:Manifest"}')
+
+        self.assertEqual(len(setex_calls), 1)
+        self.assertEqual(setex_calls[0], MANIFEST_CACHE_PREFIX + self.article.id)
+
+    def test_cached_manifest_skips_manifest_factory(self):
+        mock_r, store = self._mock_redis()
+        from scan_explorer_service.views.manifest import MANIFEST_CACHE_PREFIX
+        store[MANIFEST_CACHE_PREFIX + self.article.id] = ('{"@type":"sc:Manifest"}', time.monotonic() + 3600)
+
+        with patch('scan_explorer_service.views.manifest.manifest_factory') as mock_factory:
+            url = url_for("manifest.get_manifest", id=self.article.id)
+            r = self.client.get(url)
+            self.assertStatus(r, 200)
+            mock_factory.create_manifest.assert_not_called()
+
+    def test_404_not_cached(self):
+        mock_r, store = self._mock_redis()
+        from scan_explorer_service.views.manifest import MANIFEST_CACHE_PREFIX
+
+        url = url_for("manifest.get_manifest", id='nonexistent')
+        r = self.client.get(url)
+        self.assertStatus(r, 404)
+        self.assertNotIn(MANIFEST_CACHE_PREFIX + 'nonexistent', store)
+
+    def test_redis_unavailable_falls_through(self):
+        import scan_explorer_service.views.manifest as m
+        m._redis_client = None
+
+        with patch('scan_explorer_service.views.manifest.redis.from_url', side_effect=Exception("connection refused")):
+            url = url_for("manifest.get_manifest", id=self.article.id)
+            r = self.client.get(url)
+            self.assertStatus(r, 200)
+            data = json.loads(r.data)
+            self.assertEqual(data['@type'], 'sc:Manifest')
+
+
+class TestPdfEarlyLimitCheck(TestCaseDatabase):
+
+    def create_app(self):
+        from scan_explorer_service.app import create_app
+        return create_app(**{
+            'SQLALCHEMY_DATABASE_URI': self.postgresql_url,
+            'SQLALCHEMY_ECHO': False,
+            'TESTING': True,
+            'PROPAGATE_EXCEPTIONS': True,
+            'TRAP_BAD_REQUEST_ERRORS': True,
+            'PRESERVE_CONTEXT_ON_EXCEPTION': False,
+            'IMAGE_PDF_MEMORY_LIMIT': 100 * 1024 * 1024,
+            'IMAGE_PDF_PAGE_LIMIT': 100,
+        })
+
+    def setUp(self):
+        Base.metadata.drop_all(bind=self.app.db.engine)
+        Base.metadata.create_all(bind=self.app.db.engine)
+
+        self.collection = Collection(type='type', journal='journal', volume='volume')
+        self.app.db.session.add(self.collection)
+        self.app.db.session.commit()
+
+    def test_over_limit_returns_400_immediately(self):
+        response = self.client.get(url_for('proxy.pdf_save',
+                                           id=self.collection.id,
+                                           page_start=1,
+                                           page_end=150))
+        self.assertEqual(response.status_code, 400)
+        data = json.loads(response.data)
+        self.assertIn('exceeds limit', data['Message'])
+
+    def test_exactly_at_limit_passes(self):
+        with patch('scan_explorer_service.views.image_proxy.fetch_images') as mock_fi, \
+             patch('scan_explorer_service.views.image_proxy.img2pdf.convert') as mock_conv:
+            mock_fi.return_value = [b'data']
+            mock_conv.return_value = b'pdf'
+            response = self.client.get(url_for('proxy.pdf_save',
+                                               id=self.collection.id,
+                                               page_start=1,
+                                               page_end=100))
+            self.assertEqual(response.status_code, 200)
+
+    def test_one_over_limit_returns_400(self):
+        response = self.client.get(url_for('proxy.pdf_save',
+                                           id=self.collection.id,
+                                           page_start=1,
+                                           page_end=101))
+        self.assertEqual(response.status_code, 400)
+
+    def test_no_page_end_passes_limit_check(self):
+        with patch('scan_explorer_service.views.image_proxy.fetch_images') as mock_fi, \
+             patch('scan_explorer_service.views.image_proxy.img2pdf.convert') as mock_conv:
+            mock_fi.return_value = [b'data']
+            mock_conv.return_value = b'pdf'
+            response = self.client.get(url_for('proxy.pdf_save',
+                                               id=self.collection.id,
+                                               page_start=1))
+            self.assertEqual(response.status_code, 200)
+
+    @patch('scan_explorer_service.views.image_proxy.get_item')
+    def test_over_limit_does_not_touch_db(self, mock_get_item):
+        response = self.client.get(url_for('proxy.pdf_save',
+                                           id=self.collection.id,
+                                           page_start=1,
+                                           page_end=200))
+        self.assertEqual(response.status_code, 400)
+        mock_get_item.assert_not_called()
+
+    def test_inverted_page_range_returns_empty_pdf(self):
+        with patch('scan_explorer_service.views.image_proxy.fetch_images') as mock_fi, \
+             patch('scan_explorer_service.views.image_proxy.img2pdf.convert') as mock_conv:
+            mock_fi.return_value = []
+            mock_conv.return_value = b'pdf'
+            response = self.client.get(url_for('proxy.pdf_save',
+                                               id=self.collection.id,
+                                               page_start=10,
+                                               page_end=5))
+            self.assertIn(response.status_code, [200, 400])
+
+
+class TestParallelFetchImages(TestCaseDatabase):
+
+    def create_app(self):
+        from scan_explorer_service.app import create_app
+        return create_app(**{
+            'SQLALCHEMY_DATABASE_URI': self.postgresql_url,
+            'SQLALCHEMY_ECHO': False,
+            'TESTING': True,
+            'PROPAGATE_EXCEPTIONS': True,
+            'TRAP_BAD_REQUEST_ERRORS': True,
+            'PRESERVE_CONTEXT_ON_EXCEPTION': False,
+            'IMAGE_PDF_MEMORY_LIMIT': 100 * 1024 * 1024,
+            'IMAGE_PDF_PAGE_LIMIT': 100,
+        })
+
+    def setUp(self):
+        Base.metadata.drop_all(bind=self.app.db.engine)
+        Base.metadata.create_all(bind=self.app.db.engine)
+
+        self.collection = Collection(type='type', journal='journal', volume='volume')
+        self.app.db.session.add(self.collection)
+        self.app.db.session.commit()
+
+        self.article = Article(bibcode='1988ApJ...333..341R',
+                               collection_id=self.collection.id)
+        self.app.db.session.add(self.article)
+        self.app.db.session.commit()
+
+        pages = []
+        for i in range(5):
+            p = Page(name=f'page{i}', collection_id=self.collection.id,
+                     volume_running_page_num=i + 1)
+            p.width = 100
+            p.height = 100
+            p.label = str(i + 1)
+            pages.append(p)
+        self.app.db.session.add_all(pages)
+        self.app.db.session.commit()
+
+        for p in pages:
+            self.article.pages.append(p)
+        self.app.db.session.commit()
+
+        self.pages = pages
+
+    @patch('scan_explorer_service.views.image_proxy.S3Provider')
+    def test_fetch_images_returns_all_pages(self, mock_s3_cls):
+        mock_s3 = MagicMock()
+        mock_s3.read_object_s3.return_value = b'image_data'
+        mock_s3_cls.return_value = mock_s3
+
+        from scan_explorer_service.views.image_proxy import fetch_images
+        images = list(fetch_images(
+            self.app.db.session, self.collection, 1, 5, 100,
+            100 * 1024 * 1024))
+        self.assertEqual(len(images), 5)
+        self.assertTrue(all(img == b'image_data' for img in images))
+
+    @patch('scan_explorer_service.views.image_proxy.S3Provider')
+    def test_fetch_images_respects_memory_limit(self, mock_s3_cls):
+        mock_s3 = MagicMock()
+        mock_s3.read_object_s3.return_value = b'x' * 1000
+        mock_s3_cls.return_value = mock_s3
+
+        from scan_explorer_service.views.image_proxy import fetch_images
+        images = list(fetch_images(
+            self.app.db.session, self.collection, 1, 5, 100,
+            500))
+        self.assertLess(len(images), 5)
+
+    @patch('scan_explorer_service.views.image_proxy.S3Provider')
+    def test_fetch_images_skips_none_results(self, mock_s3_cls):
+        mock_s3 = MagicMock()
+        mock_s3.read_object_s3.side_effect = [b'data1', None, b'data3', b'data4', b'data5']
+        mock_s3_cls.return_value = mock_s3
+
+        from scan_explorer_service.views.image_proxy import fetch_images
+        images = list(fetch_images(
+            self.app.db.session, self.collection, 1, 5, 100,
+            100 * 1024 * 1024))
+        self.assertEqual(len(images), 4)
+
+    @patch('scan_explorer_service.views.image_proxy.S3Provider')
+    def test_single_s3provider_instance(self, mock_s3_cls):
+        mock_s3 = MagicMock()
+        mock_s3.read_object_s3.return_value = b'image_data'
+        mock_s3_cls.return_value = mock_s3
+
+        from scan_explorer_service.views.image_proxy import fetch_images
+        list(fetch_images(
+            self.app.db.session, self.collection, 1, 5, 100,
+            100 * 1024 * 1024))
+        self.assertEqual(mock_s3_cls.call_count, 1)
+
+
+if __name__ == '__main__':
+    unittest.main()

scan_explorer_service/tests/test_proxy.py

@@ -157,9 +157,11 @@ def test_get_item(self):
                 get_item(self.app.db.session, 'non-existent-id')
             assert("ID: non-existent-id not found" in str(context.exception))
 
-    @patch('scan_explorer_service.views.image_proxy.fetch_object')
-    def test_fetch_images(self, mock_fetch_object):
-        mock_fetch_object.return_value = b'image_data'
+    @patch('scan_explorer_service.views.image_proxy.S3Provider')
+    def test_fetch_images(self, mock_s3_cls):
+        mock_s3 = MagicMock()
+        mock_s3.read_object_s3.return_value = b'image_data'
+        mock_s3_cls.return_value = mock_s3
         item = self.article
         page_start = 1
         page_end = 2
@@ -169,7 +171,7 @@ def test_fetch_images(self, mock_fetch_object):
         gen = fetch_images(self.app.db.session, item, page_start, page_end, page_limit, memory_limit)
         images = list(gen)
         self.assertEqual(images, [b'image_data', b'image_data'])
-        mock_fetch_object.assert_called()
+        mock_s3.read_object_s3.assert_called()
 
     @patch('scan_explorer_service.utils.s3_utils.S3Provider.read_object_s3')
     def test_fetch_object(self, mock_read_object_s3):