import mimetypes from typing import Optional from datetime import datetime, timedelta from azure.storage.blob import ( BlobServiceClient, generate_blob_sas, generate_container_sas, ContentSettings, UserDelegationKey, ) from azure.core.credentials import AzureNamedKeyCredential from azure.storage.queue import QueueClient, BinaryBase64EncodePolicy import chardet from .env_helper import EnvHelper from azure.identity import DefaultAzureCredential def connection_string(account_name: str, account_key: str): return f"DefaultEndpointsProtocol=https;AccountName={account_name};AccountKey={account_key};EndpointSuffix=core.windows.net" def create_queue_client(): env_helper: EnvHelper = EnvHelper() if env_helper.AZURE_AUTH_TYPE == "rbac": return QueueClient( account_url=f"https://{env_helper.AZURE_BLOB_ACCOUNT_NAME}.queue.core.windows.net/", queue_name=env_helper.DOCUMENT_PROCESSING_QUEUE_NAME, credential=DefaultAzureCredential(), message_encode_policy=BinaryBase64EncodePolicy(), ) else: return QueueClient.from_connection_string( conn_str=connection_string( env_helper.AZURE_BLOB_ACCOUNT_NAME, env_helper.AZURE_BLOB_ACCOUNT_KEY ), queue_name=env_helper.DOCUMENT_PROCESSING_QUEUE_NAME, message_encode_policy=BinaryBase64EncodePolicy(), ) class AzureBlobStorageClient: def __init__( self, account_name: Optional[str] = None, account_key: Optional[str] = None, container_name: Optional[str] = None, ): env_helper: EnvHelper = EnvHelper() self.auth_type = env_helper.AZURE_AUTH_TYPE self.account_name = account_name or env_helper.AZURE_BLOB_ACCOUNT_NAME self.container_name = container_name or env_helper.AZURE_BLOB_CONTAINER_NAME self.endpoint = env_helper.AZURE_STORAGE_ACCOUNT_ENDPOINT if self.auth_type == "rbac": self.account_key = None self.blob_service_client = BlobServiceClient( account_url=self.endpoint, credential=DefaultAzureCredential() ) self.user_delegation_key = self.request_user_delegation_key( blob_service_client=self.blob_service_client ) else: self.account_key = account_key or env_helper.AZURE_BLOB_ACCOUNT_KEY self.blob_service_client = BlobServiceClient( self.endpoint, credential=AzureNamedKeyCredential( name=self.account_name, key=self.account_key ), ) self.user_delegation_key = None def request_user_delegation_key( self, blob_service_client: BlobServiceClient ) -> UserDelegationKey: # Get a user delegation key that's valid for 1 day delegation_key_start_time = datetime.utcnow() delegation_key_expiry_time = delegation_key_start_time + timedelta(days=1) user_delegation_key = blob_service_client.get_user_delegation_key( key_start_time=delegation_key_start_time, key_expiry_time=delegation_key_expiry_time, ) return user_delegation_key def file_exists(self, file_name): blob_client = self.blob_service_client.get_blob_client( container=self.container_name, blob=file_name ) return blob_client.exists() def upload_file( self, bytes_data, file_name, content_type: Optional[str] = None, metadata: Optional[dict[str, str]] = None, ): # Create a blob client using the local file name as the name for the blob blob_client = self.blob_service_client.get_blob_client( container=self.container_name, blob=file_name ) content_settings = ContentSettings(content_type=content_type) if content_type is None: content_type = mimetypes.MimeTypes().guess_type(file_name)[0] charset = ( f"; charset={chardet.detect(bytes_data)['encoding']}" if content_type == "text/plain" else "" ) content_type = content_type if content_type is not None else "text/plain" content_settings = ContentSettings(content_type=content_type + charset) # Upload the created file blob_client.upload_blob( bytes_data, overwrite=True, content_settings=content_settings, metadata=metadata, ) # Generate a SAS URL to the blob and return it, if auth_type is rbac, account_key is None, if not, user_delegation_key is None. return ( blob_client.url + "?" + generate_blob_sas( self.account_name, self.container_name, file_name, user_delegation_key=self.user_delegation_key, account_key=self.account_key, permission="r", expiry=datetime.utcnow() + timedelta(hours=3), ) ) def download_file(self, file_name): blob_client = self.blob_service_client.get_blob_client( container=self.container_name, blob=file_name ) return blob_client.download_blob().readall() def delete_file(self, file_name): """ Deletes a file from the Azure Blob Storage container. Args: file_name (str): The name of the file to delete. Returns: None """ blob_client = self.blob_service_client.get_blob_client( container=self.container_name, blob=file_name ) if blob_client.exists(): blob_client.delete_blob() def delete_files(self, files, integrated_vectorization: bool): """ Deletes files from the Azure Blob Storage container. Args: files (list[str]): The names of the files to delete. Returns: None """ for filename, ids in files.items(): if not integrated_vectorization: filename = filename.split("/")[-1] self.delete_file(filename) def get_all_files(self): # Get all files in the container from Azure Blob Storage container_client = self.blob_service_client.get_container_client( self.container_name ) blob_list = container_client.list_blobs(include="metadata") # sas = generate_blob_sas(account_name, container_name, blob.name,account_key=account_key, permission="r", expiry=datetime.utcnow() + timedelta(hours=3)) sas = generate_container_sas( self.account_name, self.container_name, user_delegation_key=self.user_delegation_key, account_key=self.account_key, permission="r", expiry=datetime.utcnow() + timedelta(hours=3), ) files = [] converted_files = {} for blob in blob_list: if not blob.name.startswith("converted/"): files.append( { "filename": blob.name, "converted": ( blob.metadata.get("converted", "false") == "true" if blob.metadata else False ), "embeddings_added": ( blob.metadata.get("embeddings_added", "false") == "true" if blob.metadata else False ), "fullpath": f"{self.endpoint}{self.container_name}/{blob.name}?{sas}", "converted_filename": ( blob.metadata.get("converted_filename", "") if blob.metadata else "" ), "converted_path": "", } ) else: converted_files[blob.name] = ( f"{self.endpoint}{self.container_name}/{blob.name}?{sas}" ) for file in files: converted_filename = file.pop("converted_filename", "") if converted_filename in converted_files: file["converted"] = True file["converted_path"] = converted_files[converted_filename] return files def upsert_blob_metadata(self, file_name, metadata): blob_client = self.blob_service_client.get_blob_client( container=self.container_name, blob=file_name ) # Read metadata from the blob blob_metadata = blob_client.get_blob_properties().metadata # Update metadata blob_metadata.update(metadata) # Add metadata to the blob blob_client.set_blob_metadata(metadata=blob_metadata) def get_container_sas(self): # Generate a SAS URL to the container and return it return "?" + generate_container_sas( account_name=self.account_name, container_name=self.container_name, user_delegation_key=self.user_delegation_key, account_key=self.account_key, permission="r", expiry=datetime.utcnow() + timedelta(days=365 * 5), ) def get_blob_sas(self, file_name): # Generate a SAS URL to the blob and return it return ( f"{self.endpoint}{self.container_name}/{file_name}" + "?" + generate_blob_sas( account_name=self.account_name, container_name=self.container_name, blob_name=file_name, user_delegation_key=self.user_delegation_key, account_key=self.account_key, permission="r", expiry=datetime.utcnow() + timedelta(hours=1), ) )