Add az login support (#301)

https://user-images.githubusercontent.com/40952/207650170-0e752502-3065-4733-b7b7-8818ee93db5b.mov

Co-authored-by: Serge Smertin <259697+nfx@users.noreply.github.com>

Author: fjakobs

packages/databricks-sdk-js/src/api-client.ts

@@ -7,6 +7,8 @@ import {ExposedLoggers, Utils, withLogContext} from "./logging";
 // eslint-disable-next-line @typescript-eslint/no-unused-vars
 import {context} from "./context";
 import {Context} from "./context";
+import retry, {RetriableError} from "./retries/retries";
+import Time, {TimeUnits} from "./retries/Time";
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const sdkVersion = require("../package.json").version;
@@ -109,38 +111,75 @@ export class ApiClient {
             }
         }
 
-        let response;
+        const response = await retry<
+            Awaited<Awaited<ReturnType<typeof fetch>>["response"]>
+        >({
+            timeout: new Time(10, TimeUnits.seconds),
+            fn: async () => {
+                let response;
+                try {
+                    const {abort, response: responsePromise} = await fetch(
+                        url.toString(),
+                        options
+                    );
+                    if (context?.cancellationToken?.onCancellationRequested) {
+                        context?.cancellationToken?.onCancellationRequested(
+                            abort
+                        );
+                    }
+                    response = await responsePromise;
+                } catch (e: any) {
+                    const err =
+                        e.code && e.code === "ENOTFOUND"
+                            ? new HttpError(
+                                  `Can't connect to ${url.toString()}`,
+                                  500
+                              )
+                            : e;
+                    throw logAndReturnError(url, options, "", err, context);
+                }
+
+                switch (response.status) {
+                    case 500:
+                    case 429:
+                        throw new RetriableError();
+
+                    default:
+                        break;
+                }
+                return response;
+            },
+        });
 
+        let responseText!: string;
         try {
-            const {abort, response: responsePromise} = await fetch(
-                url.toString(),
-                options
-            );
-            if (context?.cancellationToken?.onCancellationRequested) {
-                context?.cancellationToken?.onCancellationRequested(abort);
-            }
-            response = await responsePromise;
+            const responseBody = await response.arrayBuffer();
+            responseText = new TextDecoder().decode(responseBody);
         } catch (e: any) {
-            const err =
-                e.code && e.code === "ENOTFOUND"
-                    ? new HttpError(`Can't connect to ${url.toString()}`, 500)
-                    : e;
-            throw logAndReturnError(url, options, response, err, context);
+            logAndReturnError(url, options, "", e, context);
+            throw new ApiClientResponseError(
+                `Can't parse response from ${url.toString()}`,
+                ""
+            );
         }
 
         // throw error if the URL is incorrect and we get back an HTML page
         if (response.headers.get("content-type")?.match("text/html")) {
-            throw logAndReturnError(
-                url,
-                options,
-                response,
-                new HttpError(`Can't connect to ${url.toString()}`, 404),
-                context
-            );
-        }
+            // When the AAD tenent is not configured correctly, the response is a HTML page with a title like this:
+            // "Error 400 io.jsonwebtoken.IncorrectClaimException: Expected iss claim to be: https://sts.windows.net/aaaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaa/, but was: https://sts.windows.net/bbbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbb/."
+            const m = responseText.match(/<title>(Error \d+.*?)<\/title>/);
+            let error: HttpError;
+            if (m) {
+                error = new HttpError(m[1], response.status);
+            } else {
+                error = new HttpError(
+                    `Can't connect to ${url.toString()}`,
+                    response.status
+                );
+            }
 
-        const responseBody = await response.arrayBuffer();
-        const responseText = new TextDecoder().decode(responseBody);
+            throw logAndReturnError(url, options, response, error, context);
+        }
 
         // TODO proper error handling
         if (!response.ok) {
@@ -150,33 +189,40 @@ export class ApiClient {
             throw logAndReturnError(url, options, responseText, err, context);
         }
 
+        let responseJson: any;
         try {
-            response = JSON.parse(responseText);
+            responseJson = JSON.parse(responseText);
         } catch (e) {
             logAndReturnError(url, options, responseText, e, context);
-            new ApiClientResponseError(responseText, response);
+            throw new ApiClientResponseError(responseText, responseJson);
         }
 
-        if ("error" in response) {
-            logAndReturnError(url, options, response, response.error, context);
-            throw new ApiClientResponseError(response.error, response);
+        if ("error" in responseJson) {
+            logAndReturnError(
+                url,
+                options,
+                responseJson,
+                responseJson.error,
+                context
+            );
+            throw new ApiClientResponseError(responseJson.error, responseJson);
         }
 
-        if ("error_code" in response) {
+        if ("error_code" in responseJson) {
             const message =
-                response.message || `HTTP error ${response.error_code}`;
+                responseJson.message || `HTTP error ${responseJson.error_code}`;
             throw logAndReturnError(
                 url,
                 options,
-                response,
-                new HttpError(message, response.error_code),
+                responseJson,
+                new HttpError(message, responseJson.error_code),
                 context
             );
         }
         context?.logger?.debug(url.toString(), {
             request: options,
-            response: response,
+            response: responseJson,
         });
-        return response as any;
+        return responseJson;
     }
 }

packages/databricks-sdk-js/src/auth/Token.ts

@@ -0,0 +1,13 @@
+import {Credentials} from "./types";
+
+export class Token implements Credentials {
+    constructor(
+        public readonly host: URL,
+        public readonly token: string,
+        public readonly expiry: number
+    ) {}
+
+    public isValid(): boolean {
+        return this.expiry > Date.now() + 10_000;
+    }
+}

packages/databricks-sdk-js/src/auth/configFile.test.ts

@@ -5,6 +5,7 @@ import {
     HostParsingError,
     isConfigFileParsingError,
     loadConfigFile,
+    Profile,
     resolveConfigFilePath,
     TokenParsingError,
 } from "./configFile";
@@ -103,6 +104,30 @@ token = dapitest54321
         });
     });
 
+    it("should support profiles with username/password where user name is 'token'", async () => {
+        await withFile(async ({path}) => {
+            await writeFile(
+                path,
+                `[WITH_USERNAME]
+host = https://cloud.databricks.com/
+username = token
+password = dapitest54321
+
+[WITH_TOKEN]
+host = https://staging.cloud.databricks.com/
+token = dapitest54321
+`
+            );
+
+            const profiles = await loadConfigFile(path);
+            assert.equal(Object.keys(profiles).length, 2);
+            for (const profile of Object.values(profiles)) {
+                assert(!(profile instanceof TokenParsingError));
+                assert.equal((profile as Profile).token, "dapitest54321");
+            }
+        });
+    });
+
     it("should load all valid profiles and return errors for rest", async () => {
         await withFile(async ({path}) => {
             await writeFile(

packages/databricks-sdk-js/src/auth/configFile.ts

@@ -67,13 +67,21 @@ function getProfileOrError(
         return new HostParsingError(String(e));
     }
 
-    if (config.token === undefined) {
-        return new TokenParsingError('"token" it not defined');
+    if (config.token !== undefined) {
+        return {
+            host: host,
+            token: config.token,
+        };
     }
-    return {
-        host: host,
-        token: config.token,
-    };
+
+    if (config.username === "token" && config.password !== undefined) {
+        return {
+            host: host,
+            token: config.password,
+        };
+    }
+
+    return new TokenParsingError('"token" it not defined');
 }
 
 export async function loadConfigFile(filePath?: string): Promise<Profiles> {

packages/databricks-sdk-js/src/auth/fromAzureCli.integ.ts

@@ -0,0 +1,19 @@
+import {CurrentUserService, ApiClient} from "..";
+import {sleep} from "../test/IntegrationTestSetup";
+import {fromAzureCli} from "./fromAzureCli";
+
+// we can't run this test in CI because it requires Azure CLI to be installed
+// and logged in on the machine
+describe.skip(__filename, function () {
+    this.timeout(15_000);
+
+    it("should login with Azure CLI", async () => {
+        const client = new ApiClient("test", "0.1", fromAzureCli());
+
+        const scimApi = new CurrentUserService(client);
+        await scimApi.me();
+
+        await sleep(1200);
+        await scimApi.me();
+    });
+});

packages/databricks-sdk-js/src/auth/fromAzureCli.ts

@@ -0,0 +1,55 @@
+import * as child_process from "node:child_process";
+import {promisify} from "node:util";
+import {refreshableCredentialProvider} from "./refreshableCredentialProvider";
+import {Token} from "./Token";
+import {CredentialProvider, CredentialsProviderError} from "./types";
+
+const execFile = promisify(child_process.execFile);
+
+// Resource ID of the Azure application we need to log in.
+const azureDatabricksLoginAppID = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d";
+
+/**
+ * Credentials provider that uses Azure CLI to get a token.
+ *
+ * If host is not passed in then it will be read from DATABRICKS_HOST environment variable.
+ */
+export const fromAzureCli = (host?: URL): CredentialProvider => {
+    if (!host) {
+        const hostEnv = process.env["DATABRICKS_HOST"];
+        if (!hostEnv) {
+            throw new CredentialsProviderError(
+                "Can't find DATABRICKS_HOST environment variables"
+            );
+        }
+
+        host = new URL(hostEnv);
+    }
+
+    return refreshableCredentialProvider(async () => {
+        let stdout = "";
+        try {
+            ({stdout} = await execFile("az", [
+                "account",
+                "get-access-token",
+                "--resource",
+                azureDatabricksLoginAppID,
+            ]));
+        } catch (e: any) {
+            if (e.code === "ENOENT") {
+                throw new CredentialsProviderError(
+                    "Can't find 'az' command. Please install Azure CLI: https://docs.microsoft.com/en-us/cli/azure/install-azure-cli'"
+                );
+            } else {
+                throw e;
+            }
+        }
+
+        const azureToken = JSON.parse(stdout);
+        return new Token(
+            host!,
+            azureToken.accessToken,
+            new Date(azureToken.expiresOn).getTime()
+        );
+    });
+};

packages/databricks-sdk-js/src/auth/fromChain.ts

@@ -1,5 +1,5 @@
 import {CredentialProvider, CredentialsProviderError} from "./types";
-import {fromEnv} from "./fromEnv";
+import {fromToken} from "./fromToken";
 import {fromConfigFile} from "./fromConfigFile";
 
 export const fromChain = (
@@ -26,4 +26,4 @@ export const fromChain = (
     };
 };
 
-export const fromDefaultChain = fromChain([fromEnv(), fromConfigFile()]);
+export const fromDefaultChain = fromChain([fromToken(), fromConfigFile()]);

packages/databricks-sdk-js/src/auth/fromEnv.test.ts renamed to packages/databricks-sdk-js/src/auth/fromToken.test.ts

@@ -1,6 +1,6 @@
 import assert from "node:assert";
 import {CredentialsProviderError} from "./types";
-import {fromEnv} from "./fromEnv";
+import {fromToken} from "./fromToken";
 
 describe(__filename, () => {
     let origEnv: any;
@@ -17,15 +17,26 @@ describe(__filename, () => {
         process.env["DATABRICKS_HOST"] = "https://cloud.databricks.com/";
         process.env["DATABRICKS_TOKEN"] = "dapitest1234";
 
-        const provider = fromEnv();
+        const provider = fromToken();
+        const credentials = await provider();
+
+        assert.equal(credentials.host.href, "https://cloud.databricks.com/");
+        assert.equal(credentials.token, "dapitest1234");
+    });
+
+    it("should alow for passing in values", async () => {
+        const provider = fromToken(
+            new URL("https://cloud.databricks.com/"),
+            "dapitest1234"
+        );
         const credentials = await provider();
 
         assert.equal(credentials.host.href, "https://cloud.databricks.com/");
         assert.equal(credentials.token, "dapitest1234");
     });
 
     it("should throw if environment variables are not set", async () => {
-        const provider = fromEnv();
+        const provider = fromToken();
 
         assert.rejects(async () => {
             await provider();

packages/databricks-sdk-js/src/auth/fromEnv.ts renamed to packages/databricks-sdk-js/src/auth/fromToken.ts

@@ -10,15 +10,23 @@ function strip(strToStrip: string, str: string) {
     return str.split(strToStrip).join("");
 }
 
-export const fromEnv = (): CredentialProvider => {
+export const fromToken = (host?: URL, token?: string): CredentialProvider => {
     return async () => {
-        const host = process.env["DATABRICKS_HOST"];
-        const token = process.env["DATABRICKS_TOKEN"];
+        token = token || process.env["DATABRICKS_TOKEN"];
+        if (!host) {
+            const hostEnv = process.env["DATABRICKS_HOST"];
+            if (!hostEnv) {
+                throw new CredentialsProviderError(
+                    "Can't find DATABRICKS_HOST environment variable"
+                );
+            }
+            host = new URL(getValidHost(strip("'", hostEnv)));
+        }
 
         if (host && token) {
             return {
                 token: strip("'", token),
-                host: new URL(getValidHost(strip("'", host))),
+                host,
             };
         }