Enable Unity Catalog (and Docs) in Databricks Remote SSH mode

packages/databricks-vscode/package.json

@@ -564,12 +564,12 @@
                 {
                     "id": "unityCatalogView",
                     "name": "Unity Catalog",
-                    "when": "databricks.context.activated && databricks.context.loggedIn"
+                    "when": "databricks.context.activated && (databricks.context.loggedIn || databricks.context.remoteMode)"
                 },
                 {
                     "id": "databricksDocsView",
                     "name": "Documentation",
-                    "when": "databricks.context.activated && !databricks.context.remoteMode"
+                    "when": "databricks.context.activated"
                 }
             ]
         },

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

@@ -1,24 +1,110 @@
 /* eslint-disable @typescript-eslint/naming-convention */
 
+import assert from "assert";
 import {Disposable} from "vscode";
+import {anything, instance, mock, reset, verify, when} from "ts-mockito";
+import {WorkspaceClient} from "@databricks/sdk-experimental";
+import {ConnectionManager} from "./ConnectionManager";
+import {ConfigModel} from "./models/ConfigModel";
+import {CliWrapper} from "../cli/CliWrapper";
+import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager";
+import {CustomWhenContext} from "../vscode-objs/CustomWhenContext";
+import {Telemetry} from "../telemetry";
+import {AuthProvider} from "./auth/AuthProvider";
 
 describe(__filename, () => {
     let disposables: Array<Disposable>;
 
+    let mockCli: CliWrapper;
+    let mockConfigModel: ConfigModel;
+    let mockWorkspaceFolderManager: WorkspaceFolderManager;
+    let mockCustomWhenContext: CustomWhenContext;
+    let mockAuthProvider: AuthProvider;
+    let mockWorkspaceClient: WorkspaceClient;
+
+    function buildConnectionManager(): ConnectionManager {
+        return new ConnectionManager(
+            instance(mockCli),
+            instance(mockConfigModel),
+            instance(mockWorkspaceFolderManager),
+            instance(mockCustomWhenContext),
+            new Telemetry()
+        );
+    }
+
     beforeEach(() => {
         disposables = [];
+        mockCli = mock(CliWrapper);
+        mockConfigModel = mock(ConfigModel);
+        mockWorkspaceFolderManager = mock(WorkspaceFolderManager);
+        mockCustomWhenContext = mock(CustomWhenContext);
+        mockAuthProvider = mock<AuthProvider>();
+        mockWorkspaceClient = mock(WorkspaceClient);
+
+        // DatabricksWorkspace.load() reads the org id from a header on the
+        // currentUser.me() response and (best-effort) the workspace conf.
+        when(mockWorkspaceClient.currentUser).thenReturn({
+            me: async () =>
+                ({
+                    "userName": "test@databricks.com",
+                    "x-databricks-org-id": "1234",
+                }) as any,
+        } as any);
+        when(mockWorkspaceClient.apiClient).thenReturn(undefined as any);
+        when(mockAuthProvider.getWorkspaceClient()).thenResolve(
+            instance(mockWorkspaceClient)
+        );
+        when(mockAuthProvider.host).thenReturn(
+            new URL("https://test.databricks.com")
+        );
     });
 
     afterEach(() => {
         disposables.forEach((d) => d.dispose());
+        reset(mockConfigModel);
+    });
+
+    it("connectFromEnvironment connects using the injected auth provider", async () => {
+        const cm = buildConnectionManager();
+        disposables.push(cm);
+
+        await cm.connectFromEnvironment(instance(mockAuthProvider));
+
+        assert.equal(cm.state, "CONNECTED");
+        assert.ok(cm.workspaceClient);
+        assert.ok(cm.databricksWorkspace);
+        assert.equal(
+            cm.databricksWorkspace?.host.toString(),
+            "https://test.databricks.com/"
+        );
+        verify(mockCustomWhenContext.setLoggedIn(true)).atLeast(1);
     });
 
-    // TODO
-    // login
-    // logout
-    // configure
-    // attach cluster
-    // detach cluster
-    // attach workspace
-    // detach workspace
+    it("connectFromEnvironment does not touch the config model (no bundle coupling)", async () => {
+        const cm = buildConnectionManager();
+        disposables.push(cm);
+
+        await cm.connectFromEnvironment(instance(mockAuthProvider));
+
+        verify(mockConfigModel.set(anything(), anything())).never();
+        verify(mockConfigModel.setAuthProvider(anything())).never();
+    });
+
+    it("connectFromEnvironment disconnects and rethrows on failure", async () => {
+        when(mockAuthProvider.getWorkspaceClient()).thenReject(
+            new Error("no credentials")
+        );
+        const cm = buildConnectionManager();
+        disposables.push(cm);
+
+        await assert.rejects(
+            () => cm.connectFromEnvironment(instance(mockAuthProvider)),
+            /no credentials/
+        );
+
+        assert.equal(cm.state, "DISCONNECTED");
+        assert.equal(cm.workspaceClient, undefined);
+        assert.equal(cm.databricksWorkspace, undefined);
+        verify(mockCustomWhenContext.setLoggedIn(false)).atLeast(1);
+    });
 });

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

@@ -1,4 +1,5 @@
 import {
+    Config,
     WorkspaceClient,
     ApiClient,
     logging,
@@ -18,7 +19,12 @@ import {DatabricksWorkspace} from "./DatabricksWorkspace";
 import {CustomWhenContext} from "../vscode-objs/CustomWhenContext";
 import {ConfigModel} from "./models/ConfigModel";
 import {onError, withOnErrorHandler} from "../utils/onErrorDecorator";
-import {AuthProvider, ProfileAuthProvider} from "./auth/AuthProvider";
+import {
+    AuthProvider,
+    PersonalAccessTokenAuthProvider,
+    ProfileAuthProvider,
+} from "./auth/AuthProvider";
+import {normalizeHost} from "../utils/urlUtils";
 import {Mutex} from "../locking";
 import {MetadataService} from "./auth/MetadataService";
 import {Events, Telemetry} from "../telemetry";
@@ -40,6 +46,7 @@ export type ConnectionState = "CONNECTED" | "CONNECTING" | "DISCONNECTED";
 export class ConnectionManager implements Disposable {
     private disposables: Disposable[] = [];
     private _state: ConnectionState = "DISCONNECTED";
+    private _connectionError?: string;
     private loginLogoutMutex: Mutex = new Mutex();
     private savedAuthMutex: Mutex = new Mutex();
     private configureLoginMutex: Mutex = new Mutex();
@@ -220,6 +227,16 @@ export class ConnectionManager implements Disposable {
         return this._state;
     }
 
+    /**
+     * The error message from the most recent failed connection attempt, if any.
+     * Cleared on a successful connection. Used to surface why an
+     * environment-based connection (remote mode) failed instead of showing an
+     * empty view.
+     */
+    get connectionError(): string | undefined {
+        return this._connectionError;
+    }
+
     get cluster(): Cluster | undefined {
         return this._clusterManager?.cluster;
     }
@@ -261,6 +278,67 @@ export class ConnectionManager implements Disposable {
         }
     }
 
+    /**
+     * Connect using the host and token that the SDK resolves from the ambient
+     * environment (e.g. the DATABRICKS_HOST / DATABRICKS_TOKEN variables that
+     * the Databricks Remote SSH session injects). Only PAT credentials are
+     * supported here - the remote environment always provides a token, so we
+     * fail fast if one isn't present rather than attempting other auth types.
+     *
+     * Unlike the normal login flow this does not depend on a bundle/config
+     * project (host + target) and skips all sync/cluster/config machinery. It's
+     * used in Databricks Remote SSH sessions where only Unity Catalog is
+     * surfaced and credentials come from the environment.
+     */
+    async connectFromEnvironment(authProvider?: AuthProvider): Promise<void> {
+        await this.loginLogoutMutex.synchronise(async () => {
+            // We intentionally inline the connect/disconnect steps here rather
+            // than delegating to _connect()/disconnect(): both of those acquire
+            // loginLogoutMutex, which is non-reentrant, so calling them while we
+            // already hold it would deadlock. We also deliberately skip the
+            // sync/cluster/config-project machinery they run, since remote mode
+            // only needs a workspace client for the Unity Catalog view.
+            this._connectionError = undefined;
+            this.updateState("CONNECTING");
+            try {
+                // The authProvider is only injected by tests; in production it
+                // is resolved from the SDK's default credential chain.
+                if (authProvider === undefined) {
+                    const config = new Config({});
+                    await config.ensureResolved();
+                    if (config.host === undefined) {
+                        throw new Error(
+                            "No Databricks host found in the environment"
+                        );
+                    }
+                    if (config.token === undefined) {
+                        throw new Error(
+                            "No Databricks token found in the environment"
+                        );
+                    }
+                    authProvider = new PersonalAccessTokenAuthProvider(
+                        normalizeHost(config.host),
+                        config.token,
+                        this.cli
+                    );
+                }
+                this._workspaceClient = await authProvider.getWorkspaceClient();
+                this._databricksWorkspace = await DatabricksWorkspace.load(
+                    this._workspaceClient,
+                    authProvider
+                );
+                this.updateState("CONNECTED");
+            } catch (e) {
+                this._workspaceClient = undefined;
+                this._databricksWorkspace = undefined;
+                this._connectionError =
+                    e instanceof Error ? e.message : String(e);
+                this.updateState("DISCONNECTED");
+                throw e;
+            }
+        });
+    }
+
     private async loginWithSavedAuth(source: AutoLoginSource) {
         if (this.savedAuthMutex.locked) {
             return;