firebase · TomasMorton · Aug 30, 2024 · Jul 17, 2024 · Jul 17, 2024 · Jul 17, 2024
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,2 +1,3 @@
 - Improve Firebase Data Connect postgres security by granting fine grained SQL privileges to the users the need it. (#7578)
 - Remove `dataconnect:sql:migrate` command hard dependency on 'roles/cloudsql.admin'. (#7578)
+- Add support for setting the encryption configuration of restored firestore databases (#7483)
diff --git a/src/commands/firestore-databases-restore.ts b/src/commands/firestore-databases-restore.ts
@@ -7,14 +7,25 @@ import { logger } from "../logger";
 import { requirePermissions } from "../requirePermissions";
 import { Emulators } from "../emulator/types";
 import { warnEmulatorNotSupported } from "../emulator/commandUtils";
-import { FirestoreOptions } from "../firestore/options";
+import { EncryptionType, FirestoreOptions } from "../firestore/options";
 import { PrettyPrint } from "../firestore/pretty-print";
 import { FirebaseError } from "../error";
 
 export const command = new Command("firestore:databases:restore")
   .description("Restore a Firestore database in your Firebase project.")
   .option("-d, --database <databaseID>", "ID of the database to restore into")
   .option("-b, --backup <backup>", "Backup from which to restore")
+  .option(
+    "-e, --encryption-type <encryptionType>",
+    `Encryption method of the restored database; one of ${EncryptionType.USE_SOURCE_ENCRYPTION} (default), ` +
+      `${EncryptionType.CUSTOMER_MANAGED_ENCRYPTION}, ${EncryptionType.GOOGLE_DEFAULT_ENCRYPTION}`,
+  )
+  // TODO(b/356137854): Remove allowlist only message once feature is public GA.
+  .option(
+    "-k, --kms-key-name <kmsKeyName>",
+    "Resource ID of the Cloud KMS key to encrypt the restored database. This " +
+      "feature is allowlist only in initial launch.",
+  )
   .before(requirePermissions, ["datastore.backups.restoreDatabase"])
   .before(warnEmulatorNotSupported, Emulators.FIRESTORE)
   .action(async (options: FirestoreOptions) => {
@@ -32,10 +43,33 @@ export const command = new Command("firestore:databases:restore")
     }
     const backupName = options.backup;
 
+    let encryptionConfig: types.EncryptionConfig | undefined = undefined;
+    switch (options.encryptionType) {
+      case EncryptionType.GOOGLE_DEFAULT_ENCRYPTION:
+        throwIfKmsKeyNameIsSet(options.kmsKeyName);
+        encryptionConfig = { googleDefaultEncryption: {} };
+        break;
+      case EncryptionType.USE_SOURCE_ENCRYPTION:
+        throwIfKmsKeyNameIsSet(options.kmsKeyName);
+        encryptionConfig = { useSourceEncryption: {} };
+        break;
+      case EncryptionType.CUSTOMER_MANAGED_ENCRYPTION:
+        encryptionConfig = {
+          customerManagedEncryption: { kmsKeyName: getKmsKeyOrThrow(options.kmsKeyName) },
+        };
+        break;
+      case undefined:
+        throwIfKmsKeyNameIsSet(options.kmsKeyName);
+        break;
+      default:
+        throw new FirebaseError(`Invalid value for flag --encryption-type. ${helpCommandText}`);
+    }
+
     const databaseResp: types.DatabaseResp = await api.restoreDatabase(
       options.project,
       databaseId,
       backupName,
+      encryptionConfig,
     );
 
     if (options.json) {
@@ -55,4 +89,22 @@ export const command = new Command("firestore:databases:restore")
     }
 
     return databaseResp;
+
+    function throwIfKmsKeyNameIsSet(kmsKeyName: string | undefined): void {
+      if (kmsKeyName) {
+        throw new FirebaseError(
+          "--kms-key-name can only be set when specifying an --encryption-type " +
+            `of ${EncryptionType.CUSTOMER_MANAGED_ENCRYPTION}.`,
+        );
+      }
+    }
+
+    function getKmsKeyOrThrow(kmsKeyName: string | undefined): string {
+      if (kmsKeyName) return kmsKeyName;
+
+      throw new FirebaseError(
+        "--kms-key-name must be provided when specifying an --encryption-type " +
+          `of ${EncryptionType.CUSTOMER_MANAGED_ENCRYPTION}.`,
+      );
+    }
   });
diff --git a/src/firestore/api-types.ts b/src/firestore/api-types.ts
@@ -160,6 +160,7 @@ export interface DatabaseResp {
 export interface RestoreDatabaseReq {
   databaseId: string;
   backup: string;
+  encryptionConfig?: EncryptionConfig;
 }
 
 export enum RecurrenceType {
@@ -171,3 +172,14 @@ export interface CmekConfig {
   kmsKeyName: string;
   activeKeyVersion?: string[];
 }
+
+type UseGoogleDefaultEncryption = { googleDefaultEncryption: Record<string, never> };
+type UseSourceEncryption = { useSourceEncryption: Record<string, never> };
+type UseCustomerManagedEncryption = { customerManagedEncryption: CustomerManagedEncryptionOptions };
+type CustomerManagedEncryptionOptions = {
+  kmsKeyName: string;
+};
+export type EncryptionConfig =
+  | UseCustomerManagedEncryption
+  | UseSourceEncryption
+  | UseGoogleDefaultEncryption;
diff --git a/src/firestore/api.ts b/src/firestore/api.ts
@@ -695,16 +695,19 @@ export class FirestoreApi {
    * @param project the Firebase project id.
    * @param databaseId the ID of the Firestore Database to be restored into
    * @param backupName Name of the backup from which to restore
+   * @param encryptionConfig the encryption configuration of the restored database
    */
   async restoreDatabase(
     project: string,
     databaseId: string,
     backupName: string,
+    encryptionConfig?: types.EncryptionConfig,
   ): Promise<types.DatabaseResp> {
     const url = `/projects/${project}/databases:restore`;
     const payload: types.RestoreDatabaseReq = {
       databaseId,
       backup: backupName,
+      encryptionConfig: encryptionConfig,
     };
     const options = { queryParams: { databaseId: databaseId } };
     const res = await this.apiClient.post<

diff --git a/src/firestore/options.ts b/src/firestore/options.ts
@@ -18,7 +18,6 @@ export interface FirestoreOptions extends Options {
   type?: types.DatabaseType;
   deleteProtection?: types.DatabaseDeleteProtectionStateOption;
   pointInTimeRecoveryEnablement?: types.PointInTimeRecoveryEnablementOption;
-  kmsKeyName?: string;
 
   // backup schedules
   backupSchedule?: string;
@@ -28,4 +27,14 @@ export interface FirestoreOptions extends Options {
 
   // backups
   backup?: string;
+
+  // CMEK
+  encryptionType?: EncryptionType;
+  kmsKeyName?: string;
+}
+
+export enum EncryptionType {
+  CUSTOMER_MANAGED_ENCRYPTION = "CUSTOMER_MANAGED_ENCRYPTION",
+  USE_SOURCE_ENCRYPTION = "USE_SOURCE_ENCRYPTION",
+  GOOGLE_DEFAULT_ENCRYPTION = "GOOGLE_DEFAULT_ENCRYPTION",
 }