[Ananya/Sunandhita] Add. Implementation for GCP connector

Author: Sunandhita B

jb-lib/lib/file_storage/__init__.py

@@ -1,4 +1,5 @@
 from .azure import AzureAsyncStorage, AzureSyncStorage
 from .local import LocalAsyncStorage, LocalSyncStorage
+from .gcp import GcpAsyncStorage, GcpSyncStorage
 from .storage import Storage
 from .handler import StorageHandler

jb-lib/lib/file_storage/gcp/__init__.py

@@ -0,0 +1,2 @@
+from .gcp_storage import GcpAsyncStorage
+from .gcp_sync_storage import GcpSyncStorage

jb-lib/lib/file_storage/gcp/gcp_storage.py

@@ -0,0 +1,101 @@
+import os
+from typing import Union, Optional
+from datetime import datetime, timedelta, timezone
+import logging
+from google.cloud import storage
+import aiofiles
+
+logger = logging.getLogger("storage")
+
+class GcpAsyncStorage:
+    __client__ = None
+    tmp_folder = "/tmp/jb_files"
+
+    def __init__(self):
+        logger.info("Initializing GCP Storage")
+        
+        project_id = 'indian-legal-bert'
+        self.__bucket_name__ = 'jugalbandi'
+        os.environ['GOOGLE_APPLICATION_CREDENTIALS']='/Users/sunandhitab/Downloads/indian-legal-bert-72a5c6f931f1.json'
+        if not project_id or not self.__bucket_name__:
+            print(project_id, self.__bucket_name__)
+            raise ValueError(
+                "GCPAsyncStorage client not initialized. Missing project_id or bucket_name"
+            )
+
+        self.__client__ = storage.Client(project=project_id)
+        os.makedirs(self.tmp_folder, exist_ok=True)
+
+    async def write_file(
+        self,
+        file_path: str,
+        file_content: Union[str, bytes],
+        mime_type: Optional[str] = None,
+    ):
+        if not self.__client__:
+            raise Exception("GCPAsyncStorage client not initialized")
+
+        blob_name = file_path
+        bucket = self.__client__.bucket(self.__bucket_name__)
+        blob = bucket.blob(blob_name)
+
+        # Determine MIME type if not provided
+        if mime_type is None:
+            mime_type = (
+                "audio/mpeg" if file_path.lower().endswith(".mp3") else "application/octet-stream"
+            )
+
+        # Upload the blob
+        await asyncio.to_thread(blob.upload_from_string, file_content, content_type=mime_type)
+
+    async def _download_file_to_temp_storage(
+        self, file_path: Union[str, os.PathLike]
+    ) -> Union[str, os.PathLike]:
+        if not self.__client__:
+            raise Exception("GCPAsyncStorage client not initialized")
+
+        blob_name = file_path
+        bucket = self.__client__.bucket(self.__bucket_name__)
+        blob = bucket.blob(blob_name)
+
+        tmp_file_path = os.path.join(self.tmp_folder, file_path)
+
+        # Create directory if it doesn't exist
+        os.makedirs(os.path.dirname(tmp_file_path), exist_ok=True)
+
+        async with aiofiles.open(tmp_file_path, 'wb') as my_blob:
+            await asyncio.to_thread(blob.download_to_file, my_blob)
+
+        return tmp_file_path
+
+    def public_url(self, file_path: str) -> str:
+        if not self.__client__:
+            raise Exception("GCPAsyncStorage client not initialized")
+
+        blob_name = file_path
+        bucket = self.__client__.bucket(self.__bucket_name__)
+        blob = bucket.blob(blob_name)
+
+        # Generate a signed URL that expires in 1 day
+        url = blob.generate_signed_url(
+            version="v4",
+            expiration=timedelta(days=1),
+            method="GET"
+        )
+
+        return url
+
+# Example usage
+async def main():
+    storage = GcpAsyncStorage()
+    await storage.write_file('example.txt', 'Hello, World!')
+    tmp_path = await storage._download_file_to_temp_storage('example.txt')
+    print(f"File downloaded to: {tmp_path}")
+    url = storage.public_url('example.txt')
+    print(f"Public URL: {url}")
+
+if __name__ == "__main__":
+    import asyncio
+    from dotenv import load_dotenv
+    load_dotenv()
+    asyncio.run(main())

jb-lib/lib/file_storage/gcp/gcp_sync_storage.py

@@ -0,0 +1,92 @@
+import os
+from typing import Union, Optional
+from datetime import timedelta
+import logging
+from google.cloud import storage
+
+logger = logging.getLogger("storage")
+
+class GcpSyncStorage:
+    __client__ = None
+    tmp_folder = "/tmp/jb_files"
+
+    def __init__(self):
+        logger.info("Initializing GCP Storage")
+
+        project_id = 'indian-legal-bert'
+        self.__bucket_name__ = 'jugalbandi'
+        os.environ['GOOGLE_APPLICATION_CREDENTIALS']='/Users/sunandhitab/Downloads/indian-legal-bert-72a5c6f931f1.json'
+
+        if not project_id or not self.__bucket_name__:
+            raise ValueError(
+                "GCPStorage client not initialized. Missing project_id or bucket_name"
+            )
+
+        self.__client__ = storage.Client(project=project_id)
+        os.makedirs(self.tmp_folder, exist_ok=True)
+
+    def write_file(
+        self,
+        file_path: str,
+        file_content: Union[str, bytes],
+        mime_type: Optional[str] = None,
+    ):
+        if not self.__client__:
+            raise Exception("GCPStorage client not initialized")
+
+        blob_name = file_path
+        bucket = self.__client__.bucket(self.__bucket_name__)
+        blob = bucket.blob(blob_name)
+
+        if mime_type is None:
+            mime_type = (
+                "audio/mpeg" if file_path.lower().endswith(".mp3") else "application/octet-stream"
+            )
+
+        # Use synchronous method to upload
+        blob.upload_from_string(file_content, content_type=mime_type)
+
+    def download_file_to_temp_storage(
+        self, file_path: Union[str, os.PathLike]
+    ) -> Union[str, os.PathLike]:
+        if not self.__client__:
+            raise Exception("GCPStorage client not initialized")
+
+        blob_name = file_path
+        bucket = self.__client__.bucket(self.__bucket_name__)
+        blob = bucket.blob(blob_name)
+
+        tmp_file_path = os.path.join(self.tmp_folder, file_path)
+        os.makedirs(os.path.dirname(tmp_file_path), exist_ok=True)
+
+        # Download the file to the temporary location
+        with open(tmp_file_path, 'wb') as my_blob:
+            blob.download_to_file(my_blob)
+
+        return tmp_file_path
+
+    def public_url(self, file_path: str) -> str:
+        if not self.__client__:
+            raise Exception("GCPStorage client not initialized")
+
+        blob_name = file_path
+        bucket = self.__client__.bucket(self.__bucket_name__)
+        blob = bucket.blob(blob_name)
+
+        return blob.generate_signed_url(
+            version="v4",
+            expiration=timedelta(days=1),
+            method="GET"
+        )
+
+# Example usage
+def main():
+    storage = GcpSyncStorage()
+    storage.write_file('example.txt', 'Hello, World!')
+    tmp_path = storage.download_file_to_temp_storage('example.txt')
+    print(f"File downloaded to: {tmp_path}")
+    url = storage.public_url('example.txt')
+    print(f"Public URL: {url}")
+
+if __name__ == "__main__":
+    main()

jb-lib/lib/file_storage/registry.py

@@ -2,13 +2,16 @@
 from .storage import SyncStorage, AsyncStorage
 from .local import LocalAsyncStorage, LocalSyncStorage
 from .azure import AzureAsyncStorage, AzureSyncStorage
+from .gcp import GcpAsyncStorage,GcpSyncStorage
 
 STORAGE_REGISTRY: Dict[str, Type[AsyncStorage]] = {
     "local": LocalAsyncStorage,
     "azure": AzureAsyncStorage,
+    "gcp": GcpAsyncStorage,
 }
 
 SYNC_STORAGE_REGISTRY: Dict[str, Type[SyncStorage]] = {
     "local": LocalSyncStorage,
     "azure": AzureSyncStorage,
+    "gcp": GcpSyncStorage,
 }

jb-lib/tests/file_storage/gcp/test_gcp_storage.py

@@ -0,0 +1,97 @@
+import os
+import asyncio
+import unittest
+import pytest
+from unittest.mock import patch, MagicMock, AsyncMock
+from unittest import mock
+from google.cloud import storage
+from google.cloud.storage import Blob
+from lib.file_storage.gcp.gcp_storage import GcpAsyncStorage 
+
+class TestGCPAsyncStorage(unittest.TestCase):
+
+    @mock.patch('google.cloud.storage.Client')
+    def setUp(self, mock_storage_client):
+        os.environ['GCP_PROJECT_ID'] = 'test-project-id'
+        os.environ['GCP_STORAGE_BUCKET'] = 'test-bucket'
+
+        # Initialize the GCPAsyncStorage instance
+        self.storage = GcpAsyncStorage()
+
+        # Mock the bucket and blob
+        self.mock_bucket = mock_storage_client.return_value.bucket.return_value
+        self.mock_blob = self.mock_bucket.blob.return_value
+
+    @patch("google.cloud.storage.Blob")
+    @pytest.mark.asyncio
+    async def test_write_file(mock_blob):
+        with patch("lib.file_storage.gcp.gcp_storage.os.getenv") as mock_getenv, patch("google.cloud.storage.Client") as mock_storage_client:
+            mock_getenv.side_effect = lambda key: {
+                "GCP_STORAGE_BUCKET_NAME": "test_bucket",
+                "GCP_STORAGE_PROJECT": "fake_project",
+            }.get(key, None)
+
+            mock_storage_client_instance = MagicMock()
+            mock_storage_client.return_value = mock_storage_client_instance
+            
+            mock_blob_instance = MagicMock(spec=Blob)
+            mock_blob_instance.upload_from_string = AsyncMock()
+            mock_storage_client_instance.bucket.return_value.blob.return_value = mock_blob_instance
+
+            storage = GcpAsyncStorage()
+            print("Writing file")
+
+            await storage.write_file("test.txt", b"content")
+            print("File written")
+
+            mock_storage_client_instance.bucket.assert_called_once_with("test_bucket")
+            mock_storage_client_instance.bucket.return_value.blob.assert_called_once_with("test.txt")
+            mock_blob_instance.upload_from_string.assert_called_once_with(b"content")
+
+
+    @mock.patch('aiofiles.open', new_callable=mock.AsyncMock)
+    async def test_download_file_to_temp_storage(self, mock_aiofiles_open):
+        file_path = 'test.txt'
+        mock_file = mock.MagicMock()  # Mock the file object returned by aiofiles
+
+        mock_aiofiles_open.return_value.__aenter__.return_value = mock_file
+
+        # Set up the download method to do nothing when called
+        self.mock_blob.download_to_file = mock.AsyncMock()
+
+         # Call the async method using asyncio.run
+        async def run_test():
+            tmp_path = await self.storage._download_file_to_temp_storage(file_path)
+            return tmp_path
+
+        tmp_path = asyncio.run(run_test())
+
+        # Check that the correct temporary file path was returned
+        self.assertEqual(tmp_path, os.path.join(self.storage.tmp_folder, file_path))
+
+        # Assert that the download_to_file method was called with the mock file object
+        self.mock_blob.download_to_file.assert_called_once_with(mock_file)
+
+
+    async def test_public_url(self):
+        file_path = 'test.txt'
+
+        url = self.storage.public_url(file_path)
+
+        # Assert the signed URL generation was called
+        self.mock_blob.generate_signed_url.assert_called_once_with(
+            version="v4",
+            expiration=mock.ANY,
+            method="GET"
+        )
+
+        # Assert that the URL is generated correctly
+        self.assertIsInstance(url, str)
+
+    @mock.patch('os.makedirs')
+    def tearDown(self, mock_makedirs):
+        # Clean up any resources after tests
+        pass
+
+if __name__ == '__main__':
+    unittest.main()