Fix Scanner error: token limit exceeded

.gitignore

@@ -32,6 +32,7 @@ api/openapi-spec/api/v0/api.swagger.json
 .goreleaser.yaml
 .env
 test/stress/stress
+**/debug.test*
 
 gen/
 # helm charts deps

cmd/kwil-admin/cmds/snapshot/create.go

@@ -50,6 +50,7 @@ WHERE pg_stat_activity.datname = 'kwild'
 
 func createCmd() *cobra.Command {
 	var snapshotDir, dbName, dbUser, dbPass, dbHost, dbPort string
+	var maxRowSize int
 	cmd := &cobra.Command{
 		Use:     "create",
 		Short:   "Creates a snapshot of the database.",
@@ -61,7 +62,7 @@ func createCmd() *cobra.Command {
 				return fmt.Errorf("failed to expand snapshot directory path: %w", err)
 			}
 
-			return pgDump(cmd.Context(), dbName, dbUser, dbPass, dbHost, dbPort, snapshotDir)
+			return pgDump(cmd.Context(), dbName, dbUser, dbPass, dbHost, dbPort, maxRowSize, snapshotDir)
 		},
 	}
 
@@ -71,6 +72,7 @@ func createCmd() *cobra.Command {
 	cmd.Flags().StringVar(&dbPass, "password", "", "Password for the database user")
 	cmd.Flags().StringVar(&dbHost, "host", "localhost", "Host of the database")
 	cmd.Flags().StringVar(&dbPort, "port", "5432", "Port of the database")
+	cmd.Flags().IntVar(&maxRowSize, "max-row-size", 4*1024*1024, "Maximum row size to read from pg_dump (default: 4MB). Adjust this accordingly if you encounter 'bufio.Scanner: token too long' error.")
 
 	return cmd
 }
@@ -86,7 +88,7 @@ func expandPath(path string) (string, error) {
 	return filepath.Abs(path)
 }
 
-func pgDump(ctx context.Context, dbName, dbUser, dbPass, dbHost, dbPort, snapshotDir string) (err error) {
+func pgDump(ctx context.Context, dbName, dbUser, dbPass, dbHost, dbPort string, maxRowSize int, snapshotDir string) (err error) {
 	// Check if the snapshot directory exists, if not create it
 	err = os.MkdirAll(snapshotDir, 0755)
 	if err != nil {
@@ -170,7 +172,10 @@ func pgDump(ctx context.Context, dbName, dbUser, dbPass, dbHost, dbPort, snapsho
 	var totalBytes int64
 
 	// Pass the output of pg_dump through scanner to sanitize it
+	buf := make([]byte, maxRowSize)
 	scanner := bufio.NewScanner(pgDumpOutput)
+	scanner.Buffer(buf, maxRowSize)
+
 	for scanner.Scan() {
 		line := scanner.Text()
 

cmd/kwild/config/config.go

@@ -101,6 +101,7 @@ type SnapshotConfig struct {
 	RecurringHeight uint64 `mapstructure:"recurring_height"`
 	MaxSnapshots    uint64 `mapstructure:"max_snapshots"`
 	SnapshotDir     string `mapstructure:"snapshot_dir"`
+	MaxRowSize      int    `mapstructure:"max_row_size"`
 }
 
 type ChainRPCConfig struct {
@@ -579,6 +580,7 @@ func DefaultConfig() *KwildConfig {
 				RecurringHeight: 14400, // 1 day at 6s block time
 				MaxSnapshots:    3,
 				SnapshotDir:     SnapshotDirName,
+				MaxRowSize:      4 * 1024 * 1024,
 			},
 			GenesisState: "",
 		},

cmd/kwild/config/default_config.toml

@@ -117,6 +117,8 @@ recurring_height = 10000
 # Maximum number of snapshots to store
 max_snapshots = 3
 
+# Max row size that can be parsed by the snapshot store
+max_row_size = 4194304
 #######################################################################
 ###                    Logging Config Options                       ###
 #######################################################################

cmd/kwild/config/test_data/config.toml

@@ -110,6 +110,7 @@ recurring_height = 10000
 # Maximum number of snapshots to store
 max_snapshots = 3
 
+max_row_size = 4194304
 #######################################################################
 ###                    Logging Config Options                       ###
 #######################################################################

cmd/kwild/server/build.go

@@ -564,6 +564,7 @@ func buildSnapshotter(d *coreDependencies) *statesync.SnapshotStore {
 		SnapshotDir:     cfg.Snapshots.SnapshotDir,
 		RecurringHeight: cfg.Snapshots.RecurringHeight,
 		MaxSnapshots:    int(cfg.Snapshots.MaxSnapshots),
+		MaxRowSize:      cfg.Snapshots.MaxRowSize,
 	}
 
 	ss, err := statesync.NewSnapshotStore(snapshotCfg, dbCfg, *d.log.Named("snapshot-store"))

internal/statesync/snapshotter.go

@@ -43,13 +43,15 @@ const (
 type Snapshotter struct {
 	dbConfig    *DBConfig
 	snapshotDir string
+	maxRowSize  int
 	log         log.Logger
 }
 
-func NewSnapshotter(cfg *DBConfig, dir string, logger log.Logger) *Snapshotter {
+func NewSnapshotter(cfg *DBConfig, dir string, MaxRowSize int, logger log.Logger) *Snapshotter {
 	return &Snapshotter{
 		dbConfig:    cfg,
 		snapshotDir: dir,
+		maxRowSize:  MaxRowSize,
 		log:         logger,
 	}
 }
@@ -190,7 +192,10 @@ func (s *Snapshotter) sanitizeDump(height uint64, format uint32) ([]byte, error)
 	defer outputFile.Close()
 
 	// Scanner to read the dump file line by line
+	buf := make([]byte, s.maxRowSize)
 	scanner := bufio.NewScanner(dumpInst1)
+	scanner.Buffer(buf, s.maxRowSize)
+
 	var inCopyBlock bool
 	var lineHashes []hashedLine
 	var offset int64

internal/statesync/snapshotter_test.go

@@ -14,13 +14,14 @@ import (
 const (
 	dump1File = "test_data/dump1.sql"
 	dump2File = "test_data/dump2.sql"
+	tokenSize = 64 * 1024 // 64KB
 )
 
 func TestSanitizeLogicalDump(t *testing.T) {
 	dir := t.TempDir()
 	logger := log.NewStdOut(log.DebugLevel)
 	// Create a snapshotter
-	snapshotter := NewSnapshotter(nil, dir, logger)
+	snapshotter := NewSnapshotter(nil, dir, tokenSize, logger)
 
 	// Create snapshot directory
 	height := uint64(1)

internal/statesync/store.go

@@ -57,6 +57,7 @@ type SnapshotConfig struct {
 	SnapshotDir     string
 	MaxSnapshots    int
 	RecurringHeight uint64
+	MaxRowSize      int
 }
 
 type SnapshotStore struct {
@@ -80,7 +81,7 @@ type DBSnapshotter interface {
 }
 
 func NewSnapshotStore(cfg *SnapshotConfig, dbCfg *DBConfig, logger log.Logger) (*SnapshotStore, error) {
-	snapshotter := NewSnapshotter(dbCfg, cfg.SnapshotDir, logger)
+	snapshotter := NewSnapshotter(dbCfg, cfg.SnapshotDir, cfg.MaxRowSize, logger)
 	ss := &SnapshotStore{
 		cfg:         cfg,
 		snapshots:   make(map[uint64]*Snapshot),