From 39ccd99efbe12e7497e0780edaba6f2047a3be6e Mon Sep 17 00:00:00 2001
From: Benjamin Wang <benjamin.ahrtr@gmail.com>
Date: Thu, 16 Jan 2025 13:24:38 +0000
Subject: [PATCH] Minor refactor on the etcdutl migrate command

Signed-off-by: Benjamin Wang <benjamin.ahrtr@gmail.com>
---
 etcdutl/etcdutl/migrate_command.go | 55 ++++++++++++++++++------------
 1 file changed, 33 insertions(+), 22 deletions(-)

diff --git a/etcdutl/etcdutl/migrate_command.go b/etcdutl/etcdutl/migrate_command.go
index a265c7baae7..0a114268524 100644
--- a/etcdutl/etcdutl/migrate_command.go
+++ b/etcdutl/etcdutl/migrate_command.go
@@ -74,8 +74,9 @@ func (o *migrateOptions) AddFlags(cmd *cobra.Command) {
 
 func (o *migrateOptions) Config() (*migrateConfig, error) {
 	c := &migrateConfig{
-		force: o.force,
-		lg:    GetLogger(),
+		force:   o.force,
+		dataDir: o.dataDir,
+		lg:      GetLogger(),
 	}
 	var err error
 	dotCount := strings.Count(o.targetVersion, ".")
@@ -90,39 +91,43 @@ func (o *migrateOptions) Config() (*migrateConfig, error) {
 		return nil, fmt.Errorf(`target version %q not supported. Minimal "3.5"`, storageVersionToString(c.targetVersion))
 	}
 
-	dbPath := datadir.ToBackendFileName(o.dataDir)
-	c.be = backend.NewDefaultBackend(GetLogger(), dbPath)
+	return c, nil
+}
+
+type migrateConfig struct {
+	lg            *zap.Logger
+	targetVersion *semver.Version
+	walVersion    schema.WALVersion
+	dataDir       string
+	force         bool
+}
 
-	walPath := datadir.ToWALDir(o.dataDir)
-	walSnap, err := getLatestWALSnap(c.lg, o.dataDir)
+func (c *migrateConfig) finalize() error {
+	walPath := datadir.ToWALDir(c.dataDir)
+	walSnap, err := getLatestWALSnap(c.lg, c.dataDir)
 	if err != nil {
-		return nil, fmt.Errorf("failed to get the lastest snapshot: %w", err)
+		return fmt.Errorf("failed to get the lastest snapshot: %w", err)
 	}
 	w, err := wal.OpenForRead(c.lg, walPath, walSnap)
 	if err != nil {
-		return nil, fmt.Errorf(`failed to open wal: %w`, err)
+		return fmt.Errorf(`failed to open wal: %w`, err)
 	}
 	defer w.Close()
 	c.walVersion, err = wal.ReadWALVersion(w)
 	if err != nil {
-		return nil, fmt.Errorf(`failed to read wal: %w`, err)
+		return fmt.Errorf(`failed to read wal: %w`, err)
 	}
 
-	return c, nil
-}
-
-type migrateConfig struct {
-	lg            *zap.Logger
-	be            backend.Backend
-	targetVersion *semver.Version
-	walVersion    schema.WALVersion
-	force         bool
+	return nil
 }
 
 func migrateCommandFunc(c *migrateConfig) error {
-	defer c.be.Close()
-	tx := c.be.BatchTx()
-	current, err := schema.DetectSchemaVersion(c.lg, c.be.ReadTx())
+	dbPath := datadir.ToBackendFileName(c.dataDir)
+	be := backend.NewDefaultBackend(GetLogger(), dbPath)
+	defer be.Close()
+
+	tx := be.BatchTx()
+	current, err := schema.DetectSchemaVersion(c.lg, be.ReadTx())
 	if err != nil {
 		c.lg.Error("failed to detect storage version. Please make sure you are using data dir from etcd v3.5 and older")
 		return err
@@ -131,6 +136,12 @@ func migrateCommandFunc(c *migrateConfig) error {
 		c.lg.Info("storage version up-to-date", zap.String("storage-version", storageVersionToString(&current)))
 		return nil
 	}
+
+	if err = c.finalize(); err != nil {
+		c.lg.Error("Failed to finalize config", zap.Error(err))
+		return err
+	}
+
 	err = schema.Migrate(c.lg, tx, c.walVersion, *c.targetVersion)
 	if err != nil {
 		if !c.force {
@@ -139,7 +150,7 @@ func migrateCommandFunc(c *migrateConfig) error {
 		c.lg.Info("normal migrate failed, trying with force", zap.Error(err))
 		migrateForce(c.lg, tx, c.targetVersion)
 	}
-	c.be.ForceCommit()
+	be.ForceCommit()
 	return nil
 }