From ffe37f841532a2e17f681f85108cd76e2be6bb56 Mon Sep 17 00:00:00 2001
From: Tom <54514587+GAtom22@users.noreply.github.com>
Date: Tue, 13 Aug 2024 07:03:21 -0300
Subject: [PATCH] feat(confix): allow customization of migration plan (#21202)

Co-authored-by: Hieu Vu <72878483+hieuvubk@users.noreply.github.com>
(cherry picked from commit 1d7f891ea96f69c390aee99cee93fc05f0b9833f)
---
 tools/confix/CHANGELOG.md  |  4 ++++
 tools/confix/migrations.go | 22 +++++++++++++++-------
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/tools/confix/CHANGELOG.md b/tools/confix/CHANGELOG.md
index e044193f3029..531daaf7b5c1 100644
--- a/tools/confix/CHANGELOG.md
+++ b/tools/confix/CHANGELOG.md
@@ -31,6 +31,10 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ## [Unreleased]
 
+### Features
+
+* (confix) [#21202](https://github.com/cosmos/cosmos-sdk/pull/21202) Allow customization of migration `PlanBuilder`.
+
 ## [v0.1.1](https://github.com/cosmos/cosmos-sdk/releases/tag/tools/confix/v0.1.1) - 2023-12-11
 
 * [#18496](https://github.com/cosmos/cosmos-sdk/pull/18496) Remove invalid non SDK config from app.toml migration templates.
diff --git a/tools/confix/migrations.go b/tools/confix/migrations.go
index 737ed6b822cb..4a500898bc04 100644
--- a/tools/confix/migrations.go
+++ b/tools/confix/migrations.go
@@ -21,14 +21,18 @@ const (
 // MigrationMap defines a mapping from a version to a transformation plan.
 type MigrationMap map[string]func(from *tomledit.Document, to, planType string) (transform.Plan, *tomledit.Document)
 
+// loadDestConfigFile is the function signature to load the destination version
+// configuration toml file.
+type loadDestConfigFile func(to, planType string) (*tomledit.Document, error)
+
 var Migrations = MigrationMap{
 	"v0.45": NoPlan, // Confix supports only the current supported SDK version. So we do not support v0.44 -> v0.45.
-	"v0.46": PlanBuilder,
-	"v0.47": PlanBuilder,
-	"v0.50": PlanBuilder,
-	"v0.52": PlanBuilder,
+	"v0.46": defaultPlanBuilder,
+	"v0.47": defaultPlanBuilder,
+	"v0.50": defaultPlanBuilder,
+	"v0.52": defaultPlanBuilder,
 	"v2":    V2PlanBuilder,
-	// "v0.xx.x": PlanBuilder, // add specific migration in case of configuration changes in minor versions
+	// "v0.xx.x": defaultPlanBuilder, // add specific migration in case of configuration changes in minor versions
 }
 
 type v2KeyChangesMap map[string][]string
@@ -53,12 +57,16 @@ var v2KeyChanges = v2KeyChangesMap{
 	// Add other key mappings as needed
 }
 
+func defaultPlanBuilder(from *tomledit.Document, to, planType string) (transform.Plan, *tomledit.Document) {
+	return PlanBuilder(from, to, planType, LoadLocalConfig)
+}
+
 // PlanBuilder is a function that returns a transformation plan for a given diff between two files.
-func PlanBuilder(from *tomledit.Document, to, planType string) (transform.Plan, *tomledit.Document) {
+func PlanBuilder(from *tomledit.Document, to, planType string, loadFn loadDestConfigFile) (transform.Plan, *tomledit.Document) {
 	plan := transform.Plan{}
 	deletedSections := map[string]bool{}
 
-	target, err := LoadLocalConfig(to, planType)
+	target, err := loadFn(to, planType)
 	if err != nil {
 		panic(fmt.Errorf("failed to parse file: %w. This file should have been valid", err))
 	}