基于Service Account使用Google Drive

在Google中创建Service Account，进而利用Drive管理文件。

一、配置工作

1.创建Project: 在Google Dashboard创建一个名为Drive的Project

2.启用API: 在API Library中搜索并启用Google Drive API

3.创建Service Account: 在Credentials页面新建一个Service Account

4.创建key: 选择新建好的Service Account，为其创建一个JSON形式的key

JSON形式的key会自动下载到本地，此时假设key的路径是/path/to/key.json。

二、使用教程

1.配置环境: 基于Python3.8创建一个虚拟环境并安装以下依赖

pip install google-api-python-client==2.116.0

2.编写脚本: 使用key的路径/path/to/key.json来初始化GoogleDrive类即可

# -*- coding: utf-8 -*-
import os
import sys
import logging
import traceback
from types import TracebackType
from google.oauth2 import service_account
from googleapiclient.discovery import build
from googleapiclient.http import MediaFileUpload, MediaIoBaseDownload

logging.basicConfig(level=logging.INFO)


class GoogleDrive:
    _instance = None

    def __new__(cls, *args, **kwargs) -> None:
        """
        Implement singleton mode.

        Returns:
            None
        """
        if not cls._instance:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self, key_path: str = "") -> None:
        """
        Initialize an instance of the GoogleDrive class.

        Args:
            key_path (str): The path of the Service Account key. Defaults to "".

        Returns:
            None
        """
        self._key_path = key_path
        self._drive_service = None
        self._init()

    def __enter__(self) -> 'GoogleDrive':
        """
        Context manager method for entering the context.

        Returns:
            GoogleDrive: The current instance of the GoogleDrive class.
        """
        return self

    def __exit__(self, exc_type: type, exc_val: BaseException, exc_tb: TracebackType) -> None:
        """
        Context manager method for exiting the context.

        Args:
            exc_type (type): The type of the exception (if any) that occurred within the context.
            exc_val (BaseException): The exception object (if any) that occurred within the context.
            exc_tb (TracebackType): The traceback object (if any) associated with the exception.

        Returns:
            None
        """
        if exc_type:
            logging.error(f"an exception of type {exc_type} occurred: {exc_val}")

        if exc_tb:
            logging.error("".join(traceback.format_tb(exc_tb)))

    def _init(self) -> None:
        """
        Initialize the drive_service.

        Returns:
            None
        """
        credentials = service_account.Credentials.from_service_account_file(
            filename=self._key_path,
            scopes=[
                "https://www.googleapis.com/auth/drive",
                "https://www.googleapis.com/auth/drive.metadata"
            ]
        )
        self._drive_service = build(serviceName="drive", version="v3", credentials=credentials)

    def _create_folder(self, folder_name: str, parent_folder_id: str = None) -> str:
        """
        Create a new folder in Google Drive.

        Args:
            folder_name (str): The name of the folder to create.
            parent_folder_id (str): The ID of the parent folder. Default is None (create in root).

        Returns:
            str: The ID of the newly created folder, or None if creation failed.
        """
        file_metadata = {
            "name": folder_name,
            "mimeType": "application/vnd.google-apps.folder",
            "parents": [parent_folder_id]
        }

        response = self._drive_service.files().create(body=file_metadata, fields="id").execute()
        folder_id = response.get("id")
        if folder_id is not None:
            logging.info(f"created folder {folder_name}")
        else:
            logging.error(f"failed to create folder, file will be uploaded to root, response: {response}")
        return folder_id

    def _get_id(self, name: str) -> str:
        """
        Get the ID of a folder/file by its name.

        Args:
            name (str): The name of the folder/file to search for.

        Returns:
            str: The ID of the folder/file if found, otherwise empty string.
        """
        response = self._drive_service.files().list(
            q=f""" name="{name}" """).execute()

        files = response.get("files", [])
        if files:
            if len(files) > 1:
                logging.error(f"too many results for name: {name}")
                sys.exit(1)
            result_id = files[0].get("id")
        else:
            logging.error(f"no search result for name {name}")
            result_id = ""

        return result_id

    def upload_file(self, file_path: str, file_name: str = None, folder_name: str = None) -> bool:
        """
        Upload a file to Google Drive.

        Args:
            file_path (str): Path to the file to upload.
            file_name (str): Name of the file on Google Drive. Default is None (use raw filename).
            folder_name (str): Name of the folder on Google Drive. Default is None (use tmp folder).

        Returns:
            bool: True if the file was uploaded successfully, False otherwise.
        """
        file_metadata = {}
        if file_name is None:
            file_name = os.path.basename(file_path)
        file_metadata.update({"name": file_name})

        if folder_name is None:
            folder_name = "tmp"

        folder_id = self._get_id(name=folder_name)
        if not folder_id:
            folder_id = self._create_folder(folder_name=folder_name)

        file_metadata.update({"parents": [folder_id]})

        media = MediaFileUpload(file_path, resumable=True)
        response = self._drive_service.files().create(body=file_metadata, media_body=media, fields="id").execute()

        if response.get("id") is not None:
            logging.info(f"""uploaded {file_path} to folder {folder_name}, id: {response.get("id")}""")
            return True
        else:
            logging.error(f"failed to upload, response: {response}")
            return False

    def download_file(self, file_name: str, destination_path: str = None) -> bool:
        """
        Download a file from Google Drive.

        Args:
            file_name (str): The name of the file to download.
            destination_path (str): The local path to save the downloaded file.

        Returns:
            bool: True if the file was downloaded successfully, False otherwise.
        """
        file_id = self._get_id(name=file_name)
        if not file_id:
            return False

        response = self._drive_service.files().get_media(fileId=file_id)

        if destination_path is None:
            destination_path = os.path.join(os.path.dirname(__file__), file_name)

        with open(destination_path, "wb") as file:
            downloader = MediaIoBaseDownload(file, response)
            done = False
            while not done:
                status, done = downloader.next_chunk()
                logging.info(f"progress: {int(status.progress() * 100)}%")

        logging.info(f"downloaded file {file_name} to {destination_path}")
        return True

    def delete_file(self, file_name: str) -> bool:
        """
        Delete a file from Google Drive.

        Args:
            file_name (str): The name of the file to delete.

        Returns:
            bool: True if the file was deleted successfully, False otherwise.
        """
        file_id = self._get_id(name=file_name)
        if not file_id:
            return False

        self._drive_service.files().delete(fileId=file_id).execute()

        logging.info(f"deleted file {file_name} from Google Drive")
        return True

3.管理文件: 类的upload_file()/download_file()/delete_file()方法可以对文件进行上传/下载/删除操作

if __name__ == "__main__":
    # files will be managed in Google Drive of the service account
    gd = GoogleDrive(key_path="/path/to/key.json")
    # upload file
    gd.upload_file(file_path="/path/to/test.txt")
    # download file
    gd.download_file(file_name="test.txt")
    # delete file
    gd.delete_file(file_name="test.txt")

注: 脚本管理的Drive空间是属于Service Account的。

---

至此教程结束。
有错误或者改进的地方请各位积极指出！