Add AZURE_REQUEST_OPTIONS setting (#1179)

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.

Author: sdherr

docs/backends/azure.rst

@@ -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``

storages/backends/azure_storage.py

@@ -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:
@@ -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"),
@@ -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):
@@ -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):
@@ -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
 
@@ -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)
 
@@ -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()
             )
         ]
 

tests/test_azure.py

@@ -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
+            )