app_drive_upload.py

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()