File size: 9,183 Bytes
dc4e6da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 | """
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)
|