Keyvault-admin typespec v2 lro restore and backup LRO customization in v3

sdk/keyvault/keyvault-admin/src/backupClient.ts

@@ -11,20 +11,21 @@ import {
   KeyVaultSelectiveKeyRestoreResult,
 } from "./backupClientModels.js";
 // import { KeyVaultAdminPollOperationState } from "./lro/keyVaultAdminPoller.js";
-import { KeyVaultBackupOperationState } from "./lro/backup/operation.js";
+import { KeyVaultBackupOperationState, KeyVaultBackupPoller } from "./lro/backup/poller.js";
 // import { KeyVaultBackupPoller } from "./lro/backup/poller.js";
 import { KeyVaultClient } from "./generated/keyVaultClient.js";
-import { KeyVaultRestoreOperationState } from "./lro/restore/operation.js";
 // import { KeyVaultRestorePoller } from "./lro/restore/poller.js";
-import { KeyVaultSelectiveKeyRestoreOperationState } from "./lro/selectiveKeyRestore/operation.js";
 // import { KeyVaultSelectiveKeyRestorePoller } from "./lro/selectiveKeyRestore/poller.js";
 import { LATEST_API_VERSION } from "./constants.js";
-import { PollerLike } from "@azure/core-lro";
 import { TokenCredential } from "@azure/core-auth";
 import { keyVaultAuthenticationPolicy } from "@azure/keyvault-common";
 import { logger } from "./log.js";
 import { bearerTokenAuthenticationPolicyName } from "@azure/core-rest-pipeline";
 import { mappings } from "./mappings.js";
+import { PollerLike } from "./lro/keyVaultAdminPoller.js";
+import { KeyVaultRestoreOperationState, KeyVaultRestorePoller } from "./lro/restore/poller.js";
+import { SimplePollerLike, wrapPoller } from "./lro/shim.js";
+import { KeyVaultSelectiveKeyRestoreOperationState, selectiveKeyRestoreOperation } from "./lro/selectiveKeyRestore/operation.js";
 
 // TODO: discuss no longer exporting the state at the top level as a bugfix
 // export {
@@ -182,18 +183,15 @@ export class KeyVaultBackupClient {
     const options =
       typeof sasTokenOrOptions === "string" ? optionsWhenSasTokenSpecified : sasTokenOrOptions;
 
-    const poller = this.client.fullBackup(
-      {
-        storageResourceUri: blobStorageUri,
-        token: sasToken,
-        useManagedIdentity: sasToken === undefined,
-      },
-      options,
-    );
-
-    await poller.submitted();
-
-    // @ts-expect-error TODO: Poller's type assumes no result, but we need a KeyVaultBackupResult
+    const poller = new KeyVaultBackupPoller({
+      blobStorageUri,
+      sasToken,
+      client: this.client,
+      operationOptions: options,
+      resumeFrom: options.resumeFrom,
+      intervalInMs: options.intervalInMs,
+    });
+    await poller.poll();
     return poller;
   }
 
@@ -282,21 +280,15 @@ export class KeyVaultBackupClient {
       typeof sasTokenOrOptions === "string" ? optionsWhenSasTokenSpecified : sasTokenOrOptions;
 
     const folderParts = mappings.folderUriParts(folderUri);
-    const poller = this.client.fullRestoreOperation(
-      {
-        folderToRestore: folderParts.folderName,
-        sasTokenParameters: {
-          storageResourceUri: folderParts.folderUri,
-          token: sasToken,
-          useManagedIdentity: sasToken === undefined,
-        },
-      },
-      options,
-    );
-
-    await poller.submitted();
-
-    // @ts-expect-error TODO: startTime should be required in the spec
+    const poller = new KeyVaultRestorePoller({
+      ...folderParts,
+      sasToken,
+      client: this.client,
+      operationOptions: options,
+      resumeFrom: options.resumeFrom,
+      intervalInMs: options.intervalInMs,
+    });
+    await poller.poll();
     return poller;
   }
 
@@ -339,7 +331,7 @@ export class KeyVaultBackupClient {
     sasToken: string,
     options?: KeyVaultBeginSelectiveKeyRestoreOptions,
   ): Promise<
-    PollerLike<KeyVaultSelectiveKeyRestoreOperationState, KeyVaultSelectiveKeyRestoreResult>
+    SimplePollerLike<KeyVaultSelectiveKeyRestoreOperationState<KeyVaultSelectiveKeyRestoreResult>, KeyVaultSelectiveKeyRestoreResult>
   >;
 
   /**
@@ -380,7 +372,7 @@ export class KeyVaultBackupClient {
     folderUri: string,
     options?: KeyVaultBeginSelectiveKeyRestoreOptions,
   ): Promise<
-    PollerLike<KeyVaultSelectiveKeyRestoreOperationState, KeyVaultSelectiveKeyRestoreResult>
+    SimplePollerLike<KeyVaultSelectiveKeyRestoreOperationState<KeyVaultSelectiveKeyRestoreResult>, KeyVaultSelectiveKeyRestoreResult>
   >;
 
   public async beginSelectiveKeyRestore(
@@ -389,27 +381,27 @@ export class KeyVaultBackupClient {
     sasTokenOrOptions: string | KeyVaultBeginSelectiveKeyRestoreOptions = {},
     optionsWhenSasTokenSpecified: KeyVaultBeginSelectiveKeyRestoreOptions = {},
   ): Promise<
-    PollerLike<KeyVaultSelectiveKeyRestoreOperationState, KeyVaultSelectiveKeyRestoreResult>
+    SimplePollerLike<KeyVaultSelectiveKeyRestoreOperationState<KeyVaultSelectiveKeyRestoreResult>, KeyVaultSelectiveKeyRestoreResult>
   > {
     const sasToken = typeof sasTokenOrOptions === "string" ? sasTokenOrOptions : undefined;
     const options =
       typeof sasTokenOrOptions === "string" ? optionsWhenSasTokenSpecified : sasTokenOrOptions;
-
-    const folderParts = mappings.folderUriParts(folderUri);
-    const poller = this.client.selectiveKeyRestoreOperation(
-      keyName,
-      {
-        folder: folderParts.folderName,
-        sasTokenParameters: {
-          storageResourceUri: folderParts.folderUri,
-          token: sasToken,
-          useManagedIdentity: sasToken === undefined,
+    const folderUriParts = mappings.folderUriParts(folderUri);
+    const poller = await wrapPoller(
+      selectiveKeyRestoreOperation(this.client["_client"], keyName, {
+        ...options,
+        restoreBlobDetails: {
+          folder: folderUriParts.folderName,
+          sasTokenParameters: {
+            storageResourceUri: folderUri,
+            token: sasToken,
+            useManagedIdentity: !sasToken,
+          },
         },
       },
-      options,
+      ),
     );
-    await poller.submitted();
-    // @ts-expect-error TODO: startTime should be required in the spec
+    await poller.poll();
     return poller;
   }
 }

sdk/keyvault/keyvault-admin/src/backupClientModels.ts

@@ -39,19 +39,19 @@ export interface KeyVaultBackupPollerOptions extends OperationOptions {
  * An interface representing the optional parameters that can be
  * passed to {@link beginBackup}
  */
-export interface KeyVaultBeginBackupOptions extends KeyVaultBackupPollerOptions {}
+export interface KeyVaultBeginBackupOptions extends KeyVaultBackupPollerOptions { }
 
 /**
  * An interface representing the optional parameters that can be
  * passed to {@link beginRestore}
  */
-export interface KeyVaultBeginRestoreOptions extends KeyVaultBackupPollerOptions {}
+export interface KeyVaultBeginRestoreOptions extends KeyVaultBackupPollerOptions { }
 
 /**
  * An interface representing the optional parameters that can be
  * passed to {@link beginSelectiveKeyRestore}
  */
-export interface KeyVaultBeginSelectiveKeyRestoreOptions extends KeyVaultBackupPollerOptions {}
+export interface KeyVaultBeginSelectiveKeyRestoreOptions extends KeyVaultBackupPollerOptions { }
 
 /**
  * An interface representing the result of a backup operation.
@@ -95,10 +95,10 @@ export interface KeyVaultSelectiveKeyRestoreResult {
   /**
    * The start time of the selective key restore operation.
    */
-  startTime: Date;
+  startTime: string;
 
   /**
    * The end time of the selective key restore operation.
    */
-  endTime?: Date;
+  endTime?: string;
 }