Automatically enable serverless when serverless_compute_id=auto (#1550)

## Changes

Automatically enable serverless when serverless_compute_id=auto

## Tests

Manually, e2e tests in the followup

Author: ilia-db

package.json

@@ -14,7 +14,7 @@
     },
     "scripts": {
         "install:sdk": "yarn fetch:sdk && chmod a+x scripts/install_sdk.sh && scripts/install_sdk.sh",
-        "fetch:sdk": "rm -rf vendor && mkdir vendor && gh release download -R databricks/databricks-sdk-js v0.7.0 -p '*.tgz' && mv *.tgz vendor/databricks-sdk.tgz",
+        "fetch:sdk": "rm -rf vendor && mkdir vendor && gh release download -R databricks/databricks-sdk-js v0.8.0 -p '*.tgz' && mv *.tgz vendor/databricks-sdk.tgz",
         "test": "yarn workspaces foreach run test",
         "test:integ": "yarn workspaces foreach run test:integ",
         "build": "yarn workspaces foreach run build",

packages/databricks-vscode/src/configuration/ConnectionCommands.ts

@@ -135,16 +135,9 @@ export class ConnectionCommands implements Disposable {
             quickPick.title =
                 typeof title === "string" ? title : "Select Cluster";
             quickPick.keepScrollPosition = true;
-            quickPick.busy = true;
-
             quickPick.busy = true;
             quickPick.canSelectMany = false;
             const items: QuickPickItem[] = [
-                {
-                    label: "$(cloud) Serverless",
-                    detail: `Run files as Workflows or use Databricks Connect without a dedicated cluster`,
-                    alwaysShow: false,
-                },
                 {
                     label: "$(repo-create) Create New Cluster",
                     detail: `Open Databricks in the browser and create a new cluster`,
@@ -155,6 +148,15 @@ export class ConnectionCommands implements Disposable {
                     kind: QuickPickItemKind.Separator,
                 },
             ];
+            if (
+                this.connectionManager.databricksWorkspace?.isServerlessEnabled
+            ) {
+                items.unshift({
+                    label: "$(cloud) Serverless",
+                    detail: `Run files as Workflows or use Databricks Connect without a dedicated cluster`,
+                    alwaysShow: false,
+                });
+            }
             quickPick.items = items;
 
             this.clusterModel.refresh();

packages/databricks-vscode/src/configuration/ConnectionManager.ts

@@ -150,8 +150,17 @@ export class ConnectionManager implements Disposable {
     }
 
     private async updateServerless() {
+        if (this.configModel.target === undefined) {
+            return;
+        }
         const serverless = await this.configModel.get("serverless");
-        if (serverless) {
+        const computeId = this.apiClient?.config.serverlessComputeId;
+        const autoEnable =
+            serverless === undefined &&
+            !this.cluster &&
+            computeId === "auto" &&
+            this.databricksWorkspace?.isServerlessEnabled;
+        if (serverless || autoEnable) {
             await this.enableServerless();
         } else {
             await this.disableServerless();

packages/databricks-vscode/src/configuration/DatabricksWorkspace.test.ts

@@ -16,22 +16,26 @@ describe(__filename, () => {
         const wsConf = {
             enableProjectTypeInWorkspace: "true",
             enableWorkspaceFilesystem: "dbr11.0+",
+            enableServerless: true,
         } as const;
         const authProvider = new ProfileAuthProvider(
             new URL("https://fabian.databricks.com"),
             "DEFAULT",
             instance(mock(CliWrapper))
         );
         const dbWorkspace: DatabricksWorkspace = new DatabricksWorkspace(
+            "workspace-id",
             authProvider,
             user,
             wsConf
         );
 
         assert(dbWorkspace.host.toString() === host.toString());
         assert.equal(dbWorkspace.userName, user.userName);
+        assert.equal(dbWorkspace.id, "workspace-id");
         assert(dbWorkspace.isFilesInReposEnabled);
         assert(dbWorkspace.isFilesInReposEnabled);
+        assert(dbWorkspace.isServerlessEnabled);
 
         assert(
             dbWorkspace.supportFilesInReposForCluster({

packages/databricks-vscode/src/configuration/DatabricksWorkspace.ts

@@ -6,11 +6,23 @@ import {Loggers} from "../logger";
 import {AuthProvider} from "./auth/AuthProvider";
 import {RemoteUri} from "../sync/SyncDestination";
 
+type ServerlessEnablementResponse = {
+    setting?: {
+        value?: {
+            // eslint-disable-next-line @typescript-eslint/naming-convention
+            preview_enablement_val?: {
+                enabled: boolean;
+            };
+        };
+    };
+};
+
 export class DatabricksWorkspace {
     constructor(
+        public readonly id: string,
         private _authProvider: AuthProvider,
         private me: iam.User,
-        private wsConf: WorkspaceConfProps
+        private wsConf: WorkspaceConfProps & {enableServerless?: boolean}
     ) {}
 
     /**
@@ -56,6 +68,10 @@ export class DatabricksWorkspace {
         return this.wsConf.enableWorkspaceFilesystem !== "false";
     }
 
+    get isServerlessEnabled(): boolean {
+        return this.wsConf.enableServerless === true;
+    }
+
     supportFilesInReposForCluster(cluster: Cluster): boolean {
         if (!this.isReposEnabled || !this.isFilesInReposEnabled) {
             return false;
@@ -90,27 +106,36 @@ export class DatabricksWorkspace {
         @context ctx?: Context
     ) {
         const me = await client.currentUser.me(ctx);
+        // Workspace ID is returned in a header, which is then copied to the response object.
+        const id = (me as any)["x-databricks-org-id"] as string;
 
         const wsConfApi = new WorkspaceConf(client.apiClient);
-        let state: WorkspaceConfProps = {
+        let state: WorkspaceConfProps & {enableServerless?: boolean} = {
             enableProjectTypeInWorkspace: "true",
             enableWorkspaceFilesystem: "true",
         };
         try {
-            state = {
-                ...state,
-                ...(await wsConfApi.getStatus(
-                    [
-                        "enableProjectTypeInWorkspace",
-                        "enableWorkspaceFilesystem",
-                    ],
-                    ctx
-                )),
-            };
+            const confStatus = await wsConfApi.getStatus(
+                ["enableProjectTypeInWorkspace", "enableWorkspaceFilesystem"],
+                ctx
+            );
+            state = {...state, ...confStatus};
         } catch (e) {
             ctx?.logger?.error("Can't fetch workspace confs", e);
         }
+        try {
+            const serverlessEnablement = (await client.apiClient.request(
+                `/api/2.0/settings-api/workspace/${id}/serverless_job_nb`,
+                "GET"
+            )) as ServerlessEnablementResponse;
+            const enableServerless =
+                serverlessEnablement?.setting?.value?.preview_enablement_val
+                    ?.enabled === true;
+            state = {...state, enableServerless};
+        } catch (e) {
+            ctx?.logger?.error("Can't detect serverless support", e);
+        }
 
-        return new DatabricksWorkspace(authProvider, me, state);
+        return new DatabricksWorkspace(id, authProvider, me, state);
     }
 }

packages/databricks-vscode/src/telemetry/constants.ts

@@ -239,6 +239,7 @@ export class MetadataTypes {
     [Metadata.USER]: EventType<{
         hashedUserName: string;
         host: string;
+        workspaceId: string;
         authType: AuthType;
     }> = {
         hashedUserName: {
@@ -248,6 +249,9 @@ export class MetadataTypes {
             comment:
                 "The hostname of the workspace that the user is connected to",
         },
+        workspaceId: {
+            comment: "The id of the workspace",
+        },
         authType: {
             comment: "The kind of authentication used by the user",
         },

packages/databricks-vscode/src/telemetry/index.test.ts

@@ -77,6 +77,7 @@ describe(__filename, () => {
 
     it("sets user metadata correctly after logged in", async () => {
         const ws = mock(DatabricksWorkspace);
+        when(ws.id).thenReturn("workspace-id");
         when(ws.userName).thenReturn("[email protected]");
         when(ws.host).thenReturn(new URL("https://my.databricks.com"));
         const cm = mock(ConnectionManager);
@@ -110,6 +111,7 @@ describe(__filename, () => {
                 "$2b$07$97f780c0d2a202b5f54aeOJ9MKnLacIH6dXoDUpQLfJXKjhYqxJLW",
             "user.host": "my.databricks.com",
             "user.authType": "azure-cli",
+            "user.workspaceId": "workspace-id",
         });
         assert.deepEqual(metrics, {
             "event.duration": 100,

packages/databricks-vscode/src/telemetry/index.ts

@@ -57,6 +57,7 @@ async function getDatabricksWorkspaceMetadata(
     return {
         hashedUserName: hashedUserName,
         host: databricksWorkspace.host.host,
+        workspaceId: databricksWorkspace.id,
     };
 }
 

vendor/databricks-sdk.tgz

@@ -284,7 +284,7 @@ __metadata:
   linkType: hard
 
 "@databricks/databricks-sdk@file:../../vendor/databricks-sdk.tgz::locator=%40databricks%2Fdatabricks-vscode-types%40workspace%3Apackages%2Fdatabricks-vscode-types":
-  version: 0.7.0
+  version: 0.8.0
   resolution: "@databricks/databricks-sdk@file:../../vendor/databricks-sdk.tgz::locator=%40databricks%2Fdatabricks-vscode-types%40workspace%3Apackages%2Fdatabricks-vscode-types"
   dependencies:
     "@types/ini": ^4.1.0
@@ -294,12 +294,12 @@ __metadata:
     ini: ^4.1.2
     reflect-metadata: ^0.2.1
     semver: ^7.6.0
-  checksum: fc3186dc4dda2a06e01c8e7e20b5f76ec6d1b4379b65deb80cacfd506427f38841af890858bd3335c3ab76f8a7255c4f4ce295f6470d7d0905f236b211fa617b
+  checksum: be59ab37615aa63780fdb841cdd6199f930710d409e29233abc2c66d8e5a7ac9928497bdf2f38075970e89e6ade0d9ae37ea39efcf721bd85f0c484c535f46d0
   languageName: node
   linkType: hard
 
 "@databricks/databricks-sdk@file:../../vendor/databricks-sdk.tgz::locator=databricks%40workspace%3Apackages%2Fdatabricks-vscode":
-  version: 0.7.0
+  version: 0.8.0
   resolution: "@databricks/databricks-sdk@file:../../vendor/databricks-sdk.tgz::locator=databricks%40workspace%3Apackages%2Fdatabricks-vscode"
   dependencies:
     "@types/ini": ^4.1.0
@@ -309,7 +309,7 @@ __metadata:
     ini: ^4.1.2
     reflect-metadata: ^0.2.1
     semver: ^7.6.0
-  checksum: fc3186dc4dda2a06e01c8e7e20b5f76ec6d1b4379b65deb80cacfd506427f38841af890858bd3335c3ab76f8a7255c4f4ce295f6470d7d0905f236b211fa617b
+  checksum: be59ab37615aa63780fdb841cdd6199f930710d409e29233abc2c66d8e5a7ac9928497bdf2f38075970e89e6ade0d9ae37ea39efcf721bd85f0c484c535f46d0
   languageName: node
   linkType: hard
 
@@ -1448,11 +1448,11 @@ __metadata:
   linkType: hard
 
 "@types/node@npm:^18.19.23":
-  version: 18.19.64
-  resolution: "@types/node@npm:18.19.64"
+  version: 18.19.75
+  resolution: "@types/node@npm:18.19.75"
   dependencies:
     undici-types: ~5.26.4
-  checksum: e7680215b03c9bee8a33947f03d06048e8e460f23b1b7b29c45350cf437faa5f8fcb7d8c3eb8dfec8427923e7a184df42bc710c1b6252b4852e3ed7064c6228f
+  checksum: 8f780a51a2dc1d597b104a2c33a2ab175fda492defbaab2cbb5fd098124d867ebed7cda0eaf80309bba25fbd4e41b4a04c91b7f14b50547b4377a39ccf749a7c
   languageName: node
   linkType: hard
 
@@ -2220,6 +2220,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"agent-base@npm:^7.1.2":
+  version: 7.1.3
+  resolution: "agent-base@npm:7.1.3"
+  checksum: 87bb7ee54f5ecf0ccbfcba0b07473885c43ecd76cb29a8db17d6137a19d9f9cd443a2a7c5fd8a3f24d58ad8145f9eb49116344a66b107e1aeab82cf2383f4753
+  languageName: node
+  linkType: hard
+
 "agentkeepalive@npm:^4.2.1":
   version: 4.3.0
   resolution: "agentkeepalive@npm:4.3.0"
@@ -5398,12 +5405,13 @@ __metadata:
   linkType: hard
 
 "gcp-metadata@npm:^6.1.0":
-  version: 6.1.0
-  resolution: "gcp-metadata@npm:6.1.0"
+  version: 6.1.1
+  resolution: "gcp-metadata@npm:6.1.1"
   dependencies:
-    gaxios: ^6.0.0
+    gaxios: ^6.1.1
+    google-logging-utils: ^0.0.2
     json-bigint: ^1.0.0
-  checksum: 55de8ae4a6b7664379a093abf7e758ae06e82f244d41bd58d881a470bf34db94c4067ce9e1b425d9455b7705636d5f8baad844e49bb73879c338753ba7785b2b
+  checksum: 7dffe884fd718482b559a841da469dc30e766a4b86c71cda96bed42579763d58f28328238b2eb424c29ba10ef45d4bb8a6586441921734f01012b55bbea79711
   languageName: node
   linkType: hard
 
@@ -5731,16 +5739,23 @@ __metadata:
   linkType: hard
 
 "google-auth-library@npm:^9.10.0":
-  version: 9.14.2
-  resolution: "google-auth-library@npm:9.14.2"
+  version: 9.15.1
+  resolution: "google-auth-library@npm:9.15.1"
   dependencies:
     base64-js: ^1.3.0
     ecdsa-sig-formatter: ^1.0.11
     gaxios: ^6.1.1
     gcp-metadata: ^6.1.0
     gtoken: ^7.0.0
     jws: ^4.0.0
-  checksum: 64b3a6c1b1b14f1c891dbcfb850bc4db63dc8fae17e70197636244d00c83b539ac3da8688aae0bd1f09c884fc538d203945ae751edbabf666b41066385d86e30
+  checksum: 1c7660b7d21504a58b6b7780b0ca56f07a4d2b68d2ca14e5c4beaedd45cc4898baa8df1cfaf4dac15413a8db3cecc00e84662d8a327f9b9488ff2df7d8d23c84
+  languageName: node
+  linkType: hard
+
+"google-logging-utils@npm:^0.0.2":
+  version: 0.0.2
+  resolution: "google-logging-utils@npm:0.0.2"
+  checksum: 270de74cde8abe0a6639b4bebbb6db2c7d8a495cf14b23779450d3d010e759f3ffcb7eb200b1bcb05da1897fa63f3e336df2274f02fc69d3bb388019fa2b3134
   languageName: node
   linkType: hard
 
@@ -6086,12 +6101,12 @@ __metadata:
   linkType: hard
 
 "https-proxy-agent@npm:^7.0.1":
-  version: 7.0.5
-  resolution: "https-proxy-agent@npm:7.0.5"
+  version: 7.0.6
+  resolution: "https-proxy-agent@npm:7.0.6"
   dependencies:
-    agent-base: ^7.0.2
+    agent-base: ^7.1.2
     debug: 4
-  checksum: 2e1a28960f13b041a50702ee74f240add8e75146a5c37fc98f1960f0496710f6918b3a9fe1e5aba41e50f58e6df48d107edd9c405c5f0d73ac260dabf2210857
+  checksum: b882377a120aa0544846172e5db021fa8afbf83fea2a897d397bd2ddd8095ab268c24bc462f40a15f2a8c600bf4aa05ce52927f70038d4014e68aefecfa94e8d
   languageName: node
   linkType: hard
 
@@ -9435,11 +9450,11 @@ __metadata:
   linkType: hard
 
 "semver@npm:^7.6.0":
-  version: 7.6.3
-  resolution: "semver@npm:7.6.3"
+  version: 7.7.1
+  resolution: "semver@npm:7.7.1"
   bin:
     semver: bin/semver.js
-  checksum: 4110ec5d015c9438f322257b1c51fe30276e5f766a3f64c09edd1d7ea7118ecbc3f379f3b69032bacf13116dc7abc4ad8ce0d7e2bd642e26b0d271b56b61a7d8
+  checksum: 586b825d36874007c9382d9e1ad8f93888d8670040add24a28e06a910aeebd673a2eb9e3bf169c6679d9245e66efb9057e0852e70d9daa6c27372aab1dda7104
   languageName: node
   linkType: hard