""" app/tasks/upload_to_google_drive.py """ import os import json import logging from googleapiclient.discovery import build from googleapiclient.http import MediaFileUpload from google.oauth2.service_account import Credentials from google.oauth2.credentials import Credentials as OAuthCredentials from google_auth_oauthlib.flow import Flow from google.auth.transport.requests import Request from google.auth.exceptions import RefreshError from app.config import settings from app.tasks.retry_config import BaseTaskWithRetry from app.celery_app import celery logger = logging.getLogger(__name__) def get_drive_service_oauth(): """ Get Google Drive service using OAuth credentials. Uses saved refresh token to get a new access token. """ try: # Check for required OAuth settings if not (settings.google_drive_client_id and settings.google_drive_client_secret and settings.google_drive_refresh_token): logger.error("Google Drive OAuth credentials not fully configured") return None # Create credentials object from refresh token credentials = OAuthCredentials( None, # No access token initially, will be refreshed refresh_token=settings.google_drive_refresh_token, token_uri="https://oauth2.googleapis.com/token", client_id=settings.google_drive_client_id, client_secret=settings.google_drive_client_secret, # Use only drive.file scope scopes=['https://www.googleapis.com/auth/drive.file'] ) # Refresh the access token credentials.refresh(Request()) # Build and return the service service = build('drive', 'v3', credentials=credentials) return service except RefreshError as e: logger.error(f"Failed to refresh Google Drive token: {str(e)}") raise except Exception as e: logger.error(f"Failed to authenticate with Google Drive OAuth: {str(e)}") return None def get_google_drive_service(): """ Authenticate with Google Drive API using service account credentials and return an authorized service object. """ try: # Check if we should use OAuth instead of service account if getattr(settings, 'google_drive_use_oauth', False): return get_drive_service_oauth() # Load service account credentials from settings if not settings.google_drive_credentials_json: logger.error("Google Drive credentials not configured") return None credentials_dict = json.loads(settings.google_drive_credentials_json) credentials = Credentials.from_service_account_info( credentials_dict, scopes=['https://www.googleapis.com/auth/drive'] ) # Delegate to user if specified if settings.google_drive_delegate_to: credentials = credentials.with_subject(settings.google_drive_delegate_to) # Build and return the service service = build('drive', 'v3', credentials=credentials) return service except Exception as e: logger.error(f"Failed to authenticate with Google Drive: {str(e)}") return None def extract_metadata_from_file(file_path): """ Try to extract metadata from a file using several methods: 1. Check for a .json metadata file with the same name Returns a dictionary of metadata or empty dict if not found """ # Check for separate metadata JSON file metadata_path = os.path.splitext(file_path)[0] + '.json' if os.path.exists(metadata_path): try: with open(metadata_path, 'r', encoding='utf-8') as f: metadata = json.load(f) logger.info(f"Loaded metadata from external JSON file: {metadata_path}") return metadata except Exception as e: logger.warning(f"Failed to load metadata from JSON file: {str(e)}") return {} def truncate_property_value(key, value, max_bytes=100): """ Truncate a property value to ensure the key+value stays under the byte limit. Google Drive has a 124 byte limit for property key-value pairs. We reserve ~20-24 bytes for the key and leave ~100 bytes for the value. """ # Convert to string if not already str_value = str(value) # Calculate current size of key and value in bytes key_bytes = len(key.encode('utf-8')) value_bytes = len(str_value.encode('utf-8')) total_bytes = key_bytes + value_bytes # If under limit, return original value if total_bytes <= max_bytes: return str_value # Calculate how many bytes we need to trim from value # Leave a small buffer to be safe bytes_to_trim = total_bytes - max_bytes + 4 # Iteratively truncate the string until it's under the byte limit while len(str_value.encode('utf-8')) > value_bytes - bytes_to_trim: str_value = str_value[:-1] # Add ellipsis to indicate truncation if str_value != str(value): str_value = str_value[:-3] + "..." return str_value @celery.task(base=BaseTaskWithRetry) def upload_to_google_drive(file_path: str, include_metadata=True): """Uploads a file to Google Drive in the configured folder with optional metadata.""" if not os.path.exists(file_path): raise FileNotFoundError(f"File not found: {file_path}") # Extract filename from path filename = os.path.basename(file_path) # Extract metadata if available metadata = {} if include_metadata: metadata = extract_metadata_from_file(file_path) try: # Get Google Drive service service = get_google_drive_service() if not service: raise Exception("Failed to initialize Google Drive service") # Prepare the file metadata file_metadata = { 'name': filename, } # If folder ID is specified, set parent folder if settings.google_drive_folder_id: file_metadata['parents'] = [settings.google_drive_folder_id] # Add custom properties if metadata exists if metadata: # Google Drive properties must be strings and can't be nested objects file_metadata['properties'] = {} # Only add a few important top-level metadata fields as properties # Skip nested objects and long values to avoid the 124-byte limit safe_properties = {} for key, value in metadata.items(): # Skip nested structures completely - they'll be in the description if isinstance(value, (dict, list)): continue # Try to add simple values with truncation if needed try: truncated_value = truncate_property_value(key, value) safe_properties[key] = truncated_value except Exception as e: logger.warning(f"Skipping metadata property {key}: {str(e)}") # Only use the safe properties file_metadata['properties'] = safe_properties # Add minimal appProperties file_metadata['appProperties'] = { 'docuelevate': 'true' } # Add metadata to file description for better visibility in Google Drive UI # Description has much higher size limits than properties formatted_json = json.dumps(metadata, indent=2) file_metadata['description'] = f"Document Metadata:\n\n```json\n{formatted_json}\n```" logger.debug(f"Adding metadata to Google Drive file: {json.dumps(file_metadata['properties'])}") # Upload file with metadata media = MediaFileUpload( file_path, mimetype='application/pdf', resumable=True ) file = service.files().create( body=file_metadata, media_body=media, fields='id,name,webViewLink,properties,appProperties,description' ).execute() # Log success details file_id = file.get('id') web_view_link = file.get('webViewLink') logger.info(f"Successfully uploaded {filename} to Google Drive with ID: {file_id}") logger.info(f"File accessible at: {web_view_link}") result = { "status": "Completed", "file_path": file_path, "google_drive_file_id": file_id, "google_drive_web_link": web_view_link } # Add metadata info to result if included if metadata: result["metadata_included"] = True return result except Exception as e: error_msg = f"Failed to upload {filename} to Google Drive: {str(e)}" logger.error(error_msg) raise Exception(error_msg)