Revised singleton class to allow for one instance per different configuration

Author: seshubaws

aws_lambda_powertools/utilities/data_masking/base.py

@@ -1,5 +1,5 @@
 import json
-from typing import Union, Optional
+from typing import Optional, Union
 
 from aws_lambda_powertools.utilities.data_masking.provider import BaseProvider
 
@@ -94,7 +94,7 @@ def _apply_action_to_fields(self, data: Union[dict, str], fields, action, **prov
         else:
             raise TypeError(
                 "Unsupported data type. The 'data' parameter must be a dictionary or a JSON string "
-                "representation of a dictionary."
+                "representation of a dictionary.",
             )
 
         for field in fields:

aws_lambda_powertools/utilities/data_masking/providers/aws_encryption_sdk.py

@@ -8,8 +8,9 @@
     LocalCryptoMaterialsCache,
     StrictAwsKmsMasterKeyProvider,
 )
-from aws_lambda_powertools.utilities.data_masking.provider import BaseProvider
+
 from aws_lambda_powertools.shared.user_agent import register_feature_to_botocore_session
+from aws_lambda_powertools.utilities.data_masking.provider import BaseProvider
 
 
 class ContextMismatchError(Exception):
@@ -18,25 +19,28 @@ def __init__(self, key):
         self.key = key
 
 
-class SingletonMeta(type):
-    """Metaclass to cache class instances to optimize encryption"""
+class Singleton:
+    _instances: Dict[Any, "AwsEncryptionSdkProvider"] = {}
 
-    _instances: Dict["AwsEncryptionSdkProvider", Any] = {}
+    def __new__(cls, *args, **kwargs):
+        # Generate a unique key based on the configuration
+        # Create a tuple by iterating through the values in kwargs, sorting them,
+        # and then adding them to the tuple.
+        config_key = tuple(v for value in kwargs.values() for v in sorted(value))
 
-    def __call__(cls, *args, **provider_options):
-        if cls not in cls._instances:
-            instance = super().__call__(*args, **provider_options)
-            cls._instances[cls] = instance
-        return cls._instances[cls]
+        if config_key not in cls._instances:
+            cls._instances[config_key] = super(Singleton, cls).__new__(cls, *args)
+            print("in if class instances:", cls._instances)
+        return cls._instances[config_key]
 
 
 CACHE_CAPACITY: int = 100
-MAX_ENTRY_AGE_SECONDS: float = 300.0
-MAX_MESSAGES: int = 200
+MAX_CACHE_AGE_SECONDS: float = 300.0
+MAX_MESSAGES_ENCRYPTED: int = 200
 # NOTE: You can also set max messages/bytes per data key
 
 
-class AwsEncryptionSdkProvider(BaseProvider):
+class AwsEncryptionSdkProvider(BaseProvider, Singleton):
     """
     The AwsEncryptionSdkProvider is to be used as a Provider for the Datamasking class.
 
@@ -57,8 +61,8 @@ def __init__(
         keys: List[str],
         client: Optional[EncryptionSDKClient] = None,
         local_cache_capacity: Optional[int] = CACHE_CAPACITY,
-        max_cache_age_seconds: Optional[float] = MAX_ENTRY_AGE_SECONDS,
-        max_messages: Optional[int] = MAX_MESSAGES,
+        max_cache_age_seconds: Optional[float] = MAX_CACHE_AGE_SECONDS,
+        max_messages_encrypted: Optional[int] = MAX_MESSAGES_ENCRYPTED,
     ):
         self.client = client or EncryptionSDKClient()
         self.keys = keys
@@ -68,7 +72,7 @@ def __init__(
             master_key_provider=self.key_provider,
             cache=self.cache,
             max_age=max_cache_age_seconds,
-            max_messages_encrypted=max_messages,
+            max_messages_encrypted=max_messages_encrypted,
         )
 
     def encrypt(self, data: Union[bytes, str], **provider_options) -> str:
@@ -112,7 +116,9 @@ def decrypt(self, data: str, **provider_options) -> bytes:
         expected_context = provider_options.pop("encryption_context", {})
 
         ciphertext, decryptor_header = self.client.decrypt(
-            source=ciphertext_decoded, key_provider=self.key_provider, **provider_options
+            source=ciphertext_decoded,
+            key_provider=self.key_provider,
+            **provider_options,
         )
 
         for key, value in expected_context.items():

tests/e2e/data_masking/conftest.py

@@ -1,4 +1,5 @@
 import pytest
+
 from tests.e2e.data_masking.infrastructure import DataMaskingStack
 
 

tests/e2e/data_masking/handlers/basic_handler.py

@@ -1,6 +1,6 @@
+from aws_lambda_powertools import Logger
 from aws_lambda_powertools.utilities.data_masking.base import DataMasking
 from aws_lambda_powertools.utilities.data_masking.providers.aws_encryption_sdk import AwsEncryptionSdkProvider
-from aws_lambda_powertools import Logger
 
 logger = Logger()
 

tests/e2e/data_masking/infrastructure.py

@@ -1,5 +1,6 @@
 import aws_cdk.aws_kms as kms
 from aws_cdk import CfnOutput, Duration
+
 from tests.e2e.utils.infrastructure import BaseInfrastructure
 
 

tests/e2e/data_masking/test_data_masking.py

@@ -1,13 +1,15 @@
 import json
 from uuid import uuid4
-from aws_encryption_sdk.exceptions import DecryptKeyError
+
 import pytest
-from tests.e2e.utils import data_fetcher
+from aws_encryption_sdk.exceptions import DecryptKeyError
+
 from aws_lambda_powertools.utilities.data_masking.base import DataMasking
 from aws_lambda_powertools.utilities.data_masking.providers.aws_encryption_sdk import (
     AwsEncryptionSdkProvider,
     ContextMismatchError,
 )
+from tests.e2e.utils import data_fetcher
 
 
 @pytest.fixture
@@ -93,7 +95,6 @@ def test_encryption_no_context_fail(data_masker):
         data_masker.decrypt(encrypted_data, encryption_context={"this": "is_secure"})
 
 
-# TODO: metaclass?
 @pytest.mark.xdist_group(name="data_masking")
 def test_encryption_decryption_key_mismatch(data_masker, kms_key2_arn):
     # GIVEN an instantiation of DataMasking with the AWS encryption provider with a certain key
@@ -109,6 +110,21 @@ def test_encryption_decryption_key_mismatch(data_masker, kms_key2_arn):
         data_masker_key2.decrypt(encrypted_data)
 
 
+def test_encryption_provider_singleton(data_masker, kms_key1_arn, kms_key2_arn):
+    data_masker_2 = DataMasking(provider=AwsEncryptionSdkProvider(keys=[kms_key1_arn]))
+    assert data_masker.provider is data_masker_2.provider
+
+    # WHEN encrypting and then decrypting the encrypted data
+    encrypted_data = data_masker.encrypt("string")
+    decrypted_data = data_masker_2.decrypt(encrypted_data)
+
+    # THEN the result is the original input data
+    assert decrypted_data == bytes("string", "utf-8")
+
+    data_masker_3 = DataMasking(provider=AwsEncryptionSdkProvider(keys=[kms_key2_arn]))
+    assert data_masker_2.provider is not data_masker_3.provider
+
+
 @pytest.mark.xdist_group(name="data_masking")
 def test_encryption_in_logs(data_masker, basic_handler_fn, basic_handler_fn_arn):
     # GIVEN an instantiation of DataMasking with the AWS encryption provider
@@ -132,6 +148,7 @@ def test_encryption_in_logs(data_masker, basic_handler_fn, basic_handler_fn_arn)
         decrypted_data = data_masker.decrypt(encrypted_data)
         assert decrypted_data == value
 
+
 # NOTE: This test is failing currently, need to find a fix for building correct dependencies
 @pytest.mark.xdist_group(name="data_masking")
 def test_encryption_in_handler(basic_handler_fn_arn, kms_key1_arn):

tests/functional/data_masking/test_aws_encryption_sdk.py

@@ -1,12 +1,9 @@
-from unittest.mock import patch
-
 import pytest
 
 from aws_lambda_powertools.utilities.data_masking.base import DataMasking
 from aws_lambda_powertools.utilities.data_masking.providers.aws_encryption_sdk import AwsEncryptionSdkProvider
 from tests.unit.data_masking.setup import *
 
-
 AWS_SDK_KEY = "arn:aws:kms:us-west-2:683517028648:key/269301eb-81eb-4067-ac72-98e8e49bf2b3"
 
 

tests/unit/data_masking/test_data_masking.py

@@ -1,8 +1,10 @@
 import json
+
 import pytest
 from itsdangerous.url_safe import URLSafeSerializer
-from aws_lambda_powertools.utilities.data_masking.constants import DATA_MASKING_STRING
+
 from aws_lambda_powertools.utilities.data_masking.base import DataMasking
+from aws_lambda_powertools.utilities.data_masking.constants import DATA_MASKING_STRING
 from aws_lambda_powertools.utilities.data_masking.provider import BaseProvider
 
 
@@ -102,7 +104,7 @@ def test_mask_dict(data_masker):
         "a": {
             "1": {"None": "hello", "four": "world"},
             "b": {"3": {"4": "goodbye", "e": "world"}},
-        }
+        },
     }
 
     # WHEN mask is called with no fields argument
@@ -118,15 +120,18 @@ def test_mask_dict_with_fields(data_masker):
         "a": {
             "1": {"None": "hello", "four": "world"},
             "b": {"3": {"4": "goodbye", "e": "world"}},
-        }
+        },
     }
 
     # WHEN mask is called with a list of fields specified
     masked_string = data_masker.mask(data, fields=["a.1.None", "a.b.3.4"])
 
     # THEN the result is only the specified fields are masked
     assert masked_string == {
-        "a": {"1": {"None": DATA_MASKING_STRING, "four": "world"}, "b": {"3": {"4": DATA_MASKING_STRING, "e": "world"}}}
+        "a": {
+            "1": {"None": DATA_MASKING_STRING, "four": "world"},
+            "b": {"3": {"4": DATA_MASKING_STRING, "e": "world"}},
+        },
     }
 
 
@@ -137,16 +142,19 @@ def test_mask_json_dict_with_fields(data_masker):
             "a": {
                 "1": {"None": "hello", "four": "world"},
                 "b": {"3": {"4": "goodbye", "e": "world"}},
-            }
-        }
+            },
+        },
     )
 
     # WHEN mask is called with a list of fields specified
     masked_json_string = data_masker.mask(data, fields=["a.1.None", "a.b.3.4"])
 
     # THEN the result is only the specified fields are masked
     assert masked_json_string == {
-        "a": {"1": {"None": DATA_MASKING_STRING, "four": "world"}, "b": {"3": {"4": DATA_MASKING_STRING, "e": "world"}}}
+        "a": {
+            "1": {"None": DATA_MASKING_STRING, "four": "world"},
+            "b": {"3": {"4": DATA_MASKING_STRING, "e": "world"}},
+        },
     }
 
 
@@ -180,7 +188,7 @@ def test_encrypt_decrypt_bool(custom_data_masker):
     decrypted_data = custom_data_masker.decrypt(encrypted_data)
 
     # THEN the result is the original input data
-    assert decrypted_data == True
+    assert decrypted_data is True
 
 
 def test_encrypt_decrypt_none(custom_data_masker):
@@ -191,7 +199,7 @@ def test_encrypt_decrypt_none(custom_data_masker):
     decrypted_data = custom_data_masker.decrypt(encrypted_data)
 
     # THEN the result is the original input data
-    assert decrypted_data == None
+    assert decrypted_data is None
 
 
 def test_encrypt_decrypt_str(custom_data_masker):
@@ -223,7 +231,7 @@ def test_dict_encryption_with_fields(custom_data_masker):
         "a": {
             "1": {"None": "hello", "four": "world"},
             "b": {"3": {"4": "goodbye", "e": "world"}},
-        }
+        },
     }
 
     # WHEN encrypting and decrypting the data with a list of fields
@@ -242,8 +250,8 @@ def test_json_encryption_with_fields(custom_data_masker):
             "a": {
                 "1": {"None": "hello", "four": "world"},
                 "b": {"3": {"4": "goodbye", "e": "world"}},
-            }
-        }
+            },
+        },
     )
 
     # WHEN encrypting and decrypting a json representation of a dictionary with a list of fields
@@ -330,7 +338,7 @@ def test_parsing_nonexistent_fields(data_masker):
         "3": {
             "1": {"None": "hello", "four": "world"},
             "4": {"33": {"5": "goodbye", "e": "world"}},
-        }
+        },
     }
 
     # WHEN attempting to pass in fields that do not exist in the input data
@@ -347,7 +355,7 @@ def test_parsing_nonstring_fields(data_masker):
         "3": {
             "1": {"None": "hello", "four": "world"},
             "4": {"33": {"5": "goodbye", "e": "world"}},
-        }
+        },
     }
 
     # WHEN attempting to pass in a list of fields that are not strings
@@ -365,7 +373,7 @@ def test_parsing_nonstring_keys_and_fields(data_masker):
         3: {
             "1": {"None": "hello", "four": "world"},
             4: {"33": {"5": "goodbye", "e": "world"}},
-        }
+        },
     }
     masked = data_masker.mask(data, fields=[3.4])