From 7744a3dc506c14c85819d35ce156308d4058d671 Mon Sep 17 00:00:00 2001
From: Arjun Naik <arjunrn@users.noreply.github.com>
Date: Tue, 18 Apr 2023 14:01:49 +0200
Subject: [PATCH] Set Resource when GetAttribute is present

This change basically the SetResource function to continue when
the `GetAttribute` operation is present. If the `ReadOne` operation is
present then that is preferred.
---
 pkg/generate/code/set_resource.go             | 12 +++---
 pkg/generate/code/set_resource_test.go        | 43 +++++++++++++++++++
 .../models/apis/sns/0000-00-00/generator.yaml |  1 +
 .../models/apis/sqs/0000-00-00/generator.yaml |  3 ++
 4 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/pkg/generate/code/set_resource.go b/pkg/generate/code/set_resource.go
index b865dc5e..afeb4d5a 100644
--- a/pkg/generate/code/set_resource.go
+++ b/pkg/generate/code/set_resource.go
@@ -955,12 +955,14 @@ func SetResourceIdentifiers(
 ) string {
 	op := r.Ops.ReadOne
 	if op == nil {
-		if r.Ops.GetAttributes != nil {
-			// TODO(RedbackThomson): Support attribute maps for resource identifiers
-			return ""
+		switch {
+		case r.Ops.GetAttributes != nil:
+			// If single lookups can only be done with GetAttributes
+			op = r.Ops.GetAttributes
+		case r.Ops.ReadMany != nil:
+			// If single lookups can only be done using ReadMany
+			op = r.Ops.ReadMany
 		}
-		// If single lookups can only be done using ReadMany
-		op = r.Ops.ReadMany
 	}
 	inputShape := op.InputRef.Shape
 	if inputShape == nil {
diff --git a/pkg/generate/code/set_resource_test.go b/pkg/generate/code/set_resource_test.go
index acb335b1..05d771a4 100644
--- a/pkg/generate/code/set_resource_test.go
+++ b/pkg/generate/code/set_resource_test.go
@@ -2971,6 +2971,49 @@ func TestSetResource_RDS_DBInstances_SetResourceIdentifiers(t *testing.T) {
 	)
 }
 
+func TestSetResource_SNS_Topics_SetResourceIdentifiers(t *testing.T) {
+	assert := assert.New(t)
+	require := require.New(t)
+
+	g := testutil.NewModelForService(t, "sns")
+
+	crd := testutil.GetCRDByName(t, g, "Topic")
+	require.NotNil(crd)
+
+	expected := `
+	if r.ko.Status.ACKResourceMetadata == nil {
+		r.ko.Status.ACKResourceMetadata = &ackv1alpha1.ResourceMetadata{}
+	}
+	r.ko.Status.ACKResourceMetadata.ARN = identifier.ARN
+`
+	assert.Equal(
+		expected,
+		code.SetResourceIdentifiers(crd.Config(), crd, "identifier", "r.ko", 1),
+	)
+}
+
+func TestSetResource_SQS_Queues_SetResourceIdentifiers(t *testing.T) {
+	assert := assert.New(t)
+	require := require.New(t)
+
+	g := testutil.NewModelForService(t, "sqs")
+
+	crd := testutil.GetCRDByName(t, g, "Queue")
+	require.NotNil(crd)
+
+	expected := `
+	if identifier.NameOrID == "" {
+		return ackerrors.MissingNameIdentifier
+	}
+	r.ko.Status.QueueURL = &identifier.NameOrID
+
+`
+	assert.Equal(
+		expected,
+		code.SetResourceIdentifiers(crd.Config(), crd, "identifier", "r.ko", 1),
+	)
+}
+
 func TestSetResource_RDS_DBSubnetGroup_SetResourceIdentifiers(t *testing.T) {
 	assert := assert.New(t)
 	require := require.New(t)
diff --git a/pkg/testdata/models/apis/sns/0000-00-00/generator.yaml b/pkg/testdata/models/apis/sns/0000-00-00/generator.yaml
index 85a6a733..f53f143f 100644
--- a/pkg/testdata/models/apis/sns/0000-00-00/generator.yaml
+++ b/pkg/testdata/models/apis/sns/0000-00-00/generator.yaml
@@ -1,5 +1,6 @@
 resources:
   Topic:
+    is_arn_primary_key: true
     unpack_attributes_map:
       set_attributes_single_attribute: true
     fields:
diff --git a/pkg/testdata/models/apis/sqs/0000-00-00/generator.yaml b/pkg/testdata/models/apis/sqs/0000-00-00/generator.yaml
index 6696876f..6baa83cb 100644
--- a/pkg/testdata/models/apis/sqs/0000-00-00/generator.yaml
+++ b/pkg/testdata/models/apis/sqs/0000-00-00/generator.yaml
@@ -38,3 +38,6 @@ resources:
       QueueArn:
         is_attribute: true
         is_read_only: true
+      QueueUrl:
+        is_read_only: true
+        is_primary_key: true