feat: take snapshot on upgrade

packages/admin/src/api/ic.api.ts

@@ -4,8 +4,11 @@ import {
   type InstallChunkedCodeParams,
   type UploadChunkParams,
   ICManagementCanister,
-  InstallCodeParams
+  InstallCodeParams,
+  list_canister_snapshots_result,
+  snapshot_id
 } from '@dfinity/ic-management';
+import type {take_canister_snapshot_result} from '@dfinity/ic-management/dist/candid/ic-management';
 import type {CanisterStatusResponse} from '@dfinity/ic-management/dist/types/types/ic-management.responses';
 import {Principal} from '@dfinity/principal';
 import {assertNonNullish} from '@junobuild/utils';
@@ -167,3 +170,36 @@ export const canisterMetadata = async ({
 
   return result.get(path);
 };
+
+export const listCanisterSnapshots = async ({
+  actor,
+  canisterId
+}: {
+  actor: ActorParameters;
+  canisterId: Principal;
+}): Promise<list_canister_snapshots_result> => {
+  const agent = await useOrInitAgent(actor);
+
+  const {listCanisterSnapshots} = ICManagementCanister.create({
+    agent
+  });
+
+  return listCanisterSnapshots({canisterId});
+};
+
+export const takeCanisterSnapshot = async ({
+  actor,
+  ...rest
+}: {
+  actor: ActorParameters;
+  canisterId: Principal;
+  snapshotId?: snapshot_id;
+}): Promise<take_canister_snapshot_result> => {
+  const agent = await useOrInitAgent(actor);
+
+  const {takeCanisterSnapshot} = ICManagementCanister.create({
+    agent
+  });
+
+  return takeCanisterSnapshot(rest);
+};

packages/admin/src/handlers/upgrade.handlers.ts

@@ -1,5 +1,11 @@
 import {fromNullable, isNullish} from '@junobuild/utils';
-import {canisterStart, canisterStatus, canisterStop} from '../api/ic.api';
+import {
+  canisterStart,
+  canisterStatus,
+  canisterStop,
+  listCanisterSnapshots,
+  takeCanisterSnapshot
+} from '../api/ic.api';
 import {SIMPLE_INSTALL_MAX_WASM_SIZE} from '../constants/upgrade.constants';
 import {UpgradeCodeUnchangedError} from '../errors/upgrade.errors';
 import {UpgradeCodeParams, UpgradeCodeProgressStep} from '../types/upgrade.types';
@@ -13,6 +19,7 @@ export const upgrade = async ({
   canisterId,
   actor,
   onProgress,
+  takeSnapshot = true,
   ...rest
 }: UpgradeCodeParams & {reset?: boolean}) => {
   // 1. We verify that the code to be installed is different from the code already deployed. If the codes are identical, we skip the installation.
@@ -25,11 +32,16 @@ export const upgrade = async ({
   await execute({fn: stop, onProgress, step: UpgradeCodeProgressStep.StoppingCanister});
 
   try {
-    // 3. Upgrading code: If the WASM is > 2MB we proceed with the chunked installation otherwise we use the original single chunk installation method.
+    // 3. We take a snapshot - create a backup - unless the dev opted-out
+    const snapshot = async () =>
+      takeSnapshot ? await createSnapshot({canisterId, actor}) : Promise.resolve();
+    await execute({fn: snapshot, onProgress, step: UpgradeCodeProgressStep.TakingSnapshot});
+
+    // 4. Upgrading code: If the WASM is > 2MB we proceed with the chunked installation otherwise we use the original single chunk installation method.
     const upgrade = async () => await upgradeCode({wasmModule, canisterId, actor, ...rest});
     await execute({fn: upgrade, onProgress, step: UpgradeCodeProgressStep.UpgradingCode});
   } finally {
-    // 4. We restart the canister to finalize the process.
+    // 5. We restart the canister to finalize the process.
     const restart = async () => await canisterStart({canisterId, actor});
     await execute({fn: restart, onProgress, step: UpgradeCodeProgressStep.RestartingCanister});
   }
@@ -111,3 +123,13 @@ const upgradeCode = async ({
   const fn = upgradeType() === 'chunked' ? upgradeChunkedCode : upgradeSingleChunkCode;
   await fn({wasmModule, canisterId, actor, ...rest});
 };
+
+const createSnapshot = async (params: Pick<UpgradeCodeParams, 'canisterId' | 'actor'>) => {
+  const snapshots = await listCanisterSnapshots(params);
+
+  // TODO: currently only one snapshot per canister is supported on the IC
+  await takeCanisterSnapshot({
+    ...params,
+    snapshotId: snapshots?.[0]?.id
+  });
+};

packages/admin/src/services/mission-control.upgrade.services.ts

@@ -12,6 +12,7 @@ import {encoreIDLUser} from '../utils/idl.utils';
  * @param {MissionControlParameters} params.missionControl - The Mission Control parameters, including the actor and mission control ID.
  * @param {Uint8Array} params.wasmModule - The WASM module to be installed during the upgrade.
  * @param {boolean} [params.preClearChunks] - Optional. Whether to force clearing chunks before uploading a chunked WASM module. Recommended if the WASM exceeds 2MB.
+ * @param {boolean} [params.takeSnapshot=true] - Optional. Whether to take a snapshot before performing the upgrade. Defaults to true.
  * @param {function} [params.onProgress] - Optional. Callback function to report progress during the upgrade process.
  * @throws {Error} Will throw an error if the mission control principal is not defined.
  * @returns {Promise<void>} Resolves when the upgrade process is complete.
@@ -21,7 +22,10 @@ export const upgradeMissionControl = async ({
   ...rest
 }: {
   missionControl: MissionControlParameters;
-} & Pick<UpgradeCodeParams, 'wasmModule' | 'preClearChunks' | 'onProgress'>): Promise<void> => {
+} & Pick<
+  UpgradeCodeParams,
+  'wasmModule' | 'preClearChunks' | 'takeSnapshot' | 'onProgress'
+>): Promise<void> => {
   const user = await getUser({missionControl});
 
   const {missionControlId, ...actor} = missionControl;

packages/admin/src/services/orbiter.upgrade.services.ts

@@ -14,6 +14,7 @@ import {encodeIDLControllers} from '../utils/idl.utils';
  * @param {Uint8Array} params.wasmModule - The WASM module to be installed during the upgrade.
  * @param {boolean} [params.reset=false] - Optional. Indicates whether to reset the Orbiter (reinstall) instead of performing an upgrade. Defaults to `false`.
  * @param {boolean} [params.preClearChunks] - Optional. Forces clearing existing chunks before uploading a chunked WASM module. Recommended if the WASM exceeds 2MB.
+ * @param {boolean} [params.takeSnapshot=true] - Optional. Whether to take a snapshot before performing the upgrade. Defaults to true.
  * @param {function} [params.onProgress] - Optional. Callback function to track progress during the upgrade process.
  * @throws {Error} Will throw an error if the Orbiter principal is not defined.
  * @returns {Promise<void>} Resolves when the upgrade process is complete.
@@ -28,7 +29,7 @@ export const upgradeOrbiter = async ({
   reset?: boolean;
 } & Pick<
   UpgradeCodeParams,
-  'wasmModule' | 'missionControlId' | 'preClearChunks' | 'onProgress'
+  'wasmModule' | 'missionControlId' | 'preClearChunks' | 'takeSnapshot' | 'onProgress'
 >): Promise<void> => {
   const {orbiterId, ...actor} = orbiter;
 

packages/admin/src/services/satellite.upgrade.services.ts

@@ -22,6 +22,7 @@ import {encodeIDLControllers} from '../utils/idl.utils';
  * @param {boolean} params.deprecatedNoScope - Indicates whether the upgrade is deprecated and has no scope.
  * @param {boolean} [params.reset=false] - Optional. Specifies whether to reset the satellite (reinstall) instead of performing an upgrade. Defaults to `false`.
  * @param {boolean} [params.preClearChunks] - Optional. Forces clearing existing chunks before uploading a chunked WASM module. Recommended if the WASM exceeds 2MB.
+ * @param {boolean} [params.takeSnapshot=true] - Optional. Whether to take a snapshot before performing the upgrade. Defaults to true.
  * @param {function} [params.onProgress] - Optional. Callback function to track progress during the upgrade process.
  * @throws {Error} Will throw an error if the satellite parameters are invalid or missing required fields.
  * @returns {Promise<void>} Resolves when the upgrade process is complete.
@@ -39,7 +40,7 @@ export const upgradeSatellite = async ({
   reset?: boolean;
 } & Pick<
   UpgradeCodeParams,
-  'wasmModule' | 'missionControlId' | 'preClearChunks' | 'onProgress'
+  'wasmModule' | 'missionControlId' | 'preClearChunks' | 'takeSnapshot' | 'onProgress'
 >): Promise<void> => {
   const {satelliteId, ...actor} = satellite;
 

packages/admin/src/types/upgrade.types.ts

@@ -5,6 +5,7 @@ import {ActorParameters} from './actor.types';
 export enum UpgradeCodeProgressStep {
   AssertingExistingCode,
   StoppingCanister,
+  TakingSnapshot,
   UpgradingCode,
   RestartingCanister
 }
@@ -24,5 +25,6 @@ export interface UpgradeCodeParams {
   arg: Uint8Array;
   mode: canister_install_mode;
   preClearChunks?: boolean;
+  takeSnapshot?: boolean;
   onProgress?: (progress: UpgradeCodeProgress) => void;
 }