scripts/gdrive_sync.py

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