From a44d0b750d90c8857cac87c4c7611b4b752e8b75 Mon Sep 17 00:00:00 2001
From: Aditya Thebe <contact@adityathebe.com>
Date: Mon, 17 Feb 2025 23:49:41 +0545
Subject: [PATCH] feat: retry playbook actions

---
 models/playbooks.go  | 25 +++++++++++++++++++++++++
 schema/playbooks.hcl |  5 +++++
 2 files changed, 30 insertions(+)

diff --git a/models/playbooks.go b/models/playbooks.go
index 3097ac5d..f55a1bc9 100644
--- a/models/playbooks.go
+++ b/models/playbooks.go
@@ -30,6 +30,7 @@ const (
 	PlaybookRunStatusRunning         PlaybookRunStatus = "running"
 	PlaybookRunStatusScheduled       PlaybookRunStatus = "scheduled"
 	PlaybookRunStatusSleeping        PlaybookRunStatus = "sleeping"
+	PlaybookRunStatusRetrying        PlaybookRunStatus = "retrying"
 	PlaybookRunStatusWaiting         PlaybookRunStatus = "waiting" // waiting for a consumer
 )
 
@@ -220,6 +221,14 @@ func (p PlaybookRun) Schedule(db *gorm.DB) error {
 	})
 }
 
+func (p PlaybookRun) Retry(db *gorm.DB, delay time.Duration) error {
+	return p.Update(db, map[string]any{
+		"status":         PlaybookRunStatusRetrying,
+		"start_time":     gorm.Expr("CASE WHEN start_time IS NULL THEN CLOCK_TIMESTAMP() ELSE start_time END"),
+		"scheduled_time": gorm.Expr(fmt.Sprintf("CLOCK_TIMESTAMP() + INTERVAL '%d SECONDS'", int(delay.Seconds()))),
+	})
+}
+
 func (p PlaybookRun) Delay(db *gorm.DB, delay time.Duration) error {
 	return p.Update(db, map[string]any{
 		"status":         PlaybookRunStatusSleeping,
@@ -291,6 +300,19 @@ func (p PlaybookRun) Assign(db *gorm.DB, agent *Agent, action string) error {
 	return p.Waiting(db)
 }
 
+func (p PlaybookRun) RetryAction(db *gorm.DB, action string, retryCount int) (*PlaybookRunAction, error) {
+	runAction := PlaybookRunAction{
+		PlaybookRunID: p.ID,
+		Name:          action,
+		Status:        PlaybookActionStatusScheduled,
+		RetryCount:    retryCount,
+	}
+	if err := db.Save(&runAction).Error; err != nil {
+		return nil, oops.Tags("db").Wrap(err)
+	}
+	return &runAction, p.Running(db)
+}
+
 func (p PlaybookRun) StartAction(db *gorm.DB, action string) (*PlaybookRunAction, error) {
 	runAction := PlaybookRunAction{
 		PlaybookRunID: p.ID,
@@ -468,6 +490,9 @@ type PlaybookRunAction struct {
 	Error         *string              `json:"error,omitempty" gorm:"default:null"`
 	IsPushed      bool                 `json:"is_pushed"`
 	AgentID       *uuid.UUID           `json:"agent_id,omitempty"`
+
+	// RetryCount represents the Nth retry of this action
+	RetryCount int `json:"attempt,omitempty" gorm:"default:NULL"`
 }
 
 func (p PlaybookRunAction) JSON() (out map[string]any) {
diff --git a/schema/playbooks.hcl b/schema/playbooks.hcl
index 1afe0b68..c1cd1b6c 100644
--- a/schema/playbooks.hcl
+++ b/schema/playbooks.hcl
@@ -319,6 +319,11 @@ table "playbook_run_actions" {
     type    = text
     default = "running"
   }
+  column "retry_count" {
+    null    = true
+    type    = integer
+    comment = "the nth retry of this action"
+  }
   column "playbook_run_id" {
     null    = true
     type    = uuid