Spaces:
Sleeping
Sleeping
File size: 1,977 Bytes
cce8120 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | 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()
|