From 24c3d5c5b25a7fb4f4f80dd5aaf5ea75a7cd2a72 Mon Sep 17 00:00:00 2001
From: Melora Hugues <dev@faercol.me>
Date: Wed, 2 Oct 2024 16:11:24 +0200
Subject: [PATCH] Add support for check-id

---
 check.go     | 53 ++++++++++++++++++++++++++++++++++++++++++++--------
 cmd/root.go  |  2 ++
 cmd/utils.go | 12 ++++++++----
 3 files changed, 55 insertions(+), 12 deletions(-)

diff --git a/check.go b/check.go
index 166f5ce..7ed2caa 100644
--- a/check.go
+++ b/check.go
@@ -1,6 +1,7 @@
 package gohealthchecks
 
 import (
+	"fmt"
 	"path"
 
 	"github.com/google/uuid"
@@ -11,28 +12,50 @@ type Check interface {
 	params() map[string]string
 }
 
-type uuidCheck uuid.UUID
+type uuidCheck struct {
+	checkID uuid.UUID
+	runID   uuid.UUID
+}
 
 func (c uuidCheck) path() string {
-	return uuid.UUID(c).String()
+	return c.checkID.String()
 }
 
 func (c uuidCheck) params() map[string]string {
-	return nil
+	if c.runID == uuid.Nil {
+		return nil
+	}
+
+	return map[string]string{
+		"rid": c.runID.String(),
+	}
 }
 
-func NewUUIDCheck(rawUUID string) (Check, error) {
-	val, err := uuid.Parse(rawUUID)
+func NewUUIDCheck(checkRawID, runRawID string) (Check, error) {
+	checkID, err := uuid.Parse(checkRawID)
 	if err != nil {
 		return nil, err
 	}
-	return uuidCheck(val), nil
+
+	runID := uuid.Nil
+	if runRawID != "" {
+		runID, err = uuid.Parse(runRawID)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return uuidCheck{
+		checkID: checkID,
+		runID:   runID,
+	}, nil
 }
 
 type slugCheck struct {
 	pingKey    string
 	slug       string
 	autoCreate bool
+	runID      uuid.UUID
 }
 
 func (c slugCheck) path() string {
@@ -47,13 +70,27 @@ func (c slugCheck) params() map[string]string {
 	}
 	res["create"] = createVal
 
+	if c.runID != uuid.Nil {
+		res["rid"] = c.runID.String()
+	}
+
 	return res
 }
 
-func NewSlugCheck(pingKey, slug string, autoCreate bool) Check {
+func NewSlugCheck(pingKey, slug string, autoCreate bool, runRawID string) (Check, error) {
+	runID := uuid.Nil
+	if runRawID != "" {
+		var err error
+		runID, err = uuid.Parse(runRawID)
+		if err != nil {
+			return nil, fmt.Errorf("invalid run ID: %w", err)
+		}
+	}
+
 	return slugCheck{
 		pingKey:    pingKey,
 		slug:       slug,
 		autoCreate: autoCreate,
-	}
+		runID:      runID,
+	}, nil
 }
diff --git a/cmd/root.go b/cmd/root.go
index 7ec9379..9f0e0ce 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -19,6 +19,7 @@ var (
 	pingKey    string
 	checkSlug  string
 	autoCreate bool
+	runID      string
 )
 
 // Execute adds all child commands to the root command and sets flags appropriately.
@@ -33,6 +34,7 @@ func Execute() {
 func init() {
 	rootCmd.PersistentFlags().StringVar(&pingHost, "host", gohealthchecks.HealthchecksIOPingHost, "healthchecks host to use")
 	rootCmd.PersistentFlags().BoolVar(&autoCreate, "auto-create", false, "auto create check if it does not exist (slug only)")
+	rootCmd.PersistentFlags().StringVar(&runID, "run-id", "", "Optional run-id to identify the job")
 	rootCmd.PersistentFlags().StringVar(&checkUUID, "uuid", "", "check UUID")
 	rootCmd.PersistentFlags().StringVar(&pingKey, "ping-key", "", "ping key")
 	rootCmd.PersistentFlags().StringVar(&checkSlug, "slug", "", "check-slug")
diff --git a/cmd/utils.go b/cmd/utils.go
index 7277eb9..7c93cc1 100644
--- a/cmd/utils.go
+++ b/cmd/utils.go
@@ -7,8 +7,8 @@ import (
 	gohealthchecks "git.faercol.me/faercol/go-healthchecks"
 )
 
-func mustParseUUIDCheck(rawVal string) gohealthchecks.Check {
-	check, err := gohealthchecks.NewUUIDCheck(rawVal)
+func mustParseUUIDCheck(checkID, runID string) gohealthchecks.Check {
+	check, err := gohealthchecks.NewUUIDCheck(checkID, runID)
 	if err != nil {
 		failf("Invalid UUID: %s\n", err)
 	}
@@ -17,10 +17,14 @@ func mustParseUUIDCheck(rawVal string) gohealthchecks.Check {
 
 func parseCheckFlags() gohealthchecks.Check {
 	var check gohealthchecks.Check
+	var err error
 	if checkUUID != "" {
-		check = mustParseUUIDCheck(checkUUID)
+		check = mustParseUUIDCheck(checkUUID, runID)
 	} else {
-		check = gohealthchecks.NewSlugCheck(pingKey, checkSlug, autoCreate)
+		check, err = gohealthchecks.NewSlugCheck(pingKey, checkSlug, autoCreate, runID)
+		if err != nil {
+			failf("Invalid UUID: %s\n", err)
+		}
 	}
 	return check
 }