Merge pull request #38 from PySport/feat/add-load-files-cache

Add load files cache

Author: koenvo

ingestify/application/dataset_store.py

@@ -4,6 +4,8 @@
 import mimetypes
 import os
 import shutil
+from contextlib import contextmanager
+import threading
 from dataclasses import asdict
 from io import BytesIO
 
@@ -47,6 +49,8 @@ def __init__(
         self.storage_compression_method = "gzip"
         self.bucket = bucket
         self.event_bus: Optional[EventBus] = None
+        # Create thread-local storage for caching
+        self._thread_local = threading.local()
 
     # def __getstate__(self):
     #     return {"file_repository": self.file_repository, "bucket": self.bucket}
@@ -58,6 +62,34 @@ def dispatch(self, event):
         if self.event_bus:
             self.event_bus.dispatch(event)
 
+    @contextmanager
+    def with_file_cache(self):
+        """Context manager to enable file caching during its scope.
+
+        Files loaded within this context will be cached and reused,
+        avoiding multiple downloads of the same file.
+
+        Example:
+            # Without caching (loads files twice)
+            analyzer1 = StatsAnalyzer(store, dataset)
+            analyzer2 = VisualizationTool(store, dataset)
+
+            # With caching (files are loaded once and shared)
+            with store.with_file_cache():
+                analyzer1 = StatsAnalyzer(store, dataset)
+                analyzer2 = VisualizationTool(store, dataset)
+        """
+        # Enable caching for this thread
+        self._thread_local.use_file_cache = True
+        self._thread_local.file_cache = {}
+
+        try:
+            yield
+        finally:
+            # Disable caching for this thread
+            self._thread_local.use_file_cache = False
+            self._thread_local.file_cache = {}
+
     def save_ingestion_job_summary(self, ingestion_job_summary):
         self.dataset_repository.save_ingestion_job_summary(ingestion_job_summary)
 
@@ -384,10 +416,21 @@ def get_stream(file_):
                     self.file_repository.load_content(storage_path=file_.storage_path)
                 )
 
-            loaded_file = LoadedFile(
-                stream_=get_stream if lazy else get_stream(file),
-                **file.model_dump(),
-            )
+            def make_loaded_file():
+                return LoadedFile(
+                    stream_=get_stream if lazy else get_stream(file),
+                    **file.model_dump(),
+                )
+
+            # Using getattr with a default value of False - simple one-liner
+            if getattr(self._thread_local, "use_file_cache", False):
+                key = (dataset.dataset_id, current_revision.revision_id, file.file_id)
+                if key not in self._thread_local.file_cache:
+                    self._thread_local.file_cache[key] = make_loaded_file()
+                loaded_file = self._thread_local.file_cache[key]
+            else:
+                loaded_file = make_loaded_file()
+
             files[file.file_id] = loaded_file
         return FileCollection(files, auto_rewind=auto_rewind)
 

ingestify/tests/test_file_cache.py

@@ -0,0 +1,98 @@
+import pytest
+from io import BytesIO
+from unittest.mock import patch
+from datetime import datetime, timezone
+
+from ingestify.main import get_engine
+from ingestify.domain import Dataset, Identifier, Revision, File
+from ingestify.domain.models.dataset.revision import RevisionSource, SourceType
+
+
+def test_file_cache(config_file):
+    """Test file caching with the with_file_cache context manager."""
+    # Get engine from the fixture
+    engine = get_engine(config_file, "main")
+    store = engine.store
+
+    # Create a timestamp for test data
+    now = datetime.now(timezone.utc)
+
+    # Create a test file
+    test_file = File(
+        file_id="test_file_id",
+        data_feed_key="test_file",
+        tag="test_tag",
+        data_serialization_format="txt",
+        storage_path="test/path",
+        storage_size=100,
+        storage_compression_method="none",
+        created_at=now,
+        modified_at=now,
+        size=100,
+        content_type="text/plain",
+        data_spec_version="v1",
+    )
+
+    # Create a test revision with the file
+    revision = Revision(
+        revision_id=1,
+        created_at=now,
+        description="Test revision",
+        modified_files=[test_file],
+        source={"source_type": SourceType.MANUAL, "source_id": "test"},
+    )
+
+    # Create a test dataset with the revision
+    dataset = Dataset(
+        bucket="test-bucket",
+        dataset_id="test-dataset",
+        name="Test Dataset",
+        state="COMPLETE",
+        identifier=Identifier(test_id=1),
+        dataset_type="test",
+        provider="test-provider",
+        metadata={},
+        created_at=now,
+        updated_at=now,
+        last_modified_at=now,
+        revisions=[revision],
+    )
+
+    # Create a simple pass-through reader function to replace the gzip reader
+    def simple_reader(stream):
+        return stream
+
+    # Mock both the file repository and the _prepare_read_stream method
+    with patch.object(
+        store.file_repository, "load_content"
+    ) as mock_load_content, patch.object(
+        store, "_prepare_read_stream"
+    ) as mock_prepare_read_stream:
+
+        # Set up the mocks
+        mock_load_content.return_value = BytesIO(b"test content")
+        mock_prepare_read_stream.return_value = (simple_reader, "")
+
+        # Test without caching - should load files twice
+        store.load_files(dataset)
+        store.load_files(dataset)
+
+        # Should have called load_content twice (without caching)
+        assert mock_load_content.call_count == 2
+
+        # Reset the mock
+        mock_load_content.reset_mock()
+
+        # Test with caching - should load files only once
+        with store.with_file_cache():
+            store.load_files(dataset)
+            store.load_files(dataset)
+
+            # Should have called load_content only once (with caching)
+            assert mock_load_content.call_count == 1
+
+        # After exiting context, caching should be disabled
+        store.load_files(dataset)
+
+        # Should have called load_content again
+        assert mock_load_content.call_count == 2