File size: 2,532 Bytes
d626f98 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | """
gdrive_storage.py — large file transfer (up to ~2GB) via a Google Drive
folder shared with a service account. Files are uploaded once, pulled
lazily by the recipient on tap, then permanently deleted from Drive —
nothing is ever shown to the user as a link; this is internal plumbing only.
"""
import os
import io
import json
from google.oauth2 import service_account
from googleapiclient.discovery import build
from googleapiclient.http import MediaIoBaseUpload, MediaIoBaseDownload
SCOPES = ["https://www.googleapis.com/auth/drive"]
FOLDER_ID = os.environ.get("GDRIVE_FOLDER_ID")
SERVICE_ACCOUNT_JSON = os.environ.get("GDRIVE_SERVICE_ACCOUNT_JSON")
_service = None
def _get_service():
global _service
if _service is not None:
return _service
if not SERVICE_ACCOUNT_JSON:
raise RuntimeError("GDRIVE_SERVICE_ACCOUNT_JSON is not configured")
info = json.loads(SERVICE_ACCOUNT_JSON)
creds = service_account.Credentials.from_service_account_info(info, scopes=SCOPES)
_service = build("drive", "v3", credentials=creds)
return _service
def upload_large_file(file_bytes: bytes, filename: str, mime_type: str) -> dict:
"""Uploads a file into the shared Drive folder. Returns Drive file metadata."""
service = _get_service()
if not FOLDER_ID:
raise RuntimeError("GDRIVE_FOLDER_ID is not configured")
media = MediaIoBaseUpload(io.BytesIO(file_bytes), mimetype=mime_type, resumable=True)
metadata = {"name": filename, "parents": [FOLDER_ID]}
file = service.files().create(
body=metadata, media_body=media, fields="id, name, size"
).execute()
return file
def download_and_delete(file_id: str):
"""Streams a file's bytes back, then permanently deletes it from Drive.
Returns (bytes, filename)."""
service = _get_service()
meta = service.files().get(fileId=file_id, fields="name").execute()
filename = meta.get("name", "file")
request = service.files().get_media(fileId=file_id)
buffer = io.BytesIO()
downloader = MediaIoBaseDownload(buffer, request)
done = False
while not done:
_, done = downloader.next_chunk()
file_bytes = buffer.getvalue()
# Reclaim space immediately once fully downloaded
service.files().delete(fileId=file_id).execute()
return file_bytes, filename
def delete_file(file_id: str):
"""Used for cleanup of files that were never downloaded (expiry job)."""
service = _get_service()
service.files().delete(fileId=file_id).execute()
|