From 1751e2babbcd5599c58ecb551305d9b52a6b96aa Mon Sep 17 00:00:00 2001
From: cupcakearmy <hi@nicco.io>
Date: Tue, 12 Apr 2022 22:34:21 +0200
Subject: [PATCH] add option to auto forget

---
 internal/location.go | 51 +++++++++++++++++++++++++++++++-------------
 1 file changed, 36 insertions(+), 15 deletions(-)

diff --git a/internal/location.go b/internal/location.go
index da589ec..9864cf1 100644
--- a/internal/location.go
+++ b/internal/location.go
@@ -24,22 +24,31 @@ const (
 
 type HookArray = []string
 
+type LocationForgetOption string
+
+const (
+	LocationForgetYes   LocationForgetOption = "yes"
+	LocationForgetNo    LocationForgetOption = "no"
+	LocationForgetPrune LocationForgetOption = "prune"
+)
+
 type Hooks struct {
-	Dir     string    `yaml:"dir"`
-	Before  HookArray `yaml:"before,omitempty"`
-	After   HookArray `yaml:"after,omitempty"`
-	Success HookArray `yaml:"success,omitempty"`
-	Failure HookArray `yaml:"failure,omitempty"`
+	Dir     string    `mapstructure:"dir"`
+	Before  HookArray `mapstructure:"before,omitempty"`
+	After   HookArray `mapstructure:"after,omitempty"`
+	Success HookArray `mapstructure:"success,omitempty"`
+	Failure HookArray `mapstructure:"failure,omitempty"`
 }
 
 type Location struct {
-	name    string   `yaml:",omitempty"`
-	From    []string `yaml:"from,omitempty"`
-	Type    string   `yaml:"type,omitempty"`
-	To      []string `yaml:"to,omitempty"`
-	Hooks   Hooks    `yaml:"hooks,omitempty"`
-	Cron    string   `yaml:"cron,omitempty"`
-	Options Options  `yaml:"options,omitempty"`
+	name         string               `mapstructure:",omitempty"`
+	From         []string             `mapstructure:"from,omitempty"`
+	Type         string               `mapstructure:"type,omitempty"`
+	To           []string             `mapstructure:"to,omitempty"`
+	Hooks        Hooks                `mapstructure:"hooks,omitempty"`
+	Cron         string               `mapstructure:"cron,omitempty"`
+	Options      Options              `mapstructure:"options,omitempty"`
+	ForgetOption LocationForgetOption `mapstructure:"forget,omitempty"`
 }
 
 func GetLocation(name string) (Location, bool) {
@@ -87,6 +96,12 @@ func (l Location) validate() error {
 			return fmt.Errorf("invalid backend `%s`", to)
 		}
 	}
+	// Check if forget type is correct
+	if l.ForgetOption != "" {
+		if l.ForgetOption != LocationForgetYes && l.ForgetOption != LocationForgetNo && l.ForgetOption != LocationForgetPrune {
+			return fmt.Errorf("invalid value for forget option: %s", l.ForgetOption)
+		}
+	}
 	return nil
 }
 
@@ -246,15 +261,21 @@ func (l Location) Backup(cron bool, specificBackend string) []error {
 
 after:
 	var commands []string
-	if len(errors) > 0 {
-		commands = l.Hooks.Failure
-	} else {
+	var isSuccess = len(errors) == 0
+	if isSuccess {
 		commands = l.Hooks.Success
+	} else {
+		commands = l.Hooks.Failure
 	}
 	if err := l.ExecuteHooks(commands, options); err != nil {
 		errors = append(errors, err)
 	}
 
+	// Forget and optionally prune
+	if isSuccess && l.ForgetOption != "" && l.ForgetOption != LocationForgetNo {
+		l.Forget(l.ForgetOption == LocationForgetPrune, false)
+	}
+
 	if len(errors) == 0 {
 		colors.Success.Println("Done")
 	}