From 918d24a80f537b53258d5c4c35ac02a3370b26f6 Mon Sep 17 00:00:00 2001 From: Jeremy <168515712+jeremy-babylonlabs@users.noreply.github.com> Date: Tue, 3 Dec 2024 19:01:40 +1100 Subject: [PATCH] feat: add MarkDelegationAsTransitioned method (#163) --- docs/docs.go | 3 ++ docs/swagger.json | 3 ++ docs/swagger.yaml | 2 + internal/v1/db/client/delegation.go | 17 +++++++++ internal/v1/db/client/interface.go | 1 + internal/v1/db/model/delegation.go | 1 + internal/v1/service/delegation.go | 2 + internal/v2/db/client/deelgation.go | 55 ++++++++++++++++++++++++++++ internal/v2/db/client/interface.go | 1 + internal/v2/queue/handler/staking.go | 6 +++ internal/v2/service/delegation.go | 11 ++++++ internal/v2/service/interface.go | 1 + tests/mocks/mock_db_client.go | 2 +- tests/mocks/mock_ordinal_client.go | 2 +- tests/mocks/mock_v1_db_client.go | 20 +++++++++- tests/mocks/mock_v2_db_client.go | 20 +++++++++- 16 files changed, 143 insertions(+), 4 deletions(-) create mode 100644 internal/v2/db/client/deelgation.go diff --git a/docs/docs.go b/docs/docs.go index e11e52bc..45117428 100644 --- a/docs/docs.go +++ b/docs/docs.go @@ -960,6 +960,9 @@ const docTemplate = `{ "is_slashed": { "type": "boolean" }, + "is_transitioned": { + "type": "boolean" + }, "staker_pk_hex": { "type": "string" }, diff --git a/docs/swagger.json b/docs/swagger.json index fc9ca867..12a0ac3e 100644 --- a/docs/swagger.json +++ b/docs/swagger.json @@ -952,6 +952,9 @@ "is_slashed": { "type": "boolean" }, + "is_transitioned": { + "type": "boolean" + }, "staker_pk_hex": { "type": "string" }, diff --git a/docs/swagger.yaml b/docs/swagger.yaml index 21c8ec2d..2e6c49cb 100644 --- a/docs/swagger.yaml +++ b/docs/swagger.yaml @@ -214,6 +214,8 @@ definitions: type: boolean is_slashed: type: boolean + is_transitioned: + type: boolean staker_pk_hex: type: string staking_tx: diff --git a/internal/v1/db/client/delegation.go b/internal/v1/db/client/delegation.go index 83ad8881..919a1259 100644 --- a/internal/v1/db/client/delegation.go +++ b/internal/v1/db/client/delegation.go @@ -156,6 +156,23 @@ func (v1dbclient *V1Database) ScanDelegationsPaginated( ) } +// MarkDelegationAsTransitioned marks an existing delegation as transitioned +func (v1dbclient *V1Database) MarkDelegationAsTransitioned(ctx context.Context, stakingTxHashHex string) error { + client := v1dbclient.Client.Database(v1dbclient.DbName).Collection(dbmodel.V1DelegationCollection) + update := bson.M{"$set": bson.M{"is_transitioned": true}} + result, err := client.UpdateOne(ctx, bson.M{"_id": stakingTxHashHex}, update) + if err != nil { + return err + } + if result.MatchedCount == 0 { + return &db.NotFoundError{ + Key: stakingTxHashHex, + Message: "Delegation not found", + } + } + return nil +} + // TransitionState updates the state of a staking transaction to a new state // It returns an NotFoundError if the staking transaction is not found or not in the eligible state to transition func (v1dbclient *V1Database) transitionState( diff --git a/internal/v1/db/client/interface.go b/internal/v1/db/client/interface.go index 163dff1f..997bde07 100644 --- a/internal/v1/db/client/interface.go +++ b/internal/v1/db/client/interface.go @@ -30,6 +30,7 @@ type V1DBClient interface { ctx context.Context, stakingTxHashHex, unbondingTxHashHex, txHex, signatureHex string, ) error FindDelegationByTxHashHex(ctx context.Context, txHashHex string) (*v1dbmodel.DelegationDocument, error) + MarkDelegationAsTransitioned(ctx context.Context, stakingTxHashHex string) error SaveTimeLockExpireCheck(ctx context.Context, stakingTxHashHex string, expireHeight uint64, txType string) error TransitionToUnbondedState( ctx context.Context, stakingTxHashHex string, eligiblePreviousState []types.DelegationState, diff --git a/internal/v1/db/model/delegation.go b/internal/v1/db/model/delegation.go index df148576..8f64bd4f 100644 --- a/internal/v1/db/model/delegation.go +++ b/internal/v1/db/model/delegation.go @@ -22,6 +22,7 @@ type DelegationDocument struct { StakingTx *TimelockTransaction `bson:"staking_tx"` // Always exist UnbondingTx *TimelockTransaction `bson:"unbonding_tx,omitempty"` IsOverflow bool `bson:"is_overflow"` + IsTransitioned bool `bson:"is_transitioned"` } type DelegationByStakerPagination struct { diff --git a/internal/v1/service/delegation.go b/internal/v1/service/delegation.go index 93e8af7c..594cb333 100644 --- a/internal/v1/service/delegation.go +++ b/internal/v1/service/delegation.go @@ -32,6 +32,7 @@ type DelegationPublic struct { IsOverflow bool `json:"is_overflow"` IsEligibleForTransition bool `json:"is_eligible_for_transition"` IsSlashed bool `json:"is_slashed"` + IsTransitioned bool `json:"is_transitioned"` } func (s *V1Service) DelegationsByStakerPk( @@ -212,6 +213,7 @@ func (s *V1Service) FromDelegationDocument( IsOverflow: d.IsOverflow, IsEligibleForTransition: isFpTransitioned && !isSlashed && s.isEligibleForTransition(d, bbnHeight), IsSlashed: isSlashed, + IsTransitioned: d.IsTransitioned, } // Add unbonding transaction if it exists diff --git a/internal/v2/db/client/deelgation.go b/internal/v2/db/client/deelgation.go new file mode 100644 index 00000000..c985e925 --- /dev/null +++ b/internal/v2/db/client/deelgation.go @@ -0,0 +1,55 @@ +package v2dbclient + +import ( + "context" + "errors" + + "github.com/babylonlabs-io/staking-api-service/internal/shared/db" + dbmodel "github.com/babylonlabs-io/staking-api-service/internal/shared/db/model" + "go.mongodb.org/mongo-driver/bson" + "go.mongodb.org/mongo-driver/mongo" +) + +func (v2dbclient *V2Database) MarkV1DelegationAsTransitioned(ctx context.Context, stakingTxHashHex string) error { + session, err := v2dbclient.Client.StartSession() + if err != nil { + return err + } + defer session.EndSession(ctx) + + transactionWork := func(sessCtx mongo.SessionContext) (interface{}, error) { + client := v2dbclient.Client.Database(v2dbclient.DbName).Collection(dbmodel.V1DelegationCollection) + filter := bson.M{"_id": stakingTxHashHex} + + var delegation interface{} + err := client.FindOne(sessCtx, filter).Decode(&delegation) + if err != nil { + if errors.Is(err, mongo.ErrNoDocuments) { + return nil, &db.NotFoundError{ + Key: stakingTxHashHex, + Message: "Delegation not found", + } + } + return nil, err + } + + update := bson.M{"$set": bson.M{"is_transitioned": true}} + result, err := client.UpdateOne(sessCtx, filter, update) + if err != nil { + return nil, err + } + + if result.MatchedCount == 0 { + return nil, &db.NotFoundError{ + Key: stakingTxHashHex, + Message: "Delegation not found", + } + } + + return nil, nil + } + + _, err = session.WithTransaction(ctx, transactionWork) + return err +} + diff --git a/internal/v2/db/client/interface.go b/internal/v2/db/client/interface.go index 61665cdf..3c3e987a 100644 --- a/internal/v2/db/client/interface.go +++ b/internal/v2/db/client/interface.go @@ -26,4 +26,5 @@ type V2DBClient interface { SubtractStakerStats( ctx context.Context, stakingTxHashHex, stakerPkHex string, amount uint64, ) error + MarkV1DelegationAsTransitioned(ctx context.Context, stakingTxHashHex string) error } diff --git a/internal/v2/queue/handler/staking.go b/internal/v2/queue/handler/staking.go index 5db9481d..9b17eb2a 100644 --- a/internal/v2/queue/handler/staking.go +++ b/internal/v2/queue/handler/staking.go @@ -20,6 +20,12 @@ func (h *V2QueueHandler) ActiveStakingHandler(ctx context.Context, messageBody s return types.NewError(http.StatusBadRequest, types.BadRequest, err) } + // Mark as v1 delegation as transitioned if it exists + if err := h.Service.MarkV1DelegationAsTransitioned(ctx, activeStakingEvent.StakingTxHashHex); err != nil { + log.Ctx(ctx).Error().Err(err).Msg("Failed to mark v1 delegation as transitioned") + return err + } + // Perform the address lookup conversion addressLookupErr := h.performAddressLookupConversion(ctx, activeStakingEvent.StakerBtcPkHex, types.Active) if addressLookupErr != nil { diff --git a/internal/v2/service/delegation.go b/internal/v2/service/delegation.go index 80608bbe..b5b5c0de 100644 --- a/internal/v2/service/delegation.go +++ b/internal/v2/service/delegation.go @@ -170,3 +170,14 @@ func (s *V2Service) SaveUnprocessableMessages(ctx context.Context, messageBody, } return nil } + +func (s *V2Service) MarkV1DelegationAsTransitioned(ctx context.Context, stakingTxHashHex string) *types.Error { + err := s.DbClients.V2DBClient.MarkV1DelegationAsTransitioned(ctx, stakingTxHashHex) + if err != nil { + if db.IsNotFoundError(err) { + return nil + } + return types.NewInternalServiceError(err) + } + return nil +} diff --git a/internal/v2/service/interface.go b/internal/v2/service/interface.go index 1897f852..8530ecad 100644 --- a/internal/v2/service/interface.go +++ b/internal/v2/service/interface.go @@ -14,6 +14,7 @@ type V2ServiceProvider interface { IsDelegationPresent(ctx context.Context, txHashHex string) (bool, *types.Error) GetDelegation(ctx context.Context, stakingTxHashHex string) (*StakerDelegationPublic, *types.Error) GetDelegations(ctx context.Context, stakerPKHex string, paginationKey string) ([]*StakerDelegationPublic, string, *types.Error) + MarkV1DelegationAsTransitioned(ctx context.Context, stakingTxHashHex string) *types.Error GetOverallStats(ctx context.Context) (*OverallStatsPublic, *types.Error) GetStakerStats(ctx context.Context, stakerPKHex string) (*StakerStatsPublic, *types.Error) ProcessAndSaveBtcAddresses(ctx context.Context, stakerPkHex string) *types.Error diff --git a/tests/mocks/mock_db_client.go b/tests/mocks/mock_db_client.go index b247612e..9560824e 100644 --- a/tests/mocks/mock_db_client.go +++ b/tests/mocks/mock_db_client.go @@ -1,4 +1,4 @@ -// Code generated by mockery v2.41.0. DO NOT EDIT. +// Code generated by mockery v2.49.1. DO NOT EDIT. package mocks diff --git a/tests/mocks/mock_ordinal_client.go b/tests/mocks/mock_ordinal_client.go index 95dff743..36c16566 100644 --- a/tests/mocks/mock_ordinal_client.go +++ b/tests/mocks/mock_ordinal_client.go @@ -1,4 +1,4 @@ -// Code generated by mockery v2.41.0. DO NOT EDIT. +// Code generated by mockery v2.49.1. DO NOT EDIT. package mocks diff --git a/tests/mocks/mock_v1_db_client.go b/tests/mocks/mock_v1_db_client.go index 2d0d206b..c234b4e6 100644 --- a/tests/mocks/mock_v1_db_client.go +++ b/tests/mocks/mock_v1_db_client.go @@ -1,4 +1,4 @@ -// Code generated by mockery v2.41.0. DO NOT EDIT. +// Code generated by mockery v2.49.1. DO NOT EDIT. package mocks @@ -500,6 +500,24 @@ func (_m *V1DBClient) InsertPkAddressMappings(ctx context.Context, stakerPkHex s return r0 } +// MarkDelegationAsTransitioned provides a mock function with given fields: ctx, stakingTxHashHex +func (_m *V1DBClient) MarkDelegationAsTransitioned(ctx context.Context, stakingTxHashHex string) error { + ret := _m.Called(ctx, stakingTxHashHex) + + if len(ret) == 0 { + panic("no return value specified for MarkDelegationAsTransitioned") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, string) error); ok { + r0 = rf(ctx, stakingTxHashHex) + } else { + r0 = ret.Error(0) + } + + return r0 +} + // Ping provides a mock function with given fields: ctx func (_m *V1DBClient) Ping(ctx context.Context) error { ret := _m.Called(ctx) diff --git a/tests/mocks/mock_v2_db_client.go b/tests/mocks/mock_v2_db_client.go index 27929f3f..4c43e41e 100644 --- a/tests/mocks/mock_v2_db_client.go +++ b/tests/mocks/mock_v2_db_client.go @@ -1,4 +1,4 @@ -// Code generated by mockery v2.41.0. DO NOT EDIT. +// Code generated by mockery v2.49.1. DO NOT EDIT. package mocks @@ -268,6 +268,24 @@ func (_m *V2DBClient) InsertPkAddressMappings(ctx context.Context, stakerPkHex s return r0 } +// MarkV1DelegationAsTransitioned provides a mock function with given fields: ctx, stakingTxHashHex +func (_m *V2DBClient) MarkV1DelegationAsTransitioned(ctx context.Context, stakingTxHashHex string) error { + ret := _m.Called(ctx, stakingTxHashHex) + + if len(ret) == 0 { + panic("no return value specified for MarkV1DelegationAsTransitioned") + } + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, string) error); ok { + r0 = rf(ctx, stakingTxHashHex) + } else { + r0 = ret.Error(0) + } + + return r0 +} + // Ping provides a mock function with given fields: ctx func (_m *V2DBClient) Ping(ctx context.Context) error { ret := _m.Called(ctx)