Skip to content

Commit

Permalink
Add Google Cloud Storage backend (#179)
Browse files Browse the repository at this point in the history
* Add Google Cloud Storage backend

* GCS - Add docs

* GCS Backend - Add tests
  • Loading branch information
AndreaGiardini authored Jul 8, 2021
1 parent 53409db commit 8e409b4
Show file tree
Hide file tree
Showing 5 changed files with 225 additions and 2 deletions.
5 changes: 5 additions & 0 deletions cogeo_mosaic/backends/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from cogeo_mosaic.backends.base import BaseBackend
from cogeo_mosaic.backends.dynamodb import DynamoDBBackend
from cogeo_mosaic.backends.file import FileBackend
from cogeo_mosaic.backends.gs import GCSBackend
from cogeo_mosaic.backends.memory import MemoryBackend
from cogeo_mosaic.backends.s3 import S3Backend
from cogeo_mosaic.backends.sqlite import SQLiteBackend
Expand All @@ -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)
Expand Down
103 changes: 103 additions & 0 deletions cogeo_mosaic/backends/gs.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
"""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.auth.exceptions import GoogleAuthError
from google.cloud.storage import Client as gcp_session
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:
gcs_bucket = self.client.bucket(bucket)
response = gcs_bucket.blob(key).download_as_bytes()
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:
gcs_bucket = self.client.bucket(bucket)
blob = gcs_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:
gcs_bucket = self.client.bucket(bucket)
blob = gcs_bucket.blob(key)
return blob.exists()
except GoogleAuthError:
return False
19 changes: 18 additions & 1 deletion docs/advanced/backends.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ Starting in version `3.0.0`, we introduced specific `backend` to abstract mosaic

- **S3Backend** (`s3://`)

- **GCSBackend** (`gs://`)

- **DynamoDBBackend** (`dynamodb://{region}/{table_name}`). If `region` is not passed, it reads the value of the `AWS_REGION` environment variable. If that environment variable does not exist, it falls back to `us-east-1`. If you choose not to pass a `region`, you still need three `/` before the table name, like so `dynamodb:///{table_name}`.

- **SQLiteBackend** (`sqlite:///{file.db}:{mosaic_name}`)
Expand All @@ -28,7 +30,7 @@ with MemoryBackend(mosaic_def=mosaicjson) as mosaic:
img = mosaic.tile(1, 1, 1)
```

### Absctract Class
### Abstract Class

All backends are built from a `BaseBackend` which is a sub-class or `rio-tiler.io.BaseReader`.

Expand All @@ -42,6 +44,9 @@ from cogeo_mosaic.backends import MosaicBackend
with MosaicBackend("s3://mybucket/amosaic.json") as mosaic:
assert isinstance(mosaic, cogeo_mosaic.backends.s3.S3Backend)

with MosaicBackend("gs://mybucket/amosaic.json") as mosaic:
assert isinstance(mosaic, cogeo_mosaic.backends.gs.GCSBackend)

with MosaicBackend("dynamodb://us-east-1/amosaic") as mosaic:
assert isinstance(mosaic, cogeo_mosaic.backends.dynamodb.DynamoDBBackend)

Expand Down Expand Up @@ -70,6 +75,18 @@ with MosaicBackend(None, mosaic_def=mosaic) as mosaic:
assert isinstance(mosaic, cogeo_mosaic.backends.memory.MemoryBackend)
```

## GCS Backend

The GCS Backend allows read and write operations from Google Cloud Storage.

When using this backend is necessary to set the appropriate roles and IAM
permissions to let cogeo-mosaic access the files. For example:

* Read-only bucket - IAM Role `roles/storage.objectViewer`. It is possible
to restrict the read-only operation to a single bucket by using the
following condition: `resource.type == "storage.googleapis.com/Object"
&& resource.name.startsWith("projects/_/buckets/mybucket")`

## STAC Backend

The STACBackend is a **read-only** backend, meaning it can't be used to write a file. This backend will POST to the input url looking for STAC items which will then be used to create the mosaicJSON in memory.
Expand Down
1 change: 1 addition & 0 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"],
Expand Down
99 changes: 98 additions & 1 deletion tests/test_backends.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
from requests.exceptions import HTTPError, RequestException
from rio_tiler.errors import PointOutsideBounds

from cogeo_mosaic.backends import MosaicBackend
from cogeo_mosaic.backends import GCSBackend, MosaicBackend
from cogeo_mosaic.backends.dynamodb import DynamoDBBackend
from cogeo_mosaic.backends.file import FileBackend
from cogeo_mosaic.backends.memory import MemoryBackend
Expand Down Expand Up @@ -307,6 +307,103 @@ def test_s3_backend(session):
session.reset_mock()


@patch("cogeo_mosaic.backends.gs.gcp_session")
def test_gs_backend(session):
"""Test GS backend."""
with open(mosaic_gz, "rb") as f:
session.return_value.bucket.return_value.blob.return_value.download_as_bytes.return_value = (
f.read()
)

session.return_value.bucket.return_value.blob.return_value.upload_from_string.return_value = (
True
)
session.return_value.bucket.return_value.blob.return_value.exists.return_value = (
False
)

with MosaicBackend("gs://mybucket/mymosaic.json.gz") as mosaic:
assert mosaic._backend_name == "Google Cloud Storage"
assert isinstance(mosaic, GCSBackend)
assert (
mosaic.mosaicid
== "24d43802c19ef67cc498c327b62514ecf70c2bbb1bbc243dda1ee075"
)
assert mosaic.quadkey_zoom == 7
assert list(mosaic.metadata.dict(exclude_none=True).keys()) == [
"mosaicjson",
"version",
"minzoom",
"maxzoom",
"quadkey_zoom",
"bounds",
"center",
]
assert mosaic.assets_for_tile(150, 182, 9) == ["cog1.tif", "cog2.tif"]
assert mosaic.assets_for_point(-73, 45) == ["cog1.tif", "cog2.tif"]
session.return_value.bucket.assert_called_once_with("mybucket")
session.return_value.bucket.return_value.blob.assert_called_once_with(
"mymosaic.json.gz"
)

session.return_value.bucket.return_value.blob.return_value.upload_from_string.assert_not_called()
session.return_value.bucket.return_value.blob.return_value.exists.assert_not_called()
session.reset_mock()

with open(mosaic_gz, "rb") as f:
session.return_value.bucket.return_value.blob.return_value.download_as_bytes.return_value = (
f.read()
)

session.return_value.bucket.return_value.blob.return_value.upload_from_string.return_value = (
True
)
session.return_value.bucket.return_value.blob.return_value.exists.return_value = (
False
)

with MosaicBackend(
"gs://mybucket/mymosaic.json.gz", mosaic_def=mosaic_content
) as mosaic:
assert isinstance(mosaic, GCSBackend)
mosaic.write()
session.return_value.bucket.return_value.blob.return_value.download_as_bytes.assert_not_called()

session.return_value.bucket.return_value.blob.return_value.exists.assert_called_once()
session.reset_mock()

session.return_value.bucket.return_value.blob.return_value.exists.return_value = (
False
)
with MosaicBackend("gs://mybucket/00000.json", mosaic_def=mosaic_content) as mosaic:
assert isinstance(mosaic, GCSBackend)
mosaic.write()
session.return_value.bucket.return_value.blob.return_value.download_as_bytes.assert_not_called()
session.return_value.bucket.return_value.blob.return_value.exists.assert_called_once()
session.reset_mock()

session.return_value.bucket.return_value.blob.return_value.exists.return_value = (
True
)
with MosaicBackend("gs://mybucket/00000.json", mosaic_def=mosaic_content) as mosaic:
assert isinstance(mosaic, GCSBackend)
with pytest.raises(MosaicExistsError):
mosaic.write()
session.return_value.client.return_value.get_object.assert_not_called()
session.return_value.bucket.return_value.blob.return_value.exists.assert_called_once()
session.reset_mock()

session.return_value.bucket.return_value.blob.return_value.exists.return_value = (
True
)
with MosaicBackend("gs://mybucket/00000.json", mosaic_def=mosaic_content) as mosaic:
assert isinstance(mosaic, GCSBackend)
mosaic.write(overwrite=True)
session.return_value.client.return_value.get_object.assert_not_called()
session.return_value.client.return_value.head_object.assert_not_called()
session.reset_mock()


class MockMeta(object):
"""Mock Meta."""

Expand Down

0 comments on commit 8e409b4

Please sign in to comment.