From 32e710700ecaa5b484d43a9c4e6db11e738388b2 Mon Sep 17 00:00:00 2001
From: Gurjot Singh <111540954+gusin13@users.noreply.github.com>
Date: Sun, 24 Nov 2024 02:05:29 +0530
Subject: [PATCH] fix params (#64)

---
 internal/clients/bbnclient/types.go |  2 ++
 internal/db/model/params.go         | 22 +++++++++---
 internal/db/params.go               | 55 +++++++++++++----------------
 3 files changed, 44 insertions(+), 35 deletions(-)

diff --git a/internal/clients/bbnclient/types.go b/internal/clients/bbnclient/types.go
index 96bc9a7..052d891 100644
--- a/internal/clients/bbnclient/types.go
+++ b/internal/clients/bbnclient/types.go
@@ -23,6 +23,7 @@ type StakingParams struct {
 	UnbondingFeeSat              int64    `bson:"unbonding_fee_sat"`
 	MinCommissionRate            string   `bson:"min_commission_rate"`
 	DelegationCreationBaseGasFee uint64   `bson:"delegation_creation_base_gas_fee"`
+	AllowListExpirationHeight    uint64   `bson:"allow_list_expiration_height"`
 }
 
 type CheckpointParams struct {
@@ -46,6 +47,7 @@ func FromBbnStakingParams(params stakingtypes.Params) *StakingParams {
 		UnbondingFeeSat:              params.UnbondingFeeSat,
 		MinCommissionRate:            params.MinCommissionRate.String(),
 		DelegationCreationBaseGasFee: params.DelegationCreationBaseGasFee,
+		AllowListExpirationHeight:    params.AllowListExpirationHeight,
 	}
 }
 
diff --git a/internal/db/model/params.go b/internal/db/model/params.go
index 69c9ec3..11d9e0c 100644
--- a/internal/db/model/params.go
+++ b/internal/db/model/params.go
@@ -1,7 +1,21 @@
 package model
 
-type GlobalParamsDocument struct {
-	Type    string      `bson:"type"`
-	Version uint32      `bson:"version"`
-	Params  interface{} `bson:"params"`
+import "github.com/babylonlabs-io/babylon-staking-indexer/internal/clients/bbnclient"
+
+// Base document for common fields
+type BaseParamsDocument struct {
+	Type    string `bson:"type"`
+	Version uint32 `bson:"version"`
+}
+
+// Specific document for staking params
+type StakingParamsDocument struct {
+	BaseParamsDocument `bson:",inline"`
+	Params             *bbnclient.StakingParams `bson:"params"`
+}
+
+// Specific document for checkpoint params
+type CheckpointParamsDocument struct {
+	BaseParamsDocument `bson:",inline"`
+	Params             *bbnclient.CheckpointParams `bson:"params"`
 }
diff --git a/internal/db/params.go b/internal/db/params.go
index 7ae80c9..2e636ef 100644
--- a/internal/db/params.go
+++ b/internal/db/params.go
@@ -23,54 +23,47 @@ const (
 func (db *Database) SaveStakingParams(
 	ctx context.Context, version uint32, params *bbnclient.StakingParams,
 ) error {
-	collection := db.client.Database(db.dbName).
-		Collection(model.GlobalParamsCollection)
-
-	filter := bson.M{
-		"type":    STAKING_PARAMS_TYPE,
-		"version": version,
-	}
+	collection := db.client.Database(db.dbName).Collection(model.GlobalParamsCollection)
 
-	update := bson.M{
-		"$setOnInsert": &model.GlobalParamsDocument{
+	doc := &model.StakingParamsDocument{
+		BaseParamsDocument: model.BaseParamsDocument{
 			Type:    STAKING_PARAMS_TYPE,
 			Version: version,
-			Params:  params,
 		},
+		Params: params,
 	}
 
-	_, err := collection.UpdateOne(ctx, filter, update, options.Update().SetUpsert(true))
-	if err != nil {
-		return fmt.Errorf("failed to save staking params: %w", err)
+	filter := bson.M{
+		"type":    STAKING_PARAMS_TYPE,
+		"version": version,
 	}
-	return nil
+	update := bson.M{"$setOnInsert": doc}
+
+	_, err := collection.UpdateOne(ctx, filter, update, options.Update().SetUpsert(true))
+	return err
 }
 
 func (db *Database) SaveCheckpointParams(
 	ctx context.Context, params *bbnclient.CheckpointParams,
 ) error {
-	collection := db.client.Database(db.dbName).
-		Collection(model.GlobalParamsCollection)
-
-	filter := bson.M{
-		"type":    CHECKPOINT_PARAMS_TYPE,
-		"version": CHECKPOINT_PARAMS_VERSION,
-	}
+	collection := db.client.Database(db.dbName).Collection(model.GlobalParamsCollection)
 
-	update := bson.M{
-		"$setOnInsert": &model.GlobalParamsDocument{
+	doc := &model.CheckpointParamsDocument{
+		BaseParamsDocument: model.BaseParamsDocument{
 			Type:    CHECKPOINT_PARAMS_TYPE,
-			Version: CHECKPOINT_PARAMS_VERSION,
-			Params:  params,
+			Version: CHECKPOINT_PARAMS_VERSION, // hardcoded as 0
 		},
+		Params: params,
 	}
 
-	_, err := collection.UpdateOne(ctx, filter, update, options.Update().SetUpsert(true))
-	if err != nil {
-		return fmt.Errorf("failed to save checkpoint params: %w", err)
+	filter := bson.M{
+		"type":    CHECKPOINT_PARAMS_TYPE,
+		"version": CHECKPOINT_PARAMS_VERSION, // hardcoded as 0
 	}
+	update := bson.M{"$setOnInsert": doc}
 
-	return nil
+	_, err := collection.UpdateOne(ctx, filter, update, options.Update().SetUpsert(true))
+	return err
 }
 
 func (db *Database) GetStakingParams(ctx context.Context, version uint32) (*bbnclient.StakingParams, error) {
@@ -82,11 +75,11 @@ func (db *Database) GetStakingParams(ctx context.Context, version uint32) (*bbnc
 		"version": version,
 	}
 
-	var params model.GlobalParamsDocument
+	var params model.StakingParamsDocument
 	err := collection.FindOne(ctx, filter).Decode(&params)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get staking params: %w", err)
 	}
 
-	return params.Params.(*bbnclient.StakingParams), nil
+	return params.Params, nil
 }