drive_sync.py

import os
import io
import json
import time
import threading
import traceback

SYNC_INTERVAL_SECONDS = 3600
DRIVE_FILE_NAME       = "gene_function_lab.db"
LOCAL_DB_PATH         = os.environ.get(
    "GENE_LAB_DB_PATH",
    "/tmp/pubmed_llm/gene_function_lab/gene_function_lab.db",
)
DRIVE_FILE_ID_ENV     = "GOOGLE_DRIVE_DB_FILE_ID"
DRIVE_FOLDER_ID_ENV   = "GOOGLE_DRIVE_FOLDER_ID"

_last_sync     = 0
_sync_lock     = threading.Lock()
_drive_file_id = None


def _get_drive_service(write=False):
    try:
        from google.oauth2 import service_account
        from googleapiclient.discovery import build
    except ImportError:
        raise ImportError("Run: pip install google-auth google-auth-httplib2 google-api-python-client")

    sa_json = os.environ.get("GOOGLE_SERVICE_ACCOUNT_JSON", "").strip()
    if not sa_json:
        raise RuntimeError("GOOGLE_SERVICE_ACCOUNT_JSON secret not set.")

    try:
        info = json.loads(sa_json)
    except json.JSONDecodeError as e:
        raise RuntimeError(f"GOOGLE_SERVICE_ACCOUNT_JSON is not valid JSON: {e}")

    if info.get("type") != "service_account":
        raise RuntimeError(f"Wrong type: {info.get('type')} — expected service_account")

    print(f"[Drive] Service account: {info.get('client_email', 'unknown')}")

    scopes = (
    ["https://www.googleapis.com/auth/drive"]
    if write else
    ["https://www.googleapis.com/auth/drive.readonly"]  # drop drive.file
    )
    creds = service_account.Credentials.from_service_account_info(info, scopes=scopes)
    return build("drive", "v3", credentials=creds, cache_discovery=False)


def _find_db_file_id(service):
    global _drive_file_id
    if _drive_file_id:
        return _drive_file_id

    configured_file_id = os.environ.get(DRIVE_FILE_ID_ENV, "").strip()
    if configured_file_id:
        _drive_file_id = configured_file_id
        print(f"[Drive] Using {DRIVE_FILE_ID_ENV}: {_drive_file_id}")
        return _drive_file_id

    folder_id = os.environ.get(DRIVE_FOLDER_ID_ENV, "").strip()
    query_parts = [f"name='{DRIVE_FILE_NAME}'", "trashed=false"]
    if folder_id:
        query_parts.append(f"'{folder_id}' in parents")
    query = " and ".join(query_parts)

    print(f"[Drive] Searching for {DRIVE_FILE_NAME}...")
    if folder_id:
        print(f"[Drive] Restricting search to folder id: {folder_id}")
    results = service.files().list(
        q=query,
        fields="files(id, name, modifiedTime, size)",
        orderBy="modifiedTime desc",
        pageSize=10,
    ).execute()

    files = results.get("files", [])
    print(f"[Drive] Found {len(files)} file(s)")

    if not files:
        raise FileNotFoundError(
            f"{DRIVE_FILE_NAME} not found. "
            "Share the gene_function_lab folder with the service account email."
        )
    if len(files) > 1:
        print("[Drive] Multiple DB files found; using most recently modified.")
        for f in files:
            print(
                f"  - id={f.get('id')} modified={f.get('modifiedTime')} "
                f"size={f.get('size')} name={f.get('name')}"
            )

    _drive_file_id = files[0]["id"]
    print(f"[Drive] File id: {_drive_file_id}")
    return _drive_file_id


def sync_db_from_drive(force=False):
    global _last_sync
    now = time.time()
    if not force and (now - _last_sync) < SYNC_INTERVAL_SECONDS:
        return False

    with _sync_lock:
        if not force and (time.time() - _last_sync) < SYNC_INTERVAL_SECONDS:
            return False
        try:
            from googleapiclient.http import MediaIoBaseDownload
            service    = _get_drive_service()
            file_id    = _find_db_file_id(service)
            request    = service.files().get_media(fileId=file_id)
            buf        = io.BytesIO()
            downloader = MediaIoBaseDownload(buf, request)
            done = False
            while not done:
                _, done = downloader.next_chunk()
            tmp = LOCAL_DB_PATH + ".tmp"
            os.makedirs(os.path.dirname(LOCAL_DB_PATH), exist_ok=True)
            with open(tmp, "wb") as f:
                f.write(buf.getvalue())
            os.replace(tmp, LOCAL_DB_PATH)
            try:
                import db as database
                database.init_db(LOCAL_DB_PATH)
                database.DB_PATH = LOCAL_DB_PATH
            except Exception as migrate_error:
                print(f"[Drive] DB migration warning after sync: {migrate_error}")
                print(traceback.format_exc())
            _last_sync = time.time()
            print(f"[Drive] Synced successfully ({len(buf.getvalue())//1024}KB)")
            return True
        except Exception as e:
            print(f"[Drive] Sync failed: {e}")
            print(traceback.format_exc())
            return False


def upload_db_to_drive(local_path=LOCAL_DB_PATH):
    try:
        from googleapiclient.http import MediaFileUpload
        service = _get_drive_service(write=True)
        file_id = _find_db_file_id(service)
        media   = MediaFileUpload(local_path, mimetype="application/octet-stream")
        service.files().update(fileId=file_id, media_body=media).execute()
        print("[Drive] Uploaded to Drive.")
        return True
    except Exception as e:
        print(f"[Drive] Upload failed: {e}")
        print(traceback.format_exc())
        return False


def write_queue_entry_to_drive(gene, requested_by="", max_papers=300):
    import db as database
    sync_db_from_drive(force=True)
    entry = database.queue_request(
        gene=gene, requested_by=requested_by,
        max_papers=max_papers, db_path=LOCAL_DB_PATH
    )
    upload_db_to_drive(LOCAL_DB_PATH)
    return entry


def start_background_sync():
    sa = os.environ.get("GOOGLE_SERVICE_ACCOUNT_JSON", "")
    print(f"[Drive] Secret present: {bool(sa)}, length: {len(sa)}")

    def _loop():
        print("[Drive] Running initial sync...")
        try:
            sync_db_from_drive(force=True)
        except Exception as e:
            print(f"[Drive] Initial sync error: {e}")
            print(traceback.format_exc())
        while True:
            time.sleep(SYNC_INTERVAL_SECONDS)
            try:
                sync_db_from_drive()
            except Exception as e:
                print(f"[Drive] Hourly sync error: {e}")
                print(traceback.format_exc())

    t = threading.Thread(target=_loop, daemon=True)
    t.start()
    print("[Drive] Background sync started.")