team-tryouts/app/google_drive.py

"""Google Drive storage for contract PDFs.

The application uses an OAuth refresh token for the owner's personal Google
account. Tokens and client secrets come only from environment variables; the
database stores opaque file IDs, never a credential or a shareable Drive URL.
"""

import io
import os

GOOGLE_DRIVE_FOLDER_ID_ENV = 'GOOGLE_DRIVE_FOLDER_ID'
GOOGLE_DRIVE_CLIENT_ID_ENV = 'GOOGLE_DRIVE_CLIENT_ID'
GOOGLE_DRIVE_CLIENT_SECRET_ENV = 'GOOGLE_DRIVE_CLIENT_SECRET'
GOOGLE_DRIVE_REFRESH_TOKEN_ENV = 'GOOGLE_DRIVE_REFRESH_TOKEN'
GOOGLE_DRIVE_SCOPE = 'https://www.googleapis.com/auth/drive.file'


class GoogleDriveStorageError(RuntimeError):
    """A configuration or API failure while storing a contract in Drive."""


def _setting(name):
    value = os.getenv(name)
    if not value:
        raise GoogleDriveStorageError(
            f'{name} must be configured for Google Drive document storage.'
        )
    return value


def _drive_service():
    """Build an authorized Drive client from the owner's refresh token."""
    try:
        from google.oauth2.credentials import Credentials
        from googleapiclient.discovery import build
    except ImportError as exc:
        raise GoogleDriveStorageError(
            'Google Drive dependencies are not installed. Install requirements.txt again.'
        ) from exc

    credentials = Credentials(
        token=None,
        refresh_token=_setting(GOOGLE_DRIVE_REFRESH_TOKEN_ENV),
        token_uri='https://oauth2.googleapis.com/token',
        client_id=_setting(GOOGLE_DRIVE_CLIENT_ID_ENV),
        client_secret=_setting(GOOGLE_DRIVE_CLIENT_SECRET_ENV),
        scopes=[GOOGLE_DRIVE_SCOPE],
    )
    return build('drive', 'v3', credentials=credentials, cache_discovery=False)


def list_folder_files():
    """Return (file_id, name) tuples for every file in the configured folder.

    Used by the backup script to enumerate Drive-stored contracts. Returns
    only files directly in the folder (not nested inside sub-folders).

    Returns:
        list[tuple[str, str]]: (file_id, name) for each file found.
    """
    try:
        folder_id = _setting(GOOGLE_DRIVE_FOLDER_ID_ENV)
        page_token = None
        results = []
        while True:
            response = (
                _drive_service()
                .files()
                .list(
                    q=f"'{folder_id}' in parents and trashed=false",
                    fields='nextPageToken, files(id, name)',
                    pageToken=page_token,
                    pageSize=1000,
                )
                .execute()
            )
            for f in response.get('files', []):
                results.append((f['id'], f['name']))
            page_token = response.get('nextPageToken')
            if not page_token:
                break
        return results
    except GoogleDriveStorageError:
        raise
    except Exception as exc:
        raise GoogleDriveStorageError('Google Drive could not list folder contents.') from exc


def upload_file(*, stream, filename, mimetype):
    """Upload a PDF to the configured owner folder and return its Drive ID."""
    try:
        from googleapiclient.http import MediaIoBaseUpload

        stream.seek(0)
        media = MediaIoBaseUpload(stream, mimetype=mimetype, resumable=True)
        response = (
            _drive_service()
            .files()
            .create(
                body={'name': filename, 'parents': [_setting(GOOGLE_DRIVE_FOLDER_ID_ENV)]},
                media_body=media,
                fields='id',
            )
            .execute()
        )
    except GoogleDriveStorageError:
        raise
    except Exception as exc:  # Google client exceptions share no stable base class.
        raise GoogleDriveStorageError('Google Drive rejected the contract upload.') from exc

    file_id = response.get('id')
    if not file_id:
        raise GoogleDriveStorageError('Google Drive did not return an uploaded file identifier.')
    return file_id


def download_file(file_id):
    """Download a Drive file into memory for Flask's authenticated response."""
    try:
        from googleapiclient.http import MediaIoBaseDownload

        destination = io.BytesIO()
        downloader = MediaIoBaseDownload(
            destination,
            _drive_service().files().get_media(fileId=file_id),
        )
        complete = False
        while not complete:
            _status, complete = downloader.next_chunk()
        return destination.getvalue()
    except GoogleDriveStorageError:
        raise
    except Exception as exc:  # Google client exceptions share no stable base class.
        raise GoogleDriveStorageError('Google Drive could not download this contract.') from exc


def delete_file(file_id):
    """Permanently delete a Drive document when its contract record is deleted."""
    try:
        _drive_service().files().delete(fileId=file_id).execute()
    except GoogleDriveStorageError:
        raise
    except Exception as exc:  # Google client exceptions share no stable base class.
        raise GoogleDriveStorageError('Google Drive could not delete this contract.') from exc