Fix: rollback protocol

internal/models/rollback.go

@@ -27,6 +27,7 @@ type Rollback interface {
 	GlobalConstants(ctx context.Context, level int64) ([]contract.GlobalConstant, error)
 	Scripts(ctx context.Context, level int64) ([]contract.Script, error)
 	DeleteScriptsConstants(ctx context.Context, scriptIds []int64, constantsIds []int64) error
+	Protocols(ctx context.Context, level int64) error
 
 	Commit() error
 	Rollback() error

internal/postgres/rollback.go

@@ -8,6 +8,7 @@ import (
 	"github.com/baking-bad/bcdhub/internal/models/bigmapdiff"
 	"github.com/baking-bad/bcdhub/internal/models/contract"
 	"github.com/baking-bad/bcdhub/internal/models/operation"
+	"github.com/baking-bad/bcdhub/internal/models/protocol"
 	"github.com/uptrace/bun"
 )
 
@@ -133,3 +134,24 @@ func (r Rollback) DeleteScriptsConstants(ctx context.Context, scriptIds []int64,
 	_, err := query.Exec(ctx)
 	return err
 }
+
+func (r Rollback) Protocols(ctx context.Context, level int64) error {
+	result, err := r.tx.NewDelete().Model((*protocol.Protocol)(nil)).Where("start_level >= ?", level).Exec(ctx)
+	if err != nil {
+		return err
+	}
+	count, err := result.RowsAffected()
+	if err != nil {
+		return err
+	}
+	if count == 0 {
+		return nil
+	}
+
+	_, err = r.tx.NewUpdate().
+		Model((*protocol.Protocol)(nil)).
+		Where("start_level < ?", level).
+		Set("end_level = 0").
+		Exec(ctx)
+	return err
+}

internal/postgres/tests/rollback_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/baking-bad/bcdhub/internal/models/bigmapdiff"
 	"github.com/baking-bad/bcdhub/internal/models/block"
 	"github.com/baking-bad/bcdhub/internal/models/contract"
+	"github.com/baking-bad/bcdhub/internal/models/protocol"
 	"github.com/baking-bad/bcdhub/internal/models/types"
 	"github.com/baking-bad/bcdhub/internal/postgres"
 	"github.com/baking-bad/bcdhub/internal/postgres/core"
@@ -67,24 +68,26 @@ func (s *RollbackTestSuite) TearDownSuite() {
 	s.Require().NoError(s.psqlContainer.Terminate(ctx))
 }
 
-func TestSuiteRollback_Run(t *testing.T) {
-	suite.Run(t, new(RollbackTestSuite))
-}
-
-func (s *RollbackTestSuite) TestDeleteAll() {
+func (s *RollbackTestSuite) SetupTest() {
 	db, err := sql.Open("postgres", s.psqlContainer.GetDSN())
 	s.Require().NoError(err)
 
 	fixtures, err := testfixtures.New(
 		testfixtures.Database(db),
 		testfixtures.Dialect("postgres"),
-		testfixtures.Files("./fixtures/blocks.yml"),
+		testfixtures.Directory("./fixtures"),
 		testfixtures.UseAlterConstraint(),
 	)
 	s.Require().NoError(err)
 	s.Require().NoError(fixtures.Load())
 	s.Require().NoError(db.Close())
+}
+
+func TestSuiteRollback_Run(t *testing.T) {
+	suite.Run(t, new(RollbackTestSuite))
+}
 
+func (s *RollbackTestSuite) TestDeleteAll() {
 	saver, err := postgres.NewRollback(s.storage.DB)
 	s.Require().NoError(err)
 
@@ -105,19 +108,6 @@ func (s *RollbackTestSuite) TestDeleteAll() {
 }
 
 func (s *RollbackTestSuite) TestStatesChangedAtLevel() {
-	db, err := sql.Open("postgres", s.psqlContainer.GetDSN())
-	s.Require().NoError(err)
-
-	fixtures, err := testfixtures.New(
-		testfixtures.Database(db),
-		testfixtures.Dialect("postgres"),
-		testfixtures.Files("./fixtures/big_map_states.yml"),
-		testfixtures.UseAlterConstraint(),
-	)
-	s.Require().NoError(err)
-	s.Require().NoError(fixtures.Load())
-	s.Require().NoError(db.Close())
-
 	saver, err := postgres.NewRollback(s.storage.DB)
 	s.Require().NoError(err)
 
@@ -133,19 +123,6 @@ func (s *RollbackTestSuite) TestStatesChangedAtLevel() {
 }
 
 func (s *RollbackTestSuite) TestLastDiff() {
-	db, err := sql.Open("postgres", s.psqlContainer.GetDSN())
-	s.Require().NoError(err)
-
-	fixtures, err := testfixtures.New(
-		testfixtures.Database(db),
-		testfixtures.Dialect("postgres"),
-		testfixtures.Files("./fixtures/big_map_diffs.yml"),
-		testfixtures.UseAlterConstraint(),
-	)
-	s.Require().NoError(err)
-	s.Require().NoError(fixtures.Load())
-	s.Require().NoError(db.Close())
-
 	saver, err := postgres.NewRollback(s.storage.DB)
 	s.Require().NoError(err)
 
@@ -168,19 +145,6 @@ func (s *RollbackTestSuite) TestLastDiff() {
 }
 
 func (s *RollbackTestSuite) TestDeleteBigMapState() {
-	db, err := sql.Open("postgres", s.psqlContainer.GetDSN())
-	s.Require().NoError(err)
-
-	fixtures, err := testfixtures.New(
-		testfixtures.Database(db),
-		testfixtures.Dialect("postgres"),
-		testfixtures.Files("./fixtures/big_map_states.yml"),
-		testfixtures.UseAlterConstraint(),
-	)
-	s.Require().NoError(err)
-	s.Require().NoError(fixtures.Load())
-	s.Require().NoError(db.Close())
-
 	saver, err := postgres.NewRollback(s.storage.DB)
 	s.Require().NoError(err)
 
@@ -199,19 +163,6 @@ func (s *RollbackTestSuite) TestDeleteBigMapState() {
 }
 
 func (s *RollbackTestSuite) TestSaveBigMapState() {
-	db, err := sql.Open("postgres", s.psqlContainer.GetDSN())
-	s.Require().NoError(err)
-
-	fixtures, err := testfixtures.New(
-		testfixtures.Database(db),
-		testfixtures.Dialect("postgres"),
-		testfixtures.Files("./fixtures/big_map_states.yml"),
-		testfixtures.UseAlterConstraint(),
-	)
-	s.Require().NoError(err)
-	s.Require().NoError(fixtures.Load())
-	s.Require().NoError(db.Close())
-
 	saver, err := postgres.NewRollback(s.storage.DB)
 	s.Require().NoError(err)
 
@@ -243,19 +194,6 @@ func (s *RollbackTestSuite) TestSaveBigMapState() {
 }
 
 func (s *RollbackTestSuite) TestGetOperations() {
-	db, err := sql.Open("postgres", s.psqlContainer.GetDSN())
-	s.Require().NoError(err)
-
-	fixtures, err := testfixtures.New(
-		testfixtures.Database(db),
-		testfixtures.Dialect("postgres"),
-		testfixtures.Files("./fixtures/operations.yml"),
-		testfixtures.UseAlterConstraint(),
-	)
-	s.Require().NoError(err)
-	s.Require().NoError(fixtures.Load())
-	s.Require().NoError(db.Close())
-
 	saver, err := postgres.NewRollback(s.storage.DB)
 	s.Require().NoError(err)
 
@@ -272,23 +210,6 @@ func (s *RollbackTestSuite) TestGetOperations() {
 }
 
 func (s *RollbackTestSuite) TestGetContractLastActions() {
-	db, err := sql.Open("postgres", s.psqlContainer.GetDSN())
-	s.Require().NoError(err)
-
-	fixtures, err := testfixtures.New(
-		testfixtures.Database(db),
-		testfixtures.Dialect("postgres"),
-		testfixtures.Files(
-			"./fixtures/operations.yml",
-			"./fixtures/accounts.yml",
-			"./fixtures/contracts.yml",
-		),
-		testfixtures.UseAlterConstraint(),
-	)
-	s.Require().NoError(err)
-	s.Require().NoError(fixtures.Load())
-	s.Require().NoError(db.Close())
-
 	saver, err := postgres.NewRollback(s.storage.DB)
 	s.Require().NoError(err)
 
@@ -310,23 +231,6 @@ func (s *RollbackTestSuite) TestGetContractLastActions() {
 }
 
 func (s *RollbackTestSuite) TestUpdateContractStats() {
-	db, err := sql.Open("postgres", s.psqlContainer.GetDSN())
-	s.Require().NoError(err)
-
-	fixtures, err := testfixtures.New(
-		testfixtures.Database(db),
-		testfixtures.Dialect("postgres"),
-		testfixtures.Files(
-			"./fixtures/operations.yml",
-			"./fixtures/accounts.yml",
-			"./fixtures/contracts.yml",
-		),
-		testfixtures.UseAlterConstraint(),
-	)
-	s.Require().NoError(err)
-	s.Require().NoError(fixtures.Load())
-	s.Require().NoError(db.Close())
-
 	saver, err := postgres.NewRollback(s.storage.DB)
 	s.Require().NoError(err)
 
@@ -347,3 +251,23 @@ func (s *RollbackTestSuite) TestUpdateContractStats() {
 	s.Require().EqualValues(0, cntrct.TxCount)
 	s.Require().EqualValues(ts.Format(time.RFC3339), cntrct.LastAction.Format(time.RFC3339))
 }
+
+func (s *RollbackTestSuite) TestProtocols() {
+	saver, err := postgres.NewRollback(s.storage.DB)
+	s.Require().NoError(err)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	err = saver.Protocols(ctx, 2)
+	s.Require().NoError(err)
+
+	err = saver.Commit()
+	s.Require().NoError(err)
+
+	var proto protocol.Protocol
+	err = s.storage.DB.NewSelect().Model(&proto).Order("id desc").Limit(1).Scan(ctx)
+	s.Require().NoError(err)
+	s.Require().EqualValues(0, proto.EndLevel)
+	s.Require().EqualValues("Ps9mPmXaRzmzk35gbAYNCAw6UXdE2qoABTHbN2oEEc1qM7CwT9P", proto.Hash)
+}

internal/rollback/big_map_states.go

@@ -0,0 +1,52 @@
+package rollback
+
+import (
+	"context"
+
+	"github.com/baking-bad/bcdhub/internal/logger"
+)
+
+func (rm Manager) rollbackBigMapState(ctx context.Context, level int64) error {
+	logger.Info().Msg("rollback big map states...")
+	states, err := rm.rollback.StatesChangedAtLevel(ctx, level)
+	if err != nil {
+		return err
+	}
+
+	for i, state := range states {
+		diff, err := rm.rollback.LastDiff(ctx, state.Ptr, state.KeyHash, false)
+		if err != nil {
+			if rm.storage.IsRecordNotFound(err) {
+				if err := rm.rollback.DeleteBigMapState(ctx, states[i]); err != nil {
+					return err
+				}
+				continue
+			}
+			return err
+		}
+		states[i].LastUpdateLevel = diff.Level
+		states[i].LastUpdateTime = diff.Timestamp
+		states[i].IsRollback = true
+
+		if len(diff.Value) > 0 {
+			states[i].Value = diff.ValueBytes()
+			states[i].Removed = false
+		} else {
+			states[i].Removed = true
+			valuedDiff, err := rm.rollback.LastDiff(ctx, state.Ptr, state.KeyHash, true)
+			if err != nil {
+				if !rm.storage.IsRecordNotFound(err) {
+					return err
+				}
+			} else {
+				states[i].Value = valuedDiff.ValueBytes()
+			}
+		}
+
+		if err := rm.rollback.SaveBigMapState(ctx, states[i]); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}