Skip more

internal/database/snapshot.go

@@ -13,7 +13,6 @@ import (
 	"os"
 	"path/filepath"
 	"sort"
-	"strings"
 
 	"gitlab.com/accumulatenetwork/accumulate/exp/ioutil"
 	"gitlab.com/accumulatenetwork/accumulate/internal/database/smt/storage"
@@ -35,6 +34,7 @@ type CollectOptions struct {
 	BuildIndex     bool
 	Predicate      func(database.Record) (bool, error)
 	DidWriteHeader func(*snapshot.Writer) error
+	KeepMessage    func(messaging.Message) (bool, error)
 
 	Metrics *CollectMetrics
 }
@@ -516,7 +516,6 @@ func collectMessageHashes(a *Account, hashes *indexing.Bucket, opts *CollectOpti
 			continue
 		}
 
-		isSig := strings.EqualFold(chain.Name, "signature")
 		c, err := a.ChainByName(chain.Name)
 		if err != nil {
 			return errors.InvalidRecord.Wrap(err)
@@ -553,25 +552,9 @@ func collectMessageHashes(a *Account, hashes *indexing.Bucket, opts *CollectOpti
 			return errors.UnknownError.WithFormat("load %s chain entries: %w", c.Name(), err)
 		}
 		for _, h := range entries {
-			err = hashes.Write(*(*[32]byte)(h), nil)
+			err = collectMessageHash(a, c, hashes, opts, *(*[32]byte)(h))
 			if err != nil {
-				return errors.UnknownError.WithFormat("record %s chain entry: %w", c.Name(), err)
-			}
-
-			if !isSig {
-				continue
-			}
-
-			msg, err := a.parent.newMessage(messageKey{*(*[32]byte)(h)}).Main().Get()
-			if err != nil {
-				slog.Error("Failed to collect message", "account", a.Url(), "hash", logging.AsHex(h), "error", err)
-				continue
-			}
-			if msg, ok := msg.(messaging.MessageForTransaction); ok {
-				err = hashes.Write(msg.GetTxID().Hash(), nil)
-				if err != nil {
-					return errors.UnknownError.WithFormat("record %s chain entry: %w", c.Name(), err)
-				}
+				return errors.UnknownError.Wrap(err)
 			}
 		}
 		if opts.Metrics != nil {
@@ -582,6 +565,36 @@ func collectMessageHashes(a *Account, hashes *indexing.Bucket, opts *CollectOpti
 	return errors.UnknownError.Wrap(err)
 }
 
+func collectMessageHash(a *Account, c *Chain2, hashes *indexing.Bucket, opts *CollectOptions, h [32]byte) error {
+	err := hashes.Write(h, nil)
+	if err != nil {
+		return errors.UnknownError.WithFormat("record %s chain entry: %w", c.Name(), err)
+	}
+
+	msg, err := a.parent.newMessage(messageKey{h}).Main().Get()
+	if err != nil {
+		slog.Error("Failed to collect message", "account", a.Url(), "hash", logging.AsHex(h), "error", err)
+		return nil
+	}
+
+	if opts.KeepMessage != nil {
+		ok, err := opts.KeepMessage(msg)
+		if err != nil {
+			return errors.UnknownError.Wrap(err)
+		}
+		if !ok {
+			return nil
+		}
+	}
+
+	forTxn, ok := msg.(messaging.MessageForTransaction)
+	if !ok {
+		return nil
+	}
+
+	return collectMessageHash(a, c, hashes, opts, forTxn.GetTxID().Hash())
+}
+
 func writeSnapshotIndex(w *snapshot.Writer, index *indexing.Bucket, opts *CollectOptions) error {
 	if !opts.BuildIndex {
 		return nil

tools/cmd/debug/snap_collect.go

@@ -19,6 +19,7 @@ import (
 	"gitlab.com/accumulatenetwork/accumulate/pkg/database"
 	"gitlab.com/accumulatenetwork/accumulate/pkg/database/keyvalue"
 	"gitlab.com/accumulatenetwork/accumulate/pkg/database/keyvalue/remote"
+	"gitlab.com/accumulatenetwork/accumulate/pkg/types/messaging"
 	"gitlab.com/accumulatenetwork/accumulate/pkg/url"
 	"gitlab.com/accumulatenetwork/accumulate/protocol"
 )
@@ -33,6 +34,8 @@ var cmdSnapCollect = &cobra.Command{
 var flagSnapCollect = struct {
 	SkipBPT    bool
 	SkipSystem bool
+	SkipSigs   bool
+	SkipTokens bool
 	Indexed    bool
 	Partition  UrlFlag
 }{}
@@ -41,6 +44,8 @@ func init() {
 	cmdSnap.AddCommand(cmdSnapCollect)
 	cmdSnapCollect.Flags().BoolVar(&flagSnapCollect.SkipBPT, "skip-bpt", false, "Skip the BPT")
 	cmdSnapCollect.Flags().BoolVar(&flagSnapCollect.SkipSystem, "skip-system", false, "Skip system accounts")
+	cmdSnapCollect.Flags().BoolVar(&flagSnapCollect.SkipSigs, "skip-signatures", false, "Skip signatures")
+	cmdSnapCollect.Flags().BoolVar(&flagSnapCollect.SkipTokens, "skip-token-txns", false, "Skip token transactions")
 	cmdSnapCollect.Flags().BoolVar(&flagSnapCollect.Indexed, "indexed", false, "Make an indexed snapshot")
 	cmdSnapCollect.Flags().Var(&flagSnapCollect.Partition, "partition", "Specify the partition instead of determining it from the database")
 }
@@ -114,6 +119,25 @@ func collectSnapshot(_ *cobra.Command, args []string) {
 			// Retain everything
 			return true, nil
 		},
+		KeepMessage: func(msg messaging.Message) (bool, error) {
+			switch msg := msg.(type) {
+			case *messaging.SignatureMessage:
+				return !flagSnapCollect.SkipSigs, nil
+			case *messaging.TransactionMessage:
+				if !flagSnapCollect.SkipTokens {
+					break
+				}
+				switch msg.Transaction.Body.Type() {
+				case protocol.TransactionTypeSendTokens,
+					protocol.TransactionTypeBurnTokens,
+					protocol.TransactionTypeAddCredits,
+					protocol.TransactionTypeSyntheticDepositTokens,
+					protocol.TransactionTypeSyntheticDepositCredits:
+					return false, nil
+				}
+			}
+			return true, nil
+		},
 	}))
 }