feat(purge): Adding purge ability (#26)

* Adding purge ability

* Adding purge ability

* Adding init purge

README.md

@@ -10,56 +10,24 @@ the raw reports to Google Cloud Storage for further processing. This also allows
 application to be running at the same time. By allowing for MySQL or MongoDB as the database backend, this allows for
 data retention on a more reliable database.
 
-## Usage
-
-### Flags
-
-```text
-Usage of ./puppet-summary:
-  -db <string>
-        Database to use (default "sqlite"). Valid options are: sqlite, mysql, mongo.
-  -gcs
-        Enable Google Cloud Storage upload.
-  -upload-token <string>
-        Enables secure upload. Requires a token to be specified. This is useful if you want to prevent any unauthorised requests to the /upload endpoint.
-  -version
-        Print version and exit.
-```
-
-### Setup
-
-#### SQLite
-
-```shell
-./puppet-summary -db sqlite
-```
-
-This will create a `puppet-summary.db` file in the current directory.
+## Setup
 
 #### MySQL
 
-```shell
-./puppet-summary -db mysql
-```
-
-For this, you will be required to specify a `MYSQL_CONNECTION` environment variable with the connection string to your
-MySQL database. For example:
+When using MySQL, you will be required to specify a `MYSQL_CONNECTION` environment variable with the connection string
+to your MySQL database. For example:
 
 ```text
 MYSQL_CONNECTION="root:Password01@tcp(localhost:3306)/puppet-summary?timeout=90s&multiStatements=true&parseTime=true"
 ```
 
 #### MongoDB
 
-```shell
-./puppet-summary -db mongo
-```
-
-For this, you will be required to specify a `MONGO_URI` environment variable with the connection URI to your MongoDB
-database. For example:
+When using MongoDB, you will be required to specify a `MONGO_URI` environment variable with the connection URI to your
+MongoDB database. For example:
 
 ```text
-mongodb+srv://user:password@host/?retryWrites=true
+MONGO_URI="mongodb+srv://user:password@host/?retryWrites=true"
 ```
 
 #### Google Cloud Storage

cmd/summary/app.go

@@ -25,11 +25,15 @@ func newApp(_ *slog.Logger, r *mux.Router) *app {
 	}
 }
 
-func (a *app) run() error {
+func (a *app) run(purgeDays int) error {
 	if err := a.init(); err != nil {
 		return fmt.Errorf("error initializing app: %w", err)
 	}
 
+	if err := setupPurge(purgeDays); err != nil {
+		return fmt.Errorf("error setting up purge: %w", err)
+	}
+
 	if err := a.srv.ListenAndServe(); err != nil {
 		return fmt.Errorf("error running server: %w", err)
 	}
@@ -50,6 +54,7 @@ func (a *app) init() error {
 
 	a.r.HandleFunc(pathMetrics, middlewareHttp(promhttp.Handler().ServeHTTP, AuthOptionInternal)).Methods(http.MethodGet)
 	a.r.HandleFunc(pathHealth, middlewareHttp(healthHandler(), AuthOptionInternal)).Methods(http.MethodGet)
+
 	a.r.NotFoundHandler = request.NotFoundHandler()
 	a.r.MethodNotAllowedHandler = request.MethodNotAllowedHandler()
 

cmd/summary/cmd_serve.go

@@ -0,0 +1,63 @@
+package main
+
+import (
+	"context"
+	"flag"
+	"log/slog"
+
+	"github.com/Jacobbrewer1/puppet-summary/pkg/logging"
+	"github.com/google/subcommands"
+)
+
+type serveCmd struct {
+	// uploadToken is the token used to authenticate requests to the upload endpoint. If empty, the endpoint is not secure.
+	uploadToken string
+
+	// autoPurge is the number of days to keep data for. If 0 (or not set), data will not be purged.
+	autoPurge int
+
+	// dbType is the type of database to use.
+	dbType string
+
+	// gcs is whether to use Google Cloud Storage.
+	gcs bool
+}
+
+func (s *serveCmd) Name() string {
+	return "serve"
+}
+
+func (s *serveCmd) Synopsis() string {
+	return "Start the web server"
+}
+
+func (s *serveCmd) Usage() string {
+	return `serve:
+  Start the web server.
+`
+}
+
+func (s *serveCmd) SetFlags(f *flag.FlagSet) {
+	f.StringVar(&s.uploadToken, "upload-token", "", "The Bearer token used to authenticate requests to the upload endpoint.")
+	f.IntVar(&s.autoPurge, "auto-purge", 0, "The number of days to keep data for. If 0 (or not set), data will not be purged.")
+	f.StringVar(&s.dbType, "db", "sqlite", "The type of database to use. Valid values are 'sqlite', 'mysql', and 'mongodb'.")
+	f.BoolVar(&s.gcs, "gcs", false, "Whether to use Google Cloud Storage.")
+}
+
+func (s *serveCmd) Execute(_ context.Context, _ *flag.FlagSet, _ ...interface{}) subcommands.ExitStatus {
+	a, err := initializeApp()
+	if err != nil {
+		slog.Error("Error initializing application", slog.String(logging.KeyError, err.Error()))
+		return subcommands.ExitFailure
+	}
+	if err := s.generateConfig(); err != nil {
+		slog.Error("Error generating configuration", slog.String(logging.KeyError, err.Error()))
+		return subcommands.ExitFailure
+	}
+	slog.Debug("Starting application")
+	if err := a.run(s.autoPurge); err != nil {
+		slog.Error("Error running application", slog.String(logging.KeyError, err.Error()))
+		return subcommands.ExitFailure
+	}
+	return subcommands.ExitSuccess
+}

cmd/summary/cmd_version.go

@@ -0,0 +1,40 @@
+package main
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"runtime"
+
+	"github.com/google/subcommands"
+)
+
+type versionCmd struct{}
+
+func (v versionCmd) Name() string {
+	return "version"
+}
+
+func (v versionCmd) Synopsis() string {
+	return "Print application version information and exit"
+}
+
+func (v versionCmd) Usage() string {
+	return `version:
+  Print application version information and exit.
+`
+}
+
+func (v versionCmd) SetFlags(f *flag.FlagSet) {}
+
+func (v versionCmd) Execute(_ context.Context, _ *flag.FlagSet, _ ...interface{}) subcommands.ExitStatus {
+	fmt.Printf(
+		"Commit: %s\nRuntime: %s %s/%s\nDate: %s\n",
+		Commit,
+		runtime.Version(),
+		runtime.GOOS,
+		runtime.GOARCH,
+		Date,
+	)
+	return subcommands.ExitSuccess
+}

cmd/summary/config.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"flag"
 	"fmt"
 	"log/slog"
 
@@ -13,21 +12,31 @@ const (
 	appName = "summary"
 )
 
-var uploadToken = flag.String("upload-token", "", "The Bearer token used to authenticate requests to the upload endpoint.")
+var (
+	uploadToken string
+)
 
-func generateConfig() error {
-	err := dataaccess.ConnectDatabase()
+func (s *serveCmd) generateConfig() error {
+	err := dataaccess.ConnectDatabase(s.dbType)
 	if err != nil {
 		return fmt.Errorf("error connecting to database: %w", err)
 	}
-	err = dataaccess.ConnectGCS()
-	if err != nil {
-		return fmt.Errorf("error connecting to GCS: %w", err)
+	if s.gcs {
+		err = dataaccess.ConnectGCS()
+		if err != nil {
+			return fmt.Errorf("error connecting to GCS: %w", err)
+		}
 	}
-	if *uploadToken != "" {
+	if s.uploadToken != "" {
 		slog.Info("Upload token set, security on upload endpoint is enabled")
+		uploadToken = s.uploadToken
 	} else {
 		slog.Info("Upload token not set, upload endpoint is not secure")
 	}
+	if s.autoPurge != 0 {
+		slog.Info(fmt.Sprintf("Auto purge set to %d days", s.autoPurge))
+	} else {
+		slog.Info("Auto purge not set, data will not be purged")
+	}
 	return nil
 }

cmd/summary/main.go

@@ -1,14 +1,12 @@
 package main
 
 import (
+	"context"
 	"flag"
-	"fmt"
-	"log"
-	"log/slog"
 	"os"
-	"runtime"
 
 	"github.com/Jacobbrewer1/puppet-summary/pkg/logging"
+	"github.com/google/subcommands"
 )
 
 // Set at linking time
@@ -17,33 +15,20 @@ var (
 	Date   string
 )
 
-var versionFlag = flag.Bool("version", false, "Print version information and exit")
-
 func main() {
-	flag.Parse()
-	if *versionFlag {
-		fmt.Printf(
-			"Commit: %s\nRuntime: %s %s/%s\nDate: %s\n",
-			Commit,
-			runtime.Version(),
-			runtime.GOOS,
-			runtime.GOARCH,
-			Date,
-		)
-		os.Exit(0)
-	}
-
-	a, err := initializeApp()
+	_, err := logging.CommonLogger(logging.NewConfig(appName))
 	if err != nil {
-		log.Fatalln(err)
-	}
-	if err := generateConfig(); err != nil {
-		slog.Error("Error generating config", slog.String(logging.KeyError, err.Error()))
-		os.Exit(1)
-	}
-	slog.Debug("Starting application")
-	if err := a.run(); err != nil {
-		slog.Error("Error running application", slog.String(logging.KeyError, err.Error()))
-		os.Exit(1)
+		panic(err)
 	}
+
+	subcommands.Register(subcommands.HelpCommand(), "")
+	subcommands.Register(subcommands.FlagsCommand(), "")
+	subcommands.Register(subcommands.CommandsCommand(), "")
+
+	subcommands.Register(new(versionCmd), "")
+	subcommands.Register(new(serveCmd), "")
+
+	flag.Parse()
+	ctx := context.Background()
+	os.Exit(int(subcommands.Execute(ctx)))
 }

cmd/summary/purge.go

@@ -0,0 +1,44 @@
+package main
+
+import (
+	"fmt"
+	"log/slog"
+	"time"
+
+	"github.com/Jacobbrewer1/puppet-summary/pkg/dataaccess"
+	"github.com/Jacobbrewer1/puppet-summary/pkg/entities"
+	"github.com/robfig/cron"
+)
+
+func setupPurge(purgeDays int) error {
+	if purgeDays == 0 {
+		slog.Info("Auto purge not set, data will not be purged")
+		return nil
+	}
+
+	c := cron.New()
+
+	// Add a new entry to the cron scheduler to purge every (autoPurge) days at 03:00.
+	if err := c.AddFunc(fmt.Sprintf("0 3 */%d * *", purgeDays), func() { purgeData(purgeDays) }); err != nil {
+		return fmt.Errorf("error adding purge job to cron scheduler: %w", err)
+	}
+
+	c.Start()
+	slog.Debug("Cron scheduler started")
+	return nil
+}
+
+func purgeData(purgeDays int) {
+	slog.Info("Purging data")
+
+	// Get the start and end dates for the purge.
+	now := time.Now()
+	from := now.AddDate(0, 0, -purgeDays)
+
+	err := dataaccess.Purge(entities.Datetime(from))
+	if err != nil {
+		slog.Error(fmt.Sprintf("Error purging data: %s", err))
+	}
+
+	slog.Info("Purging complete")
+}

cmd/summary/upload.go

@@ -15,9 +15,9 @@ import (
 
 // uploadHandler takes the uploaded file and stores it in the database.
 func uploadHandler(w http.ResponseWriter, r *http.Request) {
-	if *uploadToken != "" {
+	if uploadToken != "" {
 		// Check the token.
-		if r.Header.Get("Authorization") != "Bearer "+*uploadToken {
+		if r.Header.Get("Authorization") != "Bearer "+uploadToken {
 			slog.Warn("Invalid token")
 			w.WriteHeader(http.StatusUnauthorized)
 			if err := json.NewEncoder(w).Encode(request.NewMessage(messages.ErrUnauthorized)); err != nil {

go.mod

@@ -6,10 +6,12 @@ require (
 	cloud.google.com/go/storage v1.36.0
 	github.com/alexliesenfeld/health v0.8.0
 	github.com/go-sql-driver/mysql v1.7.1
+	github.com/google/subcommands v1.0.1
 	github.com/google/wire v0.5.0
 	github.com/gorilla/mux v1.8.1
 	github.com/mattn/go-sqlite3 v1.14.19
 	github.com/prometheus/client_golang v1.18.0
+	github.com/robfig/cron v1.2.0
 	github.com/smallfish/simpleyaml v0.1.0
 	github.com/stretchr/testify v1.8.4
 	go.mongodb.org/mongo-driver v1.13.1

go.sum

@@ -75,6 +75,7 @@ github.com/google/martian/v3 v3.3.2 h1:IqNFLAmvJOgVlpdEBiQbDc2EwKW77amAycfTuWKdf
 github.com/google/martian/v3 v3.3.2/go.mod h1:oBOf6HBosgwRXnUGWUB05QECsc6uvmMiJ3+6W4l/CUk=
 github.com/google/s2a-go v0.1.7 h1:60BLSyTrOV4/haCDW4zb1guZItoSq8foHCXrAnjBo/o=
 github.com/google/s2a-go v0.1.7/go.mod h1:50CgR4k1jNlWBu4UfS4AcfhVe1r6pdZPygJ3R8F0Qdw=
+github.com/google/subcommands v1.0.1 h1:/eqq+otEXm5vhfBrbREPCSVQbvofip6kIz+mX5TUH7k=
 github.com/google/subcommands v1.0.1/go.mod h1:ZjhPrFU+Olkh9WazFPsl27BQ4UPiG37m3yTrtFlrHVk=
 github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/google/uuid v1.5.0 h1:1p67kYwdtXjb0gL0BPiP1Av9wiZPo5A8z2cWkTZ+eyU=
@@ -110,6 +111,8 @@ github.com/prometheus/common v0.45.0 h1:2BGz0eBc2hdMDLnO/8n0jeB3oPrt2D08CekT0lne
 github.com/prometheus/common v0.45.0/go.mod h1:YJmSTw9BoKxJplESWWxlbyttQR4uaEcGyv9MZjVOJsY=
 github.com/prometheus/procfs v0.12.0 h1:jluTpSng7V9hY0O2R9DzzJHYb2xULk9VTR1V1R/k6Bo=
 github.com/prometheus/procfs v0.12.0/go.mod h1:pcuDEFsWDnvcgNzo4EEweacyhjeA9Zk3cnaOZAZEfOo=
+github.com/robfig/cron v1.2.0 h1:ZjScXvvxeQ63Dbyxy76Fj3AT3Ut0aKsyd2/tl3DTMuQ=
+github.com/robfig/cron v1.2.0/go.mod h1:JGuDeoQd7Z6yL4zQhZ3OPEVHB7fL6Ka6skscFHfmt2k=
 github.com/rogpeppe/go-internal v1.11.0 h1:cWPaGQEPrBb5/AsnsZesgZZ9yb1OQ+GOISoDNXVBh4M=
 github.com/rogpeppe/go-internal v1.11.0/go.mod h1:ddIwULY96R17DhadqLgMfk9H9tvdUzkipdSkR5nkCZA=
 github.com/smallfish/simpleyaml v0.1.0 h1:5uAZdLAiHxS9cmzkOxg7lH0dILXKTD7uRZbAhyHmyU0=