Add drive download utility functions and test

processes/util/google_integration.py

@@ -0,0 +1,44 @@
+from io import BytesIO
+import json
+import os
+from googleapiclient.discovery import build
+from googleapiclient.http import MediaIoBaseDownload
+from google.oauth2.service_account import Credentials
+
+SERVICE_ACCOUNT_FILE = os.environ['GOOGLE_SERVICE_ACCOUNT_FILE']
+
+from logger import create_log
+
+logger = create_log(__name__)
+
+def create_drive_service(service_account_info):
+    scopes = ['https://www.googleapis.com/auth/drive']
+    credentials = Credentials.from_service_account_info(service_account_info, scopes=scopes)
+
+    return build('drive', 'v3', credentials=credentials)
+
+def get_drive_file(file_id: str) -> BytesIO:
+    drive_service = create_drive_service(service_account_info=json.loads(SERVICE_ACCOUNT_FILE))
+    request = drive_service.files().get_media(fileId=file_id)
+    file = BytesIO()
+
+    try:
+        downloader = MediaIoBaseDownload(file, request)
+
+        done = False
+        while done is False:
+            status, done = downloader.next_chunk()
+
+    except HttpError as error:
+        logger.warning(f"HTTP error occurred when downloading Drive file {file_id}: {error}")
+        file = None
+    except Exception as err:
+        logger.warning(f"Unexpected {type(err)=} occurred when downloading drive file {file_id}: {error}")
+
+    return file
+
+
+
+
+
+

tests/integration/test_google_integration.py

@@ -0,0 +1,14 @@
+import os
+from load_env import load_env_file
+
+load_env_file('local-compose', file_string='config/local-compose.yaml')
+
+from processes.util.google_integration import get_drive_file
+
+def test_get_drive_file():
+    test_id = os.environ['EXAMPLE_FILE_ID']
+    file = get_drive_file(test_id)
+
+    assert file != None
+    assert file.seek(0, 2) > 1
+