Add Google Cloud Storage backend

cogeo_mosaic/backends/__init__.py

@@ -8,6 +8,7 @@
 from cogeo_mosaic.backends.file import FileBackend
 from cogeo_mosaic.backends.memory import MemoryBackend
 from cogeo_mosaic.backends.s3 import S3Backend
+from cogeo_mosaic.backends.gs import GCSBackend
 from cogeo_mosaic.backends.sqlite import SQLiteBackend
 from cogeo_mosaic.backends.stac import STACBackend
 from cogeo_mosaic.backends.web import HttpBackend
@@ -29,6 +30,10 @@ def MosaicBackend(url: str, *args: Any, **kwargs: Any) -> BaseBackend:
     elif parsed.scheme == "s3":
         return S3Backend(url, *args, **kwargs)
 
+    # `gs://{bucket}/{key}`
+    elif parsed.scheme == "gs":
+        return GCSBackend(url, *args, **kwargs)
+
     # `dynamodb://{region}/{table}:{mosaic}`
     elif parsed.scheme == "dynamodb":
         return DynamoDBBackend(url, *args, **kwargs)

cogeo_mosaic/backends/gs.py

@@ -0,0 +1,102 @@
+"""cogeo-mosaic Google Cloud Storage backend."""
+
+import json
+from typing import Any
+from urllib.parse import urlparse
+
+import attr
+from cachetools import TTLCache, cached
+from cachetools.keys import hashkey
+
+from cogeo_mosaic.backends.base import BaseBackend
+from cogeo_mosaic.backends.utils import _compress_gz_json, _decompress_gz
+from cogeo_mosaic.cache import cache_config
+from cogeo_mosaic.errors import _HTTP_EXCEPTIONS, MosaicError, MosaicExistsError
+from cogeo_mosaic.mosaic import MosaicJSON
+
+try:
+    from google.cloud.storage import Client as gcp_session
+    from google.auth.exceptions import GoogleAuthError
+except ImportError:  # pragma: nocover
+    gcp_session = None  # type: ignore
+    GoogleAuthError = None  # type: ignore
+
+
+@attr.s
+class GCSBackend(BaseBackend):
+    """GCS Backend Adapter"""
+
+    client: Any = attr.ib(default=None)
+    bucket: str = attr.ib(init=False)
+    key: str = attr.ib(init=False)
+
+    _backend_name = "Google Cloud Storage"
+
+    def __attrs_post_init__(self):
+        """Post Init: parse path and create client."""
+        assert gcp_session is not None, "'google-cloud-storage' must be installed to use GCSBackend"
+
+        parsed = urlparse(self.path)
+        self.bucket = parsed.netloc
+        self.key = parsed.path.strip("/")
+        self.client = self.client or gcp_session()
+        super().__attrs_post_init__()
+
+    def write(self, overwrite: bool = False, **kwargs: Any):
+        """Write mosaicjson document to Google Cloud Storage."""
+        if not overwrite and self._head_object(self.key, self.bucket):
+            raise MosaicExistsError("Mosaic file already exist, use `overwrite=True`.")
+
+        mosaic_doc = self.mosaic_def.dict(exclude_none=True)
+        if self.key.endswith(".gz"):
+            body = _compress_gz_json(mosaic_doc)
+        else:
+            body = json.dumps(mosaic_doc).encode("utf-8")
+
+        self._put_object(self.key, self.bucket, body, **kwargs)
+
+    @cached(
+        TTLCache(maxsize=cache_config.maxsize, ttl=cache_config.ttl),
+        key=lambda self: hashkey(self.path),
+    )
+    def _read(self) -> MosaicJSON:  # type: ignore
+        """Get mosaicjson document."""
+        body = self._get_object(self.key, self.bucket)
+
+        self._file_byte_size = len(body)
+
+        if self.key.endswith(".gz"):
+            body = _decompress_gz(body)
+
+        return MosaicJSON(**json.loads(body))
+
+    def _get_object(self, key: str, bucket: str) -> bytes:
+        try:
+            bucket = self.client.bucket(bucket)
+            response = bucket.blob(key).download_as_string()
+        except GoogleAuthError as e:
+            status_code = e.response["ResponseMetadata"]["HTTPStatusCode"]
+            exc = _HTTP_EXCEPTIONS.get(status_code, MosaicError)
+            raise exc(e.response["Error"]["Message"]) from e
+
+        return response
+
+    def _put_object(self, key: str, bucket: str, body: bytes, **kwargs) -> str:
+        try:
+            bucket = self.client.bucket(bucket)
+            blob = bucket.blob(key)
+            blob.upload_from_string(body)
+        except GoogleAuthError as e:
+            status_code = e.response["ResponseMetadata"]["HTTPStatusCode"]
+            exc = _HTTP_EXCEPTIONS.get(status_code, MosaicError)
+            raise exc(e.response["Error"]["Message"]) from e
+
+        return key
+
+    def _head_object(self, key: str, bucket: str) -> bool:
+        try:
+            bucket = self.client.bucket(bucket)
+            blob = bucket.blob(key)
+            return blob.exists()
+        except GoogleAuthError:
+            return False

setup.py

@@ -21,6 +21,7 @@
 
 extra_reqs = {
     "aws": ["boto3"],
+    "gcp": ["google-cloud-storage"],
     "test": ["pytest", "pytest-cov"],
     "dev": ["pytest", "pytest-cov", "pre-commit"],
     "docs": ["mkdocs", "mkdocs-material", "pygments", "mkapi", "mkdocs-jupyter"],