diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index f039f83b..f63fc1c1 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -23,6 +23,7 @@ jobs: pip install -r requirements/test-requirements.txt pip install -r requirements/extras/authenticator.txt pip install -r requirements/extras/pynacl.txt + pip install -r requirements/extras/gcs.txt - name: Lint run: ./scripts/lint.sh @@ -64,6 +65,7 @@ jobs: pip install -r requirements/test-requirements.txt pip install -r requirements/extras/authenticator.txt pip install -r requirements/extras/pynacl.txt + pip install -r requirements/extras/gcs.txt - name: Test with pytest, Postgres run: ./scripts/test-postgres.sh env: @@ -93,6 +95,7 @@ jobs: pip install -r requirements/test-requirements.txt pip install -r requirements/extras/authenticator.txt pip install -r requirements/extras/pynacl.txt + pip install -r requirements/extras/gcs.txt - name: Test with pytest, SQLite run: ./scripts/test-sqlite.sh - name: Upload coverage diff --git a/CHANGES.rst b/CHANGES.rst index 61f6a485..fd080dda 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -1,6 +1,79 @@ Changes ======= +Unreleased +---------- + +Added ``GCSMediaStorage``, for storing media files in Google Cloud Storage +(``pip install 'piccolo_api[gcs]'``). + +.. note:: + Signing a URL needs a private key. On GCP compute (e.g. Cloud Run) the + Application Default Credentials don't have one, so signing goes through + the IAM API instead - grant the runtime service account + ``roles/iam.serviceAccountTokenCreator`` on itself. Note that this costs + an HTTPS request per URL, so use a service account key if you're showing + a lot of files on one page. + +Added ``CloudMediaStorage``, which is the shared base class for object storage +backends. ``S3MediaStorage`` now uses it too - a new backend just has to +implement the ``_sync`` methods, and the base class takes care of running them +in an executor. + +Storage identity (``__hash__`` / ``__eq__``) moved onto ``MediaStorage``, via a +``_hash_components`` method which each backend implements. Previously +``MediaStorage`` defined ``__eq__`` without ``__hash__``, so a custom backend +had to know to supply its own ``__hash__`` or the inherited ``__eq__`` would +fail with a confusing ``TypeError``. A custom backend which does define +``__hash__`` still works exactly as before. + +``MediaStorage`` also gained a ``_run_sync`` helper, and an ``executor`` +attribute for the backends which use it. + +Fixed several bugs in ``S3MediaStorage``: + +* The ``ContentType`` of an upload was written back to ``upload_metadata``, so + it leaked into subsequent uploads. With several uploads in flight at once it + could also be applied to the wrong file. +* ``get_file_keys()`` stripped the folder name with ``str.lstrip()``, which + removes *characters* rather than a prefix - so a file called + ``poster.jpg`` in a ``movie_posters`` folder came back as ``.jpg``. +* ``bulk_delete_files()`` sent the whole list of keys on every iteration + instead of the current batch, and the batch bounds were wrong. Deleting + more than 1000 files failed. +* A ``folder_name`` with a trailing slash stored files in one place and looked + for them in another, so ``delete_unused_files()`` never found anything. +* An uppercase file extension (``photo.JPG``) was uploaded with no + ``ContentType``, so browsers downloaded it instead of displaying it. + +``S3MediaStorage`` and ``GCSMediaStorage`` now cache their client, instead of +building one per operation. ``S3MediaStorage.get_unsigned_client`` is new. + +``LocalMediaStorage.bulk_delete_files`` and +``LocalMediaStorage.get_file_keys`` were doing their file system work directly +in the event loop, rather than in the executor like the other methods. +``LocalMediaStorage`` now follows the same convention as the other backends, +with the ``_sync`` methods doing the work and the async ones wrapping them. + +``bulk_delete_files`` now ignores files which have already gone, on every +backend - a single stale key used to abandon the rest of the batch. Any other +error is raised, including the per-key errors which ``S3MediaStorage`` used to +discard - deleting from a bucket you don't have ``s3:DeleteObject`` on +previously looked like a successful clean up. + +``folder_name`` is now tidied up when it's stored, so that a trailing slash or +a repeated slash can't send files somewhere we then fail to list. On Linux and +macOS this matches the keys files were already stored under, so nothing needs +migrating. + +.. warning:: + On **Windows** the old code built keys with a backslash + (``movie_posters\\my-file.jpeg``), and they're now built with a forward + slash. If you've been running on Windows, existing files won't be found + until you rename them in the bucket. + +------------------------------------------------------------------------------- + 1.10.0 ------ diff --git a/docs/source/api_reference/index.rst b/docs/source/api_reference/index.rst index 86a80803..c9d6a52a 100644 --- a/docs/source/api_reference/index.rst +++ b/docs/source/api_reference/index.rst @@ -43,3 +43,11 @@ Media .. currentmodule:: piccolo_api.media.s3 .. autoclass:: S3MediaStorage + +.. currentmodule:: piccolo_api.media.gcs + +.. autoclass:: GCSMediaStorage + +.. currentmodule:: piccolo_api.media.cloud + +.. autoclass:: CloudMediaStorage diff --git a/piccolo_api/media/base.py b/piccolo_api/media/base.py index 86cdfe96..77b6e03d 100644 --- a/piccolo_api/media/base.py +++ b/piccolo_api/media/base.py @@ -6,12 +6,17 @@ import pathlib import string import uuid -from collections.abc import Sequence -from typing import IO, Optional, Union +from collections.abc import Callable, Sequence +from typing import IO, TYPE_CHECKING, Optional, TypeVar, Union from piccolo.apps.user.tables import BaseUser from piccolo.columns.column_types import Array, Text, Varchar +if TYPE_CHECKING: # pragma: no cover + from concurrent.futures._base import Executor + +T = TypeVar("T") + #: Pass into ``allowed_characters`` to just allow audio files. AUDIO_EXTENSIONS = ( "mp3", @@ -76,7 +81,8 @@ class MediaStorage(metaclass=abc.ABCMeta): """ If you want to implement your own custom storage backend, create a subclass - of this class. Override each abstract method. + of this class. Override each abstract method, and + :meth:`_hash_components`. Typically, just use :class:`LocalMediaStorage ` or :class:`S3MediaStorage ` instead. @@ -232,7 +238,8 @@ async def generate_file_url( @abc.abstractmethod async def get_file(self, file_key: str) -> Optional[IO]: """ - Returns the file object matching the ``file_key``. + Returns the file object matching the ``file_key``. The caller is + responsible for closing it. """ raise NotImplementedError # pragma: no cover @@ -315,7 +322,44 @@ async def delete_unused_files( ): await self.bulk_delete_files(unused_file_keys) + ########################################################################### + + #: Subclasses which do blocking work set this, and use :meth:`_run_sync` + #: to keep that work off the event loop. If it's left as ``None``, the + #: event loop's default executor is used. + executor: Optional[Executor] = None + + async def _run_sync(self, func: Callable[[], T]) -> T: + """ + Runs a blocking function in :attr:`executor`, so it doesn't block the + event loop. + """ + loop = asyncio.get_running_loop() + return await loop.run_in_executor(self.executor, func) + + ########################################################################### + + def _hash_components(self) -> tuple: + """ + The values which identify where this storage keeps its files - for + example, the bucket and folder name. Two storages with the same + components are treated as the same storage, which is how Piccolo Admin + detects columns which would overwrite each other's files. + + Each subclass must implement this. + """ + raise NotImplementedError( + f"{type(self).__name__} must implement `_hash_components`, so we " + "can tell whether two storages save to the same place." + ) + + def __hash__(self): + return hash(self._hash_components()) + def __eq__(self, value): if not isinstance(value, MediaStorage): return False + # We compare hashes rather than the components themselves, so a + # custom backend written before `_hash_components` existed - which + # would have defined `__hash__` instead - still works. return value.__hash__() == self.__hash__() diff --git a/piccolo_api/media/cloud.py b/piccolo_api/media/cloud.py new file mode 100644 index 00000000..ecf8cbaa --- /dev/null +++ b/piccolo_api/media/cloud.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +import abc +import pathlib +from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor +from typing import IO, TYPE_CHECKING, Any, Optional, Union + +from piccolo.apps.user.tables import BaseUser +from piccolo.columns.column_types import Array, Text, Varchar + +from .base import ALLOWED_CHARACTERS, ALLOWED_EXTENSIONS, MediaStorage +from .content_type import CONTENT_TYPE + +if TYPE_CHECKING: # pragma: no cover + from concurrent.futures._base import Executor + + +class CloudMediaStorage(MediaStorage): + """ + Base class for object storage backends, such as + :class:`S3MediaStorage ` and + :class:`GCSMediaStorage `. + + The cloud SDKs are blocking, so each operation has a ``_sync`` method + which does the actual work, and an async method which runs it in an + executor to keep the event loop free. This class provides the async + methods - a subclass just has to implement the ``_sync`` ones, and set + :attr:`provider_name`. + """ + + #: Identifies the storage provider. Different providers can use the same + #: bucket name, so it's part of the storage's identity. Note that this is + #: shared by all subclasses of a backend, so someone who subclasses + #: ``S3MediaStorage`` (to customise ``get_client``, for example) still + #: compares equal to a plain ``S3MediaStorage`` for the same bucket. If a + #: new backend doesn't set it, we fall back to the class name. + provider_name: str = "" + + def __init__( + self, + column: Union[Text, Varchar, Array], + bucket_name: str, + folder_name: Optional[str] = None, + connection_kwargs: Optional[dict[str, Any]] = None, + sign_urls: bool = True, + signed_url_expiry: int = 3600, + executor: Optional[Executor] = None, + allowed_extensions: Optional[Sequence[str]] = ALLOWED_EXTENSIONS, + allowed_characters: Optional[Sequence[str]] = ALLOWED_CHARACTERS, + ): + self.bucket_name = bucket_name + self.folder_name = self._normalise_folder_name(folder_name) + self.connection_kwargs = connection_kwargs or {} + self.sign_urls = sign_urls + self.signed_url_expiry = signed_url_expiry + self.executor = executor or ThreadPoolExecutor(max_workers=10) + + super().__init__( + column=column, + allowed_extensions=allowed_extensions, + allowed_characters=allowed_characters, + ) + + ########################################################################### + + @staticmethod + def _normalise_folder_name(folder_name: Optional[str]) -> Optional[str]: + """ + Tidies up the folder name, so :meth:`folder_prefix` and + :meth:`_prepend_folder_name` can't disagree about where a file lives. + + ``PurePosixPath`` collapses a trailing slash and any repeated slashes, + giving the same answer as the ``pathlib`` join we used to build keys + with - so the keys of existing files are unchanged. + """ + if not folder_name: + return None + + normalised = str(pathlib.PurePosixPath(folder_name)) + + # `PurePosixPath('.')` is '.', but a file in '.' is just at the top + # level. + return None if normalised == "." else normalised + + @property + def folder_prefix(self) -> str: + """ + The prefix which each file key is stored under - an empty string if + no ``folder_name`` was given. + """ + folder_name = self.folder_name + + if not folder_name: + return "" + + # `folder_name` only ends with a slash if it's the bucket root. + return folder_name if folder_name.endswith("/") else f"{folder_name}/" + + def _prepend_folder_name(self, file_key: str) -> str: + return f"{self.folder_prefix}{file_key}" + + def _remove_folder_name(self, object_name: str) -> str: + """ + The inverse of :meth:`_prepend_folder_name` - turns the name of a + stored object back into a file key. + """ + prefix = self.folder_prefix + return ( + object_name[len(prefix) :] # noqa: E203 + if object_name.startswith(prefix) + else object_name + ) + + ########################################################################### + + async def store_file( + self, file_name: str, file: IO, user: Optional[BaseUser] = None + ) -> str: + return await self._run_sync( + lambda: self.store_file_sync( + file_name=file_name, file=file, user=user + ) + ) + + def store_file_sync( + self, file_name: str, file: IO, user: Optional[BaseUser] = None + ) -> str: + """ + A sync version of :meth:`store_file`. + """ + file_key = self.generate_file_key(file_name=file_name, user=user) + # `CONTENT_TYPE` is keyed by lowercase extension, but a file called + # `photo.JPG` is perfectly valid. + extension = file_key.rsplit(".", 1)[-1].lower() + + self.upload_file( + file_key=file_key, + file=file, + content_type=CONTENT_TYPE.get(extension), + ) + + return file_key + + @abc.abstractmethod + def upload_file( + self, file_key: str, file: IO, content_type: Optional[str] + ): + """ + Sends the file to the storage provider. The ``content_type`` is + ``None`` if we don't recognise the file extension. + """ + raise NotImplementedError # pragma: no cover + + async def generate_file_url( + self, file_key: str, root_url: str, user: Optional[BaseUser] = None + ) -> str: + """ + This retrieves an absolute URL for the file. + """ + return await self._run_sync( + lambda: self.generate_file_url_sync( + file_key=file_key, root_url=root_url, user=user + ) + ) + + @abc.abstractmethod + def generate_file_url_sync( + self, file_key: str, root_url: str, user: Optional[BaseUser] = None + ) -> str: + """ + A sync version of :meth:`generate_file_url`. + """ + raise NotImplementedError # pragma: no cover + + async def get_file(self, file_key: str) -> Optional[IO]: + """ + Returns the file object matching the ``file_key``. + """ + return await self._run_sync( + lambda: self.get_file_sync(file_key=file_key) + ) + + @abc.abstractmethod + def get_file_sync(self, file_key: str) -> Optional[IO]: + """ + A sync version of :meth:`get_file`. + """ + raise NotImplementedError # pragma: no cover + + async def delete_file(self, file_key: str): + """ + Deletes the file object matching the ``file_key``. + """ + return await self._run_sync( + lambda: self.delete_file_sync(file_key=file_key) + ) + + @abc.abstractmethod + def delete_file_sync(self, file_key: str): + """ + A sync version of :meth:`delete_file`. + """ + raise NotImplementedError # pragma: no cover + + async def bulk_delete_files(self, file_keys: list[str]): + await self._run_sync( + lambda: self.bulk_delete_files_sync(file_keys=file_keys) + ) + + @abc.abstractmethod + def bulk_delete_files_sync(self, file_keys: list[str]): + """ + A sync version of :meth:`bulk_delete_files`. + """ + raise NotImplementedError # pragma: no cover + + async def get_file_keys(self) -> list[str]: + """ + Returns the file key for each file we have stored. + """ + return await self._run_sync(self.get_file_keys_sync) + + @abc.abstractmethod + def get_file_keys_sync(self) -> list[str]: + """ + A sync version of :meth:`get_file_keys`. + """ + raise NotImplementedError # pragma: no cover + + ########################################################################### + + def _hash_components(self) -> tuple: + return ( + self.provider_name or type(self).__name__, + self.bucket_name, + self.folder_name, + ) diff --git a/piccolo_api/media/gcs.py b/piccolo_api/media/gcs.py new file mode 100644 index 00000000..edf7c062 --- /dev/null +++ b/piccolo_api/media/gcs.py @@ -0,0 +1,272 @@ +from __future__ import annotations + +import io +import sys +import threading +from collections.abc import Sequence +from datetime import timedelta +from typing import IO, TYPE_CHECKING, Any, Optional, Union + +from piccolo.apps.user.tables import BaseUser +from piccolo.columns.column_types import Array, Text, Varchar + +from .base import ALLOWED_CHARACTERS, ALLOWED_EXTENSIONS +from .cloud import CloudMediaStorage + +if TYPE_CHECKING: # pragma: no cover + from concurrent.futures._base import Executor + + +#: Signing goes through the IAM API, which needs this scope. The storage +#: client only asks for the ``devstorage`` ones. +IAM_SCOPE = "https://www.googleapis.com/auth/cloud-platform" + + +class GCSMediaStorage(CloudMediaStorage): + + provider_name = "gcs" + + def __init__( + self, + column: Union[Text, Varchar, Array], + bucket_name: str, + folder_name: Optional[str] = None, + connection_kwargs: Optional[dict[str, Any]] = None, + sign_urls: bool = True, + signed_url_expiry: int = 3600, + executor: Optional[Executor] = None, + allowed_extensions: Optional[Sequence[str]] = ALLOWED_EXTENSIONS, + allowed_characters: Optional[Sequence[str]] = ALLOWED_CHARACTERS, + ): + """ + Stores media files in Google Cloud Storage. This is a good option when + you have lots of files to store, and don't want them stored locally on + a server. + + :param column: + The Piccolo :class:`Column ` which the + storage is for. + :param bucket_name: + Which GCS bucket the files are stored in. + :param folder_name: + The files will be stored in this folder within the bucket. GCS + buckets don't really have folders, but if ``folder`` is + ``'movie_screenshots'``, then we store the file at + ``'movie_screenshots/my-file-abc-123.jpeg'``, to simulate it being + in a folder. + :param connection_kwargs: + These kwargs are passed directly to the + :class:`google.cloud.storage.Client`. For example:: + + GCSMediaStorage( + ..., + connection_kwargs={'project': 'my-gcp-project'} + ) + :param sign_urls: + Whether to sign the URLs - by default this is ``True``, as it's + highly recommended that you store your files in a private bucket. + :param signed_url_expiry: + Files are accessed via signed URLs, which are only valid for this + number of seconds. + :param executor: + An executor, which file save operations are run in, to avoid + blocking the event loop. If not specified, we use a sensibly + configured :class:`ThreadPoolExecutor `. + :param allowed_extensions: + Which file extensions are allowed. If ``None``, then all extensions + are allowed (not recommended unless the users are trusted). + :param allowed_characters: + Which characters are allowed in the file name. By default, it's + very strict. If set to ``None`` then all characters are allowed. + + .. note:: + Signing a URL requires a private key. Locally that comes from a + service-account JSON (``GOOGLE_APPLICATION_CREDENTIALS``). On GCP + compute (e.g. Cloud Run), Application Default Credentials from the + metadata server have no private key, so we sign via the IAM + ``signBlob`` API instead - for that to work, grant the runtime + service account ``roles/iam.serviceAccountTokenCreator`` on itself. + + Note that signing this way costs an HTTPS request per URL, so a + page showing lots of files makes lots of requests. If that's a + problem, mount a service-account key so the signing can happen + locally, or use ``sign_urls=False`` with a public bucket. + """ # noqa: E501 + + try: + from google.cloud import storage # noqa + except ImportError: # pragma: no cover + sys.exit( + "Please install google-cloud-storage to use this feature " + "`pip install 'piccolo_api[gcs]'`" + ) + else: + self.storage = storage + + self._client = None + self._signing_credentials = None + # Reentrant, so nesting a cached lookup inside another is safe. + self._client_lock = threading.RLock() + + super().__init__( + column=column, + bucket_name=bucket_name, + folder_name=folder_name, + connection_kwargs=connection_kwargs, + sign_urls=sign_urls, + signed_url_expiry=signed_url_expiry, + executor=executor, + allowed_extensions=allowed_extensions, + allowed_characters=allowed_characters, + ) + + def get_client(self): + """ + Returns a GCS client. It's cached, because creating one resolves the + credentials, which on GCP compute means a call to the metadata server + - and we'd otherwise do that for every file on an admin page. Note + that signing still costs a request per URL, see :meth:`__init__`. + """ + with self._client_lock: + if self._client is None: + self._client = self.storage.Client(**self.connection_kwargs) + return self._client + + def get_bucket(self): # pragma: no cover + return self.get_client().bucket(self.bucket_name) + + def _get_blob(self, file_key: str): + return self.get_bucket().blob(self._prepend_folder_name(file_key)) + + def upload_file( + self, file_key: str, file: IO, content_type: Optional[str] + ): + self._get_blob(file_key).upload_from_file( + file, content_type=content_type + ) + + def get_signing_credentials(self): + """ + Returns credentials which can sign a URL. + + If the client's credentials have a private key then they can sign + directly. Otherwise signing goes via the IAM API, and we have to + re-scope them first - :class:`google.cloud.storage.Client` asks for + storage scopes only, but the IAM API needs ``cloud-platform``, and + Cloud Run (unlike Compute Engine) really does hand out a token + limited to whatever was asked for. + """ + from google.auth.credentials import Scoped, Signing + + # There's no public accessor for the client's credentials. + credentials = self.get_client()._credentials + + if isinstance(credentials, Signing): + return credentials + + with self._client_lock: + if self._signing_credentials is None: + self._signing_credentials = ( + credentials.with_scopes([IAM_SCOPE]) + if isinstance(credentials, Scoped) + else credentials + ) + + return self._signing_credentials + + def get_signing_kwargs(self) -> dict[str, Any]: + """ + Signing a URL needs a private key. If the credentials don't have one + (Application Default Credentials on GCP compute don't), then we have + to sign via the IAM API instead, which needs the service account's + email address and an access token. + """ + from google.auth.credentials import Signing + from google.auth.transport.requests import Request + + credentials = self.get_signing_credentials() + + if isinstance(credentials, Signing): + # We have a private key, so we can sign locally. + return {} + + if not credentials.valid: + credentials.refresh(Request()) + + service_account_email = getattr( + credentials, "service_account_email", None + ) + + if not service_account_email: + raise ValueError( + "These credentials can't be used to sign URLs - use a " + "service account, or pass `sign_urls=False`." + ) + + return { + "service_account_email": service_account_email, + "access_token": credentials.token, + } + + def generate_file_url_sync( + self, file_key: str, root_url: str, user: Optional[BaseUser] = None + ) -> str: + blob = self._get_blob(file_key) + + if not self.sign_urls: + return blob.public_url + + return blob.generate_signed_url( + version="v4", + expiration=timedelta(seconds=self.signed_url_expiry), + method="GET", + **self.get_signing_kwargs(), + ) + + ########################################################################### + + def get_file_sync(self, file_key: str) -> Optional[IO]: + # `blob.open` would avoid loading the file into memory, but it doesn't + # touch the network until it's read - so a missing file would raise + # in the caller rather than here. The other backends raise straight + # away, so we do the same. + return io.BytesIO(self._get_blob(file_key).download_as_bytes()) + + def delete_file_sync(self, file_key: str): + return self._get_blob(file_key).delete() + + def bulk_delete_files_sync(self, file_keys: list[str]): + from google.cloud.exceptions import NotFound + + if not file_keys: + return + + bucket = self.get_bucket() + + if not bucket.exists(): + # `on_error` below can't tell the 404 for a missing bucket from + # the 404 for a missing file, so without this a sweep against the + # wrong bucket would quietly report success. The other two + # backends raise in this situation. + raise NotFound(f"The bucket {self.bucket_name} doesn't exist.") + + # `on_error` is called for each file which has already gone, so one + # stale key doesn't fail the whole sweep. Every other error is still + # raised. + # + # This deliberately doesn't use `client.batch()`, which would send + # 100 deletes per request. A batch keeps only one of its failures, + # and it's the LAST one - `exception_args = exception_args or + # subresponse` walks past earlier errors, because a 4xx response is + # falsy. So a permissions error can be reported as the `NotFound` of + # a file which had already gone, and then swallowed here. + bucket.delete_blobs( + [self._prepend_folder_name(i) for i in file_keys], + on_error=lambda blob: None, + ) + + def get_file_keys_sync(self) -> list[str]: + blobs = self.get_client().list_blobs( + self.bucket_name, prefix=self.folder_prefix or None + ) + return [self._remove_folder_name(blob.name) for blob in blobs] diff --git a/piccolo_api/media/local.py b/piccolo_api/media/local.py index df6daa9e..c7129449 100644 --- a/piccolo_api/media/local.py +++ b/piccolo_api/media/local.py @@ -1,7 +1,6 @@ from __future__ import annotations -import asyncio -import functools +import contextlib import logging import os import pathlib @@ -12,7 +11,6 @@ from piccolo.apps.user.tables import BaseUser from piccolo.columns.column_types import Array, Text, Varchar -from piccolo.utils.sync import run_sync from .base import ALLOWED_CHARACTERS, ALLOWED_EXTENSIONS, MediaStorage @@ -74,77 +72,72 @@ def __init__( async def store_file( self, file_name: str, file: IO, user: Optional[BaseUser] = None ) -> str: + return await self._run_sync( + lambda: self.store_file_sync( + file_name=file_name, file=file, user=user + ) + ) + + def store_file_sync( + self, file_name: str, file: IO, user: Optional[BaseUser] = None + ) -> str: + """ + A sync version of :meth:`store_file`. + """ # If the file_name includes the entire path (e.g. /foo/bar.jpg) - we # just want bar.jpg. file_name = pathlib.Path(file_name).name file_key = self.generate_file_key(file_name=file_name, user=user) - loop = asyncio.get_running_loop() - file_permissions = self.file_permissions - - def save(): - path = os.path.join(self.media_path, file_key) - - if os.path.exists(path): - logger.error( - "A file name clash has occurred - the chances are very " - "low. Could be malicious, or a serious bug." - ) - raise IOError("Unable to save the file") + path = os.path.join(self.media_path, file_key) - with open(path, "wb") as new_file: - shutil.copyfileobj(file, new_file) - if file_permissions is not None: - os.chmod(path, file_permissions) + if os.path.exists(path): + logger.error( + "A file name clash has occurred - the chances are very " + "low. Could be malicious, or a serious bug." + ) + raise IOError("Unable to save the file") - await loop.run_in_executor(self.executor, save) + with open(path, "wb") as new_file: + shutil.copyfileobj(file, new_file) + if self.file_permissions is not None: + os.chmod(path, self.file_permissions) return file_key - def store_file_sync( - self, file_name: str, file: IO, user: Optional[BaseUser] = None - ) -> str: - """ - A sync wrapper around :meth:`store_file`. - """ - return run_sync( - self.store_file(file_name=file_name, file=file, user=user) - ) - async def generate_file_url( self, file_key: str, root_url: str, user: Optional[BaseUser] = None ) -> str: """ This retrieves an absolute URL for the file. """ - return "/".join((root_url.rstrip("/"), file_key)) + return self.generate_file_url_sync( + file_key=file_key, root_url=root_url, user=user + ) def generate_file_url_sync( self, file_key: str, root_url: str, user: Optional[BaseUser] = None ) -> str: """ - A sync wrapper around :meth:`generate_file_url`. + A sync version of :meth:`generate_file_url`. """ - return run_sync( - self.generate_file_url( - file_key=file_key, root_url=root_url, user=user - ) - ) + return "/".join((root_url.rstrip("/"), file_key)) ########################################################################### async def get_file(self, file_key: str) -> Optional[IO]: """ - Returns the file object matching the ``file_key``. + Returns the file object matching the ``file_key``. The caller is + responsible for closing it. """ - loop = asyncio.get_running_loop() - func = functools.partial(self.get_file_sync, file_key=file_key) - return await loop.run_in_executor(self.executor, func) + return await self._run_sync( + lambda: self.get_file_sync(file_key=file_key) + ) def get_file_sync(self, file_key: str) -> Optional[IO]: """ - A sync wrapper around :meth:`get_file`. + A sync version of :meth:`get_file`. """ path = os.path.join(self.media_path, file_key) return open(path, "rb") @@ -153,37 +146,55 @@ async def delete_file(self, file_key: str): """ Deletes the file object matching the ``file_key``. """ - loop = asyncio.get_running_loop() - func = functools.partial(self.delete_file_sync, file_key=file_key) - return await loop.run_in_executor(self.executor, func) + return await self._run_sync( + lambda: self.delete_file_sync(file_key=file_key) + ) def delete_file_sync(self, file_key: str): """ - A sync wrapper around :meth:`delete_file`. + A sync version of :meth:`delete_file`. """ path = os.path.join(self.media_path, file_key) os.unlink(path) async def bulk_delete_files(self, file_keys: list[str]): + await self._run_sync( + lambda: self.bulk_delete_files_sync(file_keys=file_keys) + ) + + def bulk_delete_files_sync(self, file_keys: list[str]): + """ + A sync version of :meth:`bulk_delete_files`. + """ media_path = self.media_path + + if file_keys and not os.path.isdir(media_path): + # Otherwise every delete below is quietly skipped, and a missing + # media folder looks like a successful clean up. + raise FileNotFoundError(media_path) + for file_key in file_keys: - os.unlink(os.path.join(media_path, file_key)) + # A file which has already gone shouldn't abandon the rest of the + # batch - the other backends behave this way too. + with contextlib.suppress(FileNotFoundError): + os.unlink(os.path.join(media_path, file_key)) async def get_file_keys(self) -> list[str]: """ Returns the file key for each file we have stored. """ - file_keys = [] + return await self._run_sync(self.get_file_keys_sync) + + def get_file_keys_sync(self) -> list[str]: + """ + A sync version of :meth:`get_file_keys`. + """ + file_keys: list[str] = [] for _, _, filenames in os.walk(self.media_path): file_keys.extend(filenames) break return file_keys - def __hash__(self): - return hash(("local", self.media_path)) - - def __eq__(self, value): - if not isinstance(value, LocalMediaStorage): - return False - return value.__hash__() == self.__hash__() + def _hash_components(self) -> tuple: + return ("local", self.media_path) diff --git a/piccolo_api/media/s3.py b/piccolo_api/media/s3.py index add14e2b..9468a0a9 100644 --- a/piccolo_api/media/s3.py +++ b/piccolo_api/media/s3.py @@ -1,24 +1,24 @@ from __future__ import annotations -import asyncio -import functools -import pathlib import sys -from collections.abc import Callable, Sequence -from concurrent.futures import ThreadPoolExecutor +import threading +from collections.abc import Sequence from typing import IO, TYPE_CHECKING, Any, Optional, Union from piccolo.apps.user.tables import BaseUser from piccolo.columns.column_types import Array, Text, Varchar -from .base import ALLOWED_CHARACTERS, ALLOWED_EXTENSIONS, MediaStorage -from .content_type import CONTENT_TYPE +from .base import ALLOWED_CHARACTERS, ALLOWED_EXTENSIONS +from .cloud import CloudMediaStorage if TYPE_CHECKING: # pragma: no cover from concurrent.futures._base import Executor -class S3MediaStorage(MediaStorage): +class S3MediaStorage(CloudMediaStorage): + + provider_name = "s3" + def __init__( self, column: Union[Text, Varchar, Array], @@ -130,104 +130,85 @@ def __init__( else: self.boto3 = boto3 - self.bucket_name = bucket_name self.upload_metadata = upload_metadata or {} - self.folder_name = folder_name - self.connection_kwargs = connection_kwargs or {} - self.sign_urls = sign_urls - self.signed_url_expiry = signed_url_expiry - self.executor = executor or ThreadPoolExecutor(max_workers=10) + self._client = None + self._unsigned_client = None + # Reentrant, because `get_unsigned_client` calls `get_client`. + self._client_lock = threading.RLock() super().__init__( column=column, + bucket_name=bucket_name, + folder_name=folder_name, + connection_kwargs=connection_kwargs, + sign_urls=sign_urls, + signed_url_expiry=signed_url_expiry, + executor=executor, allowed_extensions=allowed_extensions, allowed_characters=allowed_characters, ) - def get_client(self, config=None): # pragma: no cover + def get_client(self, config=None): """ Returns an S3 client. + + The default client is cached, because building one creates a boto3 + session, which parses botocore's service and endpoint data - we'd + otherwise do that for every file shown on a page. A client with a + custom ``config`` isn't cached, as we don't know what's in it. """ - session = self.boto3.session.Session() - extra_kwargs = {"config": config} if config else {} - client = session.client("s3", **self.connection_kwargs, **extra_kwargs) - return client + with self._client_lock: + if config is None and self._client is not None: + return self._client + + session = self.boto3.session.Session() + extra_kwargs = {"config": config} if config else {} + client = session.client( + "s3", **self.connection_kwargs, **extra_kwargs + ) - async def store_file( - self, file_name: str, file: IO, user: Optional[BaseUser] = None - ) -> str: - loop = asyncio.get_running_loop() + if config is None: + self._client = client - blocking_function = functools.partial( - self.store_file_sync, file_name=file_name, file=file, user=user - ) + return client - file_key = await loop.run_in_executor(self.executor, blocking_function) + def get_unsigned_client(self): + """ + Returns a client which generates unsigned URLs. Cached, for the same + reason as :meth:`get_client`. + """ + with self._client_lock: + if self._unsigned_client is None: + from botocore import UNSIGNED + from botocore.config import Config - return file_key + self._unsigned_client = self.get_client( + config=Config(signature_version=UNSIGNED) + ) - def _prepend_folder_name(self, file_key: str) -> str: - folder_name = self.folder_name - if folder_name: - return str(pathlib.Path(folder_name, file_key)) - else: - return file_key + return self._unsigned_client - def store_file_sync( - self, file_name: str, file: IO, user: Optional[BaseUser] = None - ) -> str: - """ - A sync wrapper around :meth:`store_file`. - """ - file_key = self.generate_file_key(file_name=file_name, user=user) - extension = file_key.rsplit(".", 1)[-1] - client = self.get_client() - upload_metadata: dict[str, Any] = self.upload_metadata + def upload_file( + self, file_key: str, file: IO, content_type: Optional[str] + ): + upload_metadata: dict[str, Any] = {**self.upload_metadata} - if extension in CONTENT_TYPE: - upload_metadata["ContentType"] = CONTENT_TYPE[extension] + if content_type: + upload_metadata["ContentType"] = content_type - client.upload_fileobj( + self.get_client().upload_fileobj( file, self.bucket_name, self._prepend_folder_name(file_key), ExtraArgs=upload_metadata, ) - return file_key - - async def generate_file_url( - self, file_key: str, root_url: str, user: Optional[BaseUser] = None - ) -> str: - """ - This retrieves an absolute URL for the file. - """ - loop = asyncio.get_running_loop() - - blocking_function: Callable = functools.partial( - self.generate_file_url_sync, - file_key=file_key, - root_url=root_url, - user=user, - ) - - return await loop.run_in_executor(self.executor, blocking_function) - def generate_file_url_sync( self, file_key: str, root_url: str, user: Optional[BaseUser] = None ) -> str: - """ - A sync wrapper around :meth:`generate_file_url`. - """ - if self.sign_urls: - config = None - else: - from botocore import UNSIGNED - from botocore.config import Config - - config = Config(signature_version=UNSIGNED) - - s3_client = self.get_client(config=config) + s3_client = ( + self.get_client() if self.sign_urls else self.get_unsigned_client() + ) return s3_client.generate_presigned_url( ClientMethod="get_object", @@ -240,20 +221,7 @@ def generate_file_url_sync( ########################################################################### - async def get_file(self, file_key: str) -> Optional[IO]: - """ - Returns the file object matching the ``file_key``. - """ - loop = asyncio.get_running_loop() - - func = functools.partial(self.get_file_sync, file_key=file_key) - - return await loop.run_in_executor(self.executor, func) - def get_file_sync(self, file_key: str) -> Optional[IO]: - """ - Returns the file object matching the ``file_key``. - """ s3_client = self.get_client() response = s3_client.get_object( Bucket=self.bucket_name, @@ -261,71 +229,44 @@ def get_file_sync(self, file_key: str) -> Optional[IO]: ) return response["Body"] - async def delete_file(self, file_key: str): - """ - Deletes the file object matching the ``file_key``. - """ - loop = asyncio.get_running_loop() - - func = functools.partial( - self.delete_file_sync, - file_key=file_key, - ) - - return await loop.run_in_executor(self.executor, func) - def delete_file_sync(self, file_key: str): - """ - Deletes the file object matching the ``file_key``. - """ s3_client = self.get_client() return s3_client.delete_object( Bucket=self.bucket_name, Key=self._prepend_folder_name(file_key), ) - async def bulk_delete_files(self, file_keys: list[str]): - loop = asyncio.get_running_loop() - func = functools.partial( - self.bulk_delete_files_sync, - file_keys=file_keys, - ) - await loop.run_in_executor(self.executor, func) - def bulk_delete_files_sync(self, file_keys: list[str]): s3_client = self.get_client() + # `delete_objects` rejects requests with more than 1000 keys, so we + # stay comfortably below that. batch_size = 100 - iteration = 0 - while True: - batch = file_keys[ - (iteration * batch_size) : ( # noqa: E203 - iteration + 1 * batch_size - ) - ] - if not batch: - # https://github.com/nedbat/coveragepy/issues/772 - break # pragma: no cover + for start in range(0, len(file_keys), batch_size): + batch = file_keys[start : start + batch_size] # noqa: E203 - s3_client.delete_objects( + response = s3_client.delete_objects( Bucket=self.bucket_name, Delete={ "Objects": [ { "Key": self._prepend_folder_name(file_key), } - for file_key in file_keys + for file_key in batch ], }, ) - iteration += 1 + # S3 reports a key it couldn't delete in the response body rather + # than by raising, so without this a sweep which deleted nothing + # (e.g. no `s3:DeleteObject` permission) looks like it worked. + # A key which was already gone isn't reported as an error. + errors = response.get("Errors") + if errors: + raise OSError(f"Unable to delete some files: {errors}") def get_file_keys_sync(self) -> list[str]: - """ - Returns the file key for each file we have stored. - """ s3_client = self.get_client() keys = [] @@ -337,8 +278,8 @@ def get_file_keys_sync(self) -> list[str]: if start_after: extra_kwargs["StartAfter"] = start_after - if self.folder_name: - extra_kwargs["Prefix"] = f"{self.folder_name}/" + if self.folder_prefix: + extra_kwargs["Prefix"] = self.folder_prefix response = s3_client.list_objects_v2( Bucket=self.bucket_name, @@ -356,32 +297,12 @@ def get_file_keys_sync(self) -> list[str]: # https://github.com/nedbat/coveragepy/issues/772 break # pragma: no cover - if self.folder_name: - prefix = f"{self.folder_name}/" - return [i.lstrip(prefix) for i in keys] - else: - return keys + return [self._remove_folder_name(i) for i in keys] - async def get_file_keys(self) -> list[str]: - """ - Returns the file key for each file we have stored. - """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - self.executor, self.get_file_keys_sync - ) + ########################################################################### - def __hash__(self): - return hash( - ( - "s3", - self.connection_kwargs.get("endpoint_url"), - self.bucket_name, - self.folder_name, - ) + def _hash_components(self) -> tuple: + return ( + *super()._hash_components(), + self.connection_kwargs.get("endpoint_url"), ) - - def __eq__(self, value): - if not isinstance(value, S3MediaStorage): - return False - return value.__hash__() == self.__hash__() diff --git a/pyproject.toml b/pyproject.toml index 4b33f556..0996b636 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,11 +17,20 @@ module = [ "moto", "botocore", "botocore.config", + "google.cloud", "httpx", "qrcode" ] ignore_missing_imports = true +[[tool.mypy.overrides]] +# `google` is a namespace package, so mypy resolves part of it and then +# complains about the rest. Skipping is scoped to `google` alone, so it +# doesn't quietly reduce checking elsewhere. +module = ["google.*"] +ignore_missing_imports = true +follow_imports = "skip" + [tool.coverage.run] omit = [ "*/piccolo_app.py", diff --git a/requirements/extras/gcs.txt b/requirements/extras/gcs.txt new file mode 100644 index 00000000..37e980c8 --- /dev/null +++ b/requirements/extras/gcs.txt @@ -0,0 +1 @@ +google-cloud-storage>=2.0.0 diff --git a/setup.py b/setup.py index 5528c591..371e94e0 100644 --- a/setup.py +++ b/setup.py @@ -23,6 +23,7 @@ EXTRAS = [ "authenticator", "cryptography", + "gcs", "pynacl", "s3", ] diff --git a/tests/media/test_base.py b/tests/media/test_base.py index 89c920a9..097176d1 100644 --- a/tests/media/test_base.py +++ b/tests/media/test_base.py @@ -9,6 +9,8 @@ from piccolo.columns.column_types import Array, Integer, Varchar from piccolo.table import Table +from piccolo_api.media.base import MediaStorage +from piccolo_api.media.gcs import GCSMediaStorage from piccolo_api.media.local import LocalMediaStorage from piccolo_api.media.s3 import S3MediaStorage @@ -294,9 +296,43 @@ def test_s3_media(self): ), ) + def test_gcs_media(self): + """ + Test comparing ``GCSMediaStorage``. + """ + # These should be equal, as the folder name and bucket name are the + # same. + self.assertEqual( + GCSMediaStorage( + column=Movie.poster, + bucket_name="bucker123", + folder_name="folder123", + ), + GCSMediaStorage( + column=Movie.screenshots, + bucket_name="bucker123", + folder_name="folder123", + ), + ) + + # These shouldn't be equal, as the folder names are different. + self.assertNotEqual( + GCSMediaStorage( + column=Movie.poster, + bucket_name="bucker123", + folder_name="folder123", + ), + GCSMediaStorage( + column=Movie.screenshots, + bucket_name="bucker123", + folder_name="folder456", + ), + ) + def test_mix(self): """ - Test comparing a mix of ``LocalMediaStorage`` and ``S3MediaStorage``. + Test comparing a mix of ``LocalMediaStorage``, ``S3MediaStorage`` and + ``GCSMediaStorage``. """ self.assertNotEqual( LocalMediaStorage(column=Movie.poster, media_path="/tmp/"), @@ -307,6 +343,90 @@ def test_mix(self): ), ) + # A GCS bucket and an S3 bucket can share a name, but they're not the + # same bucket. + self.assertNotEqual( + S3MediaStorage( + column=Movie.poster, + bucket_name="bucker123", + folder_name="folder123", + ), + GCSMediaStorage( + column=Movie.screenshots, + bucket_name="bucker123", + folder_name="folder123", + ), + ) + + def test_subclass_of_a_backend(self): + """ + Subclassing a backend (to customise ``get_client``, say) mustn't + change where it points - otherwise Piccolo Admin stops spotting two + columns which save to the same place. + """ + + class CustomS3MediaStorage(S3MediaStorage): + pass + + self.assertEqual( + S3MediaStorage( + column=Movie.poster, + bucket_name="bucket123", + folder_name="folder123", + ), + CustomS3MediaStorage( + column=Movie.screenshots, + bucket_name="bucket123", + folder_name="folder123", + ), + ) + + def test_custom_backend_with_its_own_hash(self): + """ + Before ``_hash_components`` existed, the only way to write a custom + backend was to define ``__hash__``. Those must keep working. + """ + + class CustomMediaStorage(MediaStorage): + def __init__(self, column, path): + self.path = path + super().__init__(column=column) + + def __hash__(self): + return hash(("custom", self.path)) + + async def store_file(self, file_name, file, user=None): ... + + async def generate_file_url( + self, file_key, root_url, user=None + ): ... + + async def get_file(self, file_key): ... + + async def delete_file(self, file_key): ... + + async def bulk_delete_files(self, file_keys): ... + + async def get_file_keys(self): ... + + same_a = CustomMediaStorage(column=Movie.poster, path="/a") + same_b = CustomMediaStorage(column=Movie.screenshots, path="/a") + different = CustomMediaStorage(column=Movie.poster, path="/b") + + self.assertEqual(same_a, same_b) + self.assertNotEqual(same_a, different) + # This is what Piccolo Admin does to spot a misconfiguration. + self.assertEqual(len({same_a, same_b}), 1) + + def test_non_storage(self): + """ + Comparing against something which isn't a storage shouldn't blow up. + """ + self.assertNotEqual( + LocalMediaStorage(column=Movie.poster, media_path="/tmp/"), + "not a storage", + ) + def test_sets(self): """ Make sure sets behave as expected. diff --git a/tests/media/test_gcs.py b/tests/media/test_gcs.py new file mode 100644 index 00000000..7e2757cc --- /dev/null +++ b/tests/media/test_gcs.py @@ -0,0 +1,438 @@ +import asyncio +import os +import uuid +from typing import Optional +from unittest import TestCase +from unittest.mock import MagicMock, patch + +from google.auth.credentials import Scoped, Signing +from google.cloud.exceptions import Forbidden, NotFound +from piccolo.columns.column_types import Varchar +from piccolo.table import Table + +from piccolo_api.media.gcs import IAM_SCOPE, GCSMediaStorage + + +class Movie(Table): + poster = Varchar() + + +class FakeBlob: + def __init__(self, name: str, client: "FakeClient"): + self.name = name + self.client = client + self.public_url = f"https://storage.example.com/{name}" + + @property + def content_type(self) -> Optional[str]: + return self.client.content_types.get(self.name) + + def upload_from_file(self, file, content_type=None): + self.client.content_types[self.name] = content_type + self.client.store[self.name] = file.read() + + def download_as_bytes(self) -> bytes: + if self.name not in self.client.store: + raise NotFound(self.name) + return self.client.store[self.name] + + def delete(self): + if self.name in self.client.forbidden: + raise Forbidden(self.name) + # The real client raises if the blob has already gone. + if self.name not in self.client.store: + raise NotFound(self.name) + self.client.store.pop(self.name) + self.client.content_types.pop(self.name, None) + + def generate_signed_url(self, version, expiration, method, **kwargs): + self.client.signing_kwargs.update(kwargs) + return f"https://storage.example.com/{self.name}?signature=abc123" + + +class FakeBucket: + def __init__(self, client: "FakeClient"): + self.client = client + + def blob(self, name: str) -> FakeBlob: + return FakeBlob(name=name, client=self.client) + + def exists(self) -> bool: + return self.client.bucket_exists + + def delete_blobs(self, blobs, on_error=None): + """ + The real one calls ``on_error`` for a blob which has already gone, + and propagates anything else. + """ + for name in blobs: + try: + self.blob(name).delete() + except NotFound: + if on_error is None: + raise + on_error(name) + + +class FakeCredentials(Scoped): + """ + Stands in for Application Default Credentials on GCP compute - i.e. no + private key, so signing has to go via the IAM API. Like the real ones, + these are scoped, and the storage client only asks for storage scopes. + """ + + def __init__(self, scopes=("https://www.googleapis.com/auth/devstorage",)): + self.valid = True + self.token = "token123" + self.service_account_email = "robot@example.iam.gserviceaccount.com" + self.refreshed = False + # `scopes` is a read-only property on `Scoped`, backed by this. + self._scopes = list(scopes) + self._default_scopes = None + + @property + def requires_scopes(self) -> bool: + return True + + def with_scopes(self, scopes, default_scopes=None): + copy = type(self)(scopes=scopes) + copy.service_account_email = self.service_account_email + return copy + + def refresh(self, request): + self.refreshed = True + + +class FakeUserCredentials(FakeCredentials): + """ + A personal Google account - no service account email, so it can't sign. + """ + + def __init__(self, scopes=()): + super().__init__(scopes=scopes) + self.service_account_email = None + + def with_scopes(self, scopes, default_scopes=None): + return self + + +class FakeSigningCredentials(Signing): + """ + A service account with a private key, which can sign locally. + """ + + valid = True + + def sign_bytes(self, message): # pragma: no cover + return b"signature" + + @property + def signer_email(self): # pragma: no cover + return "robot@example.iam.gserviceaccount.com" + + @property + def signer(self): # pragma: no cover + return None + + +class FakeClient: + def __init__(self, store: dict): + self.store = store + self.content_types: dict = {} + self.signing_kwargs: dict = {} + self.forbidden: set = set() + self.bucket_exists = True + self._credentials = FakeCredentials() + + def bucket(self, bucket_name: str) -> FakeBucket: + return FakeBucket(client=self) + + def list_blobs(self, bucket_name: str, prefix=None): + return [ + self.bucket(bucket_name).blob(name) + for name in list(self.store) + if prefix is None or name.startswith(prefix) + ] + + +class TestGCSMediaStorage(TestCase): + def setUp(self) -> None: + Movie.create_table(if_not_exists=True).run_sync() + + def tearDown(self): + Movie.alter().drop_table().run_sync() + + def get_storage( + self, + store: dict, + folder_name="movie_posters", + **kwargs, + ) -> GCSMediaStorage: + storage = GCSMediaStorage( + column=Movie.poster, + bucket_name="bucket123", + folder_name=folder_name, + **kwargs, + ) + storage.get_client = MagicMock( # type: ignore[method-assign] + return_value=FakeClient(store) + ) + return storage + + @patch("piccolo_api.media.base.uuid") + def test_store_and_retrieve(self, uuid_module: MagicMock): + """ + Store a file, then retrieve its bytes and a signed URL. + """ + uuid_module.uuid4.return_value = uuid.UUID( + "fd0125c7-8777-4976-83c1-81605d5ab155" + ) + store: dict = {} + storage = self.get_storage(store) + + path = os.path.join(os.path.dirname(__file__), "test_files/bulb.jpg") + + with open(path, "rb") as test_file: + file_key = asyncio.run( + storage.store_file(file_name="bulb.jpg", file=test_file) + ) + + with open(path, "rb") as test_file: + expected_bytes = test_file.read() + + self.assertEqual( + file_key, + "bulb-fd0125c7-8777-4976-83c1-81605d5ab155.jpg", + ) + + # It was stored under the folder prefix. + self.assertIn(f"movie_posters/{file_key}", store) + + # The content type was set from the file extension. + self.assertEqual( + storage.get_client().content_types[f"movie_posters/{file_key}"], + "image/jpeg", + ) + + # Compare against the source file, not against what we stored - + # otherwise storing the wrong bytes would pass. + self.assertEqual(store[f"movie_posters/{file_key}"], expected_bytes) + + file = asyncio.run(storage.get_file(file_key=file_key)) + assert file is not None + self.assertEqual(file.read(), expected_bytes) + + url = asyncio.run( + storage.generate_file_url(file_key=file_key, root_url="") + ) + self.assertIn("signature=", url) + + def test_public_url(self): + storage = self.get_storage({}, sign_urls=False) + + url = asyncio.run( + storage.generate_file_url(file_key="bulb.jpg", root_url="") + ) + self.assertEqual( + url, "https://storage.example.com/movie_posters/bulb.jpg" + ) + + def test_delete_file(self): + store = {"movie_posters/bulb.jpg": b"data"} + storage = self.get_storage(store) + + asyncio.run(storage.delete_file(file_key="bulb.jpg")) + + self.assertEqual(store, {}) + + def test_bulk_delete_files(self): + store = { + "movie_posters/a.jpg": b"a", + "movie_posters/b.jpg": b"b", + "movie_posters/c.jpg": b"c", + } + storage = self.get_storage(store) + + asyncio.run(storage.bulk_delete_files(file_keys=["a.jpg", "b.jpg"])) + + self.assertEqual(list(store.keys()), ["movie_posters/c.jpg"]) + + def test_bulk_delete_lots_of_files(self): + file_keys = [f"file_{i}.jpg" for i in range(250)] + store = {f"movie_posters/{i}": b"x" for i in file_keys} + storage = self.get_storage(store) + + asyncio.run(storage.bulk_delete_files(file_keys=file_keys)) + + self.assertEqual(store, {}) + + def test_bulk_delete_no_files(self): + storage = self.get_storage({}) + + asyncio.run(storage.bulk_delete_files(file_keys=[])) + + self.assertEqual(storage.get_client().store, {}) + + def test_bulk_delete_ignores_missing_files(self): + """ + A file which has already gone shouldn't abandon the rest of the batch. + """ + store = {"movie_posters/a.jpg": b"a", "movie_posters/c.jpg": b"c"} + storage = self.get_storage(store) + + asyncio.run( + storage.bulk_delete_files( + # `b.jpg` isn't there: + file_keys=["a.jpg", "b.jpg", "c.jpg"] + ) + ) + + self.assertEqual(store, {}) + + def test_bulk_delete_with_missing_bucket(self): + """ + A missing bucket gives a 404, just like a missing file - but it means + nothing was deleted, so it mustn't be swallowed. + """ + storage = self.get_storage({"movie_posters/a.jpg": b"a"}) + storage.get_client().bucket_exists = False + + with self.assertRaises(NotFound): + asyncio.run(storage.bulk_delete_files(file_keys=["a.jpg"])) + + def test_bulk_delete_surfaces_other_errors(self): + """ + Only a missing file is tolerated - a permissions error has to be + raised, otherwise a failed clean up looks like a successful one. + """ + store = {"movie_posters/a.jpg": b"a", "movie_posters/b.jpg": b"b"} + storage = self.get_storage(store) + storage.get_client().forbidden = {"movie_posters/b.jpg"} + + with self.assertRaises(Forbidden): + asyncio.run( + storage.bulk_delete_files(file_keys=["a.jpg", "b.jpg"]) + ) + + def test_signing_credentials_are_rescoped(self): + """ + The storage client only asks for storage scopes, but signing goes via + the IAM API, which needs ``cloud-platform``. On Cloud Run the metadata + server really does hand back a token limited to what was asked for, so + signing with the client's own credentials would be rejected. + """ + storage = self.get_storage({}) + + credentials = storage.get_signing_credentials() + + self.assertEqual(credentials.scopes, [IAM_SCOPE]) + # It's cached, rather than re-scoped for every file. + self.assertIs(credentials, storage.get_signing_credentials()) + + def test_client_is_cached(self): + """ + Building a client resolves credentials, so we only want to do it + once. + """ + from google.auth.credentials import AnonymousCredentials + + storage = GCSMediaStorage( + column=Movie.poster, + bucket_name="bucket123", + connection_kwargs={ + "project": "project123", + "credentials": AnonymousCredentials(), + }, + ) + + self.assertIs(storage.get_client(), storage.get_client()) + + def test_get_missing_file(self): + """ + Fetching a file which isn't there should raise straight away, rather + than handing back a file object which fails when it's read. + """ + storage = self.get_storage({}) + + with self.assertRaises(NotFound): + asyncio.run(storage.get_file(file_key="missing.jpg")) + + def test_signing_locally_with_a_private_key(self): + """ + Credentials which can sign (i.e. a service account JSON) need no help + from the IAM API. + """ + storage = self.get_storage({}) + storage.get_client()._credentials = FakeSigningCredentials() + + self.assertEqual(storage.get_signing_kwargs(), {}) + + def test_expired_credentials_are_refreshed(self): + storage = self.get_storage({}) + # The rescoped copy is what actually signs, so that's what gets + # refreshed - not the client's own credentials. + credentials = storage.get_signing_credentials() + credentials.valid = False + + storage.get_signing_kwargs() + + self.assertTrue(credentials.refreshed) + + def test_credentials_which_cant_sign(self): + """ + A personal Google account has no service account email, so it can't + sign at all - say so, rather than failing deep in the SDK. + """ + storage = self.get_storage({}) + storage.get_client()._credentials = FakeUserCredentials() + + with self.assertRaises(ValueError): + storage.get_signing_kwargs() + + def test_signed_url_uses_iam_api_without_a_private_key(self): + """ + Credentials with no private key (e.g. on Cloud Run) can't sign + locally, so we have to pass the service account email and an access + token, which makes ``generate_signed_url`` use the IAM API. + """ + store = {"movie_posters/bulb.jpg": b"data"} + storage = self.get_storage(store) + + asyncio.run( + storage.generate_file_url(file_key="bulb.jpg", root_url="") + ) + + self.assertEqual( + storage.get_client().signing_kwargs, + { + "service_account_email": ( + "robot@example.iam.gserviceaccount.com" + ), + "access_token": "token123", + }, + ) + + def test_get_file_keys(self): + store = { + "movie_posters/a.jpg": b"a", + "movie_posters/b.jpg": b"b", + } + storage = self.get_storage(store) + + keys = asyncio.run(storage.get_file_keys()) + + self.assertEqual(sorted(keys), ["a.jpg", "b.jpg"]) + + def test_no_folder(self): + """ + With no ``folder_name``, keys are stored and listed without a prefix. + """ + store = {"a.jpg": b"a", "b.jpg": b"b"} + storage = self.get_storage(store, folder_name=None) + + keys = asyncio.run(storage.get_file_keys()) + self.assertEqual(sorted(keys), ["a.jpg", "b.jpg"]) + + # Deleting hits the no-folder key path (no prefix prepended). + asyncio.run(storage.delete_file(file_key="a.jpg")) + self.assertEqual(sorted(store.keys()), ["b.jpg"]) diff --git a/tests/media/test_local.py b/tests/media/test_local.py index 04117343..de384728 100644 --- a/tests/media/test_local.py +++ b/tests/media/test_local.py @@ -112,3 +112,60 @@ def test_store_file(self, uuid_module: MagicMock): ) self.assertListEqual(os.listdir(media_path), ["file_3.txt"]) + + def get_storage(self) -> LocalMediaStorage: + media_path = os.path.join( + tempfile.gettempdir(), "piccolo-admin-media-async" + ) + + if os.path.exists(media_path): + shutil.rmtree(media_path) + + os.mkdir(media_path) + + return LocalMediaStorage(column=Movie.poster, media_path=media_path) + + def test_async_store_and_url(self): + """ + The async methods should work as well as the ``_sync`` ones. + """ + storage = self.get_storage() + + with open( + os.path.join(os.path.dirname(__file__), "test_files/bulb.jpg"), + "rb", + ) as test_file: + file_key = asyncio.run( + storage.store_file(file_name="bulb.jpg", file=test_file) + ) + + self.assertIn(file_key, os.listdir(storage.media_path)) + + url = asyncio.run( + storage.generate_file_url(file_key, root_url="/media/") + ) + self.assertEqual(url, f"/media/{file_key}") + + def test_bulk_delete_ignores_missing_files(self): + storage = self.get_storage() + + with open(os.path.join(storage.media_path, "a.txt"), "w") as f: + f.write("test") + + asyncio.run( + # `b.txt` isn't there: + storage.bulk_delete_files(file_keys=["a.txt", "b.txt"]) + ) + + self.assertListEqual(os.listdir(storage.media_path), []) + + def test_bulk_delete_with_missing_media_path(self): + """ + A media folder which has gone shouldn't look like a successful clean + up. + """ + storage = self.get_storage() + shutil.rmtree(storage.media_path) + + with self.assertRaises(FileNotFoundError): + asyncio.run(storage.bulk_delete_files(file_keys=["a.txt"])) diff --git a/tests/media/test_s3.py b/tests/media/test_s3.py index 2d4d3ddf..76a78c27 100644 --- a/tests/media/test_s3.py +++ b/tests/media/test_s3.py @@ -298,11 +298,250 @@ def test_no_folder(self, get_client: MagicMock, uuid_module: MagicMock): ) +class TestUploadMetadata(TestCase): + @patch("piccolo_api.media.s3.S3MediaStorage.get_client") + def test_content_type_not_persisted(self, get_client: MagicMock): + """ + The ``ContentType`` we add for a file must not be written back to + ``upload_metadata``, otherwise it leaks into the next upload. + """ + bucket_name = "bucket123" + + with mock_aws(): + s3 = boto3.resource("s3", region_name="us-east-1") + s3.create_bucket(Bucket=bucket_name) + + get_client.return_value = boto3.client( + "s3", region_name="us-east-1" + ) + + upload_metadata = {"ACL": "public-read"} + + storage = S3MediaStorage( + column=Movie.poster, + bucket_name=bucket_name, + upload_metadata=upload_metadata, + ) + + asyncio.run( + storage.store_file( + file_name="bulb.jpg", file=io.BytesIO(b"test") + ) + ) + + self.assertEqual(upload_metadata, {"ACL": "public-read"}) + + +class TestBulkDelete(TestCase): + @patch("piccolo_api.media.s3.S3MediaStorage.get_client") + def test_errors_are_raised(self, get_client: MagicMock): + """ + S3 reports a key it couldn't delete in the response body rather than + by raising, so without checking it a sweep which deleted nothing + looks like it worked. + """ + client = MagicMock() + client.delete_objects.return_value = { + "Deleted": [], + "Errors": [{"Key": "movie_posters/a.jpg", "Code": "AccessDenied"}], + } + get_client.return_value = client + + storage = S3MediaStorage(column=Movie.poster, bucket_name="bucket123") + + with self.assertRaises(OSError) as manager: + asyncio.run(storage.bulk_delete_files(file_keys=["a.jpg"])) + + self.assertIn("AccessDenied", str(manager.exception)) + + +class TestClientCaching(TestCase): + def get_storage(self) -> S3MediaStorage: + return S3MediaStorage( + column=Movie.poster, + bucket_name="bucket123", + connection_kwargs={ + "aws_access_key_id": "abc123", + "aws_secret_access_key": "xyz123", + "region_name": "us-east-1", + }, + ) + + def test_client_is_cached(self): + """ + Building a client creates a boto3 session, which isn't cheap. + """ + storage = self.get_storage() + self.assertIs(storage.get_client(), storage.get_client()) + + def test_unsigned_client_is_cached_separately(self): + storage = self.get_storage() + + unsigned = storage.get_unsigned_client() + + self.assertIs(unsigned, storage.get_unsigned_client()) + # The unsigned client mustn't be handed out as the default one. + self.assertIsNot(unsigned, storage.get_client()) + + def test_custom_config_is_not_cached(self): + """ + We don't know what's in a custom config, so it shouldn't displace the + default client. + """ + from botocore.config import Config + + storage = self.get_storage() + default = storage.get_client() + + custom = storage.get_client(config=Config(region_name="eu-west-1")) + + self.assertIsNot(custom, default) + self.assertIs(storage.get_client(), default) + + +class TestContentType(TestCase): + @patch("piccolo_api.media.s3.S3MediaStorage.get_client") + def test_uppercase_extension(self, get_client: MagicMock): + """ + ``CONTENT_TYPE`` is keyed by lowercase extension, but an uppercase + file extension is perfectly valid. + """ + bucket_name = "bucket123" + + with mock_aws(): + s3 = boto3.resource("s3", region_name="us-east-1") + s3.create_bucket(Bucket=bucket_name) + + get_client.return_value = boto3.client( + "s3", region_name="us-east-1" + ) + + storage = S3MediaStorage( + column=Movie.poster, bucket_name=bucket_name + ) + + file_key = asyncio.run( + storage.store_file( + file_name="Photo.JPG", file=io.BytesIO(b"test") + ) + ) + + response = get_client.return_value.head_object( + Bucket=bucket_name, Key=file_key + ) + self.assertEqual(response["ContentType"], "image/jpeg") + + +class TestGetFileKeys(TestCase): + @patch("piccolo_api.media.s3.S3MediaStorage.get_client") + def test_folder_name_with_a_trailing_slash(self, get_client: MagicMock): + """ + A trailing slash on ``folder_name`` mustn't stop us listing the files + we stored. + """ + bucket_name = "bucket123" + + with mock_aws(): + s3 = boto3.resource("s3", region_name="us-east-1") + s3.create_bucket(Bucket=bucket_name) + + get_client.return_value = boto3.client( + "s3", region_name="us-east-1" + ) + + storage = S3MediaStorage( + column=Movie.poster, + bucket_name=bucket_name, + folder_name="movie_posters/", + ) + + file_key = asyncio.run( + storage.store_file( + file_name="bulb.jpg", file=io.BytesIO(b"test") + ) + ) + + self.assertListEqual( + asyncio.run(storage.get_file_keys()), [file_key] + ) + + @patch("piccolo_api.media.s3.S3MediaStorage.get_client") + def test_folder_name_only_stripped_as_a_prefix( + self, get_client: MagicMock + ): + """ + Only the folder prefix should be removed from a key - not every + leading character which happens to appear in the folder name. + """ + bucket_name = "bucket123" + + with mock_aws(): + s3 = boto3.resource("s3", region_name="us-east-1") + s3.create_bucket(Bucket=bucket_name) + + get_client.return_value = boto3.client( + "s3", region_name="us-east-1" + ) + + storage = S3MediaStorage( + column=Movie.poster, + bucket_name=bucket_name, + folder_name="movie_posters", + ) + + # Every character of `poster` also appears in `movie_posters`. + file_key = asyncio.run( + storage.store_file( + file_name="poster.jpg", file=io.BytesIO(b"test") + ) + ) + + self.assertListEqual( + asyncio.run(storage.get_file_keys()), [file_key] + ) + + class TestFolderName(TestCase): """ Make sure the folder name is correctly added to the file key. """ + def test_leading_slash_is_preserved(self): + """ + A leading slash is a legal key prefix, and it's where existing users' + files already live - so normalising it away would strand them. + """ + storage = S3MediaStorage( + column=Movie.poster, + bucket_name="test_bucket", + folder_name="/test_folder", + ) + self.assertEqual( + storage._prepend_folder_name(file_key="abc123.jpeg"), + "/test_folder/abc123.jpeg", + ) + + def test_folder_name_is_tidied(self): + """ + A trailing slash, a repeated slash and a bare '.' all have to agree + with what ``folder_prefix`` looks for. + """ + for folder_name, expected in ( + ("test_folder/", "test_folder/abc123.jpeg"), + ("a//b", "a/b/abc123.jpeg"), + (".", "abc123.jpeg"), + ("/", "/abc123.jpeg"), + ): + storage = S3MediaStorage( + column=Movie.poster, + bucket_name="test_bucket", + folder_name=folder_name, + ) + key = storage._prepend_folder_name(file_key="abc123.jpeg") + self.assertEqual(key, expected) + # Whatever we stored has to be findable again. + self.assertEqual(storage._remove_folder_name(key), "abc123.jpeg") + def test_with_folder_name(self): storage = S3MediaStorage( column=Movie.poster,