import os import io import logging from googleapiclient.discovery import build from googleapiclient.http import MediaIoBaseDownload, MediaFileUpload from google.oauth2.credentials import Credentials from google.oauth2 import service_account logger = logging.getLogger("gap_system.gdrive") SCOPES = ['https://www.googleapis.com/auth/drive.file'] SERVICE_ACCOUNT_FILE = os.getenv("GDRIVE_SERVICE_ACCOUNT_JSON", "service_account.json") FOLDER_ID = os.getenv("GDRIVE_FOLDER_ID", "") DB_FILENAME = "gap_system.db" def get_drive_service(): if not os.path.exists(SERVICE_ACCOUNT_FILE): return None creds = service_account.Credentials.from_service_account_file( SERVICE_ACCOUNT_FILE, scopes=SCOPES ) return build('drive', 'v3', credentials=creds) def find_file(service): query = f"name='{DB_FILENAME}' and '{FOLDER_ID}' in parents and trashed=false" results = service.files().list(q=query, fields="files(id, name)").execute() files = results.get('files', []) if not files: return None return files[0]['id'] def download_db(local_path: str): """Download database from Google Drive on startup.""" service = get_drive_service() if not service or not FOLDER_ID: logger.warning("Google Drive credentials or FOLDER_ID not found, skipping sync.") return False try: file_id = find_file(service) if not file_id: logger.info("No remote gap_system.db found.") return False request = service.files().get_media(fileId=file_id) fh = io.FileIO(str(local_path), 'wb') downloader = MediaIoBaseDownload(fh, request) done = False while done is False: status, done = downloader.next_chunk() logger.info("Successfully downloaded gap_system.db from Google Drive.") return True except Exception as e: logger.error("Failed to download DB from Google Drive: %s", e) return False def upload_db(local_path: str): """Upload database to Google Drive.""" service = get_drive_service() if not service or not FOLDER_ID: return False if not os.path.exists(local_path): return False try: file_id = find_file(service) media = MediaFileUpload(local_path, mimetype='application/x-sqlite3', resumable=True) if file_id: # Update existing file service.files().update(fileId=file_id, media_body=media).execute() logger.info("Successfully updated gap_system.db on Google Drive.") else: # Create new file file_metadata = { 'name': DB_FILENAME, 'parents': [FOLDER_ID] } service.files().create(body=file_metadata, media_body=media, fields='id').execute() logger.info("Successfully uploaded gap_system.db to Google Drive.") return True except Exception as e: logger.error("Failed to upload DB to Google Drive: %s", e) return False