From a1c583f5a1e38c49263012a65d1b3f95fbf3ecf4 Mon Sep 17 00:00:00 2001
From: Kit Ewbank <Kit_Ewbank@hotmail.com>
Date: Thu, 8 Feb 2024 15:27:17 -0500
Subject: [PATCH] generate/tags: Handle 'ServiceTagsMap' in paginated ListTags.

---
 internal/generate/tags/main.go                | 44 ++++++++++---------
 .../tags/templates/v1/header_body.tmpl        |  3 ++
 .../tags/templates/v1/list_tags_body.tmpl     |  8 ++++
 3 files changed, 34 insertions(+), 21 deletions(-)

diff --git a/internal/generate/tags/main.go b/internal/generate/tags/main.go
index c3e87ec2536..2577da94100 100644
--- a/internal/generate/tags/main.go
+++ b/internal/generate/tags/main.go
@@ -116,34 +116,34 @@ func newTemplateBody(version int, kvtValues bool) *TemplateBody {
 	switch version {
 	case sdkV1:
 		return &TemplateBody{
-			"\n" + v1.GetTagBody,
-			v1.HeaderBody,
-			"\n" + v1.ListTagsBody,
-			"\n" + v1.ServiceTagsMapBody,
-			"\n" + v1.ServiceTagsSliceBody,
-			"\n" + v1.UpdateTagsBody,
-			"\n" + v1.WaitTagsPropagatedBody,
+			getTag:             "\n" + v1.GetTagBody,
+			header:             v1.HeaderBody,
+			listTags:           "\n" + v1.ListTagsBody,
+			serviceTagsMap:     "\n" + v1.ServiceTagsMapBody,
+			serviceTagsSlice:   "\n" + v1.ServiceTagsSliceBody,
+			updateTags:         "\n" + v1.UpdateTagsBody,
+			waitTagsPropagated: "\n" + v1.WaitTagsPropagatedBody,
 		}
 	case sdkV2:
 		if kvtValues {
 			return &TemplateBody{
-				"\n" + v2.GetTagBody,
-				v2.HeaderBody,
-				"\n" + v2.ListTagsBody,
-				"\n" + v2.ServiceTagsValueMapBody,
-				"\n" + v2.ServiceTagsSliceBody,
-				"\n" + v2.UpdateTagsBody,
-				"\n" + v2.WaitTagsPropagatedBody,
+				getTag:             "\n" + v2.GetTagBody,
+				header:             v2.HeaderBody,
+				listTags:           "\n" + v2.ListTagsBody,
+				serviceTagsMap:     "\n" + v2.ServiceTagsValueMapBody,
+				serviceTagsSlice:   "\n" + v2.ServiceTagsSliceBody,
+				updateTags:         "\n" + v2.UpdateTagsBody,
+				waitTagsPropagated: "\n" + v2.WaitTagsPropagatedBody,
 			}
 		}
 		return &TemplateBody{
-			"\n" + v2.GetTagBody,
-			v2.HeaderBody,
-			"\n" + v2.ListTagsBody,
-			"\n" + v2.ServiceTagsMapBody,
-			"\n" + v2.ServiceTagsSliceBody,
-			"\n" + v2.UpdateTagsBody,
-			"\n" + v2.WaitTagsPropagatedBody,
+			getTag:             "\n" + v2.GetTagBody,
+			header:             v2.HeaderBody,
+			listTags:           "\n" + v2.ListTagsBody,
+			serviceTagsMap:     "\n" + v2.ServiceTagsMapBody,
+			serviceTagsSlice:   "\n" + v2.ServiceTagsSliceBody,
+			updateTags:         "\n" + v2.UpdateTagsBody,
+			waitTagsPropagated: "\n" + v2.WaitTagsPropagatedBody,
 		}
 	default:
 		return nil
@@ -172,6 +172,7 @@ type TemplateData struct {
 	ParentNotFoundErrCode      string
 	ParentNotFoundErrMsg       string
 	RetryCreateOnNotFound      string
+	ServiceTagsMap             bool
 	SetTagsOutFunc             string
 	TagInCustomVal             string
 	TagInIDElem                string
@@ -333,6 +334,7 @@ func main() {
 		ListTagsOutTagsElem:        *listTagsOutTagsElem,
 		ParentNotFoundErrCode:      *parentNotFoundErrCode,
 		ParentNotFoundErrMsg:       *parentNotFoundErrMsg,
+		ServiceTagsMap:             *serviceTagsMap,
 		SetTagsOutFunc:             *setTagsOutFunc,
 		TagInCustomVal:             *tagInCustomVal,
 		TagInIDElem:                *tagInIDElem,
diff --git a/internal/generate/tags/templates/v1/header_body.tmpl b/internal/generate/tags/templates/v1/header_body.tmpl
index 72cc432cadf..2423875c80c 100644
--- a/internal/generate/tags/templates/v1/header_body.tmpl
+++ b/internal/generate/tags/templates/v1/header_body.tmpl
@@ -50,4 +50,7 @@ import (
 	{{- if .NamesPkg }}
     "github.com/hashicorp/terraform-provider-aws/names"
 	{{- end }}
+	{{- if and .ListTagsOpPaginated .ServiceTagsMap }}
+	"golang.org/x/exp/maps"
+	{{- end }}
 )
diff --git a/internal/generate/tags/templates/v1/list_tags_body.tmpl b/internal/generate/tags/templates/v1/list_tags_body.tmpl
index ccb4e6f7235..30b63823db9 100644
--- a/internal/generate/tags/templates/v1/list_tags_body.tmpl
+++ b/internal/generate/tags/templates/v1/list_tags_body.tmpl
@@ -22,18 +22,26 @@ func {{ .ListTagsFunc }}(ctx context.Context, conn {{ .ClientType }}, identifier
 		{{- end }}
 	}
 {{- if .ListTagsOpPaginated }}
+    {{- if .ServiceTagsMap }}
+	output := make(map[string]*string)
+	{{- else }}
 	var output []*{{ .TagPackage }}.{{ or .TagType2 .TagType }}
+    {{- end }}
 
 	err := conn.{{ .ListTagsOp }}PagesWithContext(ctx, input, func(page *{{ .TagPackage  }}.{{ .ListTagsOp }}Output, lastPage bool) bool {
 		if page == nil {
 			return !lastPage
 		}
 
+    {{ if .ServiceTagsMap }}
+		maps.Copy(output, page.{{ .ListTagsOutTagsElem }})
+	{{- else }}
 		for _, v := range page.{{ .ListTagsOutTagsElem }} {
 			if v != nil {
 				output = append(output, v)
 			}
 		}
+    {{- end }}
 
 		return !lastPage
 	})