feat(cmd): add cleanup command (#1489)

* **New Features**
* Added `cleanup` command to remove old DAG run history with
configurable retention periods
* Supports `--dry-run` flag to preview which runs would be removed
without deleting
  * Includes `--yes` flag to skip confirmation prompts
This commit is contained in:
Yota Hamada 2025-12-16 23:36:24 +09:00 committed by GitHub
parent 5d6e50df04
commit c72623154d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
14 changed files with 527 additions and 32 deletions

View File

@ -45,6 +45,7 @@ func init() {
rootCmd.AddCommand(cmd.Retry())
rootCmd.AddCommand(cmd.StartAll())
rootCmd.AddCommand(cmd.Migrate())
rootCmd.AddCommand(cmd.Cleanup())
config.Version = version
}

134
internal/cmd/cleanup.go Normal file
View File

@ -0,0 +1,134 @@
package cmd
import (
"bufio"
"fmt"
"os"
"strconv"
"strings"
"github.com/dagu-org/dagu/internal/core/execution"
"github.com/spf13/cobra"
)
// Cleanup creates and returns a cobra command for removing old DAG run history.
func Cleanup() *cobra.Command {
return NewCommand(
&cobra.Command{
Use: "cleanup [flags] <DAG name>",
Short: "Remove old DAG run history",
Long: `Remove old DAG run history for a specified DAG.
By default, removes all history except for currently active runs.
Use --retention-days to keep recent history.
Active runs are never deleted for safety.
Examples:
dagu cleanup my-workflow # Delete all history (with confirmation)
dagu cleanup --retention-days 30 my-workflow # Keep last 30 days
dagu cleanup --dry-run my-workflow # Preview what would be deleted
dagu cleanup -y my-workflow # Skip confirmation
`,
Args: cobra.ExactArgs(1),
},
cleanupFlags,
runCleanup,
)
}
var cleanupFlags = []commandLineFlag{
retentionDaysFlag,
dryRunFlag,
yesFlag,
}
func runCleanup(ctx *Context, args []string) error {
dagName := args[0]
// Parse retention days (flags are string-based in this codebase)
retentionStr, err := ctx.StringParam("retention-days")
if err != nil {
return fmt.Errorf("failed to get retention-days: %w", err)
}
retentionDays, err := strconv.Atoi(retentionStr)
if err != nil {
return fmt.Errorf("invalid retention-days value %q: must be a non-negative integer", retentionStr)
}
// Reject negative retention (clearer error than silent no-op)
if retentionDays < 0 {
return fmt.Errorf("retention-days cannot be negative (got %d)", retentionDays)
}
// Get boolean flags
dryRun, _ := ctx.Command.Flags().GetBool("dry-run")
skipConfirm, _ := ctx.Command.Flags().GetBool("yes")
// Build description message
var actionDesc string
if retentionDays == 0 {
actionDesc = fmt.Sprintf("all history for DAG %q", dagName)
} else {
actionDesc = fmt.Sprintf("history older than %d days for DAG %q", retentionDays, dagName)
}
// Build options for RemoveOldDAGRuns
var opts []execution.RemoveOldDAGRunsOption
if dryRun {
opts = append(opts, execution.WithDryRun())
}
// Dry run mode - show what would be deleted
if dryRun {
runIDs, err := ctx.DAGRunStore.RemoveOldDAGRuns(ctx, dagName, retentionDays, opts...)
if err != nil {
return fmt.Errorf("failed to check history for %q: %w", dagName, err)
}
if len(runIDs) == 0 {
fmt.Printf("Dry run: No runs to delete for DAG %q\n", dagName)
} else {
fmt.Printf("Dry run: Would delete %d run(s) for DAG %q:\n", len(runIDs), dagName)
for _, runID := range runIDs {
fmt.Printf(" - %s\n", runID)
}
}
return nil
}
// Confirmation prompt (unless --yes or --quiet)
if !skipConfirm && !ctx.Quiet {
fmt.Printf("This will delete %s.\n", actionDesc)
fmt.Println("Active runs will be preserved.")
fmt.Print("Continue? [y/N]: ")
reader := bufio.NewReader(os.Stdin)
response, err := reader.ReadString('\n')
if err != nil {
return fmt.Errorf("failed to read user input: %w", err)
}
response = strings.TrimSpace(strings.ToLower(response))
if response != "y" && response != "yes" {
fmt.Println("Cancelled.")
return nil
}
}
// Execute cleanup using the existing DAGRunStore method
runIDs, err := ctx.DAGRunStore.RemoveOldDAGRuns(ctx, dagName, retentionDays, opts...)
if err != nil {
return fmt.Errorf("failed to cleanup history for %q: %w", dagName, err)
}
if !ctx.Quiet {
if len(runIDs) == 0 {
fmt.Printf("No runs to delete for DAG %q\n", dagName)
} else {
fmt.Printf("Successfully removed %d run(s) for DAG %q\n", len(runIDs), dagName)
}
}
return nil
}

View File

@ -0,0 +1,280 @@
package cmd_test
import (
"os"
"path/filepath"
"testing"
"time"
"github.com/dagu-org/dagu/internal/cmd"
"github.com/dagu-org/dagu/internal/core"
"github.com/dagu-org/dagu/internal/core/execution"
"github.com/dagu-org/dagu/internal/test"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestCleanupCommand(t *testing.T) {
t.Run("DeletesAllHistoryWithRetentionZero", func(t *testing.T) {
th := test.SetupCommand(t)
// Create a DAG and run it to generate history
dag := th.DAG(t, `steps:
- name: "1"
command: echo "hello"
`)
// Run the DAG to create history
th.RunCommand(t, cmd.Start(), test.CmdTest{
Args: []string{"start", dag.Location},
})
// Wait for DAG to complete
dag.AssertLatestStatus(t, core.Succeeded)
// Verify history exists
dag.AssertDAGRunCount(t, 1)
// Run cleanup with --yes to skip confirmation
th.RunCommand(t, cmd.Cleanup(), test.CmdTest{
Args: []string{"cleanup", "--yes", dag.Name},
})
// Verify history is deleted
dag.AssertDAGRunCount(t, 0)
})
t.Run("PreservesRecentHistoryWithRetentionDays", func(t *testing.T) {
th := test.SetupCommand(t)
// Create a DAG and run it
dag := th.DAG(t, `steps:
- name: "1"
command: echo "hello"
`)
// Run the DAG to create history
th.RunCommand(t, cmd.Start(), test.CmdTest{
Args: []string{"start", dag.Location},
})
// Wait for DAG to complete
dag.AssertLatestStatus(t, core.Succeeded)
// Verify history exists
dag.AssertDAGRunCount(t, 1)
// Run cleanup with retention of 30 days (should keep recent history)
th.RunCommand(t, cmd.Cleanup(), test.CmdTest{
Args: []string{"cleanup", "--retention-days", "30", "--yes", dag.Name},
})
// Verify history is still there (it's less than 30 days old)
dag.AssertDAGRunCount(t, 1)
})
t.Run("DryRunDoesNotDelete", func(t *testing.T) {
th := test.SetupCommand(t)
// Create a DAG and run it
dag := th.DAG(t, `steps:
- name: "1"
command: echo "hello"
`)
// Run the DAG to create history
th.RunCommand(t, cmd.Start(), test.CmdTest{
Args: []string{"start", dag.Location},
})
// Wait for DAG to complete
dag.AssertLatestStatus(t, core.Succeeded)
// Verify history exists
dag.AssertDAGRunCount(t, 1)
// Run cleanup with --dry-run
th.RunCommand(t, cmd.Cleanup(), test.CmdTest{
Args: []string{"cleanup", "--dry-run", dag.Name},
})
// Verify history is still there (dry run should not delete)
dag.AssertDAGRunCount(t, 1)
})
t.Run("PreservesActiveRuns", func(t *testing.T) {
th := test.SetupCommand(t)
// Create a DAG that runs for a while
dag := th.DAG(t, `steps:
- name: "1"
command: sleep 30
`)
done := make(chan struct{})
go func() {
// Start the DAG
th.RunCommand(t, cmd.Start(), test.CmdTest{
Args: []string{"start", dag.Location},
})
close(done)
}()
// Wait for DAG to start running
time.Sleep(time.Millisecond * 200)
dag.AssertLatestStatus(t, core.Running)
// Try to cleanup while running (nothing to delete since only active run exists)
th.RunCommand(t, cmd.Cleanup(), test.CmdTest{
Args: []string{"cleanup", "--yes", dag.Name},
})
// Verify the running DAG is still there (should be preserved)
dag.AssertLatestStatus(t, core.Running)
// Stop the DAG
th.RunCommand(t, cmd.Stop(), test.CmdTest{
Args: []string{"stop", dag.Location},
})
<-done
})
t.Run("RejectsNegativeRetentionDays", func(t *testing.T) {
th := test.SetupCommand(t)
dag := th.DAG(t, `steps:
- name: "1"
command: echo "hello"
`)
err := th.RunCommandWithError(t, cmd.Cleanup(), test.CmdTest{
Args: []string{"cleanup", "--retention-days", "-1", dag.Name},
})
require.Error(t, err)
assert.Contains(t, err.Error(), "cannot be negative")
})
t.Run("RejectsInvalidRetentionDays", func(t *testing.T) {
th := test.SetupCommand(t)
dag := th.DAG(t, `steps:
- name: "1"
command: echo "hello"
`)
err := th.RunCommandWithError(t, cmd.Cleanup(), test.CmdTest{
Args: []string{"cleanup", "--retention-days", "abc", dag.Name},
})
require.Error(t, err)
assert.Contains(t, err.Error(), "invalid retention-days")
})
t.Run("RequiresDAGNameArgument", func(t *testing.T) {
th := test.SetupCommand(t)
err := th.RunCommandWithError(t, cmd.Cleanup(), test.CmdTest{
Args: []string{"cleanup", "--yes"},
})
require.Error(t, err)
assert.Contains(t, err.Error(), "accepts 1 arg")
})
t.Run("SucceedsForNonExistentDAG", func(t *testing.T) {
th := test.SetupCommand(t)
// Cleanup for a DAG that doesn't exist should succeed silently
th.RunCommand(t, cmd.Cleanup(), test.CmdTest{
Args: []string{"cleanup", "--yes", "non-existent-dag"},
})
})
}
func TestCleanupCommandDirectStore(t *testing.T) {
// Test cleanup using the DAGRunStore directly to verify underlying behavior
t.Run("RemoveOldDAGRunsWithStore", func(t *testing.T) {
th := test.Setup(t)
dagName := "test-cleanup-dag"
// Create old DAG runs directly in the store
oldTime := time.Date(2021, 1, 1, 0, 0, 0, 0, time.UTC)
recentTime := time.Now()
// Create a minimal DAG for the test
testDAG := &core.DAG{Name: dagName}
// Create an old run
oldAttempt, err := th.DAGRunStore.CreateAttempt(
th.Context,
testDAG,
oldTime,
"old-run-id",
execution.NewDAGRunAttemptOptions{},
)
require.NoError(t, err)
require.NoError(t, oldAttempt.Open(th.Context))
require.NoError(t, oldAttempt.Write(th.Context, execution.DAGRunStatus{
Name: dagName,
DAGRunID: "old-run-id",
Status: core.Succeeded,
}))
require.NoError(t, oldAttempt.Close(th.Context))
// Create a recent run
recentAttempt, err := th.DAGRunStore.CreateAttempt(
th.Context,
testDAG,
recentTime,
"recent-run-id",
execution.NewDAGRunAttemptOptions{},
)
require.NoError(t, err)
require.NoError(t, recentAttempt.Open(th.Context))
require.NoError(t, recentAttempt.Write(th.Context, execution.DAGRunStatus{
Name: dagName,
DAGRunID: "recent-run-id",
Status: core.Succeeded,
}))
require.NoError(t, recentAttempt.Close(th.Context))
// Manually set old file modification time
setOldModTime(t, th.Config.Paths.DAGRunsDir, dagName, "", oldTime)
// Verify both runs exist
runs := th.DAGRunStore.RecentAttempts(th.Context, dagName, 10)
require.Len(t, runs, 2)
// Remove runs older than 7 days
removedIDs, err := th.DAGRunStore.RemoveOldDAGRuns(th.Context, dagName, 7)
require.NoError(t, err)
assert.Len(t, removedIDs, 1)
// Verify old run is deleted, recent run remains
runs = th.DAGRunStore.RecentAttempts(th.Context, dagName, 10)
require.Len(t, runs, 1)
status, err := runs[0].ReadStatus(th.Context)
require.NoError(t, err)
assert.Equal(t, "recent-run-id", status.DAGRunID)
})
}
// setOldModTime sets old modification time on DAG run files
func setOldModTime(t *testing.T, baseDir, dagName, _ string, modTime time.Time) {
t.Helper()
// Find the run directory
dagRunsDir := filepath.Join(baseDir, dagName, "dag-runs")
err := filepath.Walk(dagRunsDir, func(path string, _ os.FileInfo, err error) error {
if err != nil {
return err
}
// Set mod time on all files and directories
return os.Chtimes(path, modTime, modTime)
})
// Ignore errors if directory doesn't exist
if err != nil && !os.IsNotExist(err) {
t.Logf("Warning: failed to set mod time: %v", err)
}
}

View File

@ -266,6 +266,30 @@ var (
isBool: true,
bindViper: true,
}
// retentionDaysFlag specifies the number of days to retain history.
// Records older than this will be deleted.
// If set to 0, all records (except active) will be deleted.
retentionDaysFlag = commandLineFlag{
name: "retention-days",
defaultValue: "0",
usage: "Number of days to retain history (0 = delete all, except active runs)",
}
// dryRunFlag enables preview mode without actual deletion.
dryRunFlag = commandLineFlag{
name: "dry-run",
usage: "Preview what would be deleted without actually deleting",
isBool: true,
}
// yesFlag skips the confirmation prompt.
yesFlag = commandLineFlag{
name: "yes",
shorthand: "y",
usage: "Skip confirmation prompt",
isBool: true,
}
)
type commandLineFlag struct {

View File

@ -151,9 +151,12 @@ func (m *mockDAGRunStore) FindSubAttempt(ctx context.Context, dagRun execution.D
return args.Get(0).(execution.DAGRunAttempt), args.Error(1)
}
func (m *mockDAGRunStore) RemoveOldDAGRuns(ctx context.Context, name string, retentionDays int) error {
args := m.Called(ctx, name, retentionDays)
return args.Error(0)
func (m *mockDAGRunStore) RemoveOldDAGRuns(ctx context.Context, name string, retentionDays int, opts ...execution.RemoveOldDAGRunsOption) ([]string, error) {
args := m.Called(ctx, name, retentionDays, opts)
if args.Get(0) == nil {
return nil, args.Error(1)
}
return args.Get(0).([]string), args.Error(1)
}
func (m *mockDAGRunStore) RenameDAGRuns(ctx context.Context, oldName, newName string) error {

View File

@ -33,11 +33,12 @@ type DAGRunStore interface {
FindAttempt(ctx context.Context, dagRun DAGRunRef) (DAGRunAttempt, error)
// FindSubAttempt finds a sub dag-run record by dag-run ID.
FindSubAttempt(ctx context.Context, dagRun DAGRunRef, subDAGRunID string) (DAGRunAttempt, error)
// RemoveOldDAGRuns delete dag-run records older than retentionDays
// RemoveOldDAGRuns deletes dag-run records older than retentionDays.
// If retentionDays is negative, it won't delete any records.
// If retentionDays is zero, it will delete all records for the DAG name.
// But it will not delete the records with non-final statuses (e.g., running, queued).
RemoveOldDAGRuns(ctx context.Context, name string, retentionDays int) error
// Returns a list of dag-run IDs that were removed (or would be removed in dry-run mode).
RemoveOldDAGRuns(ctx context.Context, name string, retentionDays int, opts ...RemoveOldDAGRunsOption) ([]string, error)
// RenameDAGRuns renames all run data from oldName to newName
// The name means the DAG name, renaming it will allow user to manage those runs
// with the new DAG name.
@ -102,6 +103,22 @@ func WithDAGRunID(dagRunID string) ListDAGRunStatusesOption {
}
}
// RemoveOldDAGRunsOptions contains options for removing old dag-runs
type RemoveOldDAGRunsOptions struct {
// DryRun if true, only returns the paths that would be removed without actually deleting
DryRun bool
}
// RemoveOldDAGRunsOption is a functional option for configuring RemoveOldDAGRunsOptions
type RemoveOldDAGRunsOption func(*RemoveOldDAGRunsOptions)
// WithDryRun sets the dry-run mode for removing old dag-runs
func WithDryRun() RemoveOldDAGRunsOption {
return func(o *RemoveOldDAGRunsOptions) {
o.DryRun = true
}
}
// NewDAGRunAttemptOptions contains options for creating a new run record
type NewDAGRunAttemptOptions struct {
// RootDAGRun is the root dag-run reference for this attempt.

View File

@ -67,9 +67,12 @@ func (m *mockDAGRunStore) FindSubAttempt(ctx context.Context, dagRun execution.D
return args.Get(0).(execution.DAGRunAttempt), args.Error(1)
}
func (m *mockDAGRunStore) RemoveOldDAGRuns(ctx context.Context, name string, retentionDays int) error {
args := m.Called(ctx, name, retentionDays)
return args.Error(0)
func (m *mockDAGRunStore) RemoveOldDAGRuns(ctx context.Context, name string, retentionDays int, opts ...execution.RemoveOldDAGRunsOption) ([]string, error) {
args := m.Called(ctx, name, retentionDays, opts)
if args.Get(0) == nil {
return nil, args.Error(1)
}
return args.Get(0).([]string), args.Error(1)
}
func (m *mockDAGRunStore) RenameDAGRuns(ctx context.Context, oldName, newName string) error {
@ -242,10 +245,11 @@ func TestDAGRunStoreInterface(t *testing.T) {
assert.Equal(t, mockAttempt, childFound)
// Test RemoveOldDAGRuns
store.On("RemoveOldDAGRuns", ctx, "test-dag", 30).Return(nil)
store.On("RemoveOldDAGRuns", ctx, "test-dag", 30, mock.Anything).Return([]string{"run-1", "run-2"}, nil)
err = store.RemoveOldDAGRuns(ctx, "test-dag", 30)
removedIDs, err := store.RemoveOldDAGRuns(ctx, "test-dag", 30)
assert.NoError(t, err)
assert.Equal(t, []string{"run-1", "run-2"}, removedIDs)
// Test RenameDAGRuns
store.On("RenameDAGRuns", ctx, "old-name", "new-name").Return(nil)
@ -383,19 +387,22 @@ func TestRemoveOldDAGRunsEdgeCases(t *testing.T) {
store := &mockDAGRunStore{}
// Test with negative retention days (should not delete anything)
store.On("RemoveOldDAGRuns", ctx, "test-dag", -1).Return(nil)
err := store.RemoveOldDAGRuns(ctx, "test-dag", -1)
store.On("RemoveOldDAGRuns", ctx, "test-dag", -1, mock.Anything).Return([]string(nil), nil)
removedIDs, err := store.RemoveOldDAGRuns(ctx, "test-dag", -1)
assert.NoError(t, err)
assert.Nil(t, removedIDs)
// Test with zero retention days (should delete all except non-final statuses)
store.On("RemoveOldDAGRuns", ctx, "test-dag", 0).Return(nil)
err = store.RemoveOldDAGRuns(ctx, "test-dag", 0)
store.On("RemoveOldDAGRuns", ctx, "test-dag", 0, mock.Anything).Return([]string{"run-1", "run-2"}, nil)
removedIDs, err = store.RemoveOldDAGRuns(ctx, "test-dag", 0)
assert.NoError(t, err)
assert.Equal(t, []string{"run-1", "run-2"}, removedIDs)
// Test with positive retention days
store.On("RemoveOldDAGRuns", ctx, "test-dag", 30).Return(nil)
err = store.RemoveOldDAGRuns(ctx, "test-dag", 30)
store.On("RemoveOldDAGRuns", ctx, "test-dag", 30, mock.Anything).Return([]string{"run-old"}, nil)
removedIDs, err = store.RemoveOldDAGRuns(ctx, "test-dag", 30)
assert.NoError(t, err)
assert.Equal(t, []string{"run-old"}, removedIDs)
store.AssertExpectations(t)
}

View File

@ -276,10 +276,14 @@ func (dr DataRoot) Rename(ctx context.Context, newRoot DataRoot) error {
// If retentionDays is zero, all files will be removed.
// If retentionDays is positive, only files older than the specified number of days will be removed.
// It also removes empty directories in the hierarchy.
func (dr DataRoot) RemoveOld(ctx context.Context, retentionDays int) error {
// If dryRun is true, it returns the run IDs that would be removed without actually deleting them.
// Returns a list of dag-run IDs that were removed (or would be removed in dry-run mode).
func (dr DataRoot) RemoveOld(ctx context.Context, retentionDays int, dryRun bool) ([]string, error) {
keepTime := execution.NewUTC(time.Now().AddDate(0, 0, -retentionDays))
dagRuns := dr.listDAGRunsInRange(ctx, execution.TimeInUTC{}, keepTime, &listDAGRunsInRangeOpts{})
var removedRunIDs []string
for _, r := range dagRuns {
// Enrich context with run directory for all subsequent logs in this iteration
runCtx := logger.WithValues(ctx, tag.Dir(r.baseDir))
@ -311,13 +315,22 @@ func (dr DataRoot) RemoveOld(ctx context.Context, retentionDays int) error {
if lastUpdate.After(keepTime.Time) {
continue
}
// Add run ID to removed list
removedRunIDs = append(removedRunIDs, r.dagRunID)
// In dry-run mode, skip actual deletion
if dryRun {
continue
}
if err := r.Remove(ctx); err != nil {
logger.Error(runCtx, "Failed to remove run",
tag.Error(err))
}
dr.removeEmptyDir(ctx, filepath.Dir(r.baseDir))
}
return nil
return removedRunIDs, nil
}
func (dr DataRoot) removeEmptyDir(ctx context.Context, dayDir string) {

View File

@ -182,8 +182,9 @@ func TestDataRootRemoveOld(t *testing.T) {
assert.True(t, fileutil.FileExists(dagRun2.baseDir), "dag-run 2 should exist before cleanup")
// Remove all dag-runs (retention = 0)
err := root.RemoveOld(root.Context, 0)
removedIDs, err := root.RemoveOld(root.Context, 0, false)
require.NoError(t, err)
assert.Len(t, removedIDs, 2)
// Verify all dag-runs are removed
assert.False(t, fileutil.FileExists(dagRun1.baseDir), "dag-run 1 should be removed")
@ -228,8 +229,9 @@ func TestDataRootRemoveOld(t *testing.T) {
assert.True(t, fileutil.FileExists(dagRun2.baseDir), "Recent dag-run should exist before cleanup")
// Remove dag-runs older than 7 days (should remove old but keep recent)
err := root.RemoveOld(root.Context, 7)
removedIDs, err := root.RemoveOld(root.Context, 7, false)
require.NoError(t, err)
assert.Len(t, removedIDs, 1)
// Verify old dag-run is removed but recent one is kept
assert.False(t, fileutil.FileExists(dagRun1.baseDir), "Old dag-run should be removed")
@ -274,8 +276,9 @@ func TestDataRootRemoveOld(t *testing.T) {
assert.True(t, fileutil.FileExists(dagRun2.baseDir), "dag-run 2 should exist")
// Remove all old dag-runs (retention = 0)
err := root.RemoveOld(root.Context, 0)
removedIDs, err := root.RemoveOld(root.Context, 0, false)
require.NoError(t, err)
assert.Len(t, removedIDs, 2)
// Verify dag-runs are removed
assert.False(t, fileutil.FileExists(dagRun1.baseDir), "dag-run 1 should be removed")

View File

@ -437,16 +437,22 @@ func (store *Store) FindSubAttempt(ctx context.Context, ref execution.DAGRunRef,
// If retentionDays is negative, no files will be removed.
// If retentionDays is zero, all files will be removed.
// If retentionDays is positive, only files older than the specified number of days will be removed.
func (store *Store) RemoveOldDAGRuns(ctx context.Context, dagName string, retentionDays int) error {
// Returns a list of file paths that were removed (or would be removed in dry-run mode).
func (store *Store) RemoveOldDAGRuns(ctx context.Context, dagName string, retentionDays int, opts ...execution.RemoveOldDAGRunsOption) ([]string, error) {
if retentionDays < 0 {
logger.Warn(ctx, "Negative retentionDays, no files will be removed",
slog.Int("retention-days", retentionDays),
)
return nil
return nil, nil
}
var options execution.RemoveOldDAGRunsOptions
for _, opt := range opts {
opt(&options)
}
root := NewDataRoot(store.baseDir, dagName)
return root.RemoveOld(ctx, retentionDays)
return root.RemoveOld(ctx, retentionDays, options.DryRun)
}
// RemoveDAGRun implements models.DAGRunStore.

View File

@ -123,8 +123,9 @@ func TestJSONDB(t *testing.T) {
// Remove records older than 0 days
// It should remove all records
err := th.Store.RemoveOldDAGRuns(th.Context, "test_DAG", 0)
removedIDs, err := th.Store.RemoveOldDAGRuns(th.Context, "test_DAG", 0)
require.NoError(t, err)
assert.Len(t, removedIDs, 2) // 2 non-active runs should be removed
// Verify non active attempts are removed
attempts = th.Store.RecentAttempts(th.Context, "test_DAG", 3)

View File

@ -975,7 +975,7 @@ func (a *Agent) setupDefaultRetryPlan(ctx context.Context, nodes []*runtime.Node
func (a *Agent) setupDAGRunAttempt(ctx context.Context) (execution.DAGRunAttempt, error) {
retentionDays := a.dag.HistRetentionDays
if err := a.dagRunStore.RemoveOldDAGRuns(ctx, a.dag.Name, retentionDays); err != nil {
if _, err := a.dagRunStore.RemoveOldDAGRuns(ctx, a.dag.Name, retentionDays); err != nil {
logger.Error(ctx, "DAG runs data cleanup failed", tag.Error(err))
}

View File

@ -297,9 +297,12 @@ func (m *mockDAGRunStore) FindSubAttempt(ctx context.Context, rootDAGRun executi
return args.Get(0).(execution.DAGRunAttempt), args.Error(1)
}
func (m *mockDAGRunStore) RemoveOldDAGRuns(ctx context.Context, name string, retentionDays int) error {
args := m.Called(ctx, name, retentionDays)
return args.Error(0)
func (m *mockDAGRunStore) RemoveOldDAGRuns(ctx context.Context, name string, retentionDays int, opts ...execution.RemoveOldDAGRunsOption) ([]string, error) {
args := m.Called(ctx, name, retentionDays, opts)
if args.Get(0) == nil {
return nil, args.Error(1)
}
return args.Get(0).([]string), args.Error(1)
}
func (m *mockDAGRunStore) RenameDAGRuns(ctx context.Context, oldName, newName string) error {

View File

@ -396,9 +396,12 @@ func (m *mockDAGRunStore) FindSubAttempt(ctx context.Context, dagRun execution.D
return args.Get(0).(execution.DAGRunAttempt), args.Error(1)
}
func (m *mockDAGRunStore) RemoveOldDAGRuns(ctx context.Context, name string, retentionDays int) error {
args := m.Called(ctx, name, retentionDays)
return args.Error(0)
func (m *mockDAGRunStore) RemoveOldDAGRuns(ctx context.Context, name string, retentionDays int, opts ...execution.RemoveOldDAGRunsOption) ([]string, error) {
args := m.Called(ctx, name, retentionDays, opts)
if args.Get(0) == nil {
return nil, args.Error(1)
}
return args.Get(0).([]string), args.Error(1)
}
func (m *mockDAGRunStore) RenameDAGRuns(ctx context.Context, oldName, newName string) error {