from __future__ import annotations import mimetypes from pathlib import Path from google.auth.transport.requests import Request from google.oauth2 import service_account from google.oauth2.credentials import Credentials from google_auth_oauthlib.flow import InstalledAppFlow from googleapiclient.discovery import build from googleapiclient.http import MediaFileUpload SCOPES = ["https://www.googleapis.com/auth/drive.file"] def build_drive_service( *, oauth_client_secrets: Path | None = None, oauth_token_file: Path | None = None, service_account_file: Path | None = None, use_shared_drive: bool = False, ): """Prefer OAuth (personal My Drive). Service account only works with Shared Drives.""" creds = None if oauth_token_file and oauth_token_file.exists(): creds = Credentials.from_authorized_user_file(str(oauth_token_file), SCOPES) if creds and creds.expired and creds.refresh_token: creds.refresh(Request()) oauth_token_file.write_text(creds.to_json(), encoding="utf-8") elif oauth_client_secrets and oauth_client_secrets.exists(): flow = InstalledAppFlow.from_client_secrets_file(str(oauth_client_secrets), SCOPES) creds = flow.run_local_server(port=0) if oauth_token_file: oauth_token_file.parent.mkdir(parents=True, exist_ok=True) oauth_token_file.write_text(creds.to_json(), encoding="utf-8") elif service_account_file and service_account_file.exists(): if not use_shared_drive: raise RuntimeError( "Google service accounts cannot upload to normal My Drive folders " "(no storage quota). Use OAuth once (see docs/SETUP.md) or put the " "folder in a Shared Drive and set GOOGLE_USE_SHARED_DRIVE=true." ) creds = service_account.Credentials.from_service_account_file( str(service_account_file), scopes=SCOPES ) else: raise FileNotFoundError( "No Google credentials found. Add config/google_oauth_client.json " "and run: python bin/google_auth.py" ) return build("drive", "v3", credentials=creds, cache_discovery=False) class DriveUploader: def __init__( self, parent_folder_id: str, *, oauth_client_secrets: Path | None = None, oauth_token_file: Path | None = None, service_account_file: Path | None = None, use_shared_drive: bool = False, ): if not parent_folder_id: raise ValueError("GOOGLE_DRIVE_FOLDER_ID is empty") self.service = build_drive_service( oauth_client_secrets=oauth_client_secrets, oauth_token_file=oauth_token_file, service_account_file=service_account_file, use_shared_drive=use_shared_drive, ) self.parent_folder_id = parent_folder_id self._drive_params = {"supportsAllDrives": True} if use_shared_drive else {} def create_folder(self, name: str, parent_id: str | None = None) -> str: parent = parent_id or self.parent_folder_id meta = { "name": name, "mimeType": "application/vnd.google-apps.folder", "parents": [parent], } folder = ( self.service.files() .create(body=meta, fields="id, webViewLink", **self._drive_params) .execute() ) return folder["id"] def find_folder(self, name: str, parent_id: str | None = None) -> str | None: """Return folder id if a child folder with this exact name exists.""" parent = parent_id or self.parent_folder_id # Escape single quotes in name for Drive query safe = name.replace("'", "\\'") q = ( f"name = '{safe}' and '{parent}' in parents " f"and mimeType = 'application/vnd.google-apps.folder' and trashed = false" ) resp = ( self.service.files() .list( q=q, spaces="drive", fields="files(id, name)", pageSize=5, **self._drive_params, ) .execute() ) files = resp.get("files") or [] return files[0]["id"] if files else None def get_or_create_folder(self, name: str, parent_id: str | None = None) -> str: existing = self.find_folder(name, parent_id) if existing: return existing return self.create_folder(name, parent_id) def upload_file(self, local_path: Path, parent_id: str) -> dict: mime, _ = mimetypes.guess_type(str(local_path)) mime = mime or "application/octet-stream" media = MediaFileUpload(str(local_path), mimetype=mime, resumable=True) meta = {"name": local_path.name, "parents": [parent_id]} return ( self.service.files() .create(body=meta, media_body=media, fields="id, name, webViewLink", **self._drive_params) .execute() ) def upload_tree( self, local_dir: Path, *, date_folder: str, batch_folder: str, ) -> dict: """Upload preserving structure under: parent / YYYY-MM-DD / batch_HHMMSS / ... Example: Content Root/ 2026-07-16/ batch_093933/ idea-01_topic/ linkedin.txt hero.png manifest.json """ date_id = self.get_or_create_folder(date_folder, self.parent_folder_id) batch_id = self.create_folder(batch_folder, date_id) # Map relative dir path -> Drive folder id folder_cache: dict[str, str] = {"": batch_id} uploaded = [] for path in sorted(local_dir.rglob("*")): if not path.is_file(): continue if path.name in {"frames.txt"}: continue if "frames" in path.parts and path.suffix == ".png": continue if "_video_work" in path.parts: continue rel = path.relative_to(local_dir) parent_rel = str(rel.parent) if str(rel.parent) != "." else "" parent_drive_id = self._ensure_nested_folder(parent_rel, folder_cache, batch_id) uploaded.append(self.upload_file(path, parent_drive_id)) link = f"https://drive.google.com/drive/folders/{batch_id}" date_link = f"https://drive.google.com/drive/folders/{date_id}" return { "folder_id": batch_id, "folder_url": link, "date_folder": date_folder, "date_folder_id": date_id, "date_folder_url": date_link, "batch_folder": batch_folder, "files": uploaded, } def _ensure_nested_folder(self, rel_dir: str, cache: dict[str, str], root_id: str) -> str: if not rel_dir: return root_id if rel_dir in cache: return cache[rel_dir] parts = Path(rel_dir).parts current = "" parent_id = root_id for part in parts: current = f"{current}/{part}" if current else part if current not in cache: # Always create unique idea folders under this batch (don't reuse across batches) cache[current] = self.create_folder(part, parent_id) parent_id = cache[current] return cache[rel_dir]