From 9a54a0293c45a31d585a7a497e33df92cc857463 Mon Sep 17 00:00:00 2001
From: Jorge Esteban Quilcate Otoya <jorge.quilcate@aiven.io>
Date: Wed, 20 Dec 2023 09:45:08 -0500
Subject: [PATCH] feat: add custom Azure host/port to support custom blob
 endpoint

 e.g. azurite
---
 rohmu/object_storage/azure.py     | 62 ++++++++++++++++++++++++-------
 rohmu/object_storage/config.py    |  3 ++
 test/object_storage/test_azure.py | 12 ++++++
 3 files changed, 64 insertions(+), 13 deletions(-)

diff --git a/rohmu/object_storage/azure.py b/rohmu/object_storage/azure.py
index 0a09e07d..3e32b4da 100644
--- a/rohmu/object_storage/azure.py
+++ b/rohmu/object_storage/azure.py
@@ -64,6 +64,9 @@ def __init__(
         account_key: Optional[str] = None,
         sas_token: Optional[str] = None,
         prefix: Optional[str] = None,
+        is_secure: bool = True,
+        host: Optional[str] = None,
+        port: Optional[int] = None,
         azure_cloud: Optional[str] = None,
         proxy_info: Optional[dict[str, Union[str, int]]] = None,
         notifier: Optional[Notifier] = None,
@@ -78,16 +81,13 @@ def __init__(
         self.account_key = account_key
         self.container_name = bucket_name
         self.sas_token = sas_token
-        try:
-            endpoint_suffix = ENDPOINT_SUFFIXES[azure_cloud]
-        except KeyError:
-            raise InvalidConfigurationError(f"Unknown azure cloud {repr(azure_cloud)}")
-
-        conn_str = (
-            "DefaultEndpointsProtocol=https;"
-            f"AccountName={self.account_name};"
-            f"AccountKey={self.account_key};"
-            f"EndpointSuffix={endpoint_suffix}"
+        conn_str = self.conn_string(
+            account_name=account_name,
+            account_key=account_key,
+            azure_cloud=azure_cloud,
+            host=host,
+            port=port,
+            is_secure=is_secure,
         )
         config: dict[str, Any] = {"max_block_size": MAX_BLOCK_SIZE}
         if proxy_info:
@@ -97,13 +97,13 @@ def __init__(
                 auth = f"{username}:{password}@"
             else:
                 auth = ""
-            host = proxy_info["host"]
-            port = proxy_info["port"]
+            proxy_host = proxy_info["host"]
+            proxy_port = proxy_info["port"]
             if proxy_info.get("type") == "socks5":
                 schema = "socks5"
             else:
                 schema = "http"
-            config["proxies"] = {"https": f"{schema}://{auth}{host}:{port}"}
+            config["proxies"] = {"https": f"{schema}://{auth}{proxy_host}:{proxy_port}"}
 
         self.conn: BlobServiceClient = BlobServiceClient.from_connection_string(
             conn_str=conn_str,
@@ -113,6 +113,42 @@ def __init__(
         self.container = self.get_or_create_container(self.container_name)
         self.log.debug("AzureTransfer initialized, %r", self.container_name)
 
+    @staticmethod
+    def conn_string(
+        account_name: str,
+        account_key: Optional[str],
+        azure_cloud: Optional[str],
+        host: Optional[str],
+        port: Optional[int],
+        is_secure: bool,
+    ) -> str:
+        if not host and not port:
+            try:
+                endpoint_suffix = ENDPOINT_SUFFIXES[azure_cloud]
+            except KeyError:
+                raise InvalidConfigurationError(f"Unknown azure cloud {repr(azure_cloud)}")
+
+            conn_str = (
+                "DefaultEndpointsProtocol=https;"
+                f"AccountName={account_name};"
+                f"AccountKey={account_key};"
+                f"EndpointSuffix={endpoint_suffix}"
+            )
+        else:
+            print(host)
+            print(port)
+            if not host or not port:
+                raise InvalidConfigurationError("Custom host and port must be specified together")
+
+            protocol = "https" if is_secure else "http"
+            conn_str = (
+                f"DefaultEndpointsProtocol={protocol};"
+                f"AccountName={account_name};"
+                f"AccountKey={account_key};"
+                f"BlobEndpoint={protocol}://{host}:{port}/{account_name};"
+            )
+        return conn_str
+
     def copy_file(
         self, *, source_key: str, destination_key: str, metadata: Optional[Metadata] = None, **kwargs: Any
     ) -> None:
diff --git a/rohmu/object_storage/config.py b/rohmu/object_storage/config.py
index 5bfec47c..82676ace 100644
--- a/rohmu/object_storage/config.py
+++ b/rohmu/object_storage/config.py
@@ -83,6 +83,9 @@ class AzureObjectStorageConfig(StorageModel):
     account_key: Optional[str] = Field(None, repr=False)
     sas_token: Optional[str] = Field(None, repr=False)
     prefix: Optional[str] = None
+    is_secure: bool = True
+    host: Optional[str] = None
+    port: Optional[int] = None
     azure_cloud: Optional[str] = None
     proxy_info: Optional[ProxyInfo] = None
     storage_type: Literal[StorageDriver.azure] = StorageDriver.azure
diff --git a/test/object_storage/test_azure.py b/test/object_storage/test_azure.py
index e193f53d..cad5fa56 100644
--- a/test/object_storage/test_azure.py
+++ b/test/object_storage/test_azure.py
@@ -103,3 +103,15 @@ def test_get_contents_to_fileobj_raises_error_on_invalid_byte_range(azure_module
             fileobj_to_store_to=BytesIO(),
             byte_range=(100, 10),
         )
+
+
+def test_conn_string_custom_host_port() -> None:
+    from rohmu.object_storage.azure import AzureTransfer
+
+    conn_string = AzureTransfer.conn_string(
+        account_name="test_name", account_key="test_key", azure_cloud=None, host="localhost", port=10000, is_secure=False
+    )
+    assert (
+        "DefaultEndpointsProtocol=http;AccountName=test_name;AccountKey=test_key;"
+        "BlobEndpoint=http://localhost:10000/test_name;"
+    ) == conn_string