gdrive_storage.py

"""
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()