From 94fb48c2dc022c761725cbb02184e151e04f4a8d Mon Sep 17 00:00:00 2001
From: Mechiel Lukkien <mechiel@ueber.net>
Date: Sat, 7 Dec 2024 17:00:00 +0100
Subject: [PATCH] mox retrain: make the parameter, for account, optional and
 retrain all accounts when absent

for more easily retraining all accounts. users should be retraining their
accounts with the next release, due to the fix in the previous commit.
---
 ctl.go  | 111 +++++++++++++++++++++++++++++++-------------------------
 doc.go  |   6 +--
 main.go |  12 ++++--
 3 files changed, 72 insertions(+), 57 deletions(-)

diff --git a/ctl.go b/ctl.go
index 81f73e4f4..19ba7a85e 100644
--- a/ctl.go
+++ b/ctl.go
@@ -1335,67 +1335,78 @@ func servectlcmd(ctx context.Context, ctl *ctl, shutdown func()) {
 	case "retrain":
 		/* protocol:
 		> "retrain"
-		> account
+		> account or empty
 		< "ok" or error
 		*/
 		account := ctl.xread()
-		acc, err := store.OpenAccount(log, account)
-		ctl.xcheck(err, "open account")
-		defer func() {
-			if acc != nil {
-				err := acc.Close()
-				log.Check(err, "closing account after retraining")
-			}
-		}()
-
-		// todo: can we retrain an account without holding a write lock? perhaps by writing a junkfilter to a new location, and staying informed of message changes while we go through all messages in the account?
-
-		acc.WithWLock(func() {
-			conf, _ := acc.Conf()
-			if conf.JunkFilter == nil {
-				ctl.xcheck(store.ErrNoJunkFilter, "looking for junk filter")
-			}
 
-			// Remove existing junk filter files.
-			basePath := mox.DataDirPath("accounts")
-			dbPath := filepath.Join(basePath, acc.Name, "junkfilter.db")
-			bloomPath := filepath.Join(basePath, acc.Name, "junkfilter.bloom")
-			err := os.Remove(dbPath)
-			log.Check(err, "removing old junkfilter database file", slog.String("path", dbPath))
-			err = os.Remove(bloomPath)
-			log.Check(err, "removing old junkfilter bloom filter file", slog.String("path", bloomPath))
-
-			// Open junk filter, this creates new files.
-			jf, _, err := acc.OpenJunkFilter(ctx, log)
-			ctl.xcheck(err, "open new junk filter")
+		xretrain := func(name string) {
+			acc, err := store.OpenAccount(log, name)
+			ctl.xcheck(err, "open account")
 			defer func() {
-				if jf == nil {
-					return
+				if acc != nil {
+					err := acc.Close()
+					log.Check(err, "closing account after retraining")
 				}
-				err := jf.Close()
-				log.Check(err, "closing junk filter during cleanup")
 			}()
 
-			// Read through messages with junk or nonjunk flag set, and train them.
-			var total, trained int
-			q := bstore.QueryDB[store.Message](ctx, acc.DB)
-			q.FilterEqual("Expunged", false)
-			err = q.ForEach(func(m store.Message) error {
-				total++
-				ok, err := acc.TrainMessage(ctx, log, jf, m)
-				if ok {
-					trained++
+			// todo: can we retrain an account without holding a write lock? perhaps by writing a junkfilter to a new location, and staying informed of message changes while we go through all messages in the account?
+
+			acc.WithWLock(func() {
+				conf, _ := acc.Conf()
+				if conf.JunkFilter == nil {
+					ctl.xcheck(store.ErrNoJunkFilter, "looking for junk filter")
 				}
-				return err
+
+				// Remove existing junk filter files.
+				basePath := mox.DataDirPath("accounts")
+				dbPath := filepath.Join(basePath, acc.Name, "junkfilter.db")
+				bloomPath := filepath.Join(basePath, acc.Name, "junkfilter.bloom")
+				err := os.Remove(dbPath)
+				log.Check(err, "removing old junkfilter database file", slog.String("path", dbPath))
+				err = os.Remove(bloomPath)
+				log.Check(err, "removing old junkfilter bloom filter file", slog.String("path", bloomPath))
+
+				// Open junk filter, this creates new files.
+				jf, _, err := acc.OpenJunkFilter(ctx, log)
+				ctl.xcheck(err, "open new junk filter")
+				defer func() {
+					if jf == nil {
+						return
+					}
+					err := jf.Close()
+					log.Check(err, "closing junk filter during cleanup")
+				}()
+
+				// Read through messages with junk or nonjunk flag set, and train them.
+				var total, trained int
+				q := bstore.QueryDB[store.Message](ctx, acc.DB)
+				q.FilterEqual("Expunged", false)
+				err = q.ForEach(func(m store.Message) error {
+					total++
+					ok, err := acc.TrainMessage(ctx, log, jf, m)
+					if ok {
+						trained++
+					}
+					return err
+				})
+				ctl.xcheck(err, "training messages")
+				log.Info("retrained messages", slog.Int("total", total), slog.Int("trained", trained))
+
+				// Close junk filter, marking success.
+				err = jf.Close()
+				jf = nil
+				ctl.xcheck(err, "closing junk filter")
 			})
-			ctl.xcheck(err, "training messages")
-			log.Info("retrained messages", slog.Int("total", total), slog.Int("trained", trained))
+		}
 
-			// Close junk filter, marking success.
-			err = jf.Close()
-			jf = nil
-			ctl.xcheck(err, "closing junk filter")
-		})
+		if account == "" {
+			for _, name := range mox.Conf.Accounts() {
+				xretrain(name)
+			}
+		} else {
+			xretrain(account)
+		}
 		ctl.xwriteok()
 
 	case "recalculatemailboxcounts":
diff --git a/doc.go b/doc.go
index cc262871f..ce2b22533 100644
--- a/doc.go
+++ b/doc.go
@@ -105,7 +105,7 @@ any parameters. Followed by the help and usage information for each command.
 	mox dnsbl check zone ip
 	mox dnsbl checkhealth zone
 	mox mtasts lookup domain
-	mox retrain accountname
+	mox retrain [accountname]
 	mox sendmail [-Fname] [ignoredflags] [-t] [<message]
 	mox spf check domain ip
 	mox spf lookup domain
@@ -1390,12 +1390,12 @@ should be used, and how long the policy can be cached.
 
 # mox retrain
 
-Recreate and retrain the junk filter for the account.
+Recreate and retrain the junk filter for the account or all accounts.
 
 Useful after having made changes to the junk filter configuration, or if the
 implementation has changed.
 
-	usage: mox retrain accountname
+	usage: mox retrain [accountname]
 
 # mox sendmail
 
diff --git a/main.go b/main.go
index 41b6eb9e9..57c586e77 100644
--- a/main.go
+++ b/main.go
@@ -2874,19 +2874,23 @@ should be used, and how long the policy can be cached.
 }
 
 func cmdRetrain(c *cmd) {
-	c.params = "accountname"
-	c.help = `Recreate and retrain the junk filter for the account.
+	c.params = "[accountname]"
+	c.help = `Recreate and retrain the junk filter for the account or all accounts.
 
 Useful after having made changes to the junk filter configuration, or if the
 implementation has changed.
 `
 	args := c.Parse()
-	if len(args) != 1 {
+	if len(args) > 1 {
 		c.Usage()
 	}
+	var account string
+	if len(args) == 1 {
+		account = args[0]
+	}
 
 	mustLoadConfig()
-	ctlcmdRetrain(xctl(), args[0])
+	ctlcmdRetrain(xctl(), account)
 }
 
 func ctlcmdRetrain(ctl *ctl, account string) {