feat: added upgradeAndCall fn for flexibility

push-protocol · Nov 20, 2024 · b612cdb · b612cdb
1 parent e5ff018
commit b612cdb
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/src/lib.cairo b/src/lib.cairo
@@ -25,6 +25,7 @@ pub mod PushComm {
     use openzeppelin::access::ownable::OwnableComponent;
     use openzeppelin::upgrades::UpgradeableComponent;
     use openzeppelin::upgrades::interface::IUpgradeable;
+    use openzeppelin::upgrades::interface::IUpgradeAndCall;
 
     component!(path: OwnableComponent, storage: ownable, event: OwnableEvent);
     component!(path: UpgradeableComponent, storage: upgradeable, event: UpgradeableEvent);
@@ -170,6 +171,15 @@ pub mod PushComm {
         fn upgrade(ref self: ContractState, new_class_hash: ClassHash) {
             self.ownable.assert_only_owner();
             self.upgradeable.upgrade(new_class_hash);
+        }   
+    }
+
+    #[abi(embed_v0)]
+    impl UpgradeAndCallImpl of IUpgradeAndCall<ContractState> {
+        fn upgrade_and_call(ref self: ContractState, new_class_hash: ClassHash, selector: felt252,
+            calldata: Span<felt252>) -> Span<felt252> {
+            self.ownable.assert_only_owner();
+            self.upgradeable.upgrade_and_call(new_class_hash, selector, calldata)
         }
     }