Skip to content

Commit

Permalink
Use runtime.Poller[T] for LROs (#18076)
Browse files Browse the repository at this point in the history 
* Use runtime.Poller[T] for LROs

Removed the in-box poller types.
ResumeTokens are now passed by value.
Removed some dead code.

* bump version
  • Loading branch information
@jhendrixMSFT
jhendrixMSFT authored May 18, 2022
1 parent 6c35acf commit 844c0b5
Show file tree
Hide file tree
Showing 30 changed files with 2,794 additions and 10,513 deletions.
5 changes: 4 additions & 1 deletion sdk/keyvault/azkeys/CHANGELOG.md
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
# Release History

## 0.5.2 (Unreleased)
## 0.6.0 (Unreleased)

### Features Added

Expand All @@ -9,6 +9,9 @@
* `ListDeletedKeys` to `NewListDeletedKeysPager`
* `ListPropertiesOfKeys` to `NewListPropertiesOfKeysPager`
* `ListPropertiesOfKeyVersions` to `NewListPropertiesOfKeyVersionsPager`
* Removed types `DeleteKeyPoller` and `RecoverDeletedKeyPoller`.
* Methods `BeginDeleteKey` and `BeginRecoverDeletedKey` now return a `*runtime.Poller[T]` with their respective response types.
* Option types with a `ResumeToken` field now take the token by value.

### Bugs Fixed

Expand Down
3 changes: 2 additions & 1 deletion sdk/keyvault/azkeys/README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -256,6 +256,7 @@ func main() {
```go
import (
"github.com/Azure/azure-sdk-for-go/sdk/keyvault/azkeys"
"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
)
Expand All @@ -275,7 +276,7 @@ func main() {
if err != nil {
panic(err)
}
pollResp, err := resp.PollUntilDone(context.TODO(), 1*time.Second)
pollResp, err := resp.PollUntilDone(context.TODO(), &runtime.PollUntilDoneOptions{Frequency: time.Second})
if err != nil {
panic(err)
}
Expand Down
296 changes: 46 additions & 250 deletions sdk/keyvault/azkeys/client.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,10 +8,7 @@ package azkeys

import (
"context"
"encoding/json"
"errors"
"net/http"
"time"

"github.com/Azure/azure-sdk-for-go/sdk/azcore"
"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
Expand Down Expand Up @@ -579,157 +576,45 @@ type DeleteKeyResponse struct {
DeletedKey
}

// convert interal response to DeleteKeyResponse
func deleteKeyResponseFromGenerated(g generated.KeyVaultClientDeleteKeyResponse) DeleteKeyResponse {
vaultURL, name, version := shared.ParseID(g.Key.Kid)
return DeleteKeyResponse{
DeletedKey: DeletedKey{
Properties: keyPropertiesFromGenerated(g.Attributes, g.Key.Kid, name, version, g.Managed, vaultURL, g.Tags),
Key: jsonWebKeyFromGenerated(g.Key),
RecoveryID: g.RecoveryID,
ReleasePolicy: keyReleasePolicyFromGenerated(g.ReleasePolicy),
DeletedOn: g.DeletedDate,
ScheduledPurgeDate: g.ScheduledPurgeDate,
},
}
}

// BeginDeleteKeyOptions contains the optional parameters for the Client.BeginDeleteKey method.
type BeginDeleteKeyOptions struct {
// ResumeToken is a string to rehydrate a poller for an operation that has already begun.
ResumeToken *string
}

// convert public options to generated options struct
func (b *BeginDeleteKeyOptions) toGenerated() *generated.KeyVaultClientDeleteKeyOptions {
return &generated.KeyVaultClientDeleteKeyOptions{}
}

// DeleteKeyPoller is the poller returned by the Client.StartDeleteKey operation
type DeleteKeyPoller struct {
keyName string // This is the key to Poll for in GetDeletedKey
vaultUrl string
client *generated.KeyVaultClient
deleteResponse generated.KeyVaultClientDeleteKeyResponse
deleteRawResponse *http.Response
lastResponse generated.KeyVaultClientGetDeletedKeyResponse
resumeToken string
}

// Done returns true if the LRO has reached a terminal state
func (s *DeleteKeyPoller) Done() bool {
if s.deleteRawResponse == nil {
return false
}
return s.deleteRawResponse.StatusCode == http.StatusOK
}

// Poll fetches the latest state of the LRO. It returns an HTTP response or error.(
// If the LRO has completed successfully, the poller's state is updated and the HTTP response is returned.
// If the LRO has completed with failure or was cancelled, the poller's state is updated and the error is returned.)
func (s *DeleteKeyPoller) Poll(ctx context.Context) (*http.Response, error) {
var deleteRawResponse *http.Response
ctx = runtime.WithCaptureResponse(ctx, &deleteRawResponse)
resp, err := s.client.GetDeletedKey(ctx, s.vaultUrl, s.keyName, nil)
s.deleteRawResponse = deleteRawResponse
if deleteRawResponse.StatusCode == http.StatusOK {
// Service recognizes DeletedKey, operation is done
s.lastResponse = resp
return s.deleteRawResponse, nil
}

var httpResponseErr *azcore.ResponseError
if errors.As(err, &httpResponseErr) {
if httpResponseErr.StatusCode == http.StatusNotFound {
// This is the expected result
return s.deleteRawResponse, nil
}
}
return s.deleteRawResponse, err
}

// FinalResponse returns the final response after the operations has finished
func (s *DeleteKeyPoller) FinalResponse(ctx context.Context) (DeleteKeyResponse, error) {
return deleteKeyResponseFromGenerated(s.deleteResponse), nil
}

// PollUntilDone continually calls the Poll operation until the operation is completed. In between each
// Poll is a wait determined by the t parameter.
func (s *DeleteKeyPoller) PollUntilDone(ctx context.Context, t time.Duration) (DeleteKeyResponse, error) {
for {
_, err := s.Poll(ctx)
if err != nil {
return DeleteKeyResponse{}, err
}
if s.Done() {
break
}
time.Sleep(t)
}

vaultURL, name, version := shared.ParseID(s.deleteResponse.Key.Kid)
return DeleteKeyResponse{
DeletedKey: DeletedKey{
RecoveryID: s.deleteResponse.RecoveryID,
DeletedOn: s.deleteResponse.DeletedDate,
ScheduledPurgeDate: s.deleteResponse.ScheduledPurgeDate,
ReleasePolicy: keyReleasePolicyFromGenerated(s.deleteResponse.ReleasePolicy),
Properties: keyPropertiesFromGenerated(s.deleteResponse.Attributes, s.deleteResponse.Key.Kid, name, version, s.deleteResponse.Managed, vaultURL, s.deleteResponse.Tags),
Key: jsonWebKeyFromGenerated(s.deleteResponse.Key),
},
}, nil
}

// ResumeToken returns a token for resuming polling at a later time
func (s *DeleteKeyPoller) ResumeToken() (string, error) {
return s.resumeToken, nil
ResumeToken string
}

// BeginDeleteKey deletes a key from the keyvault. Delete cannot be applied to an individual version of a key. This operation
// requires the key/delete permission. This response contains a Poller struct that can be used to Poll for a response, or the
// PollUntilDone function can be used to poll until completion. Pass nil to use the default options.
func (c *Client) BeginDeleteKey(ctx context.Context, name string, options *BeginDeleteKeyOptions) (*DeleteKeyPoller, error) {
func (c *Client) BeginDeleteKey(ctx context.Context, name string, options *BeginDeleteKeyOptions) (*runtime.Poller[DeleteKeyResponse], error) {
if options == nil {
options = &BeginDeleteKeyOptions{}
}
var resumeToken string
var delResp generated.KeyVaultClientDeleteKeyResponse
var err error
if options.ResumeToken == nil {
delResp, err = c.kvClient.DeleteKey(ctx, c.vaultURL, name, options.toGenerated())
if err != nil {
return nil, err
}

resumeTokenMarshalled, err := json.Marshal(delResp)
if err != nil {
return nil, err
}
resumeToken = string(resumeTokenMarshalled)
} else {
resumeToken = *options.ResumeToken
err = json.Unmarshal([]byte(resumeToken), &delResp)
if err != nil {
return nil, err
}
handler := beginDeleteKeyPoller{
poll: func(ctx context.Context) (*http.Response, error) {
req, err := c.kvClient.GetDeletedKeyCreateRequest(ctx, c.vaultURL, name, nil)
if err != nil {
return nil, err
}
return c.kvClient.Pl.Do(req)
},
}

getResp, err := c.kvClient.GetDeletedKey(ctx, c.vaultURL, name, nil)
var httpErr *azcore.ResponseError
if errors.As(err, &httpErr) {
if httpErr.StatusCode != http.StatusNotFound {
return nil, err
}
if options.ResumeToken != "" {
return runtime.NewPollerFromResumeToken(options.ResumeToken, c.kvClient.Pl, &runtime.NewPollerFromResumeTokenOptions[DeleteKeyResponse]{
Handler: &handler,
})
}

return &DeleteKeyPoller{
vaultUrl: c.vaultURL,
keyName: name,
client: c.kvClient,
deleteResponse: delResp,
lastResponse: getResp,
resumeToken: resumeToken,
}, nil
var rawResp *http.Response
ctx = runtime.WithCaptureResponse(ctx, &rawResp)
if _, err := c.kvClient.DeleteKey(ctx, c.vaultURL, name, nil); err != nil {
return nil, err
}

return runtime.NewPoller(rawResp, c.kvClient.Pl, &runtime.NewPollerOptions[DeleteKeyResponse]{
Handler: &handler,
})
}

// BackupKeyOptions contains the optional parameters for the Client.BackupKey method
Expand Down Expand Up @@ -779,79 +664,10 @@ func (c *Client) BackupKey(ctx context.Context, name string, options *BackupKeyO
return backupKeyResponseFromGenerated(resp), nil
}

// RecoverDeletedKeyPoller implements the RecoverDeletedKeyPoller interface
type RecoverDeletedKeyPoller struct {
keyName string
vaultUrl string
client *generated.KeyVaultClient
recoverResponse generated.KeyVaultClientRecoverDeletedKeyResponse
lastResponse generated.KeyVaultClientGetKeyResponse
lastRawResponse *http.Response
finished bool
resumeToken string
}

// Done returns true when the polling operation is completed
func (p *RecoverDeletedKeyPoller) Done() bool {
if p.lastRawResponse == nil {
return false
}
return p.finished
}

// Poll fetches the latest state of the LRO. It returns an HTTP response or error.
// If the LRO has completed successfully, the poller's state is updated and the HTTP response is returned.
// If the LRO has completed with failure or was cancelled, the poller's state is updated and the error is returned.
func (p *RecoverDeletedKeyPoller) Poll(ctx context.Context) (*http.Response, error) {
resp, err := p.client.GetKey(ctx, p.vaultUrl, p.keyName, "", nil)
if err == nil {
// Polling is finished
p.finished = true
return nil, nil
}
p.lastResponse = resp
var httpErr *azcore.ResponseError
if errors.As(err, &httpErr) {
p.lastRawResponse = httpErr.RawResponse
if httpErr.StatusCode == http.StatusOK || httpErr.StatusCode == http.StatusNotFound {
return httpErr.RawResponse, nil
} else {
return httpErr.RawResponse, err
}
}
return p.lastRawResponse, err
}

// FinalResponse returns the final response after the operations has finished
func (p *RecoverDeletedKeyPoller) FinalResponse(ctx context.Context) (RecoverDeletedKeyResponse, error) {
return recoverDeletedKeyResponseFromGenerated(p.recoverResponse), nil
}

// PollUntilDone continually calls the Poll operation until the operation is completed. In between each
// Poll is a wait determined by the t parameter.
func (p *RecoverDeletedKeyPoller) PollUntilDone(ctx context.Context, t time.Duration) (RecoverDeletedKeyResponse, error) {
for {
_, err := p.Poll(ctx)
if err != nil {
return RecoverDeletedKeyResponse{}, err
}
if p.Done() {
break
}
time.Sleep(t)
}
return recoverDeletedKeyResponseFromGenerated(p.recoverResponse), nil
}

// BeginRecoverDeletedKeyOptions contains the optional parameters for the Client.BeginRecoverDeletedKey operation
type BeginRecoverDeletedKeyOptions struct {
// ResumeToken returns a string for creating a new poller to begin polling again
ResumeToken *string
}

// Convert the publicly exposed options object to the generated version
func (b BeginRecoverDeletedKeyOptions) toGenerated() *generated.KeyVaultClientRecoverDeletedKeyOptions {
return &generated.KeyVaultClientRecoverDeletedKeyOptions{}
ResumeToken string
}

// RecoverDeletedKeyResponse is the response object for the Client.RecoverDeletedKey operation.
Expand All @@ -872,59 +688,39 @@ func recoverDeletedKeyResponseFromGenerated(g generated.KeyVaultClientRecoverDel
}
}

// ResumeToken returns a token for resuming polling at a later time
func (r *RecoverDeletedKeyPoller) ResumeToken() (string, error) {
return r.resumeToken, nil
}

// BeginRecoverDeletedKey recovers the deleted key in the specified vault to the latest version.
// This operation can only be performed on a soft-delete enabled vault. This operation requires
// the keys/recover permission. Pass nil to use the default options.
func (c *Client) BeginRecoverDeletedKey(ctx context.Context, name string, options *BeginRecoverDeletedKeyOptions) (*RecoverDeletedKeyPoller, error) {
func (c *Client) BeginRecoverDeletedKey(ctx context.Context, name string, options *BeginRecoverDeletedKeyOptions) (*runtime.Poller[RecoverDeletedKeyResponse], error) {
if options == nil {
options = &BeginRecoverDeletedKeyOptions{}
}
var resumeToken string
var recoverResp generated.KeyVaultClientRecoverDeletedKeyResponse
var err error
if options.ResumeToken == nil {
recoverResp, err = c.kvClient.RecoverDeletedKey(ctx, c.vaultURL, name, options.toGenerated())
if err != nil {
return nil, err
}

marshalled, err := json.Marshal(recoverResp)
if err != nil {
return nil, err
}
resumeToken = string(marshalled)
} else {
resumeToken = *options.ResumeToken
err = json.Unmarshal([]byte(resumeToken), &recoverResp)
if err != nil {
return nil, err
}
handler := beginRecoverDeletedKeyPoller{
poll: func(ctx context.Context) (*http.Response, error) {
req, err := c.kvClient.GetKeyCreateRequest(ctx, c.vaultURL, name, "", nil)
if err != nil {
return nil, err
}
return c.kvClient.Pl.Do(req)
},
}

var getRawResp *http.Response
ctx = runtime.WithCaptureResponse(ctx, &getRawResp)
getResp, err := c.kvClient.GetKey(ctx, c.vaultURL, name, "", nil)
var httpErr *azcore.ResponseError
if errors.As(err, &httpErr) {
if httpErr.StatusCode != http.StatusNotFound {
return nil, err
}
if options.ResumeToken != "" {
return runtime.NewPollerFromResumeToken(options.ResumeToken, c.kvClient.Pl, &runtime.NewPollerFromResumeTokenOptions[RecoverDeletedKeyResponse]{
Handler: &handler,
})
}

return &RecoverDeletedKeyPoller{
lastResponse: getResp,
keyName: name,
client: c.kvClient,
vaultUrl: c.vaultURL,
recoverResponse: recoverResp,
lastRawResponse: getRawResp,
resumeToken: resumeToken,
}, nil
var rawResp *http.Response
ctx = runtime.WithCaptureResponse(ctx, &rawResp)
if _, err := c.kvClient.RecoverDeletedKey(ctx, c.vaultURL, name, nil); err != nil {
return nil, err
}

return runtime.NewPoller(rawResp, c.kvClient.Pl, &runtime.NewPollerOptions[RecoverDeletedKeyResponse]{
Handler: &handler,
})
}

// UpdateKeyPropertiesOptions contains the optional parameters for the Client.UpdateKeyProperties method
Expand Down
Loading

0 comments on commit 844c0b5

Please sign in to comment.