Repository: thelastpickle/cassandra-medusa

Issue:
Azure Storage Support for different Cloud Contexts (e.g. AzureUSGovernment)
[Project board link](https://github.com/orgs/k8ssandra/projects/8/views/1?pane=issue&itemId=43950753)

## Overview
Azure Storage looks to be configured for only Commercial cloud on line https://github.com/thelastpickle/cassandra-medusa/blob/master/medusa/storage/azure_storage.py#L58 _windows.net_

In AzureUSGovernment cloud, the URLs end with `usgovcloudapi.net`.

The Azure SDK allows for the appropriate cloud context to be set, and with the original azure-cli implementation of storage that was possible `< 0.16.0`.


## The Ask
The connection string for Azure Blob Storage should support a conditional so that it can *either* be the current `blob.core.windows.net` for "Azure" cloud or `blob.core.usgovcloudapi.net` for "AzureUSGovernment" cloud.  
This should be configurable in the medusa configuration file, rather than hard-coded in the above-referenced file.

_I'd code this up and offer a patch, but I don't immediately have the time, and hoped that the maintainers and existing contributors could tackle that faster than I can_


## Additional Azure References
https://learn.microsoft.com/en-us/azure/azure-government/compare-azure-government-global-azure

![Screen Shot 2023-11-07 at 3 18 07 PM](https://github.com/thelastpickle/cassandra-medusa/assets/3441274/bf9aaed7-88a1-43bb-9882-b621d728dcd2)


Comments:

Patch:
diff --git a/docs/azure_blobs_setup.md b/docs/azure_blobs_setup.md
index 135b7c5..447dbc7 100644
--- a/docs/azure_blobs_setup.md
+++ b/docs/azure_blobs_setup.md
@@ -11,15 +11,7 @@ Create a new storage account or use an existing one which will be used to store
     "key": "YOUR_KEY"
 }
 ```
-If you need to set a different host for Azure (for example the host for Azure Gov is `<storageAccount>.blob.core.usgovcloudapi.net`), please ADDITIONALLY set these two fields in the JSON file (the connection string can be found with the key):
-
-```
-"host": "YOUR_HOST"
-"connection_string": "YOUR_CONNECTION_STRING"
-
-```
-
-Place this file on all Apache Cassandra™ nodes running medusa under `/etc/medusa/`and set the rigths appropriately so that onyl users running Medusa can read/modify it.
+Place this file on all Apache Cassandra™ nodes running medusa under `/etc/medusa/`and set the rights appropriately so that only users running Medusa can read/modify it.
 
 ### Create a container
 
@@ -36,3 +28,8 @@ key_file = /etc/medusa/medusa-azure-credentials
 
 Medusa should now be able to access the bucket and perform all required operations.
 
+If you need to set a different host for Azure (for example the host for Azure Gov is `<storageAccount>.blob.core.usgovcloudapi.net`), please use the `host` parameter in the `[storage]` section of `/etc/medusa/medusa.ini`:
+
+```
+"host": "usgovcloudapi.net"
+```
diff --git a/medusa/storage/azure_storage.py b/medusa/storage/azure_storage.py
index 9faee5a..9d25b79 100644
--- a/medusa/storage/azure_storage.py
+++ b/medusa/storage/azure_storage.py
@@ -49,15 +49,25 @@ class AzureStorage(AbstractStorage):
         self.account_name = self.credentials.named_key.name
         self.bucket_name = config.bucket_name
 
+        self.azure_blob_service_url = self._make_blob_service_url(self.account_name, config)
+
         # disable chatty loggers
         logging.getLogger('azure.core.pipeline.policies.http_logging_policy').setLevel(logging.WARNING)
         logging.getLogger('chardet.universaldetector').setLevel(logging.WARNING)
 
         super().__init__(config)
 
+    def _make_blob_service_url(self, account_name, config):
+        domain = 'windows.net' if config.host is None else config.host
+        if config.port is None:
+            url = f"https://{account_name}.blob.core.{domain}/"
+        else:
+            url = f"https://{account_name}.blob.core.{domain}:{config.port}/"
+        return url
+
     def connect(self):
         self.azure_blob_service = BlobServiceClient(
-            account_url=f"https://{self.account_name}.blob.core.windows.net/",
+            account_url=self.azure_blob_service_url,
             credential=self.credentials
         )
         self.azure_container_client = self.azure_blob_service.get_container_client(self.bucket_name)


Test patch:
diff --git a/tests/storage/abstract_storage_test.py b/tests/storage/abstract_storage_test.py
index 6051142..5662483 100644
--- a/tests/storage/abstract_storage_test.py
+++ b/tests/storage/abstract_storage_test.py
@@ -18,7 +18,13 @@ import unittest
 from medusa.storage.abstract_storage import AbstractStorage
 
 
-class S3StorageTest(unittest.TestCase):
+class AttributeDict(dict):
+    __slots__ = ()
+    __getattr__ = dict.__getitem__
+    __setattr__ = dict.__setitem__
+
+
+class AbstractStorageTest(unittest.TestCase):
 
     def test_convert_human_friendly_size_to_bytes(self):
         self.assertEqual(50, AbstractStorage._human_size_to_bytes('50B'))
diff --git a/tests/storage/azure_storage_test.py b/tests/storage/azure_storage_test.py
new file mode 100644
index 0000000..7809eac
--- /dev/null
+++ b/tests/storage/azure_storage_test.py
@@ -0,0 +1,79 @@
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import tempfile
+import unittest
+
+from medusa.storage.azure_storage import AzureStorage
+from tests.storage.abstract_storage_test import AttributeDict
+
+
+class AzureStorageTest(unittest.TestCase):
+
+    credentials_file_content = """
+    {
+      "storage_account": "medusa-unit-test",
+      "key": "randomString=="
+    }
+    """
+
+    def test_make_connection_url(self):
+        with tempfile.NamedTemporaryFile() as credentials_file:
+            credentials_file.write(self.credentials_file_content.encode())
+            credentials_file.flush()
+            config = AttributeDict({
+                'region': 'region-from-config',
+                'storage_provider': 'azure_blobs',
+                'key_file': credentials_file.name,
+                'bucket_name': 'bucket-from-config',
+                'concurrent_transfers': '1',
+                'host': None,
+                'port': None,
+            })
+            azure_storage = AzureStorage(config)
+            self.assertEqual(
+                'https://medusa-unit-test.blob.core.windows.net/',
+                azure_storage.azure_blob_service_url
+            )
+
+    def test_make_connection_url_with_custom_host(self):
+        with tempfile.NamedTemporaryFile() as credentials_file:
+            credentials_file.write(self.credentials_file_content.encode())
+            credentials_file.flush()
+            config = AttributeDict({
+                'region': 'region-from-config',
+                'storage_provider': 'azure_blobs',
+                'key_file': credentials_file.name,
+                'bucket_name': 'bucket-from-config',
+                'concurrent_transfers': '1',
+                'host': 'custom.host.net',
+                'port': None,
+            })
+            azure_storage = AzureStorage(config)
+            self.assertEqual(
+                'https://medusa-unit-test.blob.core.custom.host.net/',
+                azure_storage.azure_blob_service_url
+            )
+
+    def test_make_connection_url_with_custom_host_port(self):
+        with tempfile.NamedTemporaryFile() as credentials_file:
+            credentials_file.write(self.credentials_file_content.encode())
+            credentials_file.flush()
+            config = AttributeDict({
+                'region': 'region-from-config',
+                'storage_provider': 'azure_blobs',
+                'key_file': credentials_file.name,
+                'bucket_name': 'bucket-from-config',
+                'concurrent_transfers': '1',
+                'host': 'custom.host.net',
+                'port': 123,
+            })
+            azure_storage = AzureStorage(config)
+            self.assertEqual(
+                'https://medusa-unit-test.blob.core.custom.host.net:123/',
+                azure_storage.azure_blob_service_url
+            )
diff --git a/tests/storage/google_storage_test.py b/tests/storage/google_storage_test.py
index 53bfd6d..891b49f 100644
--- a/tests/storage/google_storage_test.py
+++ b/tests/storage/google_storage_test.py
@@ -22,7 +22,7 @@ from pathlib import Path
 from medusa.storage.google_storage import _group_by_parent, _is_in_folder
 
 
-class RestoreNodeTest(unittest.TestCase):
+class GoogleStorageTest(unittest.TestCase):
 
     def test_is_in_folder(self):
         folder = Path('foo/bar')
diff --git a/tests/storage/s3_storage_test.py b/tests/storage/s3_storage_test.py
index 08e98c6..11a5daa 100644
--- a/tests/storage/s3_storage_test.py
+++ b/tests/storage/s3_storage_test.py
@@ -21,12 +21,7 @@ import tempfile
 from unittest.mock import patch, MagicMock
 
 from medusa.storage.s3_base_storage import S3BaseStorage
-
-
-class AttributeDict(dict):
-    __slots__ = ()
-    __getattr__ = dict.__getitem__
-    __setattr__ = dict.__setitem__
+from tests.storage.abstract_storage_test import AttributeDict
 
 
 class S3StorageTest(unittest.TestCase):