From a550068008ff8fdf167bfd1dc095e875d13f3381 Mon Sep 17 00:00:00 2001
From: VargaElod23 <vargaelod23@gmail.com>
Date: Wed, 2 Aug 2023 23:26:24 +0300
Subject: [PATCH] simplifies code

---
 .../storage/pebble/migrations/0_dag_removeSender.go  | 12 ++++++------
 internal/storage/pebble/migrations/manager.go        | 11 +++++++----
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/internal/storage/pebble/migrations/0_dag_removeSender.go b/internal/storage/pebble/migrations/0_dag_removeSender.go
index 1a70b6a..5bbf254 100644
--- a/internal/storage/pebble/migrations/0_dag_removeSender.go
+++ b/internal/storage/pebble/migrations/0_dag_removeSender.go
@@ -39,19 +39,19 @@ func (m *RemoveSenderMigration) Apply(s *pebble.Storage) error {
 				Timestamp:        o.Timestamp,
 				TransactionCount: o.TransactionCount,
 			}
-			batch.AddToBatchFullKey(&dag, key)
+			err = batch.AddToBatchFullKey(&dag, key)
+
+			if err != nil {
+				log.WithFields(log.Fields{"migration": m.id, "error": err}).Fatal("Error adding Dag to batch")
+			}
 
 			last_key = key
 			count++
-			if count == DAG_BATCH_THRESHOLD {
-				return true
-			}
-			return false
+			return count == DAG_BATCH_THRESHOLD
 		})
 
 		if count < DAG_BATCH_THRESHOLD {
 			batch.CommitBatch()
-			done = true
 			break
 		}
 	}
diff --git a/internal/storage/pebble/migrations/manager.go b/internal/storage/pebble/migrations/manager.go
index f703b1e..700e9ad 100644
--- a/internal/storage/pebble/migrations/manager.go
+++ b/internal/storage/pebble/migrations/manager.go
@@ -35,18 +35,21 @@ func (m *Manager) IsApplied(migration Migration) bool {
 	return err == nil
 }
 
-func (m *Manager) ApplyAll() error {
+func (m *Manager) ApplyAll() (err error) {
 	log.Info("Migration Manager: Running migrations")
 	for _, migration := range m.migrations {
 		isApplied := m.IsApplied(migration)
 		if !isApplied {
 			log.WithFields(log.Fields{"migration": migration.GetId()}).Info("Running migration")
-			err := migration.Apply(m.s)
+			err = migration.Apply(m.s)
 			if err != nil {
-				return err
+				return
 			}
 			b := m.s.NewBatch()
-			b.AddToBatchFullKey(migration.GetId(), []byte(migration_prefix+migration.GetId()))
+			err = b.AddToBatchFullKey(migration.GetId(), []byte(migration_prefix+migration.GetId()))
+			if err != nil {
+				return
+			}
 			b.CommitBatch()
 			log.WithFields(log.Fields{"migration": migration.GetId()}).Info("Applied migration")
 		} else {