Commit 753f5ec5 authored by Jeremie Hallal's avatar Jeremie Hallal
Browse files

update get_dask_storage_parameters

parent 2029e90c
Pipeline #46780 passed with stage
in 32 seconds
......@@ -2,7 +2,6 @@ from typing import Optional
import adlfs
import fsspec
from dask.distributed import WorkerPlugin
from osdu.core.api.storage.dask_storage_parameters import DaskStorageParameters
from osdu.core.api.storage.tenant import Tenant
......@@ -18,25 +17,16 @@ class AzureBlobFileSystemWithDefaultCredentials(adlfs.AzureBlobFileSystem):
def __init__(self, *args, **kwargs):
has_credential = (
"credential" in kwargs or "account_key" in kwargs
"credential" in kwargs or "account_key" in kwargs or "connection_string" in kwargs
)
if not has_credential:
kwargs["credential"] = AzureAioBlobStorage._get_credentials()
super().__init__(*args, **kwargs)
class AzureWorkerPlugin(WorkerPlugin):
""" worker plugin enables custom code to run at different stages of the Workers' lifecycle
At startup, we wrap the azure blob storage to add the authentication.
"""
def setup(self, worker):
fsspec.register_implementation("abfs", AzureBlobFileSystemWithDefaultCredentials)
fsspec.register_implementation("az", AzureBlobFileSystemWithDefaultCredentials)
return super().setup(worker)
fsspec.register_implementation("abfs", AzureBlobFileSystemWithDefaultCredentials)
fsspec.register_implementation("az", AzureBlobFileSystemWithDefaultCredentials)
def register_azure_fsspec():
fsspec.register_implementation("abfs", AzureBlobFileSystemWithDefaultCredentials)
fsspec.register_implementation("az", AzureBlobFileSystemWithDefaultCredentials)
async def get_dask_storage_parameters(tenant: Tenant, directory: Optional[str] = None) -> DaskStorageParameters:
......@@ -48,5 +38,4 @@ async def get_dask_storage_parameters(tenant: Tenant, directory: Optional[str] =
return DaskStorageParameters(protocol='az',
base_directory=base_directory,
storage_options=storage_options,
worker_plugin=(AzureWorkerPlugin, "AzureWorkerPlugin"))
register_fsspec_implementation=register_azure_fsspec)
......@@ -4,4 +4,3 @@ azure-keyvault
cachetools
fsspec
adlfs
dask[distributed]==2021.4.1
# osdu core lib main python
--extra-index-url \
https://community.opengroup.org/api/v4/projects/465/packages/pypi/simple/
osdu-core-lib-python==1.1.0.dev339911
osdu-core-lib-python==1.1.0.dev395936
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment