Decoupling LRO from core-client

src/coreClientLRO.ts

@@ -0,0 +1,303 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import {
+  OperationArguments,
+  OperationSpec,
+  OperationResponseMap,
+  FullOperationResponse
+} from "@azure/core-client";
+import {
+  FinalStateVia,
+  GetLROState,
+  InitializePollerState,
+  LRO,
+  LROConfig,
+  LROMode,
+  LROResult,
+  LROState,
+  PollingOperation,
+  createGetLROState,
+  terminalStates
+} from "./lro";
+
+export type SendOperationFn<T> = (
+  args: OperationArguments,
+  spec: OperationSpec
+) => Promise<LROResult<T>>;
+
+export function createPollingMethod<TResult>(
+  sendOperationFn: SendOperationFn<TResult>,
+  getLROState: GetLROState<TResult>,
+  args: OperationArguments,
+  spec: OperationSpec,
+  mode?: LROMode
+): PollingOperation<TResult> {
+  /**
+   * Polling calls will always return a status object i.e. {"status": "success"}
+   * these intermediate responses are not described in the swagger so we need to
+   * pass custom mappers at runtime.
+   * This function replaces all the existing mappers to be able to deserialize a status object
+   * @param responses Original set of responses defined in the operation
+   */
+  function getCompositeMappers(responses: {
+    [responseCode: string]: OperationResponseMap;
+  }): {
+    [responseCode: string]: OperationResponseMap;
+  } {
+    return Object.keys(responses).reduce((acc, statusCode) => {
+      return {
+        ...acc,
+        [statusCode]: {
+          ...responses[statusCode],
+          bodyMapper: {
+            type: {
+              name: "Composite",
+              modelProperties: {
+                status: {
+                  serializedName: "status",
+                  type: {
+                    name: "String"
+                  }
+                }
+              }
+            }
+          }
+        }
+      };
+    }, {} as { [responseCode: string]: OperationResponseMap });
+  }
+  let response: LROState<TResult> | undefined = undefined;
+  const customerCallback = args?.options?.onResponse;
+  const updatedArgs = {
+    ...args,
+    options: {
+      ...args.options,
+      onResponse: (
+        rawResponse: FullOperationResponse,
+        flatResponse: unknown
+      ): void => {
+        response = getLROState(
+          {
+            statusCode: rawResponse.status,
+            body: rawResponse.parsedBody,
+            headers: rawResponse.headers.toJSON()
+          },
+          flatResponse as TResult
+        );
+        if (response.done) {
+          customerCallback?.(rawResponse, flatResponse);
+        }
+      }
+    }
+  };
+  // Make sure we don't send any body to the get request
+  const { requestBody, responses, ...restSpec } = spec;
+  if (mode === "AzureAsync") {
+    return async (path?: string) => {
+      await sendOperationFn(updatedArgs, {
+        ...restSpec,
+        responses: getCompositeMappers(responses),
+        httpMethod: "GET",
+        ...(path && { path })
+      });
+      return response!;
+    };
+  }
+  return async (path?: string) => {
+    await sendOperationFn(updatedArgs, {
+      ...restSpec,
+      responses: responses,
+      httpMethod: "GET",
+      ...(path && { path })
+    });
+    return response!;
+  };
+}
+
+/**
+ * We need to selectively deserialize our responses, only deserializing if we
+ * are in a final LRO response, not deserializing any polling non-terminal responses
+ */
+export function shouldDeserializeLRO(finalStateVia?: string) {
+  let initialOperationInfo: LROResponseInfo | undefined;
+  let isInitialRequest = true;
+
+  return (response: FullOperationResponse) => {
+    if (response.status < 200 || response.status >= 300) {
+      return true;
+    }
+
+    if (!initialOperationInfo) {
+      initialOperationInfo = getLROData(response);
+    } else {
+      isInitialRequest = false;
+    }
+
+    if (
+      initialOperationInfo.azureAsyncOperation ||
+      initialOperationInfo.operationLocation
+    ) {
+      return (
+        !isInitialRequest &&
+        isAsyncOperationFinalResponse(
+          response,
+          initialOperationInfo,
+          finalStateVia
+        )
+      );
+    }
+
+    if (initialOperationInfo.location) {
+      return isLocationFinalResponse(response);
+    }
+
+    if (initialOperationInfo.requestMethod === "PUT") {
+      return isBodyPollingFinalResponse(response);
+    }
+
+    return true;
+  };
+}
+
+function isAsyncOperationFinalResponse(
+  response: FullOperationResponse,
+  initialOperationInfo: LROResponseInfo,
+  finalStateVia?: string
+): boolean {
+  const status: string = response.parsedBody?.status || "Succeeded";
+  if (!terminalStates.includes(status.toLowerCase())) {
+    return false;
+  }
+
+  if (initialOperationInfo.requestMethod === "DELETE") {
+    return true;
+  }
+
+  if (
+    initialOperationInfo.requestMethod === "PUT" &&
+    finalStateVia &&
+    finalStateVia.toLowerCase() === "azure-asyncoperation"
+  ) {
+    return true;
+  }
+
+  if (
+    initialOperationInfo.requestMethod !== "PUT" &&
+    !initialOperationInfo.location
+  ) {
+    return true;
+  }
+
+  return false;
+}
+
+function isLocationFinalResponse(response: FullOperationResponse): boolean {
+  return response.status !== 202;
+}
+
+function isBodyPollingFinalResponse(response: FullOperationResponse): boolean {
+  const provisioningState: string =
+    response.parsedBody?.properties?.provisioningState || "Succeeded";
+
+  if (terminalStates.includes(provisioningState.toLowerCase())) {
+    return true;
+  }
+
+  return false;
+}
+
+interface LROResponseInfo {
+  requestMethod: string;
+  azureAsyncOperation?: string;
+  operationLocation?: string;
+  location?: string;
+}
+
+function getLROData(result: FullOperationResponse): LROResponseInfo {
+  return {
+    azureAsyncOperation: result.headers.get("azure-asyncoperation"),
+    operationLocation: result.headers.get("operation-location"),
+    location: result.headers.get("location"),
+    requestMethod: result.request.method
+  };
+}
+
+export function getSpecPath(spec: OperationSpec): string {
+  if (spec.path) {
+    return spec.path;
+  } else {
+    throw Error("Bad spec: request path is not found!");
+  }
+}
+
+export class CoreClientLRO<T> implements LRO<T> {
+  constructor(
+    private sendOperationFn: SendOperationFn<T>,
+    private args: OperationArguments,
+    private spec: OperationSpec,
+    private finalStateVia?: FinalStateVia,
+    public requestPath: string = spec.path!,
+    public requestMethod: string = spec.httpMethod
+  ) {}
+  public async sendInitialRequest(
+    initializeState: InitializePollerState
+  ): Promise<LROResult<T>> {
+    const { onResponse, ...restOptions } = this.args.options || {};
+    return this.sendOperationFn(
+      {
+        ...this.args,
+        options: {
+          ...restOptions,
+          onResponse: (
+            rawResponse: FullOperationResponse,
+            flatResponse: unknown
+          ) => {
+            const isCompleted = initializeState(
+              {
+                statusCode: rawResponse.status,
+                body: rawResponse.parsedBody,
+                headers: rawResponse.headers.toJSON()
+              },
+              flatResponse
+            );
+            if (isCompleted) {
+              onResponse?.(rawResponse, flatResponse);
+            }
+          }
+        }
+      },
+      this.spec
+    );
+  }
+
+  public async sendPollRequest(
+    config: LROConfig,
+    path?: string
+  ): Promise<LROState<T>> {
+    const getLROState = createGetLROState(this, config, this.finalStateVia);
+    return createPollingMethod(
+      this.sendOperationFn,
+      getLROState,
+      this.args,
+      this.spec,
+      config.mode
+    )(path);
+  }
+  public async retrieveAzureAsyncResource(path?: string): Promise<LROState<T>> {
+    const updatedArgs = { ...this.args };
+    if (updatedArgs.options) {
+      (updatedArgs.options as any).shouldDeserialize = true;
+    }
+    return createPollingMethod(
+      this.sendOperationFn,
+      (rawResponse, flatResponse) => ({
+        rawResponse,
+        flatResponse,
+        done: true
+      }),
+      updatedArgs,
+      this.spec
+    )(path);
+  }
+}

src/generators/LROGenerator.ts

@@ -13,8 +13,9 @@ export async function generateLROFiles(
   if (!hasAnyLRO(clientDetails.operationGroups)) {
     return;
   }
-
-  const lroDir = joinPath(__dirname, "..", "..", "..", "src", "lro");
+  const baseTargetPath = srcPath || "";
+  const srcDir = joinPath(__dirname, "..", "..", "..", "src");
+  const lroDir = joinPath(srcDir, "lro");
   const lroFiles = await promises.readdir(lroDir);
 
   for (let i = 0; i < lroFiles.length; i++) {
@@ -23,11 +24,22 @@ export async function generateLROFiles(
     const fileContent = await promises.readFile(filePath, "utf-8");
 
     project.createSourceFile(
-      joinPath(srcPath || "", "lro", file),
+      joinPath(baseTargetPath, "lro", file),
       fileContent,
       { overwrite: true }
     );
   }
+  const fileContent = await promises.readFile(
+    joinPath(srcDir, "coreClientLRO.ts"),
+    "utf-8"
+  );
+  project.createSourceFile(
+    joinPath(baseTargetPath, "coreClientLRO.ts"),
+    fileContent,
+    {
+      overwrite: true
+    }
+  );
 }
 
 function hasAnyLRO(operationGroups: OperationGroupDetails[]) {

src/generators/operationGenerator.ts

@@ -739,13 +739,15 @@ function addOperationOverloads(
 function compileOperationOptionsToRequestOptionsBase(
   options: string,
   isLRO: boolean,
-  finalStateVia: string
+  finalStateVia?: string
 ): string {
   const { useCoreV2 } = getAutorestOptions();
   // In LRO we have a couple extra properties to add that's why we use
   // the private getOperationOptions function instead of the one in core-http
   return isLRO
-    ? `this.getOperationOptions(${options}, "${finalStateVia}")`
+    ? `this.getOperationOptions(${options}${
+        finalStateVia === undefined ? "" : ", ${finalStateVia}"
+      })`
     : !useCoreV2
     ? `coreHttp.operationOptionsToRequestOptionsBase(options || {})`
     : `options || {}`;
@@ -963,16 +965,19 @@ function writeLROOperationBody(
         }
       };
       const flatResponse = await directSendOperation(updatedArgs, spec);
-      return { flatResponse, rawResponse: currentRawResponse! };
+      return { flatResponse, rawResponse: {
+        statusCode: currentRawResponse!.status,
+        body: currentRawResponse!.parsedBody,
+        headers: currentRawResponse!.headers.toJSON()
+      }};
   }`;
 
   methodDeclaration.addStatements([
     sendOperationStatement,
+    `const lro = new CoreClientLRO(sendOperation,${operationParamsName},
+      ${operationSpecName},${finalStateStr})`,
     `return new LROPoller({intervalInMs: options?.updateIntervalInMs},
-      ${operationParamsName},
-      ${operationSpecName},
-      sendOperation,
-      ${finalStateStr}
+      lro
     );`
   ]);
 
@@ -1323,9 +1328,13 @@ function addImports(
 
   if (hasLROOperation(operationGroupDetails)) {
     operationGroupFile.addImportDeclaration({
-      namedImports: ["LROPoller", "shouldDeserializeLRO"],
+      namedImports: ["LROPoller"],
       moduleSpecifier: `../lro`
     });
+    operationGroupFile.addImportDeclaration({
+      namedImports: ["CoreClientLRO", "shouldDeserializeLRO"],
+      moduleSpecifier: `../coreClientLRO`
+    });
     operationGroupFile.addImportDeclaration({
       namedImports: ["PollerLike", "PollOperationState"],
       moduleSpecifier: "@azure/core-lro"