gdrive.py

import os
import json

from pydrive2.auth import GoogleAuth
from pydrive2.drive import GoogleDrive
from oauth2client.service_account import ServiceAccountCredentials


def login_with_service_account():
    scope = ["https://www.googleapis.com/auth/drive"]
    gauth = GoogleAuth()
    gauth.auth_method = 'service'
    cred_dict = json.load(open('evreal_cred.json', 'r'))
    cred_dict['private_key_id'] = os.environ['PRIVATE_KEY_ID']
    cred_dict['private_key'] = os.environ['PRIVATE_KEY'].replace('\\n', '\n')
    gauth.credentials = ServiceAccountCredentials.from_json_keyfile_dict(cred_dict, scope)
    return gauth


gauth = login_with_service_account()
drive = GoogleDrive(gauth)


def get_file_id_inside_folder(drive, folder_id, file_name):
    file_list = drive.ListFile({'q': "'{}' in parents and trashed=false".format(folder_id)}).GetList()
    for file in file_list:
        if file['title'] == file_name:
            return file['id']
    return None


def get_file_id_for_file_path(drive, file_path):
    # file_parts = os.path.split(file_path)
    path = os.path.normpath(file_path)
    file_parts = path.split(os.sep)
    file_id = '1-5zmRu3Od9s2dYu1S01u4aC2ihAmo1K2'
    for file_part in file_parts:
        if file_part == '':
            continue
        file_id = get_file_id_inside_folder(drive, file_id, file_part)
        if file_id is None:
            return None

    return file_id


def download_file(file_path, local_path):
    file_id = get_file_id_for_file_path(drive, file_path)
    if file_id is None:
        return False
    file = drive.CreateFile({'id': file_id})
    os.makedirs(os.path.dirname(local_path), exist_ok=True)
    file.GetContentFile(local_path)
    return True


# test if main
if __name__ == "__main__":
    drive_file_path = "0000_000000/ECD_calibration/videos/recon.mp4"
    # drive_file_path = "test.txt"
    local_file_path = os.path.join("temp_data", drive_file_path)
    # gauth = login_with_service_account()
    # drive = GoogleDrive(gauth)
    os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
    download_file(drive_file_path, local_file_path)