From 9d72ac00e50c2673aabf6586bc90e52d675d92f7 Mon Sep 17 00:00:00 2001
From: Andrew Starr-Bochicchio <a.starr.b@gmail.com>
Date: Thu, 7 Nov 2024 15:51:39 -0500
Subject: [PATCH] builder: create tag if it does not exist.

---
 builder/digitalocean/builder_acc_test.go |  8 ++++--
 builder/digitalocean/step_snapshot.go    | 35 ++++++++++++++++++++----
 go.mod                                   |  2 +-
 3 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/builder/digitalocean/builder_acc_test.go b/builder/digitalocean/builder_acc_test.go
index c285438..7af5eee 100644
--- a/builder/digitalocean/builder_acc_test.go
+++ b/builder/digitalocean/builder_acc_test.go
@@ -11,6 +11,7 @@ import (
 
 	"github.com/digitalocean/godo"
 	"github.com/digitalocean/packer-plugin-digitalocean/version"
+	"github.com/google/uuid"
 	"github.com/hashicorp/packer-plugin-sdk/acctest"
 	"github.com/hashicorp/packer-plugin-sdk/useragent"
 	"golang.org/x/oauth2"
@@ -22,7 +23,7 @@ func TestBuilderAcc_basic(t *testing.T) {
 	}
 	acctest.TestPlugin(t, &acctest.PluginTestCase{
 		Name:     "test-digitalocean-builder-basic",
-		Template: fmt.Sprintf(testBuilderAccBasic, "ubuntu-20-04-x64"),
+		Template: fmt.Sprintf(testBuilderAccBasic, "ubuntu-20-04-x64", uuid.New().String()),
 		Check: func(buildCommand *exec.Cmd, logfile string) error {
 			if buildCommand.ProcessState != nil {
 				if buildCommand.ProcessState.ExitCode() != 0 {
@@ -152,7 +153,7 @@ func makeTemplateWithImageId(t *testing.T) string {
 			t.Fatalf("failed to retrieve image ID: %s", err)
 		}
 
-		return fmt.Sprintf(testBuilderAccBasic, image.ID)
+		return fmt.Sprintf(testBuilderAccBasic, image.ID, uuid.New().String())
 	}
 
 	return ""
@@ -167,7 +168,8 @@ const (
 		"size": "s-1vcpu-1gb",
 		"image": "%v",
 		"ssh_username": "root",
-		"tags": ["tag1","tag2"]
+		"tags": ["tag1","tag2"],
+		"snapshot_tags": ["%s"]
 	}]
 }
 `
diff --git a/builder/digitalocean/step_snapshot.go b/builder/digitalocean/step_snapshot.go
index 1eca4fa..1edeb80 100644
--- a/builder/digitalocean/step_snapshot.go
+++ b/builder/digitalocean/step_snapshot.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"log"
+	"net/http"
 	"strconv"
 	"time"
 
@@ -81,16 +82,40 @@ func (s *stepSnapshot) Run(ctx context.Context, state multistep.StateBag) multis
 	}
 
 	if len(c.SnapshotTags) > 0 {
+		var finalErr error
 		for _, tag := range c.SnapshotTags {
-			_, err = client.Tags.TagResources(context.TODO(), tag, &godo.TagResourcesRequest{Resources: []godo.Resource{{ID: strconv.Itoa(imageId), Type: "image"}}})
+			tagReq := &godo.TagResourcesRequest{
+				Resources: []godo.Resource{
+					{ID: strconv.Itoa(imageId), Type: "image"},
+				},
+			}
+
+			resp, err := client.Tags.TagResources(context.TODO(), tag, tagReq)
 			if err != nil {
-				err := fmt.Errorf("Error Tagging Image: %s", err)
-				state.Put("error", err)
-				ui.Error(err.Error())
-				return multistep.ActionHalt
+				// If the tag doesn't exist, create it and try again.
+				if resp != nil && resp.StatusCode == http.StatusNotFound {
+					log.Printf("Tag '%s' not found; creating it...", tag)
+					_, _, err = client.Tags.Create(context.TODO(), &godo.TagCreateRequest{Name: tag})
+					if err != nil {
+						finalErr = fmt.Errorf("Error creating tag: %s", err)
+						break
+					}
+					_, err = client.Tags.TagResources(context.TODO(), tag, tagReq)
+					if err != nil {
+						finalErr = fmt.Errorf("Error tagging image: %s", err)
+						break
+					}
+				}
 			}
+
 			ui.Say(fmt.Sprintf("Added snapshot tag: %s...", tag))
 		}
+		if finalErr != nil {
+			err := fmt.Errorf("Error tagging image: %s", err)
+			state.Put("error", err)
+			ui.Error(err.Error())
+			return multistep.ActionHalt
+		}
 	}
 
 	if len(c.SnapshotRegions) > 0 {
diff --git a/go.mod b/go.mod
index 1416452..1af46a6 100644
--- a/go.mod
+++ b/go.mod
@@ -7,6 +7,7 @@ toolchain go1.22.4
 require (
 	github.com/aws/aws-sdk-go v1.44.114
 	github.com/digitalocean/godo v1.122.0
+	github.com/google/uuid v1.4.0
 	github.com/hashicorp/hcl/v2 v2.19.1
 	github.com/hashicorp/packer-plugin-sdk v0.5.4
 	github.com/mitchellh/mapstructure v1.5.0
@@ -39,7 +40,6 @@ require (
 	github.com/golang/protobuf v1.5.3 // indirect
 	github.com/google/go-querystring v1.1.0 // indirect
 	github.com/google/s2a-go v0.1.7 // indirect
-	github.com/google/uuid v1.4.0 // indirect
 	github.com/googleapis/enterprise-certificate-proxy v0.3.2 // indirect
 	github.com/googleapis/gax-go/v2 v2.12.0 // indirect
 	github.com/hashicorp/consul/api v1.25.1 // indirect