From 3d80d9d92dadf6755f8860ffa8c50715754169aa Mon Sep 17 00:00:00 2001
From: Marri Harish <harishmarri551@gmail.com>
Date: Mon, 26 Aug 2024 15:50:15 +0530
Subject: [PATCH] fix: add wasm state-sync snapshot extension (#184)

* fix: enable wasm snapshot extension

* add snapshot and pruning cmds
---
 app/app.go                   | 10 ++++++++++
 cmd/omniflixhubd/cmd/root.go | 11 +++++++----
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/app/app.go b/app/app.go
index c63235b..4af05de 100644
--- a/app/app.go
+++ b/app/app.go
@@ -285,6 +285,16 @@ func NewOmniFlixApp(
 	app.SetPrecommiter(app.PreCommitter)
 	app.SetPrepareCheckStater(app.PrepareCheckStater)
 
+	// Register snapshot extensions to enable state-sync for wasm.
+	if manager := app.SnapshotManager(); manager != nil {
+		err := manager.RegisterExtensions(
+			wasmkeeper.NewWasmSnapshotter(app.CommitMultiStore(), &app.WasmKeeper),
+		)
+		if err != nil {
+			panic(fmt.Errorf("failed to register snapshot extension: %s", err))
+		}
+	}
+
 	if loadLatest {
 		if err := app.LoadLatestVersion(); err != nil {
 			tmos.Exit(fmt.Sprintf("failed to load latest version: %s", err))
diff --git a/cmd/omniflixhubd/cmd/root.go b/cmd/omniflixhubd/cmd/root.go
index 99781dc..bf74b08 100644
--- a/cmd/omniflixhubd/cmd/root.go
+++ b/cmd/omniflixhubd/cmd/root.go
@@ -27,7 +27,9 @@ import (
 	"github.com/cosmos/cosmos-sdk/client"
 	"github.com/cosmos/cosmos-sdk/client/flags"
 	"github.com/cosmos/cosmos-sdk/client/keys"
+	"github.com/cosmos/cosmos-sdk/client/pruning"
 	"github.com/cosmos/cosmos-sdk/client/rpc"
+	"github.com/cosmos/cosmos-sdk/client/snapshot"
 	runtimeservices "github.com/cosmos/cosmos-sdk/runtime/services"
 	"github.com/cosmos/cosmos-sdk/server"
 	servertypes "github.com/cosmos/cosmos-sdk/server/types"
@@ -130,6 +132,9 @@ func initCometBftConfig() *tmcfg.Config {
 }
 
 func initRootCmd(rootCmd *cobra.Command, encodingConfig params.EncodingConfig, tempApp *app.OmniFlixApp) {
+	ac := appCreator{
+		encCfg: encodingConfig,
+	}
 	rootCmd.AddCommand(
 		genutilcli.InitCmd(tempApp.ModuleBasics, app.DefaultNodeHome),
 		tmcli.NewCompletionCmd(rootCmd, true),
@@ -150,12 +155,10 @@ func initRootCmd(rootCmd *cobra.Command, encodingConfig params.EncodingConfig, t
 		genutilcli.ValidateGenesisCmd(tempApp.ModuleBasics),
 		tmcli.NewCompletionCmd(rootCmd, true),
 		addDebugCommands(debug.Cmd()),
+		pruning.Cmd(ac.newApp, app.DefaultNodeHome),
+		snapshot.Cmd(ac.newApp),
 	)
 
-	ac := appCreator{
-		encCfg: encodingConfig,
-	}
-
 	server.AddCommands(rootCmd, app.DefaultNodeHome, ac.newApp, ac.appExport, addModuleInitFlags)
 
 	// add keybase, auxiliary RPC, query, and tx child commands