106 lines
3.8 KiB
Python
106 lines
3.8 KiB
Python
"""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 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
|