import io import json import os from pathlib import Path from google.oauth2 import service_account from googleapiclient.discovery import build from googleapiclient.http import MediaIoBaseUpload SCOPES = [ "https://www.googleapis.com/auth/drive" ] class DriveUploader: def __init__(self): service_json = os.environ["GDRIVE_SERVICE_ACCOUNT_JSON"] folder_id = os.environ["GDRIVE_FOLDER_ID"] self.folder_id = folder_id credentials = service_account.Credentials.from_service_account_info( json.loads(service_json), scopes=SCOPES, ) self.service = build( "drive", "v3", credentials=credentials, cache_discovery=False, ) def upload(self, artifact: Path): with artifact.open("rb") as fp: media = MediaIoBaseUpload( io.BytesIO(fp.read()), mimetype="application/octet-stream", resumable=False, ) metadata = { "name": artifact.name, "parents": [ self.folder_id ], } created = ( self.service.files() .create( body=metadata, media_body=media, fields="id,name,webViewLink,webContentLink", ) .execute() ) self.service.permissions().create( fileId=created["id"], body={ "type": "anyone", "role": "reader", }, ).execute() return { "id": created["id"], "name": created["name"], "download_url": ( f"https://drive.google.com/uc?export=download&id={created['id']}" ), "view_url": created.get("webViewLink"), } drive_uploader = DriveUploader()