refactored pre_and_post_upgrade

.github/workflows/test.yml

@@ -36,7 +36,6 @@
 # "examples/null_example",
 # "examples/optional_types",
 # "examples/outgoing_http_requests",
-# "examples/pre_and_post_upgrade",
 # "examples/rejections",
 # "examples/robust_imports",
 # "examples/run_time_errors",
@@ -113,6 +112,7 @@ jobs:
                     "examples/list_of_lists",
                     "examples/management_canister",
                     "examples/manual_reply",
+                    "examples/pre_and_post_upgrade",
                     "examples/primitive_types",
                     "examples/principal",
                     "examples/query",

examples/pre_and_post_upgrade/dfx.json

@@ -4,9 +4,10 @@
             "type": "custom",
             "build": "npx azle pre_and_post_upgrade",
             "root": "src",
-            "ts": "src/pre_and_post_upgrade.ts",
-            "candid": "src/pre_and_post_upgrade.did",
-            "wasm": ".azle/pre_and_post_upgrade/pre_and_post_upgrade.wasm.gz",
+            "ts": "src/index.ts",
+            "candid": "src/index.did",
+            "wasm": ".azle/pre_and_post_upgrade/pre_and_post_upgrade.wasm",
+            "gzip": true,
             "declarations": {
                 "output": "test/dfx_generated/pre_and_post_upgrade",
                 "node_compatibility": true

examples/pre_and_post_upgrade/src/index.did

@@ -0,0 +1,6 @@
+type rec_1 = record {key:text; value:nat64};
+type rec_2 = record {key:text; value:nat64};
+service: () -> {
+    setEntry: (rec_1) -> ();
+    getEntries: () -> (vec rec_2) query;
+}

examples/pre_and_post_upgrade/src/index.ts

@@ -0,0 +1,74 @@
+import {
+    init,
+    nat64,
+    postUpgrade,
+    preUpgrade,
+    query,
+    Record,
+    Service,
+    StableBTreeMap,
+    text,
+    update,
+    Vec,
+    Void
+} from 'azle';
+
+const Entry = Record({
+    key: text,
+    value: nat64
+});
+
+let stableStorage = StableBTreeMap(text, Vec(Entry), 0);
+
+let entries: {
+    [key: string]: nat64;
+} = {};
+
+export default Service({
+    init: init([], () => {
+        console.log('init');
+
+        stableStorage.insert('entries', []);
+    }),
+    postUpgrade: postUpgrade([], () => {
+        console.log('postUpgrade');
+
+        const stableEntriesOpt = stableStorage.get('entries');
+
+        const stableEntries =
+            stableEntriesOpt.length === 0
+                ? stableEntriesOpt
+                : stableEntriesOpt[0];
+
+        entries = stableEntries.reduce((result, entry) => {
+            return {
+                ...result,
+                [entry.key]: entry.value
+            };
+        }, {});
+    }),
+    preUpgrade: preUpgrade(() => {
+        console.log('preUpgrade');
+
+        stableStorage.insert(
+            'entries',
+            Object.entries(entries).map((entry) => {
+                return {
+                    key: entry[0],
+                    value: entry[1]
+                };
+            })
+        );
+    }),
+    setEntry: update([Entry], Void, (entry) => {
+        entries[entry.key] = entry.value;
+    }),
+    getEntries: query([], Vec(Entry), () => {
+        return Object.entries(entries).map((entry) => {
+            return {
+                key: entry[0],
+                value: entry[1]
+            };
+        });
+    })
+});

src/lib_functional/candid/reference/service.ts

@@ -68,6 +68,7 @@ export function Service<T extends ServiceOptions>(
         // }
     }, {});
 
+    // TODO Once types have names we should deduplicate the init and post_upgrade param types
     const candidTypes = Object.values(serviceOptions).reduce(
         (acc: string[], canisterMethodInfo) => {
             return [...acc, ...canisterMethodInfo.candidTypes];
@@ -85,6 +86,26 @@ export function Service<T extends ServiceOptions>(
                   name: initOption[0]
               };
 
+    const postUpgradeOption = Object.entries(serviceOptions).find(
+        ([key, value]) => value.mode === 'postUpgrade'
+    );
+    const postUpgrade =
+        postUpgradeOption === undefined
+            ? undefined
+            : {
+                  name: postUpgradeOption[0]
+              };
+
+    const preUpgradeOption = Object.entries(serviceOptions).find(
+        ([key, value]) => value.mode === 'preUpgrade'
+    );
+    const preUpgrade =
+        preUpgradeOption === undefined
+            ? undefined
+            : {
+                  name: preUpgradeOption[0]
+              };
+
     const heartbeatOption = Object.entries(serviceOptions).find(
         ([key, value]) => value.mode === 'heartbeat'
     );
@@ -199,6 +220,8 @@ export function Service<T extends ServiceOptions>(
 `;
 
     returnFunction.init = init;
+    returnFunction.postUpgrade = postUpgrade;
+    returnFunction.preUpgrade = preUpgrade;
     returnFunction.heartbeat = heartbeat;
     returnFunction.inspect_message = inspectMessage;
     returnFunction.queries = queries;

src/lib_functional/canister_methods/index.ts

@@ -9,11 +9,20 @@ import {
 export * from './heartbeat';
 export * from './init';
 export * from './inspect_message';
+export * from './post_upgrade';
+export * from './pre_upgrade';
 export * from './query';
 export * from './update';
 
 export type CanisterMethodInfo<T extends ReadonlyArray<any>, K> = {
-    mode: 'query' | 'update' | 'init' | 'heartbeat' | 'inspectMessage';
+    mode:
+        | 'query'
+        | 'update'
+        | 'init'
+        | 'heartbeat'
+        | 'inspectMessage'
+        | 'postUpgrade'
+        | 'preUpgrade';
     async: boolean;
     callback?: (...args: any) => any;
     candid: string;
@@ -53,6 +62,11 @@ export function executeMethod(
         return;
     }
 
+    if (mode === 'preUpgrade') {
+        callback();
+        return;
+    }
+
     const decoded = IDL.decode(paramCandid[0] as any, args[0]);
 
     const myDecodedObject = paramCandid[0].map((idl: any, index: any) => {

src/lib_functional/canister_methods/post_upgrade.ts

@@ -0,0 +1,48 @@
+import {
+    handleRecursiveParams,
+    handleRecursiveReturn,
+    newTypesToStingArr
+} from '../../lib_new/method_decorators';
+import { Callback, CanisterMethodInfo, executeMethod } from '.';
+import { TypeMapping } from '../candid';
+import { Void } from '../../lib_new';
+
+export function postUpgrade<
+    const Params extends ReadonlyArray<any>,
+    GenericCallback extends Callback<Params, Void>
+>(
+    paramsIdls: Params,
+    callback?: Awaited<ReturnType<GenericCallback>> extends TypeMapping<Void>
+        ? GenericCallback
+        : never
+): CanisterMethodInfo<Params, Void> {
+    const paramCandid = handleRecursiveParams(paramsIdls as any);
+    const returnCandid = handleRecursiveReturn(Void as any, paramCandid[2]);
+
+    const finalCallback =
+        callback === undefined
+            ? undefined
+            : (...args: any[]) => {
+                  executeMethod(
+                      'postUpgrade',
+                      paramCandid,
+                      returnCandid,
+                      args,
+                      callback,
+                      paramsIdls as any,
+                      Void,
+                      false
+                  );
+              };
+
+    return {
+        mode: 'postUpgrade',
+        callback: finalCallback,
+        candid: paramCandid[1].join(', '),
+        candidTypes: newTypesToStingArr(returnCandid[2]),
+        paramsIdls: paramsIdls as any,
+        returnIdl: Void,
+        async: false,
+        guard: undefined
+    };
+}

src/lib_functional/canister_methods/pre_upgrade.ts

@@ -0,0 +1,31 @@
+import { isAsync } from '../../lib_new/method_decorators';
+import { CanisterMethodInfo, executeMethod } from '.';
+import { Void } from '../../lib_new';
+
+export function preUpgrade(
+    callback: () => void | Promise<void>
+): CanisterMethodInfo<[], Void> {
+    const finalCallback = (...args: any[]) => {
+        executeMethod(
+            'preUpgrade',
+            undefined,
+            undefined,
+            args,
+            callback,
+            [],
+            Void,
+            false
+        );
+    };
+
+    return {
+        mode: 'preUpgrade',
+        callback: finalCallback,
+        candid: '',
+        candidTypes: [],
+        paramsIdls: [],
+        returnIdl: Void,
+        async: isAsync(callback),
+        guard: undefined
+    };
+}