Update LRO generation

packages/autorest.typescript/package.json

@@ -60,7 +60,7 @@
     "@azure/core-client": "^1.6.1",
     "@azure/core-http": "^1.2.4",
     "@azure/core-http-compat": "^1.2.0",
-    "@azure/core-lro": "^2.4.0",
+    "@azure/core-lro": "^2.5.0",
     "@azure/core-paging": "^1.2.0",
     "@azure/core-rest-pipeline": "^1.8.0",
     "@azure/core-tracing": "^1.0.0",

packages/autorest.typescript/src/autorestSession.ts

@@ -41,6 +41,7 @@ export interface AutorestOptions {
   coreHttpCompatMode?: boolean;
   dependencyInfo?: DependencyInfo;
   lenientModelDeduplication?: boolean;
+  useLegacyLro?: boolean;
 }
 
 let host: AutorestExtensionHost;

packages/autorest.typescript/src/generators/LROGenerator.ts

@@ -15,9 +15,17 @@ export async function generateLroFiles(
   }
   const LroClassFile = "lroImpl.ts";
   const baseTargetPath = srcPath || "";
-  const srcDir = joinPath(__dirname, "..", "..", "..", "src");
+  const staticDir = joinPath(
+    __dirname,
+    "..",
+    "..",
+    "..",
+    "src",
+    "generators",
+    "static"
+  );
   const fileContent = await promises.readFile(
-    joinPath(srcDir, LroClassFile),
+    joinPath(staticDir, LroClassFile),
     "utf-8"
   );
   project.createSourceFile(

packages/autorest.typescript/src/generators/clientFileGenerator.ts

@@ -32,6 +32,7 @@ import { ParameterDetails } from "../models/parameterDetails";
 import { EndpointDetails } from "../transforms/urlTransforms";
 import { PackageDetails } from "../models/packageDetails";
 import { getSecurityInfoFromModel } from "../utils/schemaHelpers";
+import { createLroImports } from "../utils/lroHelpers";
 
 type OperationDeclarationDetails = { name: string; typeName: string };
 
@@ -41,7 +42,8 @@ export function generateClient(clientDetails: ClientDetails, project: Project) {
     hideClients,
     srcPath,
     packageDetails,
-    coreHttpCompatMode
+    coreHttpCompatMode,
+    useLegacyLro
   } = getAutorestOptions();
   const { addCredentials } = getSecurityInfoFromModel(clientDetails.security);
   const hasMappers = !!clientDetails.mappers.length;
@@ -135,11 +137,11 @@ export function generateClient(clientDetails: ClientDetails, project: Project) {
 
   if (hasInlineOperations && hasLro) {
     clientFile.addImportDeclaration({
-      namedImports: ["PollerLike", "PollOperationState", "LroEngine"],
+      namedImports: createLroImports(useLegacyLro),
       moduleSpecifier: "@azure/core-lro"
     });
     clientFile.addImportDeclaration({
-      namedImports: ["LroImpl"],
+      namedImports: ["createLroSpec"],
       moduleSpecifier: `./lroImpl`
     });
   }

packages/autorest.typescript/src/generators/operationGenerator.ts

@@ -50,6 +50,7 @@ import {
 } from "./utils/pagingOperations";
 import { calculateMethodName } from "./utils/operationsUtils";
 import { getAutorestOptions } from "../autorestSession";
+import { createLroImports, createLroType } from "../utils/lroHelpers";
 
 /**
  * Function that writes the code for all the operations.
@@ -448,8 +449,10 @@ function getReturnType(
     modelNames
   );
 
+  const { useLegacyLro } = getAutorestOptions();
+
   return operation.isLro
-    ? `Promise<PollerLike<PollOperationState<${responseName}>,${responseName}>>`
+    ? createLroType({ responseName, useLegacyLro })
     : `Promise<${responseName}>`;
 }
 
@@ -900,7 +903,7 @@ function writeLroOperationBody(
   isInline?: boolean,
   isTracingEnabled = false
 ) {
-  const { useCoreV2 } = getAutorestOptions();
+  const { useCoreV2, useLegacyLro } = getAutorestOptions();
   const spanName = `${clientDetails.className}.${methodDeclaration.getName()}`;
   const client = isInline ? "" : ".client";
   const sendRequestStatement = `this${client}.sendOperationRequest(args, spec)`;
@@ -929,7 +932,7 @@ function writeLroOperationBody(
       spanName
     )}
     };
-    const sendOperation = async (args: coreClient.OperationArguments, spec: coreClient.OperationSpec) => {
+    const sendOperationFn = async (args: coreClient.OperationArguments, spec: coreClient.OperationSpec) => {
       let currentRawResponse: coreClient.FullOperationResponse | undefined = undefined;
       const providedCallback = args.options?.onResponse;
       const callback: coreClient.RawResponseCallback = (
@@ -952,21 +955,29 @@ function writeLroOperationBody(
         headers: currentRawResponse!.headers.toJSON()
       }};
   }`;
+
+  const commonOptions = `intervalInMs: options?.updateIntervalInMs${
+    lroResourceLocationConfig
+      ? `, ${
+          useLegacyLro ? "lroResourceLocationConfig" : "resourceLocationConfig"
+        }: "${lroResourceLocationConfig.toLowerCase()}"`
+      : ""
+  }`;
   methodDeclaration.addStatements([
     sendOperationStatement,
-    `const lro = new LroImpl(sendOperation,${operationParamsName},
-      ${operationSpecName})`,
-    `const poller = new LroEngine(lro,{ resumeFrom: options?.resumeFrom, intervalInMs: options?.updateIntervalInMs${
-      lroResourceLocationConfig
-        ? `, lroResourceLocationConfig: "${lroResourceLocationConfig.toLowerCase()}"`
-        : ""
-    } });`,
+    `const lro = createLroSpec({sendOperationFn, args: ${operationParamsName},
+      spec: ${operationSpecName}})`,
+    `const poller = ${
+      useLegacyLro
+        ? `new LroEngine(lro, { resumeFrom: options?.resumeFrom, ${commonOptions} })`
+        : `await createHttpPoller<${responseName}, OperationState<${responseName}>>(lro, { restoreFrom: options?.resumeFrom, ${commonOptions} })`
+    };`,
     "await poller.poll();",
     "return poller;"
   ]);
 
   methodDeclaration.setReturnType(
-    `Promise<PollerLike<PollOperationState<${responseName}>,${responseName}>>`
+    createLroType({ responseName, useLegacyLro })
   );
 }
 
@@ -1228,7 +1239,7 @@ function addImports(
   operationGroupFile: SourceFile,
   clientDetails: ClientDetails
 ) {
-  const { useCoreV2 } = getAutorestOptions();
+  const { useCoreV2, useLegacyLro } = getAutorestOptions();
 
   const { className, mappers } = clientDetails;
   addPagingEsNextRef(operationGroupDetails.operations, operationGroupFile);
@@ -1295,11 +1306,11 @@ function addImports(
 
   if (hasLroOperation(operationGroupDetails)) {
     operationGroupFile.addImportDeclaration({
-      namedImports: ["PollerLike", "PollOperationState", "LroEngine"],
+      namedImports: createLroImports(useLegacyLro),
       moduleSpecifier: "@azure/core-lro"
     });
     operationGroupFile.addImportDeclaration({
-      namedImports: ["LroImpl"],
+      namedImports: ["createLroSpec"],
       moduleSpecifier: `../lroImpl`
     });
   }

packages/autorest.typescript/src/generators/operationInterfaceGenerator.ts

@@ -29,6 +29,7 @@ import {
 } from "./utils/pagingOperations";
 import { calculateMethodName } from "./utils/operationsUtils";
 import { getAutorestOptions } from "../autorestSession";
+import { createLroType } from "../utils/lroHelpers";
 
 /**
  * Function that writes the code for all the operations.
@@ -112,8 +113,10 @@ function getReturnType(
     modelNames
   );
 
+  const { useLegacyLro } = getAutorestOptions();
+
   return operation.isLro
-    ? `Promise<PollerLike<PollOperationState<${responseName}>,${responseName}>>`
+    ? createLroType({ responseName, useLegacyLro })
     : `Promise<${responseName}>`;
 }
 
@@ -237,7 +240,7 @@ function addImports(
   operationGroupFile: SourceFile,
   clientDetails: ClientDetails
 ) {
-  const { useCoreV2 } = getAutorestOptions();
+  const { useCoreV2, useLegacyLro } = getAutorestOptions();
   addPagingEsNextRef(operationGroupDetails.operations, operationGroupFile);
   addTracingOperationImports(operationGroupFile);
   addPagingImports(operationGroupDetails.operations, operationGroupFile);
@@ -260,7 +263,9 @@ function addImports(
 
   if (hasLroOperation(operationGroupDetails)) {
     operationGroupFile.addImportDeclaration({
-      namedImports: ["PollerLike", "PollOperationState"],
+      namedImports: useLegacyLro
+        ? ["PollerLike", "PollOperationState"]
+        : ["SimplePollerLike", "OperationState"],
       moduleSpecifier: "@azure/core-lro"
     });
   }

packages/autorest.typescript/src/generators/samples/rlcSamples.ts.hbs

@@ -34,7 +34,7 @@ async function {{name}}() {
   }
   {{else if this.isLRO}}
   const initialResponse = await client.path({{this.pathParamNames}}).{{method}}({{methodParamNames}});
-  const poller = getLongRunningPoller(client, initialResponse);
+  const poller = await getLongRunningPoller(client, initialResponse);
   const result = await poller.pollUntilDone();
   {{else}}
   const result = await client.path({{this.pathParamNames}}).{{method}}({{methodParamNames}});

packages/autorest.typescript/src/generators/static/lroImpl.ts

@@ -0,0 +1,34 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+import { AbortSignalLike } from "@azure/abort-controller";
+import { LongRunningOperation, LroResponse } from "@azure/core-lro";
+
+export function createLroSpec<T>(inputs: {
+  sendOperationFn: (args: any, spec: any) => Promise<LroResponse<T>>;
+  args: Record<string, unknown>;
+  spec: {
+    readonly requestBody?: unknown;
+    readonly path?: string;
+    readonly httpMethod: string;
+  } & Record<string, any>;
+}): LongRunningOperation<T> {
+  const { args, spec, sendOperationFn } = inputs;
+  return {
+    requestMethod: spec.httpMethod,
+    requestPath: spec.path!,
+    sendInitialRequest: () => sendOperationFn(args, spec),
+    sendPollRequest: (
+      path: string,
+      options?: { abortSignal?: AbortSignalLike }
+    ) => {
+      const { requestBody, ...restSpec } = spec;
+      return sendOperationFn(args, {
+        ...restSpec,
+        httpMethod: "GET",
+        path,
+        abortSignal: options?.abortSignal
+      });
+    }
+  };
+}

packages/autorest.typescript/src/generators/static/packageFileGenerator.ts

@@ -86,7 +86,7 @@ function regularAutorestPackage(
       node: ">=14.0.0"
     },
     dependencies: {
-      ...(hasLro && { "@azure/core-lro": "^2.2.0" }),
+      ...(hasLro && { "@azure/core-lro": "^2.5.0" }),
       ...(hasLro && { "@azure/abort-controller": "^1.0.0" }),
       ...(hasAsyncIterators && { "@azure/core-paging": "^1.2.0" }),
       ...(!useCoreV2 && { "@azure/core-http": "^2.0.0" }),

packages/autorest.typescript/src/utils/autorestOptions.ts

@@ -11,6 +11,7 @@ export async function extractAutorestOptions(): Promise<AutorestOptions> {
   const host = getHost();
   const useCoreV2 = await getUseCoreV2(host);
   const restLevelClient = await getRestLevelClient(host);
+  const useLegacyLro = await getUseLegacyLro(host);
   const rlcShortcut = await getHasShortcutMethods(host);
   const azureArm = await getIsAzureArm(host);
   const addCredentials = await getAddCredentials(host);
@@ -73,10 +74,16 @@ export async function extractAutorestOptions(): Promise<AutorestOptions> {
     azureSdkForJs,
     productDocLink,
     coreHttpCompatMode,
-    dependencyInfo
+    dependencyInfo,
+    useLegacyLro
   };
 }
 
+async function getUseLegacyLro(host: AutorestExtensionHost): Promise<boolean> {
+  const useLegacyLroOption = await host.getValue("use-legacy-lro");
+  return useLegacyLroOption === null ? false : Boolean(useLegacyLroOption);
+}
+
 async function getHasShortcutMethods(
   host: AutorestExtensionHost
 ): Promise<boolean> {

packages/autorest.typescript/src/utils/lroHelpers.ts

@@ -0,0 +1,17 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+export const createLroImports = (useLegacyLro?: boolean) =>
+  useLegacyLro
+    ? ["PollerLike", "PollOperationState", "LroEngine"]
+    : ["SimplePollerLike", "OperationState", "createHttpPoller"];
+
+export function createLroType(inputs: {
+  useLegacyLro?: boolean;
+  responseName: string;
+}): string {
+  const { responseName, useLegacyLro } = inputs;
+  return useLegacyLro
+    ? `Promise<PollerLike<PollOperationState<${responseName}>,${responseName}>>`
+    : `Promise<SimplePollerLike<OperationState<${responseName}>,${responseName}>>`;
+}

packages/autorest.typescript/test/integration/generated/datafactory/package.json

@@ -6,7 +6,7 @@
   "version": "1.0.0-preview1",
   "engines": { "node": ">=14.0.0" },
   "dependencies": {
-    "@azure/core-lro": "^2.2.0",
+    "@azure/core-lro": "^2.5.0",
     "@azure/abort-controller": "^1.0.0",
     "@azure/core-paging": "^1.2.0",
     "@azure/core-client": "^1.7.0",

packages/autorest.typescript/test/integration/generated/datafactory/src/lroImpl.ts

@@ -1,26 +1,34 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+import { AbortSignalLike } from "@azure/abort-controller";
 import { LongRunningOperation, LroResponse } from "@azure/core-lro";
 
-export class LroImpl<T> implements LongRunningOperation<T> {
-  constructor(
-    private sendOperationFn: (args: any, spec: any) => Promise<LroResponse<T>>,
-    private args: Record<string, unknown>,
-    private spec: {
-      readonly requestBody?: unknown;
-      readonly path?: string;
-      readonly httpMethod: string;
-    } & Record<string, any>,
-    public requestPath: string = spec.path!,
-    public requestMethod: string = spec.httpMethod
-  ) {}
-  public async sendInitialRequest(): Promise<LroResponse<T>> {
-    return this.sendOperationFn(this.args, this.spec);
-  }
-  public async sendPollRequest(path: string): Promise<LroResponse<T>> {
-    const { requestBody, ...restSpec } = this.spec;
-    return this.sendOperationFn(this.args, {
-      ...restSpec,
-      path,
-      httpMethod: "GET"
-    });
-  }
+export function createLroSpec<T>(inputs: {
+  sendOperationFn: (args: any, spec: any) => Promise<LroResponse<T>>;
+  args: Record<string, unknown>;
+  spec: {
+    readonly requestBody?: unknown;
+    readonly path?: string;
+    readonly httpMethod: string;
+  } & Record<string, any>;
+}): LongRunningOperation<T> {
+  const { args, spec, sendOperationFn } = inputs;
+  return {
+    requestMethod: spec.httpMethod,
+    requestPath: spec.path!,
+    sendInitialRequest: () => sendOperationFn(args, spec),
+    sendPollRequest: (
+      path: string,
+      options?: { abortSignal?: AbortSignalLike }
+    ) => {
+      const { requestBody, ...restSpec } = spec;
+      return sendOperationFn(args, {
+        ...restSpec,
+        httpMethod: "GET",
+        path,
+        abortSignal: options?.abortSignal
+      });
+    }
+  };
 }