Fix and optimize babylon migration

baking-bad · Sep 29, 2023 · 3c37305 · 3c37305
1 parent ba37f66
commit 3c37305
Show file tree

Hide file tree

Showing 8 changed files with 18 additions and 76 deletions.
diff --git a/internal/models/bigmapdiff/repository.go b/internal/models/bigmapdiff/repository.go
@@ -12,7 +12,6 @@ type Repository interface {
 	GetForAddress(ctx context.Context, address string) ([]BigMapState, error)
 	Count(ctx context.Context, ptr int64) (int, error)
 	Current(ctx context.Context, keyHash string, ptr int64) (BigMapState, error)
-	CurrentByContract(ctx context.Context, contract string) ([]BigMapState, error)
 	Previous(ctx context.Context, diffs []BigMapDiff) ([]BigMapDiff, error)
 	GetStats(ctx context.Context, ptr int64) (Stats, error)
 	Keys(ctx context.Context, reqCtx GetContext) (states []BigMapState, err error)

diff --git a/internal/models/interface.go b/internal/models/interface.go
@@ -53,7 +53,7 @@ type Transaction interface {
 	ToJakarta(ctx context.Context) error
 	JakartaVesting(ctx context.Context, contract *contract.Contract) error
 	JakartaUpdateNonDelegator(ctx context.Context, contract *contract.Contract) error
-	BabylonBigMapStates(ctx context.Context, state *bigmapdiff.BigMapState) error
+	DeleteBigMapStatesByContract(ctx context.Context, contract string) ([]bigmapdiff.BigMapState, error)
 
 	Commit() error
 	Rollback() error

diff --git a/internal/models/mock/bigmapdiff/mock.go b/internal/models/mock/bigmapdiff/mock.go
diff --git a/internal/parsers/migrations/babylon.go b/internal/parsers/migrations/babylon.go
@@ -118,7 +118,7 @@ func (p *Babylon) getUpdates(ctx context.Context, script noderpc.Script, contrac
 		}
 	}
 
-	keys, err := p.bmdRepo.CurrentByContract(ctx, contract.Account.Address)
+	keys, err := tx.DeleteBigMapStatesByContract(ctx, contract.Account.Address)
 	if err != nil {
 		return err
 	}
@@ -128,7 +128,9 @@ func (p *Babylon) getUpdates(ctx context.Context, script noderpc.Script, contrac
 
 	for i := range keys {
 		keys[i].Ptr = newPtr
-		if err := tx.BabylonBigMapStates(ctx, &keys[i]); err != nil {
+		keys[i].ID = 0
+
+		if err := tx.BigMapStates(ctx, &keys[i]); err != nil {
 			return err
 		}
 	}

diff --git a/internal/postgres/bigmapdiff/storage.go b/internal/postgres/bigmapdiff/storage.go
@@ -182,14 +182,6 @@ func (storage *Storage) GetStats(ctx context.Context, ptr int64) (stats bigmapdi
 	return
 }
 
-// CurrentByContract -
-func (storage *Storage) CurrentByContract(ctx context.Context, contract string) (keys []bigmapdiff.BigMapState, err error) {
-	err = storage.DB.NewSelect().Model(&keys).
-		Where("contract = ?", contract).
-		Scan(ctx)
-	return
-}
-
 // Keys -
 func (storage *Storage) Keys(ctx context.Context, req bigmapdiff.GetContext) (states []bigmapdiff.BigMapState, err error) {
 	err = storage.buildGetContextForState(req).Scan(ctx, &states)

diff --git a/internal/postgres/core/transaction.go b/internal/postgres/core/transaction.go
@@ -218,11 +218,6 @@ func (t Transaction) ToJakarta(ctx context.Context) error {
 	return err
 }
 
-func (t Transaction) BabylonBigMapStates(ctx context.Context, state *bigmapdiff.BigMapState) error {
-	_, err := t.tx.NewUpdate().Model(state).WherePK().Column("ptr").Exec(ctx)
-	return err
-}
-
 func (t Transaction) UpdateStats(ctx context.Context, stats stats.Stats) error {
 	_, err := t.tx.NewInsert().
 		Model(&stats).
@@ -269,3 +264,12 @@ func (t Transaction) TicketBalances(ctx context.Context, balances ...*ticket.Bal
 		Exec(ctx)
 	return err
 }
+
+func (t Transaction) DeleteBigMapStatesByContract(ctx context.Context, contract string) (states []bigmapdiff.BigMapState, err error) {
+	_, err = t.tx.NewDelete().
+		Model((*bigmapdiff.BigMapState)(nil)).
+		Where("contract = ?", contract).
+		Returning("*").
+		Exec(ctx, &states)
+	return
+}
diff --git a/internal/postgres/tests/big_map_diffs_test.go b/internal/postgres/tests/big_map_diffs_test.go
@@ -147,15 +147,6 @@ func (s *StorageTestSuite) TestBigMapDiffsGetStats() {
 	s.Require().EqualValues("KT1NSpRTVR4MUwx64XCADXDUmpMGQw5yVNK1", stats.Contract)
 }
 
-func (s *StorageTestSuite) TestBigMapDiffsCurrentByContract() {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
-	defer cancel()
-
-	diff, err := s.bigMapDiffs.CurrentByContract(ctx, "KT1NSpRTVR4MUwx64XCADXDUmpMGQw5yVNK1")
-	s.Require().NoError(err)
-	s.Require().Len(diff, 4)
-}
-
 func (s *StorageTestSuite) TestBigMapDiffsKeys() {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 	defer cancel()

diff --git a/internal/postgres/tests/transaction_test.go b/internal/postgres/tests/transaction_test.go
@@ -319,28 +319,21 @@ func (s *StorageTestSuite) TestToJakarta() {
 	s.Require().EqualValues(14, newContract.JakartaID)
 }
 
-func (s *StorageTestSuite) TestBabylonBigMapStates() {
+func (s *StorageTestSuite) TestDeleteBigMapStates() {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 	defer cancel()
 
 	tx, err := core.NewTransaction(ctx, s.storage.DB)
 	s.Require().NoError(err)
 
-	err = tx.BabylonBigMapStates(ctx, &bigmapdiff.BigMapState{
-		ID:       3,
-		Ptr:      1000,
-		KeyHash:  "expruDuAZnFKqmLoisJqUGqrNzXTvw7PJM2rYk97JErM5FHCerQqgn",
-		Contract: "KT1Pz65ssbPF7Zv9Dh7ggqUkgAYNSuJ9iia7",
-	})
+	states, err := tx.DeleteBigMapStatesByContract(ctx, "KT1Pz65ssbPF7Zv9Dh7ggqUkgAYNSuJ9iia7")
 	s.Require().NoError(err)
 
 	err = tx.Commit()
 	s.Require().NoError(err)
+	s.Require().Len(states, 9)
 
-	var state bigmapdiff.BigMapState
-	err = s.storage.DB.NewSelect().Model(&state).Where("id = 3").Scan(ctx)
-	s.Require().NoError(err)
-	s.Require().EqualValues(1000, state.Ptr)
+	s.Require().Equal("KT1Pz65ssbPF7Zv9Dh7ggqUkgAYNSuJ9iia7", states[0].Contract)
 }
 
 func (s *StorageTestSuite) TestUpdateStats() {