migration 5

Author: imabhichow

DynamoDbEncryption/runtimes/python/src/aws_dbesdk_dynamodb/internaldafny/extern/InternalLegacyOverride.py

@@ -34,18 +34,19 @@ def Build(config: DynamoDbItemEncryptorConfig_DynamoDbItemEncryptorConfig):
 
         legacy_override = config.legacyOverride.value
 
+        maybe_encryptor = aws_dbesdk_dynamodb.smithygenerated.aws_cryptography_dbencryptionsdk_dynamodb.dafny_to_smithy.aws_cryptography_dbencryptionsdk_dynamodb_LegacyDynamoDbEncryptorReference(
+            legacy_override.encryptor
+        )
+
         # Precondition: The encryptor MUST be a DynamoDBEncryptor
         if not _HAS_LEGACY_DDBEC:
             return InternalLegacyOverride.CreateBuildFailure(
                 InternalLegacyOverride.CreateError("Could not find aws-dynamodb-encryption-python installation")
             )
-        if not isinstance(legacy_override.encryptor, EncryptedClient):
+        if not isinstance(maybe_encryptor, EncryptedClient):
             return InternalLegacyOverride.CreateBuildFailure(
                 InternalLegacyOverride.CreateError("Legacy encryptor is not supported")
             )
-        maybe_encryptor = aws_dbesdk_dynamodb.smithygenerated.aws_cryptography_dbencryptionsdk_dynamodb.dafny_to_smithy.aws_cryptography_dbencryptionsdk_dynamodb_LegacyDynamoDbEncryptorReference(
-            legacy_override.encryptor
-        )
 
         # Preconditions: MUST be able to create valid encryption context
         maybe_encryption_context = InternalLegacyOverride.legacyEncryptionContext(config)
@@ -59,7 +60,9 @@ def Build(config: DynamoDbItemEncryptorConfig_DynamoDbItemEncryptorConfig):
 
         # Create and return the legacy override instance
         legacy_instance = InternalLegacyOverride()
-        legacy_instance.encryptor = maybe_encryptor
+        # Store both the original encryptor and the wrapped reference
+        legacy_instance.original_encryptor = legacy_override.encryptor  # Store the original EncryptedClient
+        legacy_instance.encryptor = maybe_encryptor  # Store the wrapped LegacyDynamoDbEncryptorReference
         legacy_instance.policy = legacy_override.policy
         # Access the value property, not calling it as a function
         legacy_instance.encryption_context = maybe_encryption_context.value
@@ -200,9 +203,11 @@ def DecryptItem(self, input):
                     InternalLegacyOverride.CreateError("Item does not contain required legacy fields")
                 )
 
-            # Use the encryptor to decrypt the item using the instance attributes
-            decrypted_item = self.encryptor.decrypt_item(
-                encrypted_item, actions=self.attribute_actions, encryption_context=self.encryption_context
+            # Use the original EncryptedClient to decrypt the item
+            decrypted_item = self.original_encryptor.decrypt_item(
+                encrypted_item,
+                attribute_actions=self.attribute_actions,
+                encryption_context=self.encryption_context
             )
 
             # Create the output with the decrypted item
@@ -215,7 +220,8 @@ def DecryptItem(self, input):
 
     def __init__(self):
         super().__init__()
-        self.encryptor = None
+        self.encryptor = None  # LegacyDynamoDbEncryptorReference
+        self.original_encryptor = None  # Original EncryptedClient with encrypt_item/decrypt_item methods
         self.attribute_actions = None
         self.encryption_context = None
         self.policy = None