Preserve error type and response metadata

aws · Dec 13, 2022 · 87493d8 · 87493d8
1 parent 1740546
commit 87493d8
Show file tree

Hide file tree

Showing 6 changed files with 79 additions and 19 deletions.
diff --git a/private/model/api/api.go b/private/model/api/api.go
@@ -643,7 +643,7 @@ func newClient(cfg aws.Config, handlers request.Handlers, partitionID, endpoint,
 		{{- $_ := $.AddSDKImport "private/protocol" }}
 		svc.Handlers.UnmarshalError.PushBackNamed(
 			{{-  if .Metadata.AWSQueryCompatible }}
-			protocol.NewUnmarshalErrorHandler({{ .ProtocolPackage }}.NewUnmarshalTypedErrorWithOptions(exceptionFromCode, {{ .ProtocolPackage }}.WithQueryCompatibility())).NamedHandler(),
+			protocol.NewUnmarshalErrorHandler({{ .ProtocolPackage }}.NewUnmarshalTypedErrorWithOptions(exceptionFromCode, {{ .ProtocolPackage }}.WithQueryCompatibility(queryExceptionFromCode))).NamedHandler(),
 			{{- else }}
 			protocol.NewUnmarshalErrorHandler({{ .ProtocolPackage }}.NewUnmarshalTypedError(exceptionFromCode)).NamedHandler(),
 			{{- end}}
@@ -940,6 +940,13 @@ const (
 			"{{ $s.ErrorName }}": newError{{ $s.ShapeName }},
 		{{- end }}
 	}
+	{{- if .Metadata.AWSQueryCompatible }}
+	var queryExceptionFromCode = map[string]func(protocol.ResponseMetadata, string)error {
+		{{- range $_, $s := $.ShapeListErrors }}
+			"{{ $s.ErrorName }}": newQueryCompatibleError{{ $s.ShapeName }},
+		{{- end }}
+	}
+	{{- end }}
 {{- end }}
 `))
 

diff --git a/private/model/api/codegentest/service/awsquerycompatible/api.go b/private/model/api/codegentest/service/awsquerycompatible/api.go
diff --git a/private/model/api/codegentest/service/awsquerycompatible/errors.go b/private/model/api/codegentest/service/awsquerycompatible/errors.go
diff --git a/private/model/api/codegentest/service/awsquerycompatible/service.go b/private/model/api/codegentest/service/awsquerycompatible/service.go
diff --git a/private/model/api/shape.go b/private/model/api/shape.go
@@ -783,6 +783,9 @@ type {{ $.ShapeName }} struct {
 	{{- if $.Exception }}
 		{{- $_ := $.API.AddSDKImport "private/protocol" }}
 		RespMetadata protocol.ResponseMetadata` + "`json:\"-\" xml:\"-\"`" + `
+	{{- if $.API.Metadata.AWSQueryCompatible }}
+		code string
+	{{- end }}
 	{{- end }}
 
 	{{- if $.OutputEventStreamAPI }}
@@ -922,8 +925,22 @@ func newError{{ $.ShapeName }}(v protocol.ResponseMetadata) error {
 	}
 }
 
+{{- if $.API.Metadata.AWSQueryCompatible }}
+func newQueryCompatibleError{{ $.ShapeName }}(v protocol.ResponseMetadata, code string) error {
+	return &{{ $.ShapeName }}{
+		RespMetadata: v,
+		code: code,
+	}
+}
+{{- end }}
+
 // Code returns the exception type name.
 func (s *{{ $.ShapeName }}) Code() string {
+	{{- if $.API.Metadata.AWSQueryCompatible }}
+	if s.code != "" {
+		return s.code
+	}
+	{{- end }}
 	return "{{ $.ErrorName }}"
 }
 

diff --git a/private/protocol/jsonrpc/unmarshal_error.go b/private/protocol/jsonrpc/unmarshal_error.go
@@ -21,15 +21,15 @@ const (
 // for both typed and untyped errors.
 type UnmarshalTypedError struct {
 	exceptions map[string]func(protocol.ResponseMetadata) error
-	parseQueryError bool
+	queryExceptions map[string]func(protocol.ResponseMetadata, string) error
 }
 
 // NewUnmarshalTypedError returns an UnmarshalTypedError initialized for the
 // set of exception names to the error unmarshalers
 func NewUnmarshalTypedError(exceptions map[string]func(protocol.ResponseMetadata) error) *UnmarshalTypedError {
 	return &UnmarshalTypedError{
 		exceptions: exceptions,
-		parseQueryError: false,
+		queryExceptions: map[string]func(protocol.ResponseMetadata, string) error{},
 	}
 }
 
@@ -41,9 +41,9 @@ func NewUnmarshalTypedErrorWithOptions(exceptions map[string]func(protocol.Respo
 	return unmarshaledError
 }
 
-func WithQueryCompatibility() func(*UnmarshalTypedError) {
+func WithQueryCompatibility(queryExceptions map[string]func(protocol.ResponseMetadata, string) error) func(*UnmarshalTypedError) {
 	return func(typedError *UnmarshalTypedError) {
-		typedError.parseQueryError = true
+		typedError.queryExceptions = queryExceptions
 	}
 }
 
@@ -70,30 +70,33 @@ func (u *UnmarshalTypedError) UnmarshalError(
 	code := codeParts[len(codeParts)-1]
 	msg := jsonErr.Message
 
-	queryCodeHeader := resp.Header.Get(awsQueryError)
-	if queryCodeHeader != "" && u.parseQueryError {
-		queryCodeParts := strings.Split(queryCodeHeader, ";")
-		if queryCodeParts != nil && len(queryCodeParts) == 2 {
-			return awserr.NewRequestFailure(
-				awserr.New(queryCodeParts[0], msg, nil),
-				respMeta.StatusCode,
-				respMeta.RequestID,
-				), nil
-		}
-	}
+	queryCodeParts := queryCodeParts(resp, u)
 
 	if fn, ok := u.exceptions[code]; ok {
 		// If exception code is know, use associated constructor to get a value
 		// for the exception that the JSON body can be unmarshaled into.
-		v := fn(respMeta)
+		var v error
+		queryErrFn, queryExceptionsFound := u.queryExceptions[code]
+		if queryCodeParts != nil && len(queryCodeParts) == 2 && queryExceptionsFound {
+			v = queryErrFn(respMeta, queryCodeParts[0])
+		} else {
+			v = fn(respMeta)
+		}
 		err := jsonutil.UnmarshalJSONCaseInsensitive(v, body)
 		if err != nil {
 			return nil, err
 		}
-
 		return v, nil
 	}
 
+	if queryCodeParts != nil && len(queryCodeParts) == 2 {
+		return awserr.NewRequestFailure(
+			awserr.New(queryCodeParts[0], msg, nil),
+			respMeta.StatusCode,
+			respMeta.RequestID,
+			), nil
+	}
+
 	// fallback to unmodeled generic exceptions
 	return awserr.NewRequestFailure(
 		awserr.New(code, msg, nil),
@@ -102,6 +105,15 @@ func (u *UnmarshalTypedError) UnmarshalError(
 	), nil
 }
 
+func queryCodeParts(resp *http.Response, u *UnmarshalTypedError) []string {
+	queryCodeHeader := resp.Header.Get(awsQueryError)
+	var queryCodeParts []string
+	if queryCodeHeader != "" && len(u.queryExceptions) > 0 {
+		queryCodeParts = strings.Split(queryCodeHeader, ";")
+	}
+	return queryCodeParts
+}
+
 // UnmarshalErrorHandler is a named request handler for unmarshaling jsonrpc
 // protocol request errors
 var UnmarshalErrorHandler = request.NamedHandler{