""" Google Drive integration for uploading generated documents. Accepts OAuth tokens directly from frontend (no backend OAuth flow). """ import io import pathlib from typing import Optional from google.oauth2.credentials import Credentials from googleapiclient.discovery import build from googleapiclient.http import MediaFileUpload, MediaIoBaseUpload from googleapiclient.errors import HttpError from google.auth.transport.requests import Request from datetime import datetime, timedelta from .config import settings class GoogleDriveClient: """Google Drive API client for file uploads using frontend-provided tokens""" def __init__(self, access_token: str, refresh_token: Optional[str] = None): """ Initialize Google Drive client with OAuth tokens from frontend. Args: access_token: Google OAuth access token (provided by frontend) refresh_token: Google OAuth refresh token (optional, for token renewal) Raises: ValueError: If token is invalid or expired """ self.access_token = access_token self.refresh_token = refresh_token self.credentials = self._create_credentials() self.service = build('drive', 'v3', credentials=self.credentials) def _create_credentials(self) -> Credentials: """Create credentials object from provided tokens""" # Validate refresh token requirements if self.refresh_token: # If refresh_token is provided, we need client credentials for auto-refresh if not settings.GOOGLE_CLIENT_ID or not settings.GOOGLE_CLIENT_SECRET: raise ValueError( "GOOGLE_CLIENT_ID and GOOGLE_CLIENT_SECRET must be set in .env " "to support token refresh. Either:\n" "1. Set these environment variables, OR\n" "2. Ensure the access token doesn't expire during processing (get fresh token)" ) credentials = Credentials( token=self.access_token, refresh_token=self.refresh_token, token_uri='https://oauth2.googleapis.com/token', client_id=settings.GOOGLE_CLIENT_ID, client_secret=settings.GOOGLE_CLIENT_SECRET, scopes=['https://www.googleapis.com/auth/drive.file'] ) else: # No refresh token - token must be valid for entire operation credentials = Credentials( token=self.access_token, scopes=['https://www.googleapis.com/auth/drive.file'] ) # Try to refresh if expired upfront (only if refresh_token available) if credentials.expired and credentials.refresh_token: try: print(f"[Google Drive] Token expired, refreshing...") credentials.refresh(Request()) print(f"[Google Drive] Token refreshed successfully") except Exception as e: raise ValueError( f"Failed to refresh Google Drive token: {str(e)}. " "User needs to re-authenticate." ) elif credentials.expired: raise ValueError( "Google Drive token has expired and no refresh token provided. " "User needs to re-authenticate with a fresh token." ) return credentials def upload_file( self, file_path: pathlib.Path, filename: Optional[str] = None, folder_name: str = "DocGenie Documents", mime_type: str = "application/zip" ) -> str: """ Upload a file to user's Google Drive. Args: file_path: Path to local file to upload filename: Name for file in Google Drive (default: use file_path name) folder_name: Name of folder to create/use in Drive mime_type: MIME type of the file Returns: Google Drive file URL (shareable link) Raises: HttpError: If upload fails """ try: # Get or create folder folder_id = self._get_or_create_folder(folder_name) # Prepare file metadata file_metadata = { 'name': filename or file_path.name, 'parents': [folder_id] } # Upload file media = MediaFileUpload( str(file_path), mimetype=mime_type, resumable=True ) file = self.service.files().create( body=file_metadata, media_body=media, fields='id, webViewLink, webContentLink' ).execute() # Make file accessible (reader permissions) self._share_file(file['id']) # Return shareable link return file.get('webViewLink', file.get('webContentLink')) except HttpError as error: print(f"Google Drive upload error: {error}") raise def upload_bytes( self, file_bytes: bytes, filename: str, folder_name: str = "DocGenie Documents", mime_type: str = "application/zip" ) -> str: """ Upload bytes directly to Google Drive (without saving to disk). Args: file_bytes: File content as bytes filename: Name for file in Google Drive folder_name: Name of folder to create/use in Drive mime_type: MIME type of the file Returns: Google Drive file URL (shareable link) """ try: folder_id = self._get_or_create_folder(folder_name) file_metadata = { 'name': filename, 'parents': [folder_id] } # Create media from bytes media = MediaIoBaseUpload( io.BytesIO(file_bytes), mimetype=mime_type, resumable=True ) file = self.service.files().create( body=file_metadata, media_body=media, fields='id, webViewLink, webContentLink' ).execute() self._share_file(file['id']) return file.get('webViewLink', file.get('webContentLink')) except HttpError as error: print(f"Google Drive upload error: {error}") raise def _get_or_create_folder(self, folder_name: str) -> str: """ Get or create a folder in user's Google Drive. Args: folder_name: Name of the folder Returns: Folder ID """ # Search for existing folder query = f"name='{folder_name}' and mimeType='application/vnd.google-apps.folder' and trashed=false" results = self.service.files().list( q=query, spaces='drive', fields='files(id, name)' ).execute() folders = results.get('files', []) if folders: # Folder exists, return its ID return folders[0]['id'] # Create new folder file_metadata = { 'name': folder_name, 'mimeType': 'application/vnd.google-apps.folder' } folder = self.service.files().create( body=file_metadata, fields='id' ).execute() return folder['id'] def _share_file(self, file_id: str): """ Make file shareable (anyone with link can view). Args: file_id: Google Drive file ID """ try: permission = { 'type': 'anyone', 'role': 'reader' } self.service.permissions().create( fileId=file_id, body=permission ).execute() except HttpError as error: print(f"Warning: Could not share file {file_id}: {error}") # Don't raise - file uploaded successfully even if sharing fails def upload_to_google_drive( access_token: str, file_path: pathlib.Path, refresh_token: Optional[str] = None, filename: Optional[str] = None ) -> str: """ Convenience function to upload a file to user's Google Drive. Args: access_token: Google OAuth access token (from frontend) file_path: Path to file to upload refresh_token: Google OAuth refresh token (optional) filename: Optional custom filename Returns: Google Drive URL Raises: ValueError: If token is invalid or expired HttpError: If upload fails """ client = GoogleDriveClient(access_token=access_token, refresh_token=refresh_token) return client.upload_file(file_path, filename)