Skip to content

Commit

Permalink
Add AZURE_REQUEST_OPTIONS setting (#1179)
Browse files Browse the repository at this point in the history
This PR changes no behavior except to provide a settings "hook" that you
can use to pass-through azure cli request options. There was already one
such option being passed through, `timeout`, so I've refactored things
slighly to avoid duplication.
  • Loading branch information
sdherr committed Oct 30, 2024
1 parent f029e50 commit e2f6b83
Show file tree
Hide file tree
Showing 3 changed files with 57 additions and 6 deletions.
16 changes: 16 additions & 0 deletions docs/backends/azure.rst
Original file line number Diff line number Diff line change
Expand Up @@ -195,6 +195,22 @@ Settings
Additionally, this setting can be used to configure the client retry settings. To see how follow the
`Python retry docs <https://learn.microsoft.com/en-us/azure/storage/blobs/storage-retry-policy-python>`__.

``request_options`` or ``AZURE_REQUEST_OPTIONS``

Default: ``{}``

A dict of kwarg options to set on each request for the ``BlobServiceClient``. A partial list of options can be found
`in the client docs <https://learn.microsoft.com/en-us/python/api/overview/azure/storage-blob-readme?view=azure-python#other-client--per-operation-configuration>`__.

A no-argument callable can be used to set the value at request time. For example, if you are using django-guid
and want to pass through the request id::

from django_guid import get_guid

AZURE_REQUEST_OPTIONS = {
"client_request_id": get_guid
}

``api_version`` or ``AZURE_API_VERSION``

Default: ``None``
Expand Down
29 changes: 23 additions & 6 deletions storages/backends/azure_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ def _get_file(self):

if "r" in self._mode or "a" in self._mode:
download_stream = self._storage.client.download_blob(
self._path, timeout=self._storage.timeout
self._path, **self._storage._request_options()
)
download_stream.readinto(file)
if "r" in self._mode:
Expand Down Expand Up @@ -132,6 +132,22 @@ def __init__(self, **settings):
if not self.account_key and "AccountKey" in parsed:
self.account_key = parsed["AccountKey"]

def _request_options(self):
"""
If callables were provided in request_options, evaluate them and return
the concrete values. Include "timeout", which was a previously-supported
request option before the introduction of the request_options setting.
"""
if not self.request_options:
return {"timeout": self.timeout}
callable_allowed = ("raw_response_hook", "raw_request_hook")
options = self.request_options.copy()
options["timeout"] = self.timeout
for key, value in self.request_options.items():
if key not in callable_allowed and callable(value):
options[key] = value()
return options

def get_default_settings(self):
return {
"account_name": setting("AZURE_ACCOUNT_NAME"),
Expand All @@ -154,6 +170,7 @@ def get_default_settings(self):
"token_credential": setting("AZURE_TOKEN_CREDENTIAL"),
"api_version": setting("AZURE_API_VERSION", None),
"client_options": setting("AZURE_CLIENT_OPTIONS", {}),
"request_options": setting("AZURE_REQUEST_OPTIONS", {}),
}

def _get_service_client(self):
Expand Down Expand Up @@ -252,13 +269,13 @@ def exists(self, name):

def delete(self, name):
try:
self.client.delete_blob(self._get_valid_path(name), timeout=self.timeout)
self.client.delete_blob(self._get_valid_path(name), **self._request_options())
except ResourceNotFoundError:
pass

def size(self, name):
blob_client = self.client.get_blob_client(self._get_valid_path(name))
properties = blob_client.get_blob_properties(timeout=self.timeout)
properties = blob_client.get_blob_properties(**self._request_options())
return properties.size

def _save(self, name, content):
Expand All @@ -276,8 +293,8 @@ def _save(self, name, content):
content,
content_settings=ContentSettings(**params),
max_concurrency=self.upload_max_conn,
timeout=self.timeout,
overwrite=self.overwrite_files,
**self._request_options(),
)
return cleaned_name

Expand Down Expand Up @@ -350,7 +367,7 @@ def get_modified_time(self, name):
USE_TZ is True, otherwise returns a naive datetime in the local timezone.
"""
blob_client = self.client.get_blob_client(self._get_valid_path(name))
properties = blob_client.get_blob_properties(timeout=self.timeout)
properties = blob_client.get_blob_properties(**self._request_options())
if not setting("USE_TZ", False):
return timezone.make_naive(properties.last_modified)

Expand All @@ -372,7 +389,7 @@ def list_all(self, path=""):
return [
blob.name
for blob in self.client.list_blobs(
name_starts_with=path, timeout=self.timeout
name_starts_with=path, **self._request_options()
)
]

Expand Down
18 changes: 18 additions & 0 deletions tests/test_azure.py
Original file line number Diff line number Diff line change
Expand Up @@ -378,3 +378,21 @@ def test_client_settings(self, bsc):
bsc.assert_called_once_with(
"https://test.blob.core.windows.net", credential=None, api_version="1.3"
)

def test_lazy_evaluated_request_options(self):
foo = mock.MagicMock()
foo.side_effect = [1, 2] # return different values the two times it is called
with override_settings(AZURE_REQUEST_OPTIONS={"key1": 5, "key2": foo}):
storage = azure_storage.AzureStorage()
client_mock = mock.MagicMock()
storage._client = client_mock

_, _ = storage.listdir("")
client_mock.list_blobs.assert_called_with(
name_starts_with="", timeout=20, key1=5, key2=1
)

_, _ = storage.listdir("")
client_mock.list_blobs.assert_called_with(
name_starts_with="", timeout=20, key1=5, key2=2
)

0 comments on commit e2f6b83

Please sign in to comment.