diff --git a/restful.go b/restful.go index aec7a6c..8e31fc2 100644 --- a/restful.go +++ b/restful.go @@ -377,6 +377,12 @@ func (c *APIClient) PollUntilQueryEnd(ctx context.Context, resp *QueryResponse) data := resp.Data resp, err = c.PollQuery(ctx, resp.NextURI) if err != nil { + if errors.Is(err, context.Canceled) { + // context might be canceled due to timeout or canceled. if it's canceled, we need call + // the kill url to tell the backend it's killed. + fmt.Printf("query canceled, kill query:%s", resp.ID) + _ = c.KillQuery(context.Background(), resp) + } return nil, err } if resp.Error != nil {