Spaces:

Text-to-Document-Generation
/

Docgenie-API

Paused

Ahadhassan-2003

deploy: update HF Space

dc4e6da 6 months ago

9.18 kB

"""
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)

	"""
	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)