fix decorators for init and post upgrade

demergent-labs · Dec 30, 2024 · cc43919 · cc43919
1 parent 2052715
commit cc43919
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 67 deletions.
diff --git a/src/lib/stable/canister_methods/index.ts b/src/lib/stable/canister_methods/index.ts
@@ -1,6 +1,5 @@
 import { IDL } from '@dfinity/candid';
 
-import { Method, MethodMeta } from '../../../build/stable/utils/types';
 import { handleUncaughtError } from '../error';
 import {
     CanisterMethodMode,
@@ -14,7 +13,6 @@ export type MethodType<This, Args extends any[], Return> = (
 
 export function decoratorArgumentsHandler<This, Args extends any[], Return>(
     canisterMethodMode: CanisterMethodMode,
-    methodMetaKey: keyof MethodMeta,
     param1?: MethodType<This, Args, Return> | IDL.Type[],
     param2?: ClassMethodDecoratorContext | IDL.Type,
     param3?: { composite?: boolean; manual?: boolean }
@@ -26,7 +24,6 @@ export function decoratorArgumentsHandler<This, Args extends any[], Return>(
 
         return decoratorImplementation(
             canisterMethodMode,
-            methodMetaKey,
             originalMethod,
             context
         );
@@ -43,7 +40,6 @@ export function decoratorArgumentsHandler<This, Args extends any[], Return>(
         ): MethodType<This, Args, Return> => {
             return decoratorImplementation(
                 canisterMethodMode,
-                methodMetaKey,
                 originalMethod,
                 context,
                 paramIdlTypes,
@@ -56,7 +52,6 @@ export function decoratorArgumentsHandler<This, Args extends any[], Return>(
 
 function decoratorImplementation<This, Args extends any[], Return>(
     canisterMethodMode: CanisterMethodMode,
-    methodMetaKey: keyof MethodMeta,
     originalMethod: MethodType<This, Args, Return>,
     context: ClassMethodDecoratorContext,
     paramIdlTypes?: IDL.Type[],
@@ -68,29 +63,13 @@ function decoratorImplementation<This, Args extends any[], Return>(
     const index = globalThis._azleCanisterMethodsIndex++;
     const indexString = index.toString();
 
-    if (canisterMethodMode !== 'heartbeat') {
-        if (Array.isArray(globalThis._azleMethodMeta[methodMetaKey])) {
-            if (canisterMethodMode === 'query') {
-                (globalThis._azleMethodMeta[methodMetaKey] as Method[]).push({
-                    name,
-                    index,
-                    composite: options?.composite ?? false
-                });
-            } else {
-                (globalThis._azleMethodMeta[methodMetaKey] as Method[]).push({
-                    name,
-                    index
-                });
-            }
-        } else {
-            (globalThis._azleMethodMeta[methodMetaKey] as Method) = {
-                name,
-                index
-            };
-        }
-    }
-
     if (canisterMethodMode === 'query') {
+        globalThis._azleMethodMeta.queries?.push({
+            name,
+            index,
+            composite: options?.composite ?? false
+        });
+
         globalThis._azleCanisterMethodIdlTypes[name] = IDL.Func(
             paramIdlTypes ?? [],
             returnIdlType === undefined ? [] : [returnIdlType],
@@ -99,25 +78,36 @@ function decoratorImplementation<This, Args extends any[], Return>(
     }
 
     if (canisterMethodMode === 'update') {
+        globalThis._azleMethodMeta.updates?.push({
+            name,
+            index
+        });
+
         globalThis._azleCanisterMethodIdlTypes[name] = IDL.Func(
             paramIdlTypes ?? [],
             returnIdlType === undefined ? [] : [returnIdlType]
         );
     }
 
     if (canisterMethodMode === 'init') {
-        globalThis._azleCanisterMethodIdlTypes[name] = IDL.Func(
-            paramIdlTypes ?? [],
-            [],
-            ['init']
+        globalThis._azleMethodMeta.init = {
+            name,
+            index
+        };
+
+        globalThis._azleInitAndPostUpgradeIdlTypes.push(
+            IDL.Func(paramIdlTypes ?? [], [], ['init'])
         );
     }
 
     if (canisterMethodMode === 'postUpgrade') {
-        globalThis._azleCanisterMethodIdlTypes[name] = IDL.Func(
-            paramIdlTypes ?? [],
-            [],
-            ['post_upgrade']
+        globalThis._azleMethodMeta.post_upgrade = {
+            name,
+            index
+        };
+
+        globalThis._azleInitAndPostUpgradeIdlTypes.push(
+            IDL.Func(paramIdlTypes ?? [], [], ['post_upgrade'])
         );
     }
 

diff --git a/src/lib/stable/canister_methods/init.ts b/src/lib/stable/canister_methods/init.ts
@@ -8,11 +8,7 @@ export function init<This, Args extends any[], Return>(
 ): MethodType<This, Args, Return>;
 
 export function init<This, Args extends any[], Return>(
-    paramIdlTypes?: IDL.Type[],
-    returnIdlType?: IDL.Type,
-    options?: {
-        manual?: boolean;
-    }
+    paramIdlTypes?: IDL.Type[]
 ): (
     originalMethod: MethodType<This, Args, Return>,
     context: ClassMethodDecoratorContext
@@ -23,5 +19,5 @@ export function init<This, Args extends any[], Return>(
     param2?: ClassMethodDecoratorContext | IDL.Type,
     param3?: { manual?: boolean }
 ): any {
-    return decoratorArgumentsHandler('init', 'init', param1, param2, param3);
+    return decoratorArgumentsHandler('init', param1, param2, param3);
 }
diff --git a/src/lib/stable/canister_methods/post_upgrade.ts b/src/lib/stable/canister_methods/post_upgrade.ts
@@ -8,11 +8,7 @@ export function postUpgrade<This, Args extends any[], Return>(
 ): MethodType<This, Args, Return>;
 
 export function postUpgrade<This, Args extends any[], Return>(
-    paramIdlTypes?: IDL.Type[],
-    returnIdlType?: IDL.Type,
-    options?: {
-        manual?: boolean;
-    }
+    paramIdlTypes?: IDL.Type[]
 ): (
     originalMethod: MethodType<This, Args, Return>,
     context: ClassMethodDecoratorContext
@@ -23,11 +19,5 @@ export function postUpgrade<This, Args extends any[], Return>(
     param2?: ClassMethodDecoratorContext | IDL.Type,
     param3?: { manual?: boolean }
 ): any {
-    return decoratorArgumentsHandler(
-        'postUpgrade',
-        'post_upgrade',
-        param1,
-        param2,
-        param3
-    );
+    return decoratorArgumentsHandler('postUpgrade', param1, param2, param3);
 }
diff --git a/src/lib/stable/canister_methods/query.ts b/src/lib/stable/canister_methods/query.ts
@@ -24,11 +24,5 @@ export function query<This, Args extends any[], Return>(
     param2?: ClassMethodDecoratorContext | IDL.Type,
     param3?: { manual?: boolean }
 ): any {
-    return decoratorArgumentsHandler(
-        'query',
-        'queries',
-        param1,
-        param2,
-        param3
-    );
+    return decoratorArgumentsHandler('query', param1, param2, param3);
 }
diff --git a/src/lib/stable/canister_methods/update.ts b/src/lib/stable/canister_methods/update.ts
@@ -24,11 +24,5 @@ export function update<This, Args extends any[], Return>(
     param2?: ClassMethodDecoratorContext | IDL.Type,
     param3?: { manual?: boolean }
 ): any {
-    return decoratorArgumentsHandler(
-        'update',
-        'updates',
-        param1,
-        param2,
-        param3
-    );
+    return decoratorArgumentsHandler('update', param1, param2, param3);
 }