update documents using _id field (#20) (#21)

internal/db/dbclient.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/options"
 
@@ -112,35 +113,35 @@ func (db *Database) FindSendUnbondingDocuments(ctx context.Context) ([]model.Unb
 
 func (db *Database) updateUnbondingDocumentState(
 	ctx context.Context,
-	unbondingTxHashHex string,
+	id primitive.ObjectID,
 	newState model.UnbondingState) error {
 	client := db.Client.Database(db.DbName).Collection(model.UnbondingCollection)
-	filter := bson.M{"unbonding_tx_hash_hex": unbondingTxHashHex}
+	filter := bson.M{"_id": id}
 	update := bson.M{"$set": bson.M{"state": newState}}
 	_, err := client.UpdateOne(ctx, filter, update)
 	return err
 }
 
 func (db *Database) SetUnbondingDocumentSend(
 	ctx context.Context,
-	unbondingTxHashHex string) error {
-	return db.updateUnbondingDocumentState(ctx, unbondingTxHashHex, model.Send)
+	id primitive.ObjectID) error {
+	return db.updateUnbondingDocumentState(ctx, id, model.Send)
 }
 
 func (db *Database) SetUnbondingDocumentFailed(
 	ctx context.Context,
-	unbondingTxHashHex string) error {
-	return db.updateUnbondingDocumentState(ctx, unbondingTxHashHex, model.Failed)
+	id primitive.ObjectID) error {
+	return db.updateUnbondingDocumentState(ctx, id, model.Failed)
 }
 
 func (db *Database) SetUnbondingDocumentInputAlreadySpent(
 	ctx context.Context,
-	unbondingTxHashHex string) error {
-	return db.updateUnbondingDocumentState(ctx, unbondingTxHashHex, model.InputAlreadySpent)
+	id primitive.ObjectID) error {
+	return db.updateUnbondingDocumentState(ctx, id, model.InputAlreadySpent)
 }
 
 func (db *Database) SetUnbondingDocumentFailedToGetCovenantSignatures(
 	ctx context.Context,
-	unbondingTxHashHex string) error {
-	return db.updateUnbondingDocumentState(ctx, unbondingTxHashHex, model.FailedToGetCovenantSignatures)
+	id primitive.ObjectID) error {
+	return db.updateUnbondingDocumentState(ctx, id, model.FailedToGetCovenantSignatures)
 }

internal/db/model/unbonding.go

@@ -1,5 +1,7 @@
 package model
 
+import "go.mongodb.org/mongo-driver/bson/primitive"
+
 // TODO Double check with staking-api-service
 const (
 	UnbondingCollection   = "unbonding_queue"
@@ -17,15 +19,17 @@ const (
 )
 
 type UnbondingDocument struct {
-	StakerPkHex        string         `bson:"staker_pk_hex"`
-	FinalityPkHex      string         `bson:"finality_pk_hex"`
-	UnbondingTxSigHex  string         `bson:"unbonding_tx_sig_hex"`
-	State              UnbondingState `bson:"state"`
-	UnbondingTxHashHex string         `bson:"unbonding_tx_hash_hex"` // Unique Index
-	UnbondingTxHex     string         `bson:"unbonding_tx_hex"`
-	StakingTxHex       string         `bson:"staking_tx_hex"`
-	StakingOutputIndex uint64         `bson:"staking_output_index"`
-	StakingTimelock    uint64         `bson:"staking_timelock"`
-	StakingAmount      uint64         `bson:"staking_amount"`
-	StakingTxHashHex   string         `json:"staking_tx_hash_hex"`
+	// Can be nil only if we are inserting a document
+	ID                 *primitive.ObjectID `bson:"_id,omitempty"`
+	StakerPkHex        string              `bson:"staker_pk_hex"`
+	FinalityPkHex      string              `bson:"finality_pk_hex"`
+	UnbondingTxSigHex  string              `bson:"unbonding_tx_sig_hex"`
+	State              UnbondingState      `bson:"state"`
+	UnbondingTxHashHex string              `bson:"unbonding_tx_hash_hex"` // Unique Index
+	UnbondingTxHex     string              `bson:"unbonding_tx_hex"`
+	StakingTxHex       string              `bson:"staking_tx_hex"`
+	StakingOutputIndex uint64              `bson:"staking_output_index"`
+	StakingTimelock    uint64              `bson:"staking_timelock"`
+	StakingAmount      uint64              `bson:"staking_amount"`
+	StakingTxHashHex   string              `json:"staking_tx_hash_hex"`
 }

internal/services/expected_interfaces.go

@@ -8,6 +8,7 @@ import (
 	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/wire"
+	"go.mongodb.org/mongo-driver/bson/primitive"
 )
 
 // Minimal set of data necessary to sign unbonding transaction
@@ -85,6 +86,7 @@ type StakingTransactionData struct {
 }
 
 type UnbondingTxData struct {
+	UnbondingDocID           primitive.ObjectID
 	UnbondingTransaction     *wire.MsgTx
 	UnbondingTransactionHash *chainhash.Hash
 	UnbondingTransactionSig  *schnorr.Signature
@@ -97,13 +99,15 @@ func (u UnbondingTxData) StakingOutput() *wire.TxOut {
 }
 
 func NewUnbondingTxData(
+	id primitive.ObjectID,
 	tx *wire.MsgTx,
 	hash *chainhash.Hash,
 	sig *schnorr.Signature,
 	info *StakingInfo,
 	sd *StakingTransactionData,
 ) *UnbondingTxData {
 	return &UnbondingTxData{
+		UnbondingDocID:           id,
 		UnbondingTransaction:     tx,
 		UnbondingTransactionHash: hash,
 		UnbondingTransactionSig:  sig,

internal/services/in_mem_store.go

@@ -9,6 +9,7 @@ import (
 	"github.com/btcsuite/btcd/btcec/v2/schnorr"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/wire"
+	"go.mongodb.org/mongo-driver/bson/primitive"
 )
 
 type state string
@@ -35,8 +36,9 @@ func newUnbondingTxDataWithCounter(
 	sd *StakingTransactionData,
 	counter int,
 ) *unbondingTxDataWithCounter {
+	nullId := primitive.NilObjectID
 	return &unbondingTxDataWithCounter{
-		UnbondingTxData: *NewUnbondingTxData(tx, hash, sig, info, sd),
+		UnbondingTxData: *NewUnbondingTxData(nullId, tx, hash, sig, info, sd),
 		Counter:         counter,
 		state:           inserted,
 	}

internal/services/persistent_store.go

@@ -146,7 +146,7 @@ func documentToData(d *model.UnbondingDocument) (*UnbondingTxData, error) {
 		StakingOutputIdx:   d.StakingOutputIndex,
 	}
 
-	return NewUnbondingTxData(tx, unbondingTxHash, sig, si, sd), nil
+	return NewUnbondingTxData(*d.ID, tx, unbondingTxHash, sig, si, sd), nil
 }
 
 func (s *PersistentUnbondingStorage) AddTxWithSignature(
@@ -246,21 +246,17 @@ func (s *PersistentUnbondingStorage) GetNotProcessedUnbondingTransactions(ctx co
 }
 
 func (s *PersistentUnbondingStorage) SetUnbondingTransactionProcessed(ctx context.Context, utx *UnbondingTxData) error {
-	txHash := utx.UnbondingTransactionHash.String()
-	return s.client.SetUnbondingDocumentSend(ctx, txHash)
+	return s.client.SetUnbondingDocumentSend(ctx, utx.UnbondingDocID)
 }
 
 func (s *PersistentUnbondingStorage) SetUnbondingTransactionProcessingFailed(ctx context.Context, utx *UnbondingTxData) error {
-	txHash := utx.UnbondingTransactionHash.String()
-	return s.client.SetUnbondingDocumentFailed(ctx, txHash)
+	return s.client.SetUnbondingDocumentFailed(ctx, utx.UnbondingDocID)
 }
 
 func (s *PersistentUnbondingStorage) SetUnbondingTransactionInputAlreadySpent(ctx context.Context, utx *UnbondingTxData) error {
-	txHash := utx.UnbondingTransactionHash.String()
-	return s.client.SetUnbondingDocumentFailed(ctx, txHash)
+	return s.client.SetUnbondingDocumentInputAlreadySpent(ctx, utx.UnbondingDocID)
 }
 
 func (s *PersistentUnbondingStorage) SetUnbondingTransactionFailedToGetCovenantSignatures(ctx context.Context, utx *UnbondingTxData) error {
-	txHash := utx.UnbondingTransactionHash.String()
-	return s.client.SetUnbondingDocumentFailedToGetCovenantSignatures(ctx, txHash)
+	return s.client.SetUnbondingDocumentFailedToGetCovenantSignatures(ctx, utx.UnbondingDocID)
 }