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 root_folder_id():
    """The configured Drive folder that holds every application document.

    This is the folder the owner set up (the ``signés`` folder); new
    documents land directly in it and signed copies are filed in nested
    folders below it.
    """
    return _setting(GOOGLE_DRIVE_FOLDER_ID_ENV)


def get_or_create_folder(name, parent_id):
    """Return the Drive folder named *name* under *parent_id*, creating it.

    Drive has no ``mkdir -p``; creating a folder is a write, so a signed
    copy filed under ``<team>/<role>/<username>`` walks the chain once per
    upload.  A folder already present is reused rather than duplicated.

    Args:
        name: Folder name to find or create.
        parent_id: Drive id of the parent folder.

    Returns:
        str: The Drive id of the folder.
    """
    try:
        service = _drive_service()
        query = (
            f"name='{name}' and mimeType='application/vnd.google-apps.folder' "
            f"and '{parent_id}' in parents and trashed=false"
        )
        response = service.files().list(q=query, fields='files(id, name)', pageSize=10).execute()
        existing = response.get('files', [])
        if existing:
            return existing[0]['id']
        folder = (
            service.files()
            .create(
                body={
                    'name': name,
                    'mimeType': 'application/vnd.google-apps.folder',
                    'parents': [parent_id],
                },
                fields='id',
            )
            .execute()
        )
        return folder['id']
    except GoogleDriveStorageError:
        raise
    except Exception as exc:  # Google client exceptions share no stable base class.
        raise GoogleDriveStorageError('Google Drive could not create the document folder.') from exc


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, parents=None):
    """Upload a PDF to the configured owner folder and return its Drive ID.

    Args:
        stream: The binary stream to upload.
        filename: Name used for the uploaded file.
        mimetype: MIME type of the file.
        parents: Drive ids of the parent folders.  Defaults to the root
            folder when omitted, which is how the legacy contract flow
            still works.
    """
    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': parents or [root_folder_id()],
                },
                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