概要

Google Storageへのファイルのアップロード＆ダウンロード処理を実装する必要があったので
google-cloud-storage · PyPIの便利なラッパークラスを作りました。

Install library

% poetry add google-cloud-storage
# pip なら
# pip install google-cloud-storage

Wrapper class

import os
import mimetypes
from tempfile import SpooledTemporaryFile
from typing import Optional
from google.cloud import storage
from google.cloud.exceptions import NotFound


class GoogleCloudStorage:

    def __init__(
            self,
            project_id="your-gcp-project_id",
            bucket_name="your-bucket-name",
            location="",
            auto_create_bucket=False,
            auto_create_acl="projectPrivate",
            max_memory_size=0,
            tmp_dir="/tmp",
    ):
        self._project_id = project_id
        self._bucket_name = bucket_name
        self._location = location
        self._auto_create_bucket = auto_create_bucket
        self._auto_create_acl = auto_create_acl
        self._max_memory_size = max_memory_size
        self._tmp_dir = tmp_dir
        self._bucket = None
        self._client = None

    @property
    def client(self):
        if self._client is None:
            self._client = storage.Client(
                project=self._project_id,
            )
        return self._client

    @property
    def bucket(self):
        if self._bucket is None:
            self._bucket = self._get_or_create_bucket(self._bucket_name)
        return self._bucket

    def _get_or_create_bucket(self, name):
        try:
            return self.client.get_bucket(name)
        except NotFound:
            if self._auto_create_bucket:
                bucket = self.client.create_bucket(bucket_or_name=name, location=self._location)
                bucket.acl.save_predefined(predefined=self._auto_create_acl)
                return bucket
            raise ValueError(
                f"Bucket {name} does not exist. "
                f"Buckets can be automatically created by setting GCS_AUTO_CREATE_BUCKET to ``True``."
            )

    def exists(self, name) -> bool:
        if not name:
            try:
                _ = self.bucket
                return True
            except ValueError:
                return False
        return bool(self.bucket.get_blob(name))

    def save(self, name, content, content_type: Optional[str] = None):
        blob = self.bucket.blob(blob_name=name)
        if content_type is None:
            content_type = mimetypes.guess_type(name)[0]
        if isinstance(content, str):
            if os.path.isfile(content):
                blob.upload_from_filename(filename=content, content_type=content_type)
            else:
                blob.upload_from_string(data=content, content_type=content_type)
        else:
            blob.upload_from_file(file_obj=content, rewind=True, content_type=content_type)

    def delete(self, name):
        self.bucket.delete_blob(name)

    def open(self, name) -> SpooledTemporaryFile:
        blob = self.bucket.blob(blob_name=name)
        if not blob:
            raise IOError(f"File does not exist: {name}")
        file = SpooledTemporaryFile(max_size=self._max_memory_size, suffix=".gcs", dir=self._tmp_dir)
        blob.download_to_file(file)
        file.seek(0)
        return file

使い方

アクセス設定

• gcpのコンソールでサービスアカウントを作成しておき、ストレージのアクセス権限を付与しておきます
• サービスアカウントのキーを作成し、キーファイルをダウンロードしておきます
• 環境変数 GOOGLE_APPLICATION_CREDENTIALS にダウンロードしたキーファイルのパスを設定しておきます。

Upload

driver = GoogleCloudStorage()
filename = "path/to/gcs/test.txt"
content = "This is test."
# contentの内容で、 your-bucket-name/path/to/gcs/test.txt にテキストファイルがアップロードされます
# contentに指定するのは、ローカルのファイルオブジェクトやファイル名でもOKです
driver.save(name=filename, content=content)

ファイルの存在チェック

driver = GoogleCloudStorage()
filename = "path/to/gcs/test.txt"
file_exists: bool = driver.exists(name=filename)

Download

from tempfile import SpooledTemporaryFile

driver = GoogleCloudStorage()
filename = "path/to/gcs/test.txt"
file: SpooledTemporaryFile = driver.open(name=filename)

Delete

from tempfile import SpooledTemporaryFile

driver = GoogleCloudStorage()
filename = "path/to/gcs/test.txt"
driver.delete(name=filename)

以上です。