sa-resume-worker / app /drive_client.py
Avinashnalla7's picture
Fix Drive client query + enable heartbeat upsert
bcc0e71
Raw
History Blame Contribute Delete
2.78 kB
import io
import json
from typing import Optional, List, Dict
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
from googleapiclient.http import MediaIoBaseUpload
# Scope is enough to create/update files your app creates/has access to in Drive
DRIVE_SCOPES = ["https://www.googleapis.com/auth/drive.file"]
def build_drive_service(token_json_str: str):
"""
token_json_str = HF secret GOOGLE_TOKEN_JSON (contents of token.json)
"""
info = json.loads(token_json_str)
creds = Credentials.from_authorized_user_info(info, scopes=DRIVE_SCOPES)
return build("drive", "v3", credentials=creds, cache_discovery=False)
def list_children(service, parent_folder_id: str, page_size: int = 20) -> List[Dict]:
"""
List immediate children in a folder (files + folders).
"""
q = f"'{parent_folder_id}' in parents and trashed = false"
resp = (
service.files()
.list(
q=q,
pageSize=page_size,
fields="files(id,name,modifiedTime,mimeType)",
)
.execute()
)
return resp.get("files", [])
def find_file_by_name(service, parent_folder_id: str, filename: str) -> Optional[str]:
"""
Returns the fileId if a file exists with this exact name in the folder.
"""
# IMPORTANT: sanitize BEFORE f-string (no backslashes inside f-string expressions)
safe_name = filename.replace("'", "\\'")
q = (
f"'{parent_folder_id}' in parents "
f"and trashed = false "
f"and name = '{safe_name}'"
)
resp = (
service.files()
.list(
q=q,
pageSize=1,
fields="files(id,name)",
)
.execute()
)
files = resp.get("files", [])
return files[0]["id"] if files else None
def upsert_json_file(service, parent_folder_id: str, filename: str, payload: dict) -> str:
"""
Create the JSON file if missing, otherwise update it.
Returns the Drive fileId.
"""
data = json.dumps(payload, indent=2).encode("utf-8")
media = MediaIoBaseUpload(io.BytesIO(data), mimetype="application/json", resumable=False)
existing_id = find_file_by_name(service, parent_folder_id, filename)
if existing_id:
updated = (
service.files()
.update(
fileId=existing_id,
media_body=media,
fields="id,modifiedTime,name",
)
.execute()
)
return updated["id"]
created = (
service.files()
.create(
body={"name": filename, "parents": [parent_folder_id]},
media_body=media,
fields="id,modifiedTime,name",
)
.execute()
)
return created["id"]