feat(tracing): must specify days as flag in cleanup cmd

sensorbucket · Jan 23, 2025 · 5f673a1 · 5f673a1
1 parent 150a4ce
commit 5f673a1
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 5 deletions.
diff --git a/services/tracing/cmd/cleanDatabase.go b/services/tracing/cmd/cleanDatabase.go
@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/urfave/cli/v2"
 	"sensorbucket.nl/sensorbucket/services/tracing/tracing"
@@ -14,7 +15,7 @@ func cmdCleanDatabase(cmd *cli.Context) error {
 	}
 
 	svc := tracing.Create(db)
-	if err := svc.PeriodicCleanup(); err != nil {
+	if err := svc.PeriodicCleanup(time.Duration(cmd.Float64("days")) * time.Hour * 24); err != nil {
 		return fmt.Errorf("could not perform database cleanup: %w", err)
 	}
 

diff --git a/services/tracing/cmd/root.go b/services/tracing/cmd/root.go
@@ -12,6 +12,13 @@ var App = &cli.App{
 		{
 			Name:   "cleanup",
 			Action: cmdCleanDatabase,
+			Flags: []cli.Flag{
+				&cli.Float64Flag{
+					Name:     "days",
+					Usage:    "Keep only data since <value> ago. Fractions allowed, a day is specified as 24 hours",
+					Required: true,
+				},
+			},
 		},
 	},
 }
diff --git a/services/tracing/tracing/tracing.go b/services/tracing/tracing/tracing.go
@@ -272,12 +272,12 @@ func (svc *Service) Query(ctx context.Context, filters TraceFilter, r pagination
 	return &page, nil
 }
 
-func (svc *Service) PeriodicCleanup() error {
+func (svc *Service) PeriodicCleanup(duration time.Duration) error {
 	tx, err := svc.db.Beginx()
 	if err != nil {
 		return fmt.Errorf("")
 	}
-	n, err := tx.Exec(`DELETE FROM traces WHERE created_at < (NOW() - $1::INTERVAL)`, "1 day")
+	n, err := tx.Exec(`DELETE FROM traces WHERE created_at < (NOW() - $1::INTERVAL)`, duration)
 	if err != nil {
 		var rollbackErr error
 		if rbErr := tx.Rollback(); rbErr != nil {
@@ -290,7 +290,7 @@ func (svc *Service) PeriodicCleanup() error {
 		log.Printf("Cleaned %d traces\n", n)
 	}
 
-	n, err = tx.Exec(`DELETE FROM trace_steps WHERE queue_time < (NOW() - $1::INTERVAL)`, "1 day")
+	n, err = tx.Exec(`DELETE FROM trace_steps WHERE queue_time < (NOW() - $1::INTERVAL)`, duration)
 	if err != nil {
 		var rollbackErr error
 		if rbErr := tx.Rollback(); rbErr != nil {
@@ -303,7 +303,7 @@ func (svc *Service) PeriodicCleanup() error {
 		log.Printf("Cleaned %d trace steps\n", n)
 	}
 
-	n, err = tx.Exec(`DELETE FROM trace_ingress WHERE archived_at < (NOW() - $1::INTERVAL)`, "1 day")
+	n, err = tx.Exec(`DELETE FROM trace_ingress WHERE archived_at < (NOW() - $1::INTERVAL)`, duration)
 	if err != nil {
 		var rollbackErr error
 		if rbErr := tx.Rollback(); rbErr != nil {