Unified auth for JS SDK (#309)

Author: fjakobs

.gitignore

@@ -14,3 +14,4 @@ out/
 **/coverage/
 **/.nyc_output/
 .databricks/**
+packages/databricks-sdk-js/src/config/testdata/azure/Library

packages/databricks-sdk-js/src/WorkspaceClient.ts

@@ -0,0 +1,83 @@
+import {Config, ConfigOptions} from "./config/Config";
+import {ApiClient, ClientOptions} from "./api-client";
+import {ClustersService} from "./apis/clusters";
+import {CommandExecutionService} from "./apis/commands";
+import {DbfsService} from "./apis/dbfs";
+import {LibrariesService} from "./apis/libraries";
+import {PermissionsService} from "./apis/permissions";
+import {ReposService} from "./apis/repos";
+import {
+    CurrentUserService,
+    UsersService,
+    GroupsService,
+    ServicePrincipalsService,
+} from "./apis/scim";
+import {WorkspaceService} from "./apis/workspace";
+import {WorkspaceConfService} from "./apis/workspaceconf";
+import {JobsService} from "./apis/jobs";
+
+export class WorkspaceClient {
+    readonly config: Config;
+    readonly apiClient: ApiClient;
+
+    constructor(config: ConfigOptions | Config, options: ClientOptions = {}) {
+        if (!(config instanceof Config)) {
+            config = new Config(config);
+        }
+
+        this.config = config as Config;
+        this.apiClient = new ApiClient(this.config, options);
+    }
+
+    get clusters() {
+        return new ClustersService(this.apiClient);
+    }
+
+    get dbfs() {
+        return new DbfsService(this.apiClient);
+    }
+
+    get commands() {
+        return new CommandExecutionService(this.apiClient);
+    }
+
+    get jobs() {
+        return new JobsService(this.apiClient);
+    }
+
+    get libraries() {
+        return new LibrariesService(this.apiClient);
+    }
+
+    get repos() {
+        return new ReposService(this.apiClient);
+    }
+
+    get currentUser() {
+        return new CurrentUserService(this.apiClient);
+    }
+
+    get users() {
+        return new UsersService(this.apiClient);
+    }
+
+    get groups() {
+        return new GroupsService(this.apiClient);
+    }
+
+    get servicePrincipals() {
+        return new ServicePrincipalsService(this.apiClient);
+    }
+
+    get workspace() {
+        return new WorkspaceService(this.apiClient);
+    }
+
+    get workspaceConf() {
+        return new WorkspaceConfService(this.apiClient);
+    }
+
+    get permissions() {
+        return new PermissionsService(this.apiClient);
+    }
+}

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

@@ -2,6 +2,7 @@
 import ".";
 import assert from "node:assert";
 import {ApiClient} from "./api-client";
+import {Config} from ".";
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const sdkVersion = require("../package.json").version;
@@ -12,12 +13,20 @@ describe(__filename, () => {
     });
 
     it("should create proper user agent", () => {
-        const ua = new ApiClient({extraUserAgent: {unit: "3.4.5"}}).userAgent();
+        const ua = new ApiClient(
+            new Config({
+                authType: "pat",
+            }),
+            {
+                product: "unit",
+                productVersion: "3.4.5",
+            }
+        ).userAgent();
         assert.equal(
             ua,
             `unit/3.4.5 databricks-sdk-js/${sdkVersion} nodejs/${process.version.slice(
                 1
-            )} os/${process.platform}`
+            )} os/${process.platform} auth/pat`
         );
     });
 });

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

@@ -1,15 +1,14 @@
 /* eslint-disable @typescript-eslint/naming-convention */
 import * as https from "node:https";
 import {TextDecoder} from "node:util";
-import {fromDefaultChain} from "./auth/fromChain";
 import {fetch} from "./fetch";
 import {ExposedLoggers, Utils, withLogContext} from "./logging";
 // eslint-disable-next-line @typescript-eslint/no-unused-vars
 import {context} from "./context";
 import {Context} from "./context";
+import {Headers, Config} from "./config/Config";
 import retry, {RetriableError} from "./retries/retries";
 import Time, {TimeUnits} from "./retries/Time";
-import {CredentialProvider} from "./auth/types";
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const sdkVersion = require("../package.json").version;
@@ -43,53 +42,52 @@ function logAndReturnError(
     return error;
 }
 
-export class ApiClient {
-    private agent: https.Agent;
-    private _host?: URL;
+export type ProductVersion = `${number}.${number}.${number}`;
 
-    private credentialProvider: CredentialProvider;
-    private readonly extraUserAgent: Record<string, string>;
+export interface ClientOptions {
+    agent?: https.Agent;
+    product?: string;
+    productVersion?: ProductVersion;
+    userAgentExtra?: Record<string, string>;
+}
 
-    get host(): Promise<URL> {
-        return (async () => {
-            if (!this._host) {
-                const credentials = await this.credentialProvider();
-                this._host = credentials.host;
-            }
-            return this._host;
-        })();
+export class ApiClient {
+    private agent: https.Agent;
+    readonly product: string;
+    readonly productVersion: ProductVersion;
+    readonly userAgentExtra: Record<string, string>;
+
+    constructor(readonly config: Config, options: ClientOptions = {}) {
+        this.agent =
+            options.agent ||
+            new https.Agent({
+                keepAlive: true,
+                keepAliveMsecs: 15_000,
+                rejectUnauthorized: config.insecureSkipVerify === false,
+                timeout: (config.httpTimeoutSeconds || 5) * 1000,
+            });
+
+        this.product = options.product || "unknown";
+        this.productVersion = options.productVersion || "0.0.0";
+        this.userAgentExtra = options.userAgentExtra || {};
     }
 
-    constructor({
-        credentialProvider = fromDefaultChain,
-        extraUserAgent = {},
-    }: {
-        credentialProvider?: CredentialProvider;
-        extraUserAgent?: Record<string, string>;
-    }) {
-        this.credentialProvider = credentialProvider;
-        this.extraUserAgent = extraUserAgent;
-
-        this.agent = new https.Agent({
-            keepAlive: true,
-            keepAliveMsecs: 15_000,
-        });
+    get host(): Promise<URL> {
+        return this.config.getHost();
     }
 
     userAgent(): string {
-        const pairs: Array<string> = [];
-        for (const [key, value] of Object.entries(this.extraUserAgent)) {
-            pairs.push(`${key}/${value}`);
-        }
-
-        pairs.push(
+        const pairs = [
+            `${this.product}/${this.productVersion}`,
             `databricks-sdk-js/${sdkVersion}`,
             `nodejs/${process.version.slice(1)}`,
-            `os/${process.platform}`
-        );
+            `os/${process.platform}`,
+            `auth/${this.config.authType}`,
+        ];
 
-        // TODO: add ability of per-request extra-information,
-        // so that we can track sub-functionality, like in Terraform
+        for (const [key, value] of Object.entries(this.userAgentExtra)) {
+            pairs.push(`${key}/${value}`);
+        }
         return pairs.join(" ");
     }
 
@@ -100,15 +98,15 @@ export class ApiClient {
         payload?: any,
         @context context?: Context
     ): Promise<Record<string, unknown>> {
-        const credentials = await this.credentialProvider();
-        const headers = {
-            "Authorization": `Bearer ${credentials.token}`,
+        const headers: Headers = {
             "User-Agent": this.userAgent(),
             "Content-Type": "text/json",
         };
 
+        await this.config.authenticate(headers);
+
         // create a copy of the URL, so that we can modify it
-        const url = new URL(credentials.host.toString());
+        const url = new URL(this.config.host!);
         url.pathname = path;
 
         const options: any = {
@@ -128,7 +126,10 @@ export class ApiClient {
         const response = await retry<
             Awaited<Awaited<ReturnType<typeof fetch>>["response"]>
         >({
-            timeout: new Time(10, TimeUnits.seconds),
+            timeout: new Time(
+                this.config.retryTimeoutSeconds || 300,
+                TimeUnits.seconds
+            ),
             fn: async () => {
                 let response;
                 try {

packages/databricks-sdk-js/src/apis/commands/commands.integ.ts

@@ -1,6 +1,5 @@
 /* eslint-disable @typescript-eslint/naming-convention */
 
-import {CommandExecutionService} from "../..";
 import assert from "assert";
 
 import {IntegrationTestSetup, sleep} from "../../test/IntegrationTestSetup";
@@ -15,7 +14,7 @@ describe(__filename, function () {
     });
 
     it("should execute python with low level API", async () => {
-        const commandsApi = new CommandExecutionService(integSetup.client);
+        const commandsApi = integSetup.client.commands;
 
         const context = await commandsApi.createAndWait({
             clusterId: integSetup.cluster.id,

packages/databricks-sdk-js/src/apis/jobs/jobs.integ.ts

@@ -1,6 +1,5 @@
 /* eslint-disable @typescript-eslint/naming-convention */
 
-import {JobsService, DbfsService} from "../..";
 import assert from "assert";
 
 import {IntegrationTestSetup, sleep} from "../../test/IntegrationTestSetup";
@@ -15,9 +14,9 @@ describe(__filename, function () {
     });
 
     it("should run a notebook job", async () => {
-        const jobsService = new JobsService(integSetup.client);
+        const jobsService = integSetup.client.jobs;
 
-        const dbfsApi = new DbfsService(integSetup.client);
+        const dbfsApi = integSetup.client.dbfs;
         const jobPath = `/tmp/sdk-js-integ-${integSetup.testRunId}.py`;
 
         await dbfsApi.put({