| from azure.storage.blob import AccountSasPermissions, BlobServiceClient, ResourceTypes, generate_account_sas | from azure.storage.blob import AccountSasPermissions, BlobServiceClient, ResourceTypes, generate_account_sas | ||||
| from flask import Flask | from flask import Flask | ||||
| from extensions.ext_redis import redis_client | |||||
| from extensions.storage.base_storage import BaseStorage | from extensions.storage.base_storage import BaseStorage | ||||
| class AzureStorage(BaseStorage): | class AzureStorage(BaseStorage): | ||||
| """Implementation for azure storage. | """Implementation for azure storage. | ||||
| """ | """ | ||||
| def __init__(self, app: Flask): | def __init__(self, app: Flask): | ||||
| super().__init__(app) | super().__init__(app) | ||||
| app_config = self.app.config | app_config = self.app.config | ||||
| self.bucket_name = app_config.get('AZURE_BLOB_CONTAINER_NAME') | self.bucket_name = app_config.get('AZURE_BLOB_CONTAINER_NAME') | ||||
| sas_token = generate_account_sas( | |||||
| account_name=app_config.get('AZURE_BLOB_ACCOUNT_NAME'), | |||||
| account_key=app_config.get('AZURE_BLOB_ACCOUNT_KEY'), | |||||
| resource_types=ResourceTypes(service=True, container=True, object=True), | |||||
| permission=AccountSasPermissions(read=True, write=True, delete=True, list=True, add=True, create=True), | |||||
| expiry=datetime.now(timezone.utc).replace(tzinfo=None) + timedelta(hours=1) | |||||
| ) | |||||
| self.client = BlobServiceClient(account_url=app_config.get('AZURE_BLOB_ACCOUNT_URL'), | |||||
| credential=sas_token) | |||||
| self.account_url = app_config.get('AZURE_BLOB_ACCOUNT_URL') | |||||
| self.account_name = app_config.get('AZURE_BLOB_ACCOUNT_NAME') | |||||
| self.account_key = app_config.get('AZURE_BLOB_ACCOUNT_KEY') | |||||
| def save(self, filename, data): | def save(self, filename, data): | ||||
| blob_container = self.client.get_container_client(container=self.bucket_name) | |||||
| client = self._sync_client() | |||||
| blob_container = client.get_container_client(container=self.bucket_name) | |||||
| blob_container.upload_blob(filename, data) | blob_container.upload_blob(filename, data) | ||||
| def load_once(self, filename: str) -> bytes: | def load_once(self, filename: str) -> bytes: | ||||
| blob = self.client.get_container_client(container=self.bucket_name) | |||||
| client = self._sync_client() | |||||
| blob = client.get_container_client(container=self.bucket_name) | |||||
| blob = blob.get_blob_client(blob=filename) | blob = blob.get_blob_client(blob=filename) | ||||
| data = blob.download_blob().readall() | data = blob.download_blob().readall() | ||||
| return data | return data | ||||
| def load_stream(self, filename: str) -> Generator: | def load_stream(self, filename: str) -> Generator: | ||||
| client = self._sync_client() | |||||
| def generate(filename: str = filename) -> Generator: | def generate(filename: str = filename) -> Generator: | ||||
| blob = self.client.get_blob_client(container=self.bucket_name, blob=filename) | |||||
| blob = client.get_blob_client(container=self.bucket_name, blob=filename) | |||||
| with closing(blob.download_blob()) as blob_stream: | with closing(blob.download_blob()) as blob_stream: | ||||
| while chunk := blob_stream.readall(4096): | while chunk := blob_stream.readall(4096): | ||||
| yield chunk | yield chunk | ||||
| return generate() | return generate() | ||||
| def download(self, filename, target_filepath): | def download(self, filename, target_filepath): | ||||
| blob = self.client.get_blob_client(container=self.bucket_name, blob=filename) | |||||
| client = self._sync_client() | |||||
| blob = client.get_blob_client(container=self.bucket_name, blob=filename) | |||||
| with open(target_filepath, "wb") as my_blob: | with open(target_filepath, "wb") as my_blob: | ||||
| blob_data = blob.download_blob() | blob_data = blob.download_blob() | ||||
| blob_data.readinto(my_blob) | blob_data.readinto(my_blob) | ||||
| def exists(self, filename): | def exists(self, filename): | ||||
| blob = self.client.get_blob_client(container=self.bucket_name, blob=filename) | |||||
| client = self._sync_client() | |||||
| blob = client.get_blob_client(container=self.bucket_name, blob=filename) | |||||
| return blob.exists() | return blob.exists() | ||||
| def delete(self, filename): | def delete(self, filename): | ||||
| blob_container = self.client.get_container_client(container=self.bucket_name) | |||||
| blob_container.delete_blob(filename) | |||||
| client = self._sync_client() | |||||
| blob_container = client.get_container_client(container=self.bucket_name) | |||||
| blob_container.delete_blob(filename) | |||||
| def _sync_client(self): | |||||
| cache_key = 'azure_blob_sas_token_{}_{}'.format(self.account_name, self.account_key) | |||||
| cache_result = redis_client.get(cache_key) | |||||
| if cache_result is not None: | |||||
| sas_token = cache_result.decode('utf-8') | |||||
| else: | |||||
| sas_token = generate_account_sas( | |||||
| account_name=self.account_name, | |||||
| account_key=self.account_key, | |||||
| resource_types=ResourceTypes(service=True, container=True, object=True), | |||||
| permission=AccountSasPermissions(read=True, write=True, delete=True, list=True, add=True, create=True), | |||||
| expiry=datetime.now(timezone.utc).replace(tzinfo=None) + timedelta(hours=1) | |||||
| ) | |||||
| redis_client.set(indexing_cache_key, sas_token, ex=3000) | |||||
| return BlobServiceClient(account_url=self.account_url, credential=sas_token) |