Merge pull request #133 from PolicyEngine/350_use_caching_google_storage_client

download_file_from_gcs now uses cache.

Author: mikesmit

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: patch
+  changes:
+    fixed:
+      - downloads from google storage should now be properly cached.

policyengine/utils/google_cloud_bucket.py

@@ -1,3 +1,23 @@
+from .data.caching_google_storage_client import CachingGoogleStorageClient
+import asyncio
+from pathlib import Path
+
+_caching_client: CachingGoogleStorageClient | None = None
+
+
+def _clear_client():
+    global _caching_client
+    _caching_client = None
+
+
+def _get_client():
+    global _caching_client
+    if _caching_client is not None:
+        return _caching_client
+    _caching_client = CachingGoogleStorageClient()
+    return _caching_client
+
+
 def download_file_from_gcs(
     bucket_name: str, file_name: str, destination_path: str
 ) -> None:
@@ -12,18 +32,6 @@ def download_file_from_gcs(
     Returns:
         None
     """
-    from google.cloud import storage
-
-    # Initialize a client
-    client = storage.Client()
-
-    # Get the bucket
-    bucket = client.bucket(bucket_name)
-
-    # Create a blob object from the file name
-    blob = bucket.blob(file_name)
-
-    # Download the file to a local path
-    blob.download_to_filename(destination_path)
-
-    return destination_path
+    asyncio.run(
+        _get_client().download(bucket_name, file_name, Path(destination_path))
+    )

tests/utils/data/test_google_cloud_bucket.py

@@ -0,0 +1,39 @@
+from unittest import TestCase
+from unittest.mock import patch
+import pytest
+from pathlib import Path
+from policyengine.utils.google_cloud_bucket import (
+    download_file_from_gcs,
+    _clear_client,
+)
+
+
+class TestGoogleCloudBucket(TestCase):
+    def setUp(self):
+        _clear_client()
+
+    @patch(
+        "policyengine.utils.google_cloud_bucket.CachingGoogleStorageClient",
+        autospec=True,
+    )
+    def test_download_uses_storage_client(self, client_class):
+        client_instance = client_class.return_value
+        download_file_from_gcs(
+            "TEST_BUCKET", "TEST/FILE/NAME.TXT", "TARGET/PATH"
+        )
+        client_instance.download.assert_called_with(
+            "TEST_BUCKET", "TEST/FILE/NAME.TXT", Path("TARGET/PATH")
+        )
+
+    @patch(
+        "policyengine.utils.google_cloud_bucket.CachingGoogleStorageClient",
+        autospec=True,
+    )
+    def test_download_only_creates_client_once(self, client_class):
+        download_file_from_gcs(
+            "TEST_BUCKET", "TEST/FILE/NAME.TXT", "TARGET/PATH"
+        )
+        download_file_from_gcs(
+            "TEST_BUCKET", "TEST/FILE/NAME.TXT", "ANOTHER/PATH"
+        )
+        client_class.assert_called_once()