hanko/backend/cmd/cleanup/cleanup.go

package cleanup

import (
	"fmt"
	"github.com/spf13/cobra"
	"github.com/teamhanko/hanko/backend/config"
	"github.com/teamhanko/hanko/backend/persistence"
	"github.com/teamhanko/hanko/backend/persistence/models"
	"log"
	"sort"
	"strings"
	"time"
)

// options holds user-provided CLI options
type options struct {
	tables     []string // List of tables to clean up
	configFile string   // Path to configuration file
	pageSize   int      // The number of entities to query at once
	run        bool     // Whether to execute cleanup or simulate
}

// handlerParam holds the necessary parameters for cleanup operations
type handlerParam struct {
	table   string
	config  *config.Config
	storage persistence.Storage
	options *options
}

// handlerFunc defines the function signature for cleanup handlers
type handlerFunc func(handlerParam) error

// Table names used for cleanup operations
const (
	tableAuditLogs           = "audit_logs"
	tableFlows               = "flows"
	tableWebauthnSessionData = "webauthn_session_data"
)

// Map of table names to their respective cleanup handlers
var handler = map[string]handlerFunc{
	tableFlows: func(param handlerParam) error {
		return cleanup[models.Flow](param, param.storage.GetFlowPersister(), time.Now().UTC())
	},
	tableAuditLogs: func(param handlerParam) error {
		duration, err := time.ParseDuration(param.config.AuditLog.Retention)
		if err != nil {
			return fmt.Errorf("failed to parse the retention duration: %w", err)
		}

		return cleanup[models.AuditLog](param, param.storage.GetAuditLogPersister(), time.Now().Add(-duration).UTC())
	},
	tableWebauthnSessionData: func(param handlerParam) error {
		return cleanup[models.WebauthnSessionData](param, param.storage.GetWebauthnSessionDataPersister(), time.Now().UTC())
	},
}

// allowedTables is a list of table names that can be cleaned up
var allowedTables = func() []string {
	keys := make([]string, 0, len(handler))
	for key := range handler {
		keys = append(keys, key)
	}

	sort.Strings(keys)

	return keys
}()

// isTableAllowed checks if a given table name exists in the allowed list
func isTableAllowed(table string) bool {
	for _, allowed := range allowedTables {
		if table == allowed {
			return true
		}
	}
	return false
}

// validateTables checks if the specified table names exist in the allowed list
func validateTables(tables []string) error {
	var invalidTables []string

	for _, table := range tables {
		if !isTableAllowed(table) {
			invalidTables = append(invalidTables, table)
		}
	}

	if len(invalidTables) > 0 {
		return fmt.Errorf("invalid table name(s): %s - allowed values: %s",
			strings.Join(invalidTables, ", "), strings.Join(allowedTables, ", "))
	}

	return nil
}

// newCleanupCommand creates the Cobra command for database cleanup
func newCleanupCommand() *cobra.Command {
	opts := &options{}

	cmd := &cobra.Command{
		Use:   "cleanup",
		Short: "Cleanup the database.",
		Long:  `Cleans up the database by deleting expired entities.`,
		PreRunE: func(cmd *cobra.Command, args []string) error {
			if len(opts.tables) == 0 {
				opts.tables = allowedTables
				return nil
			}

			return validateTables(opts.tables)
		},
		RunE: func(cmd *cobra.Command, args []string) error {
			cfg, err := config.Load(&opts.configFile)
			if err != nil {
				log.Fatal(err)
			}

			storage, err := persistence.New(cfg.Database)
			if err != nil {
				log.Fatal(err)
			}

			log.Printf("Cleaning up table(s): %s...\n", strings.Join(opts.tables, ", "))

			for _, table := range opts.tables {
				param := handlerParam{
					table:   table,
					config:  cfg,
					storage: storage,
					options: opts,
				}
				err = handler[table](param)
				if err != nil {
					log.Fatal(err)
				}
			}

			log.Println("Cleanup completed.")

			if !opts.run {
				log.Println("This was a dry-run; add --run to the command to really delete the data.")
			}

			return nil
		},
	}

	cmd.Flags().StringVarP(&opts.configFile, "config", "c", config.DefaultConfigFilePath, "path to config file")
	cmd.Flags().StringSliceVarP(&opts.tables, "tables", "t", []string{}, fmt.Sprintf("specify individual tables to clean up (comma-separated) - allowed values: %s", strings.Join(allowedTables, ", ")))
	cmd.Flags().IntVarP(&opts.pageSize, "page-size", "s", 512, "the number of entities to query at once")
	cmd.Flags().BoolVar(&opts.run, "run", false, "execute the cleanup process instead of simulating")

	return cmd
}

// cleanup performs the cleanup operation for a given table and persister
func cleanup[T any](param handlerParam, persister persistence.Cleanup[T], cutoffTime time.Time) error {
	var (
		page    = 1
		deleted = 0
	)

	for {
		items, err := persister.FindExpired(cutoffTime, page, param.options.pageSize)
		if err != nil {
			return err
		}

		if len(items) > 0 {
			for _, item := range items {
				if param.options.run {
					err = persister.Delete(item)
					if err != nil {
						return err
					}
				}

				deleted++
			}

			log.Printf("Deleted %d %s in total.", deleted, param.table)

			if !param.options.run {
				page++
			}
		}

		if len(items) < param.options.pageSize {
			break
		}
	}

	return nil
}

// RegisterCommands registers the cleanup command with the parent command
func RegisterCommands(parent *cobra.Command) {
	cmd := newCleanupCommand()
	parent.AddCommand(cmd)
}