[FormRecognizer] Bring back Azure AD credential support (#8898)

Azure · May 15, 2020 · cebbff5 · cebbff5
1 parent 538727e
commit cebbff5
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 13 deletions.
diff --git a/sdk/formrecognizer/ai-form-recognizer/review/ai-form-recognizer.api.md b/sdk/formrecognizer/ai-form-recognizer/review/ai-form-recognizer.api.md
@@ -14,6 +14,7 @@ import { PipelineOptions } from '@azure/core-http';
 import { PollerLike } from '@azure/core-lro';
 import { PollOperationState } from '@azure/core-lro';
 import { RestResponse } from '@azure/core-http';
+import { TokenCredential } from '@azure/identity';
 
 // @public
 export interface AccountProperties {
@@ -199,7 +200,7 @@ export type FormPollerLike = PollerLike<PollOperationState<RecognizeFormResultRe
 
 // @public
 export class FormRecognizerClient {
-    constructor(endpointUrl: string, credential: KeyCredential, options?: FormRecognizerClientOptions);
+    constructor(endpointUrl: string, credential: TokenCredential | KeyCredential, options?: FormRecognizerClientOptions);
     beginRecognizeContent(data: FormRecognizerRequestBody, contentType?: ContentType, options?: BeginRecognizeContentOptions): Promise<ContentPollerLike>;
     beginRecognizeContentFromUrl(formFileUrl: string, options?: BeginRecognizeContentOptions): Promise<ContentPollerLike>;
     beginRecognizeCustomForms(modelId: string, data: FormRecognizerRequestBody, contentType?: ContentType, options?: BeginRecognizeFormsOptions): Promise<FormPollerLike>;
@@ -269,7 +270,7 @@ export interface FormText {
 
 // @public
 export class FormTrainingClient {
-    constructor(endpointUrl: string, credential: KeyCredential, options?: FormRecognizerClientOptions);
+    constructor(endpointUrl: string, credential: TokenCredential | KeyCredential, options?: FormRecognizerClientOptions);
     beginTraining(trainingFilesUrl: string, useTrainingLabels?: boolean, options?: BeginTrainingOptions<FormModelResponse>): Promise<PollerLike<PollOperationState<FormModelResponse>, FormModelResponse>>;
     deleteModel(modelId: string, options?: DeleteModelOptions): Promise<RestResponse>;
     readonly endpointUrl: string;
@@ -597,8 +598,8 @@ export type ValueTypes = "string" | "date" | "time" | "phoneNumber" | "number" |
 
 // Warnings were encountered during analysis:
 //
-// src/formRecognizerClient.ts:70:3 - (ae-forgotten-export) The symbol "BeginRecognizePollState" needs to be exported by the entry point index.d.ts
-// src/formTrainingClient.ts:66:3 - (ae-forgotten-export) The symbol "BeginTrainingPollState" needs to be exported by the entry point index.d.ts
+// src/formRecognizerClient.ts:73:3 - (ae-forgotten-export) The symbol "BeginRecognizePollState" needs to be exported by the entry point index.d.ts
+// src/formTrainingClient.ts:69:3 - (ae-forgotten-export) The symbol "BeginTrainingPollState" needs to be exported by the entry point index.d.ts
 
 // (No @packageDocumentation comment for this package)
 

diff --git a/sdk/formrecognizer/ai-form-recognizer/src/formRecognizerClient.ts b/sdk/formrecognizer/ai-form-recognizer/src/formRecognizerClient.ts
@@ -4,12 +4,15 @@
 import {
   createPipelineFromOptions,
   InternalPipelineOptions,
+  isTokenCredential,
+  bearerTokenAuthenticationPolicy,
   operationOptionsToRequestOptionsBase,
   AbortSignalLike,
   ServiceClientCredentials
 } from "@azure/core-http";
+import { TokenCredential } from '@azure/identity';
 import { KeyCredential } from "@azure/core-auth";
-import { SDK_VERSION } from "./constants";
+import { SDK_VERSION, DEFAULT_COGNITIVE_SCOPE } from "./constants";
 import { logger } from "./logger";
 import { createSpan } from "./tracing";
 import {
@@ -182,7 +185,7 @@ export class FormRecognizerClient {
    * @internal
    * @ignore
    */
-  private readonly credential: KeyCredential;
+  private readonly credential: TokenCredential | KeyCredential;
 
   /**
    * @internal
@@ -204,12 +207,12 @@ export class FormRecognizerClient {
    * );
    * ```
    * @param {string} endpointUrl Url to an Azure Form Recognizer service endpoint
-   * @param {KeyCredential} credential Used to authenticate requests to the service.
+   * @param {TokenCredential | KeyCredential} credential Used to authenticate requests to the service.
    * @param {FormRecognizerClientOptions} [options] Used to configure the Form Recognizer client.
    */
   constructor(
     endpointUrl: string,
-    credential: KeyCredential,
+    credential: TokenCredential | KeyCredential,
     options: FormRecognizerClientOptions = {}
   ) {
     this.endpointUrl = endpointUrl;
@@ -226,7 +229,9 @@ export class FormRecognizerClient {
       pipelineOptions.userAgentOptions.userAgentPrefix = libInfo;
     }
 
-    const authPolicy = createFormRecognizerAzureKeyCredentialPolicy(credential);
+    const authPolicy = isTokenCredential(credential)
+      ? bearerTokenAuthenticationPolicy(credential, DEFAULT_COGNITIVE_SCOPE)
+      : createFormRecognizerAzureKeyCredentialPolicy(credential);
 
     const internalPipelineOptions: InternalPipelineOptions = {
       ...pipelineOptions,

diff --git a/sdk/formrecognizer/ai-form-recognizer/src/formTrainingClient.ts b/sdk/formrecognizer/ai-form-recognizer/src/formTrainingClient.ts
@@ -6,14 +6,17 @@
 import {
   createPipelineFromOptions,
   InternalPipelineOptions,
+  isTokenCredential,
+  bearerTokenAuthenticationPolicy,
   operationOptionsToRequestOptionsBase,
   RestResponse,
   ServiceClientCredentials
 } from "@azure/core-http";
+import { TokenCredential } from '@azure/identity';
 import { KeyCredential } from "@azure/core-auth";
 import { PagedAsyncIterableIterator, PageSettings } from "@azure/core-paging";
 import "@azure/core-paging";
-import { SDK_VERSION } from "./constants";
+import { SDK_VERSION, DEFAULT_COGNITIVE_SCOPE } from "./constants";
 import { logger } from "./logger";
 import { createSpan } from "./tracing";
 import { CanonicalCode } from "@opentelemetry/api";
@@ -96,12 +99,12 @@ export class FormTrainingClient {
    * );
    * ```
    * @param {string} endpointUrl Url to an Azure Form Recognizer service endpoint
-   * @param {AzureKeyCredential} credential Used to authenticate requests to the service.
+   * @param {TokenCredential | KeyCredential} credential Used to authenticate requests to the service.
    * @param {FormRecognizerClientOptions} [options] Used to configure the client.
    */
   constructor(
     endpointUrl: string,
-    credential: KeyCredential,
+    credential: TokenCredential | KeyCredential,
     options: FormRecognizerClientOptions = {}
   ) {
     this.endpointUrl = endpointUrl;
@@ -117,7 +120,9 @@ export class FormTrainingClient {
       pipelineOptions.userAgentOptions.userAgentPrefix = libInfo;
     }
 
-    const authPolicy = createFormRecognizerAzureKeyCredentialPolicy(credential);
+    const authPolicy = isTokenCredential(credential)
+      ? bearerTokenAuthenticationPolicy(credential, DEFAULT_COGNITIVE_SCOPE)
+      : createFormRecognizerAzureKeyCredentialPolicy(credential);
 
     const internalPipelineOptions: InternalPipelineOptions = {
       ...pipelineOptions,