From 505c84380e178988d82def24e827261f7e0ffef3 Mon Sep 17 00:00:00 2001
From: wuxu92 <wuxu92@outlook.com>
Date: Wed, 9 Oct 2024 16:12:35 +0800
Subject: [PATCH] add additional query string to provisioning state poller and
 use location url instead of origin url

---
 .gitignore                                    |  1 +
 .../poller_provisioning_state.go              | 36 +++++++++++++++++--
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/.gitignore b/.gitignore
index 14b870f7927..669d7de5094 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,6 @@
 .fleet/
 .idea/
+.vscode/
 .DS_Store
 tmp/
 vendor/
diff --git a/sdk/client/resourcemanager/poller_provisioning_state.go b/sdk/client/resourcemanager/poller_provisioning_state.go
index f5c6f2e81d7..167fd86c47e 100644
--- a/sdk/client/resourcemanager/poller_provisioning_state.go
+++ b/sdk/client/resourcemanager/poller_provisioning_state.go
@@ -48,6 +48,11 @@ func provisioningStatePollerFromResponse(response *client.Response, lroIsSelfRef
 	}
 
 	resourcePath := originalUri
+	if pollingUriStr := pollingUriForLongRunningOperation(response); pollingUriStr != "" {
+		if pollingUri, err := url.ParseRequestURI(pollingUriStr); err == nil {
+			resourcePath = pollingUri.RequestURI()
+		}
+	}
 	if !lroIsSelfReference {
 		// if it's a self-reference (such as API Management's API/API Schema)
 		path, err := resourceManagerResourcePathFromUri(originalUri)
@@ -74,7 +79,8 @@ func (p *provisioningStatePoller) Poll(ctx context.Context) (*pollers.PollResult
 		},
 		HttpMethod: http.MethodGet,
 		OptionsObject: provisioningStateOptions{
-			apiVersion: p.apiVersion,
+			apiVersion:       p.apiVersion,
+			additionalParams: additonalParamsFromUri(p.resourcePath),
 		},
 		Path: p.resourcePath,
 	}
@@ -182,8 +188,11 @@ func resourceManagerResourcePathFromUri(input string) (*string, error) {
 
 var _ client.Options = provisioningStateOptions{}
 
+// For APIM API resource, the query string in the polling URI is not always the same as the original request URI.
+// These query strings are required for the polling request, otherwise, it will respond with 404 Not Found instead of a 400 Bad Request.
 type provisioningStateOptions struct {
-	apiVersion string
+	apiVersion       string
+	additionalParams map[string]string
 }
 
 func (p provisioningStateOptions) ToHeaders() *client.Headers {
@@ -196,6 +205,29 @@ func (p provisioningStateOptions) ToOData() *odata.Query {
 
 func (p provisioningStateOptions) ToQuery() *client.QueryParams {
 	q := client.QueryParams{}
+	for k, v := range p.additionalParams {
+		q.Append(k, v)
+	}
 	q.Append("api-version", p.apiVersion)
 	return &q
 }
+
+// APIM API resource need the asyncId parameter to get the real provisioning state.
+var additionalParamKeys = map[string]bool{
+	"asyncId": true,
+}
+
+func additonalParamsFromUri(uri string) map[string]string {
+	if u, err := url.ParseRequestURI(uri); err == nil {
+		if u.RawQuery != "" {
+			additionalParams := map[string]string{}
+			for k, v := range u.Query() {
+				if additionalParamKeys[k] {
+					additionalParams[k] = v[0]
+				}
+			}
+			return additionalParams
+		}
+	}
+	return nil
+}