From d48b5e8674a7931bc605630dcb8e6e97bf6af364 Mon Sep 17 00:00:00 2001 From: Ralph Slooten Date: Tue, 6 Aug 2024 17:23:28 +1200 Subject: [PATCH] Feature: Add option to control message retention by age (#338) --- cmd/root.go | 4 +++ config/config.go | 49 ++++++++++++++++++++++++++++-- internal/storage/cron.go | 64 ++++++++++++++++++++++++++++++---------- 3 files changed, 100 insertions(+), 17 deletions(-) diff --git a/cmd/root.go b/cmd/root.go index 14d3221..03b6376 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -85,6 +85,7 @@ func init() { rootCmd.Flags().StringVar(&config.Label, "label", config.Label, "Optional label identify this Mailpit instance") rootCmd.Flags().StringVar(&config.TenantID, "tenant-id", config.TenantID, "Database tenant ID to isolate data") rootCmd.Flags().IntVarP(&config.MaxMessages, "max", "m", config.MaxMessages, "Max number of messages to store") + rootCmd.Flags().StringVar(&config.MaxAge, "max-age", config.MaxAge, "Max age of messages in either (h)ours or (d)ays (eg: 3d)") rootCmd.Flags().BoolVar(&config.UseMessageDates, "use-message-dates", config.UseMessageDates, "Use message dates as the received dates") rootCmd.Flags().BoolVar(&config.IgnoreDuplicateIDs, "ignore-duplicate-ids", config.IgnoreDuplicateIDs, "Ignore duplicate messages (by Message-Id)") rootCmd.Flags().StringVar(&logger.LogFile, "log-file", logger.LogFile, "Log output to file instead of stdout") @@ -179,6 +180,9 @@ func initConfigFromEnv() { if len(os.Getenv("MP_MAX_MESSAGES")) > 0 { config.MaxMessages, _ = strconv.Atoi(os.Getenv("MP_MAX_MESSAGES")) } + if len(os.Getenv("MP_MAX_AGE")) > 0 { + config.MaxAge = os.Getenv("MP_MAX_AGE") + } if getEnabledFromEnv("MP_USE_MESSAGE_DATES") { config.UseMessageDates = true } diff --git a/config/config.go b/config/config.go index c2fdaa6..129de05 100644 --- a/config/config.go +++ b/config/config.go @@ -10,6 +10,7 @@ import ( "path" "path/filepath" "regexp" + "strconv" "strings" "github.com/axllent/mailpit/internal/auth" @@ -31,15 +32,22 @@ var ( // TenantID is an optional prefix to be applied to all database tables, // allowing multiple isolated instances of Mailpit to share a database. - TenantID = "" + TenantID string // Label to identify this Mailpit instance (optional). // This gets applied to web UI, SMTP and optional POP3 server. - Label = "" + Label string // MaxMessages is the maximum number of messages a mailbox can have (auto-pruned every minute) MaxMessages = 500 + // MaxAge is the maximum age of messages (auto-pruned every hour). + // Value can be either h for hours or d for days + MaxAge string + + // MaxAgeInHours is the maximum age of messages in hours, set with parseMaxAge() using MaxAge value + MaxAgeInHours int + // UseMessageDates sets the Created date using the message date, not the delivered date UseMessageDates bool @@ -218,6 +226,10 @@ func VerifyConfig() error { Label = tools.Normalize(Label) + if err := parseMaxAge(); err != nil { + return err + } + TenantID = tools.Normalize(TenantID) if TenantID != "" { logger.Log().Infof("[db] using tenant \"%s\"", TenantID) @@ -458,6 +470,39 @@ func VerifyConfig() error { return nil } +// Parse the --max-age value (if set) +func parseMaxAge() error { + if MaxAge == "" { + return nil + } + + re := regexp.MustCompile(`^\d+(h|d)$`) + if !re.MatchString(MaxAge) { + return fmt.Errorf("max-age must be either h for hours or d for days: %s", MaxAge) + } + + if strings.HasSuffix(MaxAge, "h") { + hours, err := strconv.Atoi(strings.TrimSuffix(MaxAge, "h")) + if err != nil { + return err + } + + MaxAgeInHours = hours + + return nil + } + + days, err := strconv.Atoi(strings.TrimSuffix(MaxAge, "d")) + if err != nil { + return err + } + + logger.Log().Debugf("[db] auto-deleting messages older than %s", MaxAge) + + MaxAgeInHours = days * 24 + return nil +} + // Parse the SMTPRelayConfigFile (if set) func parseRelayConfig(c string) error { if c == "" { diff --git a/internal/storage/cron.go b/internal/storage/cron.go index 30d1c09..884f1c3 100644 --- a/internal/storage/cron.go +++ b/internal/storage/cron.go @@ -9,6 +9,7 @@ import ( "github.com/axllent/mailpit/config" "github.com/axllent/mailpit/internal/logger" + "github.com/axllent/mailpit/internal/tools" "github.com/axllent/mailpit/server/websockets" "github.com/leporo/sqlf" ) @@ -48,34 +49,67 @@ func dbCron() { // PruneMessages will auto-delete the oldest messages if messages > config.MaxMessages. // Set config.MaxMessages to 0 to disable. func pruneMessages() { - if config.MaxMessages < 1 { + if config.MaxMessages < 1 && config.MaxAgeInHours == 0 { return } start := time.Now() - q := sqlf.Select("ID, Size"). - From(tenant("mailbox")). - OrderBy("Created DESC"). - Limit(5000). - Offset(config.MaxMessages) - ids := []string{} var prunedSize int64 var size float64 - if err := q.QueryAndClose(context.TODO(), db, func(row *sql.Rows) { - var id string - if err := row.Scan(&id, &size); err != nil { + // prune using `--max` if set + if config.MaxMessages > 0 { + q := sqlf.Select("ID, Size"). + From(tenant("mailbox")). + OrderBy("Created DESC"). + Limit(5000). + Offset(config.MaxMessages) + + if err := q.QueryAndClose(context.TODO(), db, func(row *sql.Rows) { + var id string + + if err := row.Scan(&id, &size); err != nil { + logger.Log().Errorf("[db] %s", err.Error()) + return + } + ids = append(ids, id) + prunedSize = prunedSize + int64(size) + + }); err != nil { logger.Log().Errorf("[db] %s", err.Error()) return } - ids = append(ids, id) - prunedSize = prunedSize + int64(size) + } - }); err != nil { - logger.Log().Errorf("[db] %s", err.Error()) - return + // prune using `--max-age` if set + if config.MaxAgeInHours > 0 { + // now() minus the number of hours + ts := time.Now().Add(time.Duration(-config.MaxAgeInHours) * time.Hour).UnixMilli() + + q := sqlf.Select("ID, Size"). + From(tenant("mailbox")). + Where("Created < ?", ts). + Limit(5000) + + if err := q.QueryAndClose(context.TODO(), db, func(row *sql.Rows) { + var id string + + if err := row.Scan(&id, &size); err != nil { + logger.Log().Errorf("[db] %s", err.Error()) + return + } + + if !tools.InArray(id, ids) { + ids = append(ids, id) + prunedSize = prunedSize + int64(size) + } + + }); err != nil { + logger.Log().Errorf("[db] %s", err.Error()) + return + } } if len(ids) == 0 {