Merge pull request #409 from AGhafaryy/lazyloaded-args

allowing lazy loaded arguments as arguments

Author: zprobst

nodestream/project/storage.py

@@ -1,5 +1,7 @@
 from dataclasses import dataclass, field
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Union
+
+from nodestream.file_io import LazyLoadedArgument
 
 from ..pipeline.object_storage import ObjectStore, Signer
 
@@ -9,14 +11,15 @@ class StoreConfiguration:
     name: str
     storage_type: str
     arguments: Dict[str, Any]
-    hmac_key: Optional[any] = None
+    hmac_key: Optional[Union[LazyLoadedArgument, str]] = None
 
     def initialize(self) -> ObjectStore:
-        store = ObjectStore.from_file_arguments(self.storage_type, **self.arguments)
+        resolved_args = LazyLoadedArgument.resolve_if_needed(self.arguments)
+        store = ObjectStore.from_file_arguments(self.storage_type, **resolved_args)
         if self.hmac_key:
-            return store.signed(Signer.hmac(self.hmac_key))
-        else:
-            return store
+            resolved_key = LazyLoadedArgument.resolve_if_needed(self.hmac_key)
+            return store.signed(Signer.hmac(resolved_key))
+        return store
 
     def to_file_data(self):
         return dict(
@@ -37,10 +40,15 @@ def from_file_data(data):
 
     @staticmethod
     def describe_yaml_schema():
-        from schema import Optional, Schema
+        from schema import Optional, Or, Schema
 
         return Schema(
-            {"name": str, "type": str, Optional("hmac_key"): str, Optional(str): object}
+            {
+                "name": str,
+                "type": str,
+                Optional("hmac_key"): Or(LazyLoadedArgument, str, only_one=True),
+                Optional(str): object,
+            }
         )
 
 

tests/unit/project/test_storage.py

@@ -1,6 +1,10 @@
+import base64
+from unittest.mock import Mock
+
 import pytest
 from hamcrest import assert_that, contains, equal_to, instance_of
 
+from nodestream.file_io import LazyLoadedArgument
 from nodestream.pipeline.object_storage import (
     DirectoryObjectStore,
     NullObjectStore,
@@ -88,3 +92,24 @@ def test_storage_configuration_to_file_data(store_config):
     file_data = storage_config.to_file_data()
     expected_data = {"stores": [("test_store", store_config)]}
     assert_that(file_data, equal_to(expected_data))
+
+
+def test_store_configuration_initialize_with_lazy_hmac():
+    expected_hmac = "dvHdCrVbRPp1HcmWX78Ryw=="
+    mock_lazy_hmac = Mock(spec=LazyLoadedArgument)
+    mock_lazy_hmac.get_value.return_value = expected_hmac
+    expected_arg = "testarglol"
+    mock_lazy_arg = Mock(spec=LazyLoadedArgument)
+    mock_lazy_arg.get_value.return_value = expected_arg
+
+    store_config = StoreConfiguration(
+        name="test-store",
+        storage_type="local",
+        arguments={"root": mock_lazy_arg},
+        hmac_key=mock_lazy_hmac,
+    )
+
+    store = store_config.initialize()
+    assert_that(base64.b64encode(store.signer.key).decode(), equal_to(expected_hmac))
+    assert_that(store.store.root, equal_to(expected_arg))
+    mock_lazy_hmac.get_value.assert_called_once()