Allow Azure storage to be used with an SAS URL (#28)

Co-authored-by: Nick Zaccardi <nick.zaccardi@level12.io>
Co-authored-by: Alex Lee <alex.lee@level12.io>

Author: 3 people

keg_storage/backends/azure.py

@@ -3,13 +3,15 @@
 import typing
 import urllib.parse
 from datetime import datetime
+from typing import ClassVar, List, Optional
 
 import arrow
 from azure.storage.blob import (
     BlobServiceClient,
     BlobBlock,
     ContainerClient,
     generate_blob_sas,
+    generate_container_sas,
 )
 from azure.storage.blob._models import BlobPrefix
 
@@ -47,19 +49,26 @@ class AzureWriter(AzureFile):
         2. There is no separate call to instantiate the upload. The first call to put_block will
            create the blob.
     """
-    def __init__(self, path: str, mode: base.FileMode, container_client: ContainerClient,
-                 chunk_size: int = None):
-        max_block_size = 100 * 1024 * 1024
+
+    max_block_size: ClassVar[int] = 100 * 1024 * 1024
+
+    def __init__(
+        self,
+        path: str,
+        mode: base.FileMode,
+        container_client: ContainerClient,
+        chunk_size: int = max_block_size,
+    ):
         if chunk_size is not None:
             # chunk_size cannot be larger than max_block_size due to API restrictions
-            chunk_size = min(chunk_size, max_block_size)
+            chunk_size = min(chunk_size, self.max_block_size)
         super().__init__(
             path=path,
             mode=mode,
             container_client=container_client,
             chunk_size=chunk_size,
         )
-        self.blocks = []
+        self.blocks: List[BlobBlock] = []
 
     def _gen_block_id(self) -> str:
         """
@@ -151,21 +160,43 @@ def read(self, size: int) -> bytes:
 
 
 class AzureStorage(base.StorageBackend):
-    def __init__(self, account: str, key: str, bucket: str, name: str = 'azure'):
-        super().__init__()
-        self.name = name
+    account_url: Optional[str]
+
+    def __init__(
+        self,
+        account: Optional[str] = None,
+        key: Optional[str] = None,
+        bucket: Optional[str] = None,
+        sas_container_url: Optional[str] = None,
+        name: str = 'azure',
+    ):
+        super().__init__(name)
         self.account = account
         self.key = key
         self.bucket = bucket
-        self.account_url = 'https://{}.blob.core.windows.net'.format(self.account)
+
+        if account and key and bucket:
+            self.account_url = 'https://{}.blob.core.windows.net'.format(self.account)
+            self.container_url = None
+        elif sas_container_url:
+            self.account_url = None
+            self.container_url = sas_container_url
+        else:
+            raise ValueError('Must provide either sas_container_url or account, key and bucket')
 
     def _create_service_client(self):
+        if self.account_url is None:
+            raise ValueError('Unable to construct a service client from a container SAS URL')
         return BlobServiceClient(
             account_url=self.account_url,
             credential=self.key,
         )
 
     def _create_container_client(self):
+        if self.container_url:
+            # Constructed using an SAS URL
+            return ContainerClient.from_container_url(self.container_url)
+
         service_client = self._create_service_client()
         return service_client.get_container_client(self.bucket)
 
@@ -256,6 +287,8 @@ def create_download_url(self, path: str, expire: typing.Union[arrow.Arrow, datet
     def _create_sas_url(self, path: str, sas_permissions: str,
                         expire: typing.Union[arrow.Arrow, datetime],
                         ip: typing.Optional[str] = None):
+        if not self.account_url:
+            raise ValueError('Cannot create a SAS URL without account credentials')
         path = self._clean_path(path)
         expire = expire.datetime if isinstance(expire, arrow.Arrow) else expire
         token = generate_blob_sas(
@@ -269,3 +302,19 @@ def _create_sas_url(self, path: str, sas_permissions: str,
         )
         url = urllib.parse.urljoin(self.account_url, '{}/{}'.format(self.bucket, path))
         return '{}?{}'.format(url, token)
+
+    def create_container_url(self, expire: typing.Union[arrow.Arrow, datetime],
+                             ip: typing.Optional[str] = None):
+        if not self.account_url:
+            raise ValueError('Cannot create a SAS URL without account credentials')
+        expire = expire.datetime if isinstance(expire, arrow.Arrow) else expire
+        token = generate_container_sas(
+            account_name=self.account,
+            container_name=self.bucket,
+            account_key=self.key,
+            permission='rwdl',
+            expiry=expire,
+            ip=ip
+        )
+        url = urllib.parse.urljoin(self.account_url, self.bucket)
+        return '{}?{}'.format(url, token)

keg_storage/backends/base.py

@@ -23,7 +23,7 @@ def __str__(self):
         return f'{s}b'
 
     @classmethod
-    def as_mode(cls, obj):
+    def as_mode(cls, obj) -> "FileMode":
         if isinstance(obj, cls):
             return obj
         if not isinstance(obj, str):
@@ -103,10 +103,8 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
 
 class StorageBackend:
-    name = None
-
-    def __init__(self, *args, **kwargs):
-        pass
+    def __init__(self, name: str):
+        self.name = name
 
     def list(self, path: str) -> typing.List[ListEntry]:
         """

keg_storage/backends/s3.py

@@ -181,8 +181,7 @@ def __init__(
             aws_profile=None,
             name='s3'
     ):
-        super().__init__()
-        self.name = name
+        super().__init__(name)
         self.bucket = bucket
         self.session = boto3.session.Session(
             aws_access_key_id=aws_access_key_id,

keg_storage/backends/sftp.py

@@ -54,14 +54,13 @@ def __init__(
             look_for_keys=False,
             name='sftp'
     ):
-        super().__init__()
+        super().__init__(name)
         self.host = host
         self.username = username
         self.key_filename = key_filename
         self.known_hosts_fpath = known_hosts_fpath
         self.allow_agent = allow_agent
         self.look_for_keys = look_for_keys
-        self.name = name
 
     def create_client(self):
         client = SSHClient()

keg_storage/tests/test_backend_azure.py

@@ -1,5 +1,6 @@
 import base64
 import datetime
+import re
 import string
 import urllib.parse as urlparse
 from io import BytesIO
@@ -62,17 +63,19 @@ def test_list_leading_slashes(self, m_client):
 
     def test_open_read_write(self, m_client):
         storage = self.create_storage()
-        with pytest.raises(NotImplementedError) as exc:
-            storage.open('foo', base.FileMode.read | base.FileMode.write)
-        assert str(exc.value) == 'Read+write mode not supported by the Azure backend'
+        with pytest.raises(
+            NotImplementedError,
+            match=re.escape("Read+write mode not supported by the Azure backend"),
+        ):
+            storage.open("foo", base.FileMode.read | base.FileMode.write)
 
     def test_open_bad_mode(self, m_client):
         storage = self.create_storage()
-        with pytest.raises(ValueError) as exc:
-            storage.open('foo', base.FileMode(0))
-        assert (
-            str(exc.value) == 'Unsupported mode. Accepted modes are FileMode.read or FileMode.write'
-        )
+        with pytest.raises(
+            ValueError,
+            match=re.escape("Unsupported mode. Accepted modes are FileMode.read or FileMode.write"),
+        ):
+            storage.open("foo", base.FileMode(0))
 
     def test_read_operations(self, m_client):
         storage = self.create_storage()
@@ -273,3 +276,68 @@ def test_download_url(self, m_client, expire):
         assert qs['sp'] == ['r']
         assert qs['sig']
         assert qs['sip'] == ['127.0.0.1']
+
+    def test_construct_from_sas_url(self, m_client):
+        storage = backends.AzureStorage(
+            **{"sas_container_url": "https://foo.blob.core.windows.net/test?sp=rwdl"}
+        )
+        assert storage.account_url is None
+        assert storage.container_url == "https://foo.blob.core.windows.net/test?sp=rwdl"
+
+        with pytest.raises(
+            ValueError, match="Unable to construct a service client from a container SAS URL"
+        ):
+            storage._create_service_client()
+
+    def test_construct_neither_sas_url_nor_account_info(self, m_client):
+        with pytest.raises(
+            ValueError, match="Must provide either sas_container_url or account, key and bucket"
+        ):
+            backends.AzureStorage(
+                account="foo", bucket="test",
+            )
+
+    def test_sas_create_container_url(self, m_client):
+        storage = backends.AzureStorage(
+            **{"sas_container_url": "https://foo.blob.core.windows.net/test?sp=rwdl"}
+        )
+
+        with pytest.raises(ValueError, match="Cannot create a SAS URL without account credentials"):
+            storage.create_container_url(arrow.get(2019, 1, 2, 3, 4, 5))
+
+    def test_sas_create_upload_url(self, m_client):
+        storage = backends.AzureStorage(
+            **{"sas_container_url": "https://foo.blob.core.windows.net/test?sp=rwdl"}
+        )
+
+        with pytest.raises(ValueError, match="Cannot create a SAS URL without account credentials"):
+            storage.create_upload_url("foo/bar.txt", arrow.get(2019, 1, 2, 3, 4, 5))
+
+    @pytest.mark.parametrize('expire', [
+        arrow.get(2019, 1, 2, 3, 4, 5),
+        datetime.datetime(2019, 1, 2, 3, 4, 5)
+    ])
+    def test_create_container_url(self, m_client, expire):
+        storage = self.create_storage()
+        url = storage.create_container_url(expire=expire)
+        parsed = urlparse.urlparse(url)
+        assert parsed.netloc == 'foo.blob.core.windows.net'
+        assert parsed.path == '/test'
+        qs = urlparse.parse_qs(parsed.query)
+
+        assert qs['se'] == ['2019-01-02T03:04:05Z']
+        assert qs['sp'] == ['rwdl']
+        assert qs['sig']
+        assert 'sip' not in qs
+
+        # with IP restriction
+        url = storage.create_container_url(expire=expire, ip='127.0.0.1')
+        parsed = urlparse.urlparse(url)
+        assert parsed.netloc == 'foo.blob.core.windows.net'
+        assert parsed.path == '/test'
+        qs = urlparse.parse_qs(parsed.query)
+
+        assert qs['se'] == ['2019-01-02T03:04:05Z']
+        assert qs['sp'] == ['rwdl']
+        assert qs['sig']
+        assert qs['sip'] == ['127.0.0.1']

keg_storage/tests/test_backend_s3.py

@@ -1,5 +1,6 @@
 import datetime
 import io
+import re
 from unittest import mock
 
 import arrow
@@ -166,16 +167,17 @@ def test_open_write(self, m_boto):
         assert isinstance(result, backends.s3.S3Writer)
 
     def test_open_read_write(self, m_boto):
-        s3 = backends.S3Storage('bucket', aws_region='us-east-1')
-        with pytest.raises(NotImplementedError) as exc:
-            s3.open('foo/bar', FileMode.read | FileMode.write)
-        assert str(exc.value) == 'Read+write mode not supported by the S3 backend'
-
-        with pytest.raises(ValueError) as exc:
-            s3.open('foo/bar', FileMode(0))
-        assert (
-            str(exc.value) == 'Unsupported mode. Accepted modes are FileMode.read or FileMode.write'
-        )
+        s3 = backends.S3Storage("bucket", aws_region="us-east-1")
+        with pytest.raises(
+            NotImplementedError, match=re.escape("Read+write mode not supported by the S3 backend")
+        ):
+            s3.open("foo/bar", FileMode.read | FileMode.write)
+
+        with pytest.raises(
+            ValueError,
+            match=re.escape("Unsupported mode. Accepted modes are FileMode.read or FileMode.write"),
+        ):
+            s3.open("foo/bar", FileMode(0))
 
     def test_read_operations(self, m_boto):
         s3 = backends.S3Storage('bucket', aws_region='us-east-1')
@@ -200,7 +202,7 @@ def test_read_not_found(self, m_boto):
         with pytest.raises(FileNotFoundInStorageError) as exc:
             s3.open('foo/bar', FileMode.read)
 
-        assert str(exc.value.filename) == 'foo/bar'
+        assert exc.value.filename == 'foo/bar'
         assert str(exc.value.storage_type) == 'S3Storage'
 
     def test_write_operations(self, m_boto):

keg_storage/tests/test_backend_sftp.py

@@ -94,9 +94,8 @@ def test_read_operations(self, sftp, m_sftp, m_log):
     @sftp_mocked()
     def test_read_not_permitted(self, sftp, m_sftp, m_log):
         with sftp.open('/tmp/foo.txt', FileMode.write) as file:
-            with pytest.raises(IOError) as exc:
+            with pytest.raises(IOError, match="File not opened for reading"):
                 file.read(1)
-        assert str(exc.value) == 'File not opened for reading'
 
     @sftp_mocked()
     def test_write_operations(self, sftp, m_sftp, m_log):
@@ -109,7 +108,6 @@ def test_write_operations(self, sftp, m_sftp, m_log):
 
     @sftp_mocked()
     def test_write_not_permitted(self, sftp, m_sftp, m_log):
-        with sftp.open('/tmp/foo.txt', FileMode.read) as file:
-            with pytest.raises(IOError) as exc:
-                file.write(b'')
-        assert str(exc.value) == 'File not opened for writing'
+        with sftp.open("/tmp/foo.txt", FileMode.read) as file:
+            with pytest.raises(IOError, match="File not opened for writing"):
+                file.write(b"")

keg_storage/tests/test_backends.py

@@ -1,6 +1,7 @@
 import io
 import os
 import pathlib
+import re
 
 import arrow
 import click
@@ -34,7 +35,7 @@ def close(self):
 
 class FakeBackend(keg_storage.StorageBackend):
     def __init__(self, base_dir: pathlib.Path):
-        super().__init__()
+        super().__init__("fake")
         self.base_dir = base_dir
 
     def list(self, path: str):
@@ -62,13 +63,12 @@ def delete(self, path):
 class TestStorageBackend:
 
     def test_methods_not_implemented(self):
-
-        interface = keg_storage.StorageBackend()
+        interface = keg_storage.StorageBackend("incomplete")
 
         cases = {
-            interface.list: ('path',),
-            interface.delete: ('path',),
-            interface.open: ('path', FileMode.read),
+            interface.list: ("path",),
+            interface.delete: ("path",),
+            interface.open: ("path", FileMode.read),
         }
 
         for method, args in cases.items():
@@ -259,12 +259,13 @@ def test_str(self):
     def test_as_mode(self):
         assert FileMode.as_mode(FileMode.read) == FileMode.read
 
-        assert FileMode.as_mode('r') == FileMode.read
-        assert FileMode.as_mode('w') == FileMode.write
-        assert FileMode.as_mode('rw') == FileMode.read | FileMode.write
-        assert FileMode.as_mode('wr') == FileMode.read | FileMode.write
-        assert FileMode.as_mode('rb') == FileMode.read
+        assert FileMode.as_mode("r") == FileMode.read
+        assert FileMode.as_mode("w") == FileMode.write
+        assert FileMode.as_mode("rw") == FileMode.read | FileMode.write
+        assert FileMode.as_mode("wr") == FileMode.read | FileMode.write
+        assert FileMode.as_mode("rb") == FileMode.read
 
-        with pytest.raises(ValueError) as exc:
+        with pytest.raises(
+            ValueError, match=re.escape("as_mode() accepts only FileMode or str arguments")
+        ):
             FileMode.as_mode(1)
-        assert str(exc.value) == 'as_mode() accepts only FileMode or str arguments'