122 lines
4.3 KiB
Python
122 lines
4.3 KiB
Python
"""Implements a wrapper around an Azure blob storage client that provides operations on the blob storage required by the service."""
|
|
import os
|
|
from typing import Iterable
|
|
|
|
from azure.storage.blob import BlobServiceClient, ContainerClient
|
|
|
|
from pyinfra.config import CONFIG
|
|
from pyinfra.storage.storage import StorageHandle
|
|
|
|
|
|
def get_blob_service_client(connection_string=None) -> BlobServiceClient:
|
|
"""Instantiates a minio.Minio client.
|
|
|
|
Args:
|
|
connection_string: Azure blob storage connection string.
|
|
|
|
Returns:
|
|
A minio.Minio client.
|
|
"""
|
|
connection_string = CONFIG.azure_blob_storage.connection_string if connection_string is None else connection_string
|
|
return BlobServiceClient.from_connection_string(conn_str=connection_string)
|
|
|
|
|
|
class AzureBlobStorageHandle(StorageHandle):
|
|
"""Implements a wrapper around an Azure blob storage client that provides operations on the blob storage required by
|
|
the service.
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""Initializes an AzureBlobStorageHandle"""
|
|
super().__init__()
|
|
self.client: BlobServiceClient = get_blob_service_client()
|
|
self.default_container_name = CONFIG.azure_blob_storage.container
|
|
self.backend = "azure"
|
|
|
|
def __get_container_client(self, container_name) -> ContainerClient:
|
|
|
|
if container_name is None:
|
|
container_name = self.default_container_name
|
|
|
|
container_client = self._StorageHandle__provide_container(container_name)
|
|
return container_client
|
|
|
|
def _StorageHandle__provide_container(self, container_name: str):
|
|
container_client = self.client.get_container_client(container_name)
|
|
return container_client if container_client.exists() else self.client.create_container(container_name)
|
|
|
|
def _StorageHandle__add_file(self, path, storage_path, container_name: str = None):
|
|
|
|
container_client = self.__get_container_client(container_name)
|
|
blob_client = container_client.get_blob_client(storage_path)
|
|
|
|
with open(path, "rb") as f:
|
|
blob_client.upload_blob(f, overwrite=True)
|
|
|
|
def list_files(self, container_name=None) -> Iterable[str]:
|
|
"""List all files in a container.
|
|
|
|
Args:
|
|
container_name: container to list files from.
|
|
|
|
Returns:
|
|
Iterable of filenames.
|
|
"""
|
|
return self._list_files("name", container_name=container_name)
|
|
|
|
def get_objects(self, container_name=None):
|
|
"""List all files in a container_name.
|
|
|
|
Args:
|
|
container_name: Container to list files from.
|
|
|
|
Returns:
|
|
Iterable over all objects in the container.
|
|
"""
|
|
container_client = self.__get_container_client(container_name)
|
|
blobs = container_client.list_blobs()
|
|
yield from blobs
|
|
|
|
def _StorageHandle__list_containers(self):
|
|
return self.client.list_containers()
|
|
|
|
def _StorageHandle__purge(self) -> None:
|
|
"""Deletes all files and buckets in the store."""
|
|
for container in self.client.list_containers():
|
|
self.client.delete_container(container.name)
|
|
|
|
def _StorageHandle__fget_object(self, container_name: str, object_name: str, target_path):
|
|
|
|
with open(target_path, "wb") as f:
|
|
blob_data = self.get_object(container_name, object_name)
|
|
blob_data.readinto(f)
|
|
|
|
def get_object(self, object_name: str, container_name: str = None):
|
|
|
|
if container_name is None:
|
|
container_name = self.default_container_name
|
|
|
|
container_client = self.__get_container_client(container_name)
|
|
blob_client = container_client.get_blob_client(object_name)
|
|
blob_data = blob_client.download_blob()
|
|
|
|
return blob_data
|
|
|
|
def _StorageHandle__remove_file(self, folder: str, filename: str, container_name: str = None) -> None:
|
|
"""Removes a file from the store.
|
|
|
|
Args:
|
|
folder: Folder containing file.
|
|
filename: Name of file (without folder) to remove.
|
|
container_name: container containing file.
|
|
"""
|
|
if container_name is None:
|
|
container_name = self.default_container_name
|
|
|
|
path = os.path.join(folder, filename)
|
|
|
|
container_client = self.__get_container_client(container_name)
|
|
blob_client = container_client.get_blob_client(path)
|
|
if blob_client.exists():
|
|
container_client.delete_blob(blob_client)
|