dbbackuprestore: add KMS encryption key and status message support to database restore (#303)

mdmonjurulhasan · web-flow · commit 07fcf8809081 · 2026-02-24T00:24:46.000-08:00
diff --git a/dbbackuprestore/dbbackuprestore-ali/src/main/java/com/salesforce/multicloudj/dbbackuprestore/ali/AliDBBackupRestore.java b/dbbackuprestore/dbbackuprestore-ali/src/main/java/com/salesforce/multicloudj/dbbackuprestore/ali/AliDBBackupRestore.java
@@ -269,6 +269,7 @@ private Restore convertToRestore(DescribeRestoreJobs2ResponseBody.DescribeRestor
                 .status(status)
                 .startTime(startTime)
                 .endTime(endTime)
+                .statusMessage(restoreJob.getErrorMessage())
                 .build();
     }
 
diff --git a/dbbackuprestore/dbbackuprestore-ali/src/test/java/com/salesforce/multicloudj/dbbackuprestore/ali/AliDBBackupRestoreTest.java b/dbbackuprestore/dbbackuprestore-ali/src/test/java/com/salesforce/multicloudj/dbbackuprestore/ali/AliDBBackupRestoreTest.java
@@ -208,7 +208,13 @@ void testRestoreBackupFailure() throws Exception {
         });
     }
 
-    private Restore setupGetRestoreJobTest(String restoreJobId, String status, boolean includeCompletionTime) throws Exception {
+    private Restore setupGetRestoreJobTest(String restoreJobId, String status, boolean includeCompletionTime)
+            throws Exception {
+        return setupGetRestoreJobTest(restoreJobId, status, includeCompletionTime, null);
+    }
+
+    private Restore setupGetRestoreJobTest(
+            String restoreJobId, String status, boolean includeCompletionTime, String errorMessage) throws Exception {
         Instant creationTime = Instant.now().minusSeconds(300);
         Instant completionTime = Instant.now();
 
@@ -223,6 +229,9 @@ private Restore setupGetRestoreJobTest(String restoreJobId, String status, boole
         if (includeCompletionTime) {
             restoreJob.setCompleteTime(completionTime.getEpochSecond());
         }
+        if (errorMessage != null) {
+            restoreJob.setErrorMessage(errorMessage);
+        }
 
         DescribeRestoreJobs2ResponseBody.DescribeRestoreJobs2ResponseBodyRestoreJobs restoreJobs =
                 new DescribeRestoreJobs2ResponseBody.DescribeRestoreJobs2ResponseBodyRestoreJobs();
@@ -275,6 +284,17 @@ void testGetRestore_Job_Failed() throws Exception {
         assertNotNull(restore.getEndTime());
     }
 
+    @Test
+    void testGetRestore_Job_Failed_WithStatusMessage() throws Exception {
+        String restoreJobId = "restore-job-789";
+        String failureMessage = "Snapshot not found: snapshot-789";
+        Restore restore = setupGetRestoreJobTest(restoreJobId, "FAILED", true, failureMessage);
+        assertEquals(RestoreStatus.FAILED, restore.getStatus());
+        assertEquals(failureMessage, restore.getStatusMessage());
+        assertNotNull(restore.getStartTime());
+        assertNotNull(restore.getEndTime());
+    }
+
     @Test
     void testGetRestore_Job_NotFound() throws Exception {
         String restoreJobId = "non-existent-restore";
diff --git a/dbbackuprestore/dbbackuprestore-aws/src/main/java/com/salesforce/multicloudj/dbbackuprestore/aws/AwsDBBackupRestore.java b/dbbackuprestore/dbbackuprestore-aws/src/main/java/com/salesforce/multicloudj/dbbackuprestore/aws/AwsDBBackupRestore.java
@@ -118,6 +118,11 @@ public String restoreBackup(RestoreRequest request) {
         Map<String, String> metadata = new HashMap<>();
         metadata.put("targetTableName", targetTableName);
 
+        if (StringUtils.isNotBlank(request.getKmsEncryptionKeyId())) {
+            metadata.put("encryptionType", "KMS");
+            metadata.put("kmsMasterKeyArn", request.getKmsEncryptionKeyId());
+        }
+
         StartRestoreJobRequest restoreJobRequest = StartRestoreJobRequest.builder().recoveryPointArn(request.getBackupId()).metadata(metadata).idempotencyToken(UUID.uniqueString()).iamRoleArn(iamRoleArn).resourceType("DynamoDB").build();
 
         StartRestoreJobResponse response = backupClient.startRestoreJob(restoreJobRequest);
@@ -199,6 +204,7 @@ private Restore convertToRestore(DescribeRestoreJobResponse response) {
                 .status(convertRestoreJobStatus(response.status()))
                 .startTime(response.creationDate())
                 .endTime(response.completionDate())
+                .statusMessage(response.statusMessage())
                 .build();
     }
 
diff --git a/dbbackuprestore/dbbackuprestore-aws/src/test/java/com/salesforce/multicloudj/dbbackuprestore/aws/AwsDBBackupRestoreTest.java b/dbbackuprestore/dbbackuprestore-aws/src/test/java/com/salesforce/multicloudj/dbbackuprestore/aws/AwsDBBackupRestoreTest.java
@@ -168,6 +168,37 @@ void testRestoreBackup() {
         assertEquals("DynamoDB", capturedRequest.resourceType());
         assertTrue(capturedRequest.metadata().containsKey("targetTableName"));
         assertEquals("restored-table", capturedRequest.metadata().get("targetTableName"));
+        Assertions.assertFalse(capturedRequest.metadata().containsKey("encryptionType"));
+        Assertions.assertFalse(capturedRequest.metadata().containsKey("kmsMasterKeyArn"));
+    }
+
+    @Test
+    void testRestoreBackupWithKmsEncryption() {
+        String kmsKeyArn = "arn:aws:kms:us-west-2:123456789012:key/abc-def-123";
+        RestoreRequest request = RestoreRequest.builder()
+                .backupId("arn:aws:backup:us-west-2:123456789012:recovery-point:789")
+                .targetResource("restored-table")
+                .roleId(IAM_ROLE_ARN)
+                .kmsEncryptionKeyId(kmsKeyArn)
+                .build();
+
+        StartRestoreJobResponse response = StartRestoreJobResponse.builder()
+                .restoreJobId("restore-job-456")
+                .build();
+
+        when(mockBackupClient.startRestoreJob(any(StartRestoreJobRequest.class)))
+                .thenReturn(response);
+
+        String restoreJobId = dbBackupRestore.restoreBackup(request);
+        assertEquals("restore-job-456", restoreJobId);
+
+        ArgumentCaptor<StartRestoreJobRequest> captor = ArgumentCaptor.forClass(StartRestoreJobRequest.class);
+        verify(mockBackupClient, times(1)).startRestoreJob(captor.capture());
+
+        StartRestoreJobRequest capturedRequest = captor.getValue();
+        assertEquals("restored-table", capturedRequest.metadata().get("targetTableName"));
+        assertEquals("KMS", capturedRequest.metadata().get("encryptionType"));
+        assertEquals(kmsKeyArn, capturedRequest.metadata().get("kmsMasterKeyArn"));
     }
 
     @Test
@@ -184,6 +215,10 @@ void testRestoreBackupMissingRole() {
     }
 
     private Restore setupGetRestoreJobTest(String restoreJobId, RestoreJobStatus status, boolean includeCompletionDate) {
+        return setupGetRestoreJobTest(restoreJobId, status, includeCompletionDate, null);
+    }
+
+    private Restore setupGetRestoreJobTest(String restoreJobId, RestoreJobStatus status, boolean includeCompletionDate, String statusMessage) {
         Instant creationTime = Instant.now().minusSeconds(300);
         Instant completionTime = Instant.now();
 
@@ -197,6 +232,9 @@ private Restore setupGetRestoreJobTest(String restoreJobId, RestoreJobStatus sta
         if (includeCompletionDate) {
             responseBuilder.completionDate(completionTime);
         }
+        if (statusMessage != null) {
+            responseBuilder.statusMessage(statusMessage);
+        }
 
         when(mockBackupClient.describeRestoreJob(any(DescribeRestoreJobRequest.class)))
                 .thenReturn(responseBuilder.build());
@@ -242,6 +280,18 @@ void testGetRestore_Job_Failed() {
         assertNotNull(restore.getEndTime());
     }
 
+    @Test
+    void testGetRestore_Job_Failed_WithStatusMessage() {
+        String restoreJobId = "restore-job-789";
+        String failureMessage = "Table already exists: temp-restore-table";
+        Restore restore = setupGetRestoreJobTest(restoreJobId, RestoreJobStatus.FAILED, true, failureMessage);
+
+        assertEquals(RestoreStatus.FAILED, restore.getStatus());
+        assertEquals(failureMessage, restore.getStatusMessage());
+        assertNotNull(restore.getStartTime());
+        assertNotNull(restore.getEndTime());
+    }
+
     @Test
     void testGetRestore_Job_Pending() {
         String restoreJobId = "restore-job-pending";
diff --git a/dbbackuprestore/dbbackuprestore-client/src/main/java/com/salesforce/multicloudj/dbbackuprestore/driver/Restore.java b/dbbackuprestore/dbbackuprestore-client/src/main/java/com/salesforce/multicloudj/dbbackuprestore/driver/Restore.java
@@ -51,4 +51,10 @@ public class Restore {
      * Will be null if the restore is still in progress.
      */
     private Instant endTime;
+
+    /**
+     * Status or failure message from the restore job.
+     * For failed jobs, this typically contains the provider-specific restore error reason.
+     */
+    private String statusMessage;
 }
diff --git a/dbbackuprestore/dbbackuprestore-client/src/main/java/com/salesforce/multicloudj/dbbackuprestore/driver/RestoreRequest.java b/dbbackuprestore/dbbackuprestore-client/src/main/java/com/salesforce/multicloudj/dbbackuprestore/driver/RestoreRequest.java
@@ -42,4 +42,9 @@ public class RestoreRequest {
      * is required to locate and restore from a backup.
      */
     private String vaultId;
+
+    /**
+     * KMS encryption key identifier for encrypting the restored resource.
+     */
+    private String kmsEncryptionKeyId;
 }
diff --git a/dbbackuprestore/dbbackuprestore-client/src/test/java/com/salesforce/multicloudj/dbbackuprestore/driver/RestoreRequestTest.java b/dbbackuprestore/dbbackuprestore-client/src/test/java/com/salesforce/multicloudj/dbbackuprestore/driver/RestoreRequestTest.java
@@ -16,12 +16,14 @@ void testRestoreRequestBuilder() {
                 .targetResource("restored-table")
                 .roleId("arn:aws:iam::123456789012:role/RestoreRole")
                 .vaultId("vault-abc")
+                .kmsEncryptionKeyId("arn:aws:kms:us-west-2:123456789012:key/abc-def")
                 .build();
 
         assertEquals("backup-123", request.getBackupId());
         assertEquals("restored-table", request.getTargetResource());
         assertEquals("arn:aws:iam::123456789012:role/RestoreRole", request.getRoleId());
         assertEquals("vault-abc", request.getVaultId());
+        assertEquals("arn:aws:kms:us-west-2:123456789012:key/abc-def", request.getKmsEncryptionKeyId());
     }
 
     @Test
@@ -32,11 +34,13 @@ void testRestoreRequestSetters() {
         request.setTargetResource("new-table");
         request.setRoleId("role-123");
         request.setVaultId("vault-xyz");
+        request.setKmsEncryptionKeyId("arn:aws:kms:us-west-2:123456789012:key/xyz");
 
         assertEquals("backup-789", request.getBackupId());
         assertEquals("new-table", request.getTargetResource());
         assertEquals("role-123", request.getRoleId());
         assertEquals("vault-xyz", request.getVaultId());
+        assertEquals("arn:aws:kms:us-west-2:123456789012:key/xyz", request.getKmsEncryptionKeyId());
     }
 
     @Test
@@ -45,12 +49,14 @@ void testRestoreRequestAllArgsConstructor() {
                 "backup-all",
                 "target-table",
                 "role-arn",
-                "vault-id"
+                "vault-id",
+                "arn:aws:kms:us-west-2:123456789012:key/all-args"
         );
 
         assertEquals("backup-all", request.getBackupId());
         assertEquals("target-table", request.getTargetResource());
         assertEquals("role-arn", request.getRoleId());
         assertEquals("vault-id", request.getVaultId());
+        assertEquals("arn:aws:kms:us-west-2:123456789012:key/all-args", request.getKmsEncryptionKeyId());
     }
 }
diff --git a/dbbackuprestore/dbbackuprestore-gcp-firestore/src/main/java/com/salesforce/multicloudj/dbbackuprestore/gcp/FSDBBackupRestore.java b/dbbackuprestore/dbbackuprestore-gcp-firestore/src/main/java/com/salesforce/multicloudj/dbbackuprestore/gcp/FSDBBackupRestore.java
@@ -118,6 +118,16 @@ public String restoreBackup(RestoreRequest request) {
                 .setDatabaseId(targetDBID)
                 .setBackup(request.getBackupId());
 
+        if (StringUtils.isNotBlank(request.getKmsEncryptionKeyId())) {
+            restoreBuilder.setEncryptionConfig(
+                    Database.EncryptionConfig.newBuilder()
+                            .setCustomerManagedEncryption(
+                                    Database.EncryptionConfig.CustomerManagedEncryptionOptions.newBuilder()
+                                            .setKmsKeyName(request.getKmsEncryptionKeyId())
+                                            .build())
+                            .build());
+        }
+
         // Restore is a long-running operation
         OperationFuture<Database, RestoreDatabaseMetadata> operation =
                 firestoreAdminClient.restoreDatabaseAsync(restoreBuilder.build());
@@ -191,9 +201,11 @@ private Restore convertToRestore(Operation operation) {
         Instant endTime = null;
         Instant startTime = null;
 
+        String statusMessage = null;
         if (operation.getDone()) {
             if (operation.hasError()) {
                 status = RestoreStatus.FAILED;
+                statusMessage = operation.getError().getMessage();
             } else {
                 status = RestoreStatus.COMPLETED;
             }
@@ -228,6 +240,7 @@ private Restore convertToRestore(Operation operation) {
                 .status(status)
                 .startTime(startTime)
                 .endTime(endTime)
+                .statusMessage(statusMessage)
                 .build();
     }
 
diff --git a/dbbackuprestore/dbbackuprestore-gcp-firestore/src/test/java/com/salesforce/multicloudj/dbbackuprestore/gcp/FSDBBackupRestoreTest.java b/dbbackuprestore/dbbackuprestore-gcp-firestore/src/test/java/com/salesforce/multicloudj/dbbackuprestore/gcp/FSDBBackupRestoreTest.java
@@ -161,6 +161,37 @@ void testRestoreBackup() throws Exception {
         assertEquals("projects/my-project", capturedRequest.getParent());
     }
 
+    @Test
+    void testRestoreBackup_WithCmekEncryption() throws Exception {
+        String kmsKeyName = "projects/my-project/locations/us-central1/keyRings/my-ring/cryptoKeys/my-key";
+        RestoreRequest request = RestoreRequest.builder()
+                .backupId("projects/my-project/locations/us-central1/backups/backup-789")
+                .targetResource("restored-db")
+                .kmsEncryptionKeyId(kmsKeyName)
+                .build();
+
+        @SuppressWarnings("unchecked")
+        OperationFuture<Database, RestoreDatabaseMetadata> mockOperationFuture =
+                org.mockito.Mockito.mock(OperationFuture.class);
+        when(mockOperationFuture.getName()).thenReturn("operations/restore-job-456");
+
+        when(mockFirestoreClient.restoreDatabaseAsync(any(RestoreDatabaseRequest.class)))
+                .thenReturn(mockOperationFuture);
+
+        String restoreId = dbBackupRestore.restoreBackup(request);
+
+        assertEquals("operations/restore-job-456", restoreId);
+
+        ArgumentCaptor<RestoreDatabaseRequest> captor = ArgumentCaptor.forClass(RestoreDatabaseRequest.class);
+        verify(mockFirestoreClient, times(1)).restoreDatabaseAsync(captor.capture());
+
+        RestoreDatabaseRequest capturedRequest = captor.getValue();
+        assertTrue(capturedRequest.hasEncryptionConfig());
+        assertEquals(
+                kmsKeyName,
+                capturedRequest.getEncryptionConfig().getCustomerManagedEncryption().getKmsKeyName());
+    }
+
     private Restore setupGetRestoreJobTest(String restoreId, boolean isDone, boolean hasError, boolean includeEndTime) {
         Instant startTime = Instant.now().minusSeconds(300);
         Instant endTime = Instant.now();

Original file line number	Diff line number	Diff line change
`@@ -269,6 +269,7 @@ private Restore convertToRestore(DescribeRestoreJobs2ResponseBody.DescribeRestor`
`269`	`269`	`.status(status)`
`270`	`270`	`.startTime(startTime)`
`271`	`271`	`.endTime(endTime)`
	`272`	`+ .statusMessage(restoreJob.getErrorMessage())`
`272`	`273`	`.build();`
`273`	`274`	`}`
`274`	`275`