wehub-resource-sync
b7f52be4c9
CI / Run CI (push) Has been cancelled
CI / check-backend (push) Has been cancelled
CI / check-frontend (push) Has been cancelled
CI / tests (push) Has been cancelled
CI / e2e-tests (push) Has been cancelled
Copilot Setup Steps / copilot-setup-steps (push) Has been cancelled
89 lines
2.7 KiB
Python
89 lines
2.7 KiB
Python
from typing import TYPE_CHECKING, Any, Dict, Optional, Union
|
|
|
|
from azure.storage.filedatalake import (
|
|
ContentSettings,
|
|
DataLakeFileClient,
|
|
DataLakeServiceClient,
|
|
FileSystemClient,
|
|
)
|
|
|
|
from chainlit.data.storage_clients.base import BaseStorageClient
|
|
from chainlit.logger import logger
|
|
|
|
if TYPE_CHECKING:
|
|
from azure.core.credentials import (
|
|
AzureNamedKeyCredential,
|
|
AzureSasCredential,
|
|
TokenCredential,
|
|
)
|
|
|
|
|
|
class AzureStorageClient(BaseStorageClient):
|
|
"""
|
|
Class to enable Azure Data Lake Storage (ADLS) Gen2
|
|
|
|
parms:
|
|
account_url: "https://<your_account>.dfs.core.windows.net"
|
|
credential: Access credential (AzureKeyCredential)
|
|
sas_token: Optionally include SAS token to append to urls
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
account_url: str,
|
|
container: str,
|
|
credential: Optional[
|
|
Union[
|
|
str,
|
|
Dict[str, str],
|
|
"AzureNamedKeyCredential",
|
|
"AzureSasCredential",
|
|
"TokenCredential",
|
|
]
|
|
],
|
|
sas_token: Optional[str] = None,
|
|
):
|
|
try:
|
|
self.data_lake_client = DataLakeServiceClient(
|
|
account_url=account_url, credential=credential
|
|
)
|
|
self.container_client: FileSystemClient = (
|
|
self.data_lake_client.get_file_system_client(file_system=container)
|
|
)
|
|
self.sas_token = sas_token
|
|
logger.info("AzureStorageClient initialized")
|
|
except Exception as e:
|
|
logger.warning(f"AzureStorageClient initialization error: {e}")
|
|
|
|
async def upload_file(
|
|
self,
|
|
object_key: str,
|
|
data: Union[bytes, str],
|
|
mime: str = "application/octet-stream",
|
|
overwrite: bool = True,
|
|
content_disposition: str | None = None,
|
|
) -> Dict[str, Any]:
|
|
try:
|
|
file_client: DataLakeFileClient = self.container_client.get_file_client(
|
|
object_key
|
|
)
|
|
content_settings = ContentSettings(
|
|
content_type=mime, content_disposition=content_disposition
|
|
)
|
|
file_client.upload_data(
|
|
data, overwrite=overwrite, content_settings=content_settings
|
|
)
|
|
url = (
|
|
f"{file_client.url}{self.sas_token}"
|
|
if self.sas_token
|
|
else file_client.url
|
|
)
|
|
return {"object_key": object_key, "url": url}
|
|
except Exception as e:
|
|
logger.warning(f"AzureStorageClient, upload_file error: {e}")
|
|
return {}
|
|
|
|
async def close(self) -> None:
|
|
self.container_client.close()
|
|
self.data_lake_client.close()
|