Add reclassify command (#100)

internal/app/appfx/module.go

@@ -1,6 +1,7 @@
 package appfx
 
 import (
+	"github.com/bitmagnet-io/bitmagnet/internal/app/cmd/reclassifycmd"
 	"github.com/bitmagnet-io/bitmagnet/internal/app/cmd/reindexcmd"
 	"github.com/bitmagnet-io/bitmagnet/internal/app/cmd/torrentcmd"
 	"github.com/bitmagnet-io/bitmagnet/internal/blocking/blockingfx"
@@ -44,6 +45,7 @@ func New() fx.Option {
 		versionfx.New(),
 		// cli commands:
 		fx.Provide(
+			reclassifycmd.New,
 			reindexcmd.New,
 			torrentcmd.New,
 		),

internal/app/cmd/reclassifycmd/command.go

@@ -0,0 +1,77 @@
+package reclassifycmd
+
+import (
+	"github.com/bitmagnet-io/bitmagnet/internal/boilerplate/lazy"
+	"github.com/bitmagnet-io/bitmagnet/internal/classifier/asynq/message"
+	"github.com/bitmagnet-io/bitmagnet/internal/database/dao"
+	"github.com/bitmagnet-io/bitmagnet/internal/model"
+	"github.com/bitmagnet-io/bitmagnet/internal/protocol"
+	"github.com/bitmagnet-io/bitmagnet/internal/queue/publisher"
+	"github.com/schollz/progressbar/v3"
+	"github.com/urfave/cli/v2"
+	"go.uber.org/fx"
+	"go.uber.org/zap"
+	"gorm.io/gen"
+)
+
+type Params struct {
+	fx.In
+	Dao                 lazy.Lazy[*dao.Query]
+	ClassifierPublisher lazy.Lazy[publisher.Publisher[message.ClassifyTorrentPayload]]
+	Logger              *zap.SugaredLogger
+}
+
+type Result struct {
+	fx.Out
+	Command *cli.Command `group:"commands"`
+}
+
+func New(p Params) (Result, error) {
+	return Result{Command: &cli.Command{
+		Name:  "reclassify",
+		Usage: "Queue all torrents for reclassification",
+		Flags: []cli.Flag{
+			&cli.IntFlag{
+				Name:  "batchSize",
+				Value: 100,
+			},
+		},
+		Action: func(ctx *cli.Context) error {
+			println("queueing full reclassify...")
+			d, err := p.Dao.Get()
+			if err != nil {
+				return err
+			}
+			p, err := p.ClassifierPublisher.Get()
+			if err != nil {
+				return err
+			}
+			batchSize := ctx.Int("batchSize")
+			torrentCount := int64(0)
+			if result, err := d.Torrent.WithContext(ctx.Context).Count(); err != nil {
+				return err
+			} else {
+				torrentCount = result
+			}
+			bar := progressbar.Default(torrentCount, "queuing torrents")
+			var torrentResult []*model.Torrent
+			if err := d.Torrent.WithContext(ctx.Context).FindInBatches(&torrentResult, batchSize, func(tx gen.Dao, _ int) error {
+				infoHashes := make([]protocol.ID, 0, len(torrentResult))
+				for _, c := range torrentResult {
+					infoHashes = append(infoHashes, c.InfoHash)
+				}
+				if _, err := p.Publish(ctx.Context, message.ClassifyTorrentPayload{
+					InfoHashes: infoHashes,
+				}); err != nil {
+					return err
+				}
+				_ = bar.Add(len(torrentResult))
+				return nil
+			}); err != nil {
+				return err
+			}
+			_ = bar.Finish()
+			return nil
+		},
+	}}, nil
+}

internal/app/cmd/reindexcmd/command.go

@@ -28,12 +28,19 @@ func New(p Params) (Result, error) {
 	return Result{Command: &cli.Command{
 		Name:  "reindex",
 		Usage: "Reindex all records for full text search",
+		Flags: []cli.Flag{
+			&cli.IntFlag{
+				Name:  "batchSize",
+				Value: 1000,
+			},
+		},
 		Action: func(ctx *cli.Context) error {
 			println("reindexing...")
 			d, err := p.Dao.Get()
 			if err != nil {
 				return err
 			}
+			batchSize := ctx.Int("batchSize")
 			contentCount := int64(0)
 			torrentContentCount := int64(0)
 			if result, err := d.Content.WithContext(ctx.Context).Count(); err != nil {
@@ -47,7 +54,6 @@ func New(p Params) (Result, error) {
 				torrentContentCount = result
 			}
 			contentBar := progressbar.Default(contentCount, "[1/2] reindexing content")
-			batchSize := 1000
 			tsvs := make(map[model.ContentRef]fts.Tsvector)
 			var contentResult []*model.Content
 			if err := d.Content.WithContext(ctx.Context).Preload(