allowing lazy loaded arguments as arguments

Author: Andrew Ghafari

nodestream/pipeline/extractors/credential_utils.py

@@ -1,24 +1,32 @@
 from datetime import datetime
 from time import time
-from typing import Optional
+from typing import Optional, Union
 from uuid import uuid4
 
 import boto3
 import pytz
 from botocore.credentials import RefreshableCredentials
 from botocore.session import Session
+from nodestream.file_io import LazyLoadedArgument
 
 
 class AwsClientFactory:
     def __init__(
         self,
-        assume_role_arn: Optional[str] = None,
-        assume_role_external_id: Optional[str] = None,
+        assume_role_arn: Optional[Union[LazyLoadedArgument,str]] = None,
+        assume_role_external_id: Optional[Union[LazyLoadedArgument,str]] = None,
         session_ttl: int = 3000,
         **boto_session_args
     ) -> None:
-        self.assume_role_arn = assume_role_arn
-        self.assume_role_external_id = assume_role_external_id
+        if isinstance(assume_role_arn, LazyLoadedArgument):
+            self.assume_role_arn = assume_role_arn.get_value()
+        else:
+            self.assume_role_arn = assume_role_arn
+        if isinstance(assume_role_external_id, LazyLoadedArgument):
+            self.assume_role_external_id = assume_role_external_id.get_value()
+        else:
+            self.assume_role_external_id = assume_role_external_id
+
         self.session_args = self._init_session_args(**boto_session_args)
         self.session_ttl = session_ttl
 

nodestream/project/storage.py

@@ -1,6 +1,6 @@
 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 +9,17 @@ 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)
         if self.hmac_key:
-            return store.signed(Signer.hmac(self.hmac_key))
-        else:
-            return store
+            if isinstance(self.hmac_key, LazyLoadedArgument):
+                resolved_key = self.hmac_key.get_value()
+                return store.signed(Signer.hmac(resolved_key))
+            else:
+                return store.signed(Signer.hmac(self.hmac_key))
+        return store
 
     def to_file_data(self):
         return dict(
@@ -40,7 +43,7 @@ def describe_yaml_schema():
         from schema import Optional, Schema
 
         return Schema(
-            {"name": str, "type": str, Optional("hmac_key"): str, Optional(str): object}
+            {"name": str, "type": str, Optional("hmac_key"): LazyLoadedArgument, Optional(str): object}
         )
 
 

tests/unit/pipeline/extractors/stores/aws/test_credential_utils.py

@@ -6,6 +6,9 @@
 from freezegun import freeze_time
 from hamcrest import assert_that, equal_to, has_key, not_
 
+from nodestream.file_io import LazyLoadedArgument
+from unittest.mock import Mock
+
 
 @pytest.fixture
 def client_with_role():
@@ -145,3 +148,37 @@ def test_make_client(mocker, client_without_role):
         return_value=session
     )
     client_without_role.make_client("s3")
+
+def create_mock_lazy_loaded_argument(value):
+    mock = Mock(spec=LazyLoadedArgument)
+    mock.get_value.return_value = value
+    return mock
+
+def test_init_with_lazy_loaded_role_arn():
+    from nodestream.pipeline.extractors.credential_utils import AwsClientFactory
+
+    lazy_role_arn = create_mock_lazy_loaded_argument("arn:aws:iam::123456789012:role/test")
+    client = AwsClientFactory(assume_role_arn=lazy_role_arn)
+
+    assert_that(client.assume_role_arn, equal_to("arn:aws:iam::123456789012:role/test"))
+
+def test_init_with_lazy_loaded_external_id():
+    from nodestream.pipeline.extractors.credential_utils import AwsClientFactory
+
+    lazy_external_id = create_mock_lazy_loaded_argument("test-external-id")
+    client = AwsClientFactory(assume_role_external_id=lazy_external_id)
+
+    assert_that(client.assume_role_external_id, equal_to("test-external-id"))
+
+def test_init_with_both_lazy_loaded_arguments():
+    from nodestream.pipeline.extractors.credential_utils import AwsClientFactory
+
+    lazy_role_arn = create_mock_lazy_loaded_argument("arn:aws:iam::123456789012:role/test")
+    lazy_external_id = create_mock_lazy_loaded_argument("test-external-id")
+    client = AwsClientFactory(
+        assume_role_arn=lazy_role_arn,
+        assume_role_external_id=lazy_external_id
+    )
+
+    assert_that(client.assume_role_arn, equal_to("arn:aws:iam::123456789012:role/test"))
+    assert_that(client.assume_role_external_id, equal_to("test-external-id"))

tests/unit/project/test_storage.py

@@ -1,3 +1,5 @@
+import base64
+
 import pytest
 from hamcrest import assert_that, contains, equal_to, instance_of
 
@@ -7,6 +9,10 @@
     SignedObjectStore,
 )
 from nodestream.project.storage import StorageConfiguration, StoreConfiguration
+from unittest.mock import Mock
+from hamcrest import assert_that, instance_of
+from nodestream.file_io import LazyLoadedArgument
+from nodestream.pipeline.object_storage import SignedObjectStore, DirectoryObjectStore
 
 
 @pytest.fixture
@@ -88,3 +94,23 @@ 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(tmp_path):
+    expected_hmac = "dvHdCrVbRPp1HcmWX78Ryw=="
+    mock_lazy_hmac = Mock(spec=LazyLoadedArgument)
+    mock_lazy_hmac.get_value.return_value = expected_hmac
+
+    store_config = StoreConfiguration(
+        name="test-store",
+        storage_type="local",
+        arguments={"root": tmp_path},
+        hmac_key=mock_lazy_hmac
+    )
+
+    store = store_config.initialize()
+    assert_that(
+        base64.b64encode(store.signer.key).decode(),
+        equal_to(expected_hmac)
+    )
+    mock_lazy_hmac.get_value.assert_called_once()