From fff6772a221324dd99a66dd050ccf7f8d1a4cf0c Mon Sep 17 00:00:00 2001 From: Julius Unverfehrt Date: Fri, 13 Jan 2023 08:29:09 +0100 Subject: [PATCH] add get_size method for storage --- pyinfra/storage/adapters/azure.py | 6 ++++++ pyinfra/storage/adapters/s3.py | 12 +++++++++++- 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/pyinfra/storage/adapters/azure.py b/pyinfra/storage/adapters/azure.py index a23334c..d36d774 100644 --- a/pyinfra/storage/adapters/azure.py +++ b/pyinfra/storage/adapters/azure.py @@ -42,6 +42,12 @@ class AzureStorageAdapter(object): blob_client = container_client.get_blob_client(object_name) return blob_client.exists() + def get_size(self, bucket_name, object_name) -> int: + """Get object size in bytes.""" + container_client = self.__provide_container_client(bucket_name) + blob_client = container_client.get_blob_client(object_name) + return blob_client.get_blob_properties().size + @retry(tries=3, delay=5, jitter=(1, 3)) def get_object(self, bucket_name, object_name): logger.debug(f"Downloading '{object_name}'...") diff --git a/pyinfra/storage/adapters/s3.py b/pyinfra/storage/adapters/s3.py index df33542..d8b3a2e 100644 --- a/pyinfra/storage/adapters/s3.py +++ b/pyinfra/storage/adapters/s3.py @@ -54,9 +54,19 @@ class S3StorageAdapter(object): try: self._client.stat_object(bucket_name, object_name) return True - except Exception: + except Exception as err: + logger.warning(err) return False + def get_size(self, bucket_name, object_name) -> int: + """Get object size in bytes.""" + try: + stats = self._client.stat_object(bucket_name, object_name) + return stats.size + except Exception as err: + logger.warning(err) + return 0 + @retry(tries=3, delay=5, jitter=(1, 3)) def get_object(self, bucket_name, object_name): logger.debug(f"Downloading '{object_name}'...")