From 690d76025d6114e4a26cd69c99e26914b94efa9a Mon Sep 17 00:00:00 2001
From: Niv Keidan <51288016+NivKeidan@users.noreply.github.com>
Date: Sun, 7 Jul 2024 15:07:11 +0300
Subject: [PATCH] feat(cde-2722): allow to not wait for propagation when
 deleting DNS records (#1)

allow to not wait for propagation when deleting DNS records
---
 client.go        | 54 +++++++++++++++++++++++++++++++-------------
 provider.go      | 23 +++++++++++++++++++
 provider_test.go | 59 ++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 120 insertions(+), 16 deletions(-)
 create mode 100644 provider_test.go

diff --git a/client.go b/client.go
index fabfc86..69a796e 100644
--- a/client.go
+++ b/client.go
@@ -274,11 +274,18 @@ func (p *Provider) createRecord(ctx context.Context, zoneID string, record libdn
 		HostedZoneId: aws.String(zoneID),
 	}
 
-	err := p.applyChange(ctx, createInput)
+	changeId, err := p.applyChange(ctx, createInput)
 	if err != nil {
 		return record, err
 	}
 
+	if p.WaitForPropagation {
+		err = p.waitForChangeToPropagate(ctx, changeId)
+		if err != nil {
+			return record, err
+		}
+	}
+
 	return record, nil
 }
 
@@ -315,11 +322,18 @@ func (p *Provider) updateRecord(ctx context.Context, zoneID string, record libdn
 		HostedZoneId: aws.String(zoneID),
 	}
 
-	err := p.applyChange(ctx, updateInput)
+	changeId, err := p.applyChange(ctx, updateInput)
 	if err != nil {
 		return record, err
 	}
 
+	if p.WaitForPropagation {
+		err = p.waitForChangeToPropagate(ctx, changeId)
+		if err != nil {
+			return record, err
+		}
+	}
+
 	return record, nil
 }
 
@@ -366,7 +380,7 @@ func (p *Provider) deleteRecord(ctx context.Context, zoneID string, record libdn
 		HostedZoneId: aws.String(zoneID),
 	}
 
-	err := p.applyChange(ctx, deleteInput)
+	changeId, err := p.applyChange(ctx, deleteInput)
 	if err != nil {
 		var nfe *types.InvalidChangeBatch
 		if record.Type == "TXT" && errors.As(err, &nfe) {
@@ -375,27 +389,35 @@ func (p *Provider) deleteRecord(ctx context.Context, zoneID string, record libdn
 		return record, err
 	}
 
+	if p.shouldWaitForDeletePropagation() {
+		err = p.waitForChangeToPropagate(ctx, changeId)
+		if err != nil {
+			return record, err
+		}
+	}
+
 	return record, nil
 }
 
-func (p *Provider) applyChange(ctx context.Context, input *r53.ChangeResourceRecordSetsInput) error {
+func (p *Provider) applyChange(ctx context.Context, input *r53.ChangeResourceRecordSetsInput) (*string, error) {
 	changeResult, err := p.client.ChangeResourceRecordSets(ctx, input)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	// Waiting for propagation if it's set in the provider config.
-	if p.WaitForPropagation {
-		changeInput := &r53.GetChangeInput{
-			Id: changeResult.ChangeInfo.Id,
-		}
+	return changeResult.ChangeInfo.Id, nil
+}
 
-		// Wait for the RecordSetChange status to be "INSYNC"
-		waiter := r53.NewResourceRecordSetsChangedWaiter(p.client)
-		err = waiter.Wait(ctx, changeInput, p.MaxWaitDur)
-		if err != nil {
-			return err
-		}
+func (p *Provider) waitForChangeToPropagate(ctx context.Context, changeID *string) error {
+	changeInput := &r53.GetChangeInput{
+		Id: changeID,
+	}
+
+	// Wait for the RecordSetChange status to be "INSYNC"
+	waiter := r53.NewResourceRecordSetsChangedWaiter(p.client)
+	err := waiter.Wait(ctx, changeInput, p.MaxWaitDur)
+	if err != nil {
+		return err
 	}
 
 	return nil
diff --git a/provider.go b/provider.go
index 1057713..2397cb2 100644
--- a/provider.go
+++ b/provider.go
@@ -58,6 +58,29 @@ type Provider struct {
 	// WaitForPropagation if set to true, it will wait for the record to be
 	// propagated before returning.
 	WaitForPropagation bool `json:"wait_for_propagation,omitempty"`
+
+	// WaitForDeletePropagation enables not to wait for propogation when cleaning challenge records
+	// 0 (default) will use WaitForPropagation's value (for backwards compatibility)
+	// 1 will always wait, 2 will never wait
+	WaitForDeletePropagation waitForDeletePropagationOption `json:"wait_for_delete_propogation,omitempty"`
+}
+
+type waitForDeletePropagationOption int
+
+const (
+	AlwaysWaitForDeletePropagation waitForDeletePropagationOption = 1
+	NeverWaitForDeletePropagation  waitForDeletePropagationOption = 2
+)
+
+func (p *Provider) shouldWaitForDeletePropagation() bool {
+	switch p.WaitForDeletePropagation {
+	case AlwaysWaitForDeletePropagation:
+		return true
+	case NeverWaitForDeletePropagation:
+		return false
+	default:
+		return p.WaitForPropagation
+	}
 }
 
 // GetRecords lists all the records in the zone.
diff --git a/provider_test.go b/provider_test.go
new file mode 100644
index 0000000..2308441
--- /dev/null
+++ b/provider_test.go
@@ -0,0 +1,59 @@
+package route53
+
+import (
+	"testing"
+)
+
+func TestProvider_shouldWaitForDeletePropagation(t *testing.T) {
+	type fields struct {
+		WaitForPropagation       bool
+		WaitForDeletePropagation waitForDeletePropagationOption
+	}
+	tests := []struct {
+		name   string
+		fields fields
+		want   bool
+	}{
+		{
+			name: "backwards compatability - should wait",
+			fields: fields{
+				WaitForPropagation: true,
+			},
+			want: true,
+		},
+		{
+			name: "backwards compatability - should not wait",
+			fields: fields{
+				WaitForPropagation: false,
+			},
+			want: false,
+		},
+		{
+			name: "opposite from WaitForPropagation - should wait",
+			fields: fields{
+				WaitForPropagation:       false,
+				WaitForDeletePropagation: AlwaysWaitForDeletePropagation,
+			},
+			want: true,
+		},
+		{
+			name: "opposite from WaitForPropagation - should wait",
+			fields: fields{
+				WaitForPropagation:       true,
+				WaitForDeletePropagation: NeverWaitForDeletePropagation,
+			},
+			want: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			p := &Provider{
+				WaitForPropagation:       tt.fields.WaitForPropagation,
+				WaitForDeletePropagation: tt.fields.WaitForDeletePropagation,
+			}
+			if got := p.shouldWaitForDeletePropagation(); got != tt.want {
+				t.Errorf("Provider.shouldWaitForDeletePropagation() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}