[Azure] Supply guid from django_guid as the correlationId on storage calls

storages/backends/azure_storage.py

@@ -21,6 +21,20 @@
 from storages.utils import setting
 from storages.utils import to_bytes
 
+try:
+    from django_guid import get_guid
+except ImportError:
+    def get_guid():
+        return None
+
+
+def _make_headers():
+    headers = {}
+    correlation_id = get_guid()
+    if correlation_id is not None:
+        headers["x-ms-client-request-id"] = correlation_id
+    return headers
+
 
 @deconstructible
 class AzureStorageFile(File):
@@ -44,7 +58,9 @@ 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,
+                timeout=self._storage.timeout,
+                headers=_make_headers())
             download_stream.readinto(file)
         if 'r' in self._mode:
             file.seek(0)
@@ -266,13 +282,14 @@ def delete(self, name):
         try:
             self.client.delete_blob(
                 self._get_valid_path(name),
-                timeout=self.timeout)
+                timeout=self.timeout,
+                headers=_make_headers())
         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(timeout=self.timeout, headers=_make_headers())
         return properties.size
 
     def _save(self, name, content):
@@ -291,7 +308,8 @@ def _save(self, name, content):
             content_settings=ContentSettings(**params),
             max_concurrency=self.upload_max_conn,
             timeout=self.timeout,
-            overwrite=self.overwrite_files)
+            overwrite=self.overwrite_files,
+            headers=_make_headers())
         return cleaned_name
 
     def _expire_at(self, expire):
@@ -355,7 +373,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(timeout=self.timeout, headers=_make_headers())
         if not setting('USE_TZ', False):
             return timezone.make_naive(properties.last_modified)
 
@@ -385,7 +403,8 @@ def list_all(self, path=''):
             blob.name
             for blob in self.client.list_blobs(
                 name_starts_with=path,
-                timeout=self.timeout)]
+                timeout=self.timeout,
+                headers=_make_headers())]
 
     def listdir(self, path=''):
         """

tests/test_azure.py

@@ -1,4 +1,5 @@
 import datetime
+import uuid
 from datetime import timedelta
 from unittest import mock
 
@@ -10,10 +11,21 @@
 from django.test import override_settings
 from django.utils import timezone
 from django.utils.timezone import make_aware
+from django_guid import set_guid
 
 from storages.backends import azure_storage
 
 
+def set_and_expect_guid():
+    """
+    Set a GUID via the django_guid module and expect it to be set in the headers.
+    """
+
+    guid = str(uuid.uuid4())
+    set_guid(guid)
+    return {'x-ms-client-request-id': guid}
+
+
 class AzureStorageTest(TestCase):
 
     def setUp(self, *args):
@@ -299,7 +311,8 @@ def test_container_client_params_connection_string(self):
 
     # From boto3
 
-    def test_storage_save(self):
+    @mock.patch('storages.backends.azure_storage.get_guid', return_value=None)
+    def test_storage_save(self, mocked_get_guid):
         """
         Test saving a file
         """
@@ -314,19 +327,67 @@ def test_storage_save(self):
                 content_settings='content_settings_foo',
                 max_concurrency=2,
                 timeout=20,
-                overwrite=True)
+                overwrite=True,
+                headers={})
             c_mocked.assert_called_once_with(
                 content_type='text/plain',
                 content_encoding=None,
                 cache_control=None)
             self.storage._custom_client.upload_blob.assert_not_called()
 
-    def test_storage_open_write(self):
+    def test_storage_save_with_guid(self):
+        """
+        Test saving a file
+        """
+        name = 'test storage save.txt'
+        content = ContentFile('new content')
+        headers = set_and_expect_guid()
+        with mock.patch('storages.backends.azure_storage.ContentSettings') as c_mocked:
+            c_mocked.return_value = 'content_settings_foo'
+            self.assertEqual(self.storage.save(name, content), name)
+            self.storage._client.upload_blob.assert_called_once_with(
+                name,
+                content.file,
+                content_settings='content_settings_foo',
+                max_concurrency=2,
+                timeout=20,
+                overwrite=True,
+                headers=headers)
+            c_mocked.assert_called_once_with(
+                content_type='text/plain',
+                content_encoding=None,
+                cache_control=None)
+            self.storage._custom_client.upload_blob.assert_not_called()
+
+    @mock.patch('storages.backends.azure_storage.get_guid', return_value=None)
+    def test_storage_open_write(self, mocked_get_guid):
+        """
+        Test opening a file in write mode
+        """
+        name = 'test_open_for_writïng.txt'
+        content = 'new content'
+
+        file = self.storage.open(name, 'w')
+        file.write(content)
+        written_file = file.file
+        file.close()
+        self.storage._client.upload_blob.assert_called_once_with(
+            name,
+            written_file,
+            content_settings=mock.ANY,
+            max_concurrency=2,
+            timeout=20,
+            overwrite=True,
+            headers={})
+        self.storage._custom_client.upload_blob.assert_not_called()
+
+    def test_storage_open_write_with_guid(self):
         """
         Test opening a file in write mode
         """
         name = 'test_open_for_writïng.txt'
         content = 'new content'
+        headers = set_and_expect_guid()
 
         file = self.storage.open(name, 'w')
         file.write(content)
@@ -338,26 +399,48 @@ def test_storage_open_write(self):
             content_settings=mock.ANY,
             max_concurrency=2,
             timeout=20,
-            overwrite=True)
+            overwrite=True,
+            headers=headers)
         self.storage._custom_client.upload_blob.assert_not_called()
 
-    def test_storage_exists(self):
+    @mock.patch('storages.backends.azure_storage.get_guid', return_value=None)
+    def test_storage_exists(self, mocked_get_guid):
         blob_name = "blob"
         client_mock = mock.MagicMock()
         custom_client_mock = mock.MagicMock()
         self.storage._client.get_blob_client.return_value = client_mock
         self.storage._custom_client.get_blob_client.return_value = client_mock
         self.assertTrue(self.storage.exists(blob_name))
-        self.assertEqual(client_mock.exists.call_count, 1)
+        client_mock.get_blob_properties.assert_called_once_with(headers={})
         self.assertEqual(custom_client_mock.exists.call_count, 0)
 
-    def test_delete_blob(self):
+    def test_storage_exists_with_guid(self):
+        blob_name = "blob"
+        headers = set_and_expect_guid()
+        client_mock = mock.MagicMock()
+        custom_client_mock = mock.MagicMock()
+        self.storage._client.get_blob_client.return_value = client_mock
+        self.storage._custom_client.get_blob_client.return_value = client_mock
+        self.assertTrue(self.storage.exists(blob_name))
+        client_mock.get_blob_properties.assert_called_once_with(headers=headers)
+        self.assertEqual(custom_client_mock.exists.call_count, 0)
+
+    @mock.patch('storages.backends.azure_storage.get_guid', return_value=None)
+    def test_delete_blob(self, mocked_get_guid):
         self.storage.delete("name")
         self.storage._client.delete_blob.assert_called_once_with(
-            "name", timeout=20)
+            "name", timeout=20, headers={})
         self.storage._custom_client.delete_blob.assert_not_called()
 
-    def test_storage_listdir_base(self):
+    def test_delete_blob_with_guid(self):
+        headers = set_and_expect_guid()
+        self.storage.delete("name")
+        self.storage._client.delete_blob.assert_called_once_with(
+            "name", timeout=20, headers=headers)
+        self.storage._custom_client.delete_blob.assert_not_called()
+
+    @mock.patch('storages.backends.azure_storage.get_guid', return_value=None)
+    def test_storage_listdir_base(self, mocked_get_guid):
         file_names = ["some/path/1.txt", "2.txt", "other/path/3.txt", "4.txt"]
 
         result = []
@@ -369,7 +452,7 @@ def test_storage_listdir_base(self):
 
         dirs, files = self.storage.listdir("")
         self.storage._client.list_blobs.assert_called_with(
-            name_starts_with="", timeout=20)
+            name_starts_with="", timeout=20, headers={})
         self.storage._custom_client.list_blobs.assert_not_called()
 
         self.assertEqual(len(dirs), 2)
@@ -384,20 +467,38 @@ def test_storage_listdir_base(self):
                 filename in files,
                 """ "{}" not in file list "{}".""".format(filename, files))
 
-    def test_storage_listdir_subdir(self):
-        file_names = ["some/path/1.txt", "some/2.txt"]
+    def test_storage_listdir_base_with_guid(self):
+        file_names = ["some/path/1.txt", "2.txt", "other/path/3.txt", "4.txt"]
+        headers = set_and_expect_guid()
 
         result = []
         for p in file_names:
             obj = mock.MagicMock()
             obj.name = p
             result.append(obj)
         self.storage._client.list_blobs.return_value = iter(result)
+
+        dirs, files = self.storage.listdir("")
+        self.storage._client.list_blobs.assert_called_with(
+            name_starts_with="", timeout=20, headers=headers)
         self.storage._custom_client.list_blobs.assert_not_called()
+        # Rest of functionality is tested in test_storage_listdir_base
+
+    @mock.patch('storages.backends.azure_storage.get_guid', return_value=None)
+    def test_storage_listdir_subdir(self, mocked_get_guid):
+        file_names = ["some/path/1.txt", "some/2.txt"]
+
+        result = []
+        for p in file_names:
+            obj = mock.MagicMock()
+            obj.name = p
+            result.append(obj)
+        self.storage._client.list_blobs.return_value = iter(result)
 
         dirs, files = self.storage.listdir("some/")
         self.storage._client.list_blobs.assert_called_with(
-            name_starts_with="some/", timeout=20)
+            name_starts_with="some/", timeout=20, headers={})
+        self.storage._custom_client.list_blobs.assert_not_called()
 
         self.assertEqual(len(dirs), 1)
         self.assertTrue(
@@ -409,6 +510,23 @@ def test_storage_listdir_subdir(self):
             '2.txt' in files,
             """ "2.txt" not in files list "{}".""".format(files))
 
+    def test_storage_listdir_subdir_with_guid(self):
+        file_names = ["some/path/1.txt", "some/2.txt"]
+        headers = set_and_expect_guid()
+
+        result = []
+        for p in file_names:
+            obj = mock.MagicMock()
+            obj.name = p
+            result.append(obj)
+        self.storage._client.list_blobs.return_value = iter(result)
+
+        dirs, files = self.storage.listdir("some/")
+        self.storage._client.list_blobs.assert_called_with(
+            name_starts_with="some/", timeout=20, headers=headers)
+        self.storage._custom_client.list_blobs.assert_not_called()
+        # Rest of functionality is tested in test_storage_listdir_subdir
+
     def test_size_of_file(self):
         props = BlobProperties()
         props.size = 12

tox.ini

@@ -19,6 +19,7 @@ deps =
 	django4.1: django~=4.1.0
 	djangomain: https://github.com/django/django/archive/main.tar.gz
 	cryptography
+	django-guid
 	pytest
 	pytest-cov
 	rsa