run azcopy only with storage_type BLOB (#17)

mara · Sep 28, 2023 · 4401327 · 4401327
1 parent b56669b
commit 4401327
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 8 deletions.
diff --git a/mara_storage/shell.py b/mara_storage/shell.py
@@ -74,7 +74,7 @@ def __(storage: storages.AzureStorage, file_name: str, compression: Compression
                    ) if not storage.sas else ''
 
     return (f'{azlogin_env}azcopy cp '
-            + shlex.quote(storage.build_uri(file_name))
+            + shlex.quote(storage.build_uri(file_name, storage_type='blob'))
             + ' --from-to BlobPipe'
             + (f'\\\n  | {uncompressor(compression)} - ' if compression != Compression.NONE else ''))
 
@@ -166,7 +166,7 @@ def __(storage: storages.AzureStorage, file_name: str, compression: Compression
 
     return ((f'gzip \\\n  | ' if compression == Compression.GZIP else '')
             + f'{azlogin_env}azcopy cp '
-            + shlex.quote(storage.build_uri(file_name))
+            + shlex.quote(storage.build_uri(file_name, storage_type='blob'))
             + ' --from-to PipeBlob')
 
 
@@ -251,5 +251,5 @@ def __(storage: storages.AzureStorage, file_name: str, force: bool = True, recur
                    ) if not storage.sas else ''
 
     return (f'{azlogin_env}azcopy rm '
-            + shlex.quote(storage.build_uri(file_name))
+            + shlex.quote(storage.build_uri(file_name, storage_type='blob'))
             + (' --recursive=true' if recursive else ''))
diff --git a/mara_storage/storages.py b/mara_storage/storages.py
@@ -89,6 +89,7 @@ def build_uri(self, path: str):
         """Returns a URI for a path on the storage"""
         return f"{self.base_uri}/{path}"
 
+
 class AzureStorage(Storage):
     def __init__(self, account_name: str, container_name: str, sas: str = None,
                  storage_type: str = 'blob', account_key: str = None,
@@ -123,13 +124,16 @@ def __init__(self, account_name: str, container_name: str, sas: str = None,
 
     @property
     def base_uri(self):
-        return f'https://{self.account_name}.{self.storage_type}.core.windows.net/{self.container_name}'
+        return self.build_base_uri()
 
-    def build_uri(self, path: str):
+    def build_base_uri(self, storage_type: str = None):
+        return f'https://{self.account_name}.{storage_type or self.storage_type}.core.windows.net/{self.container_name}'
+
+    def build_uri(self, path: str = None, storage_type: str = None):
         """Returns a URI for a path on the storage"""
         if path and not path.startswith('/'):
             path = '/' + path
-        return (f"{self.base_uri}{path}"
+        return (f"{self.build_base_uri(storage_type)}{path}"
                 + (f'?{self.sas}' if self.sas else ''))
 
     def connection_string(self):

diff --git a/tests/local_config.py.example b/tests/local_config.py.example
@@ -17,5 +17,6 @@ GCS_SERVICE_ACCOUNT_FILE = ''
 
 # the Azure Storage config used to perform the test
 AZ_STORAGE_ACCOUNT_NAME = ''
+AZ_STORAGE_ACCOUNT_KEY = None
 AZ_STORAGE_TYPE = 'blob'
-AZ_STORAGE_SAS = None
+AZ_STORAGE_SAS = None  # SAS Token must have access to the whole storage with permissions to create a container
diff --git a/tests/test_azcopy.py b/tests/test_azcopy.py
@@ -4,7 +4,7 @@
 from mara_storage import storages, info, shell, manage
 
 
-from .local_config import AZ_STORAGE_ACCOUNT_NAME, AZ_STORAGE_TYPE, AZ_STORAGE_SAS
+from .local_config import AZ_STORAGE_ACCOUNT_NAME, AZ_STORAGE_ACCOUNT_KEY, AZ_STORAGE_TYPE, AZ_STORAGE_SAS
 
 TEST_TOUCH_FILE_NAME = 'empty-file.txt'
 TEST_FILE_NOT_EXISTS_FILE_NAME = 'file-does-not-exist.txt'
@@ -25,6 +25,7 @@ def storage():
 
     return storages.AzureStorage(
         account_name=AZ_STORAGE_ACCOUNT_NAME,
+        account_key=AZ_STORAGE_ACCOUNT_KEY,
         container_name=container_name,
         sas=AZ_STORAGE_SAS,
         storage_type=AZ_STORAGE_TYPE)