object_storage: config models improvements

rohmu/common/models.py

@@ -1,15 +1,14 @@
 # Copyright (c) 2023 Aiven, Helsinki, Finland. https://aiven.io/
-
-
 from rohmu.common.statsd import StatsdConfig
+from rohmu.common.strenum import StrEnum
 from rohmu.notifier.interface import Notifier
 from typing import Optional
 
 import enum
 import pydantic
 
 
-class StorageOperation(str, enum.Enum):
+class StorageOperation(StrEnum):
     iter_key = "iter_key"
     copy_file = "copy_file"
     delete_key = "delete_key"
@@ -28,15 +27,22 @@ class StorageOperation(str, enum.Enum):
     multipart_aborted = "multipart_aborted"
     multipart_complete = "multipart_complete"
 
-    def __str__(self) -> str:
-        return str(self.value)
-
 
-class ProxyType(str, enum.Enum):
+class ProxyType(StrEnum):
     socks5 = "socks5"
     http = "http"
 
 
+@enum.unique
+class StorageDriver(StrEnum):
+    azure = "azure"
+    google = "google"
+    local = "local"
+    s3 = "s3"
+    sftp = "sftp"
+    swift = "swift"
+
+
 class RohmuModel(pydantic.BaseModel):
     class Config:
         # As we're keen to both export and decode json, just using
@@ -66,6 +72,7 @@ class ProxyInfo(RohmuModel):
 
 
 class StorageModel(pydantic.BaseModel):
+    storage_type: StorageDriver
     notifier: Optional[Notifier] = None
     statsd_info: Optional[StatsdConfig] = None
 

rohmu/common/statsd.py

@@ -18,21 +18,18 @@
 from __future__ import annotations
 
 from contextlib import asynccontextmanager, contextmanager
-from enum import Enum
+from rohmu.common.strenum import StrEnum
 from typing import AsyncIterator, Dict, Iterator, Optional, Union
 
 import pydantic
 import socket
 import time
 
 
-class MessageFormat(str, Enum):
+class MessageFormat(StrEnum):
     datadog = "datadog"
     telegraf = "telegraf"
 
-    def __str__(self) -> str:
-        return str(self.value)
-
 
 Tags = Dict[str, Union[int, str, None]]
 

rohmu/common/strenum.py

@@ -0,0 +1,21 @@
+"""
+rohmu - StrEnum
+
+Copyright (c) 2023 Aiven, Helsinki, Finland. https://aiven.io/
+See LICENSE for details
+"""
+from __future__ import annotations
+
+import enum
+
+
+class StrEnum(str, enum.Enum):
+    def __str__(self) -> str:
+        return str(self.value)
+
+    @classmethod
+    def of(cls, value: str) -> StrEnum | None:
+        try:
+            return cls(value)
+        except ValueError:
+            return None

rohmu/object_storage/base.py

@@ -113,7 +113,7 @@ def _should_multipart(
 
     @classmethod
     def from_model(cls, model: StorageModelT) -> BaseTransfer[StorageModelT]:
-        return cls(**model.dict(by_alias=True))
+        return cls(**model.dict(by_alias=True, exclude={"storage_type"}))
 
     def copy_file(
         self, *, source_key: str, destination_key: str, metadata: Optional[Metadata] = None, **_kwargs: Any

rohmu/object_storage/config.py

@@ -9,8 +9,9 @@
 from __future__ import annotations
 
 from enum import Enum, unique
-from rohmu.common.models import ProxyInfo, StorageModel
-from typing import Any, Dict, Final, Optional, TypeVar
+from pydantic import DirectoryPath, Field, FilePath
+from rohmu.common.models import ProxyInfo, StorageDriver, StorageModel
+from typing import Any, Dict, Final, Literal, Optional, TypeVar
 
 import platform
 
@@ -76,27 +77,30 @@ def calculate_s3_chunk_size() -> int:
 
 
 class AzureObjectStorageConfig(StorageModel):
-    bucket_name: str
+    bucket_name: Optional[str]
     account_name: str
-    account_key: Optional[str] = None
-    sas_token: Optional[str] = None
+    account_key: Optional[str] = Field(None, repr=False)
+    sas_token: Optional[str] = Field(None, repr=False)
     prefix: Optional[str] = None
     azure_cloud: Optional[str] = None
     proxy_info: Optional[ProxyInfo] = None
+    storage_type: Literal[StorageDriver.azure] = StorageDriver.azure
 
 
 class GoogleObjectStorageConfig(StorageModel):
     project_id: str
-    bucket_name: str
-    credential_file: Optional[str] = None
-    credentials: Optional[Dict[str, Any]] = None
+    bucket_name: Optional[str]
+    credential_file: Optional[FilePath] = None
+    credentials: Optional[Dict[str, Any]] = Field(None, repr=False)
     proxy_info: Optional[ProxyInfo] = None
     prefix: Optional[str] = None
+    storage_type: Literal[StorageDriver.google] = StorageDriver.google
 
 
 class LocalObjectStorageConfig(StorageModel):
-    directory: str
+    directory: DirectoryPath
     prefix: Optional[str] = None
+    storage_type: Literal[StorageDriver.local] = StorageDriver.local
 
 
 @unique
@@ -108,9 +112,9 @@ class S3AddressingStyle(Enum):
 
 class S3ObjectStorageConfig(StorageModel):
     region: str
-    bucket_name: str
+    bucket_name: Optional[str]
     aws_access_key_id: Optional[str] = None
-    aws_secret_access_key: Optional[str] = None
+    aws_secret_access_key: Optional[str] = Field(None, repr=False)
     prefix: Optional[str] = None
     host: Optional[str] = None
     port: Optional[str] = None
@@ -122,21 +126,23 @@ class S3ObjectStorageConfig(StorageModel):
     proxy_info: Optional[ProxyInfo] = None
     connect_timeout: Optional[str] = None
     read_timeout: Optional[str] = None
-    aws_session_token: Optional[str] = None
+    aws_session_token: Optional[str] = Field(None, repr=False)
+    storage_type: Literal[StorageDriver.s3] = StorageDriver.s3
 
 
 class SFTPObjectStorageConfig(StorageModel):
     server: str
     port: int
     username: str
-    password: Optional[str] = None
-    private_key: Optional[str] = None
+    password: Optional[str] = Field(None, repr=False)
+    private_key: Optional[str] = Field(None, repr=False)
     prefix: Optional[str] = None
+    storage_type: Literal[StorageDriver.sftp] = StorageDriver.sftp
 
 
 class SwiftObjectStorageConfig(StorageModel):
     user: str
-    key: str
+    key: str = Field(repr=False)
     container_name: str
     auth_url: str
     auth_version: str = "2.0"
@@ -153,3 +159,5 @@ class SwiftObjectStorageConfig(StorageModel):
     project_domain_name: Optional[str] = None
     service_type: Optional[str] = None
     endpoint_type: Optional[str] = None
+    prefix: Optional[str] = None
+    storage_type: Literal[StorageDriver.swift] = StorageDriver.swift