Spaces:
Running
Running
| import logging | |
| from googleapiclient.discovery import build | |
| from googleapiclient.http import MediaIoBaseDownload, MediaFileUpload | |
| from google.oauth2.service_account import Credentials | |
| from fastapi import HTTPException | |
| import os | |
| logger = logging.getLogger(__name__) | |
| def initialize_drive_api(): | |
| try: | |
| logger.info("Initializing Google Drive API...") | |
| creds = Credentials.from_service_account_file( | |
| SERVICE_ACCOUNT_FILE, | |
| scopes=["https://www.googleapis.com/auth/drive"] | |
| ) | |
| service = build('drive', 'v3', credentials=creds) | |
| logger.info("Google Drive API initialized successfully.") | |
| return service | |
| except Exception as e: | |
| logger.error(f"Error initializing Google Drive API: {e}") | |
| raise HTTPException(status_code=500, detail=f"Failed to initialize Google Drive API: {e}") | |
| def list_files_in_folder(service, folder_id): | |
| try: | |
| logger.info(f"Listing files in folder with ID: {folder_id}...") | |
| results = service.files().list( | |
| q=f"'{folder_id}' in parents", | |
| fields="files(id, name, mimeType)" | |
| ).execute() | |
| files = results.get('files', []) | |
| logger.info(f"Found {len(files)} files in the folder.") | |
| return files | |
| except Exception as e: | |
| logger.error(f"Error listing files: {e}") | |
| raise HTTPException(status_code=500, detail=f"Failed to list files in folder: {e}") | |
| def download_file(service, file_id, file_name, folder_name): | |
| try: | |
| subfolder_path = os.path.join(BASE_MODEL_PATH, folder_name) | |
| os.makedirs(subfolder_path, exist_ok=True) | |
| logger.info(f"Starting download for file: {file_name} with ID: {file_id}...") | |
| request = service.files().get_media(fileId=file_id) | |
| fh = io.FileIO(os.path.join(subfolder_path, file_name), 'wb') | |
| downloader = MediaIoBaseDownload(fh, request) | |
| done = False | |
| while done is False: | |
| status, done = downloader.next_chunk() | |
| logger.info(f"Download {file_name} {int(status.progress() * 100)}%.") | |
| logger.info(f"Downloaded {file_name} successfully.") | |
| except Exception as e: | |
| logger.error(f"Error downloading file {file_name}: {e}") | |
| raise HTTPException(status_code=500, detail=f"Failed to download file: {e}") | |
| def upload_file_to_google_drive(service, file_path, folder_id): | |
| try: | |
| file_name = os.path.basename(file_path) | |
| file_metadata = { | |
| 'name': file_name, | |
| 'parents': [folder_id] | |
| } | |
| media = MediaFileUpload(file_path, resumable=True) | |
| file = service.files().create( | |
| body=file_metadata, | |
| media_body=media, | |
| fields='id' | |
| ).execute() | |
| logger.info(f"File '{file_name}' uploaded successfully to Google Drive with ID: {file.get('id')}") | |
| return file.get('id') | |
| except Exception as e: | |
| logger.error(f"Error uploading file to Google Drive: {e}") | |
| raise HTTPException(status_code=500, detail=f"Failed to upload file to Google Drive: {e}") | |
| def create_folder_in_google_drive(service, folder_name, parent_folder_id): | |
| try: | |
| folder_metadata = { | |
| 'name': folder_name, | |
| 'mimeType': 'application/vnd.google-apps.folder', | |
| 'parents': [parent_folder_id] | |
| } | |
| folder = service.files().create( | |
| body=folder_metadata, | |
| fields='id' | |
| ).execute() | |
| logger.info(f"Folder '{folder_name}' created successfully with ID: {folder.get('id')}") | |
| return folder.get('id') | |
| except Exception as e: | |
| logger.error(f"Error creating folder in Google Drive: {e}") | |
| raise HTTPException(status_code=500, detail=f"Failed to create folder in Google Drive: {e}") | |
| def find_folder_by_name(service, folder_name, parent_folder_id): | |
| try: | |
| results = service.files().list( | |
| q=f"name='{folder_name}' and '{parent_folder_id}' in parents and mimeType='application/vnd.google-apps.folder'", | |
| fields="files(id, name)" | |
| ).execute |