From bd0712fa3f91efafa3ecabbdd92d4d03e76d7749 Mon Sep 17 00:00:00 2001 From: misha-db Date: Tue, 14 Jul 2026 13:21:32 +0400 Subject: [PATCH 1/5] Unity catalog in remote IDE --- packages/databricks-vscode/package.json | 2 +- .../configuration/ConnectionManager.test.ts | 102 +++++- .../src/configuration/ConnectionManager.ts | 66 +++- .../src/configuration/auth/AuthProvider.ts | 59 +++- packages/databricks-vscode/src/extension.ts | 320 +++++++++++------- .../UnityCatalogTreeDataProvider.test.ts | 42 ++- .../UnityCatalogTreeDataProvider.ts | 19 +- 7 files changed, 476 insertions(+), 134 deletions(-) diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 682c9a29c..d6d326560 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -564,7 +564,7 @@ { "id": "unityCatalogView", "name": "Unity Catalog", - "when": "databricks.context.activated && databricks.context.loggedIn" + "when": "databricks.context.activated && (databricks.context.loggedIn || databricks.context.remoteMode)" }, { "id": "databricksDocsView", diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.test.ts b/packages/databricks-vscode/src/configuration/ConnectionManager.test.ts index 57801b5b5..1dc438ea0 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionManager.test.ts +++ b/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; + 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(); + 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); + }); }); diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.ts b/packages/databricks-vscode/src/configuration/ConnectionManager.ts index 6da1e99cc..78ad60b9f 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionManager.ts +++ b/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, + EnvironmentAuthProvider, + 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,53 @@ export class ConnectionManager implements Disposable { } } + /** + * Connect using the ambient environment credentials resolved by the SDK's + * default credential chain (environment variables, metadata service, etc.). + * + * 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 { + await this.loginLogoutMutex.synchronise(async () => { + 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" + ); + } + authProvider = new EnvironmentAuthProvider( + normalizeHost(config.host), + config, + 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; diff --git a/packages/databricks-vscode/src/configuration/auth/AuthProvider.ts b/packages/databricks-vscode/src/configuration/auth/AuthProvider.ts index 268d4b2af..649da2d51 100644 --- a/packages/databricks-vscode/src/configuration/auth/AuthProvider.ts +++ b/packages/databricks-vscode/src/configuration/auth/AuthProvider.ts @@ -24,7 +24,8 @@ export type AuthType = | "databricks-cli" | "google-id" | "profile" - | "pat"; + | "pat" + | "env"; export abstract class AuthProvider { constructor( @@ -532,3 +533,59 @@ export class PersonalAccessTokenAuthProvider extends AuthProvider { }); } } + +/** + * Auth provider backed by an already-resolved SDK Config, produced by the SDK's + * default credential chain (environment variables, metadata service, etc.). + * + * Used in Databricks Remote SSH sessions where the ambient environment supplies + * credentials and there is no bundle/config project or profile to resolve. + */ +export class EnvironmentAuthProvider extends AuthProvider { + constructor( + host: URL, + private readonly config: Config, + cli: CliWrapper + ) { + super(host, "env", cli, true); + } + + describe(): string { + return "Environment"; + } + + toJSON(): Record { + return { + host: this.host.toString(), + authType: this.authType, + }; + } + + toEnv(): Record { + return { + DATABRICKS_HOST: this.host.toString(), + }; + } + + toIni(): Record | undefined { + return undefined; + } + + protected async _check(): Promise { + try { + const workspaceClient = await this.getWorkspaceClient(); + await workspaceClient.currentUser.me(); + return true; + } catch (e) { + logging.NamedLogger.getOrCreate(Loggers.Extension).error( + "Can't login using the ambient environment credentials", + e + ); + return false; + } + } + + protected _getSdkConfig(): Config { + return this.config; + } +} diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index fd7e799ec..79280edb6 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -92,6 +92,146 @@ const packageJson = require("../package.json"); const customWhenContext = new CustomWhenContext(); +/** + * Register the Unity Catalog tree view, its commands and the detail panel. + * Shared between the normal activation flow and the remote (Databricks Remote + * SSH) flow, where Unity Catalog is the only view we surface. + */ +function registerUnityCatalog( + context: ExtensionContext, + connectionManager: ConnectionManager, + stateStorage: StateStorage, + telemetry: Telemetry, + // In remote mode there is no automatic reconnection, so wire the refresh + // command to re-establish the connection before refreshing the tree. + reconnect?: () => Promise +): void { + const unityCatalogTreeDataProvider = new UnityCatalogTreeDataProvider( + connectionManager, + stateStorage, + context.extensionPath + ); + + const unityCatalogTreeView = window.createTreeView("unityCatalogView", { + treeDataProvider: unityCatalogTreeDataProvider, + }); + context.subscriptions.push( + unityCatalogTreeDataProvider, + unityCatalogTreeView, + telemetry.registerCommand( + "databricks.unityCatalog.refresh", + async () => { + if (reconnect) { + await reconnect(); + } + unityCatalogTreeDataProvider.refresh(); + } + ), + telemetry.registerCommand( + "databricks.unityCatalog.refreshNode", + (node: UnityCatalogTreeNode) => + unityCatalogTreeDataProvider.refreshNode(node) + ), + telemetry.registerCommand( + "databricks.unityCatalog.copyStorageLocation", + async (node: UnityCatalogTreeNode) => { + if ( + (node.kind === "table" || node.kind === "volume") && + node.storageLocation + ) { + await env.clipboard.writeText(node.storageLocation); + window.showInformationMessage("Copied to clipboard"); + } + } + ), + telemetry.registerCommand( + "databricks.unityCatalog.copyViewSql", + async (node: UnityCatalogTreeNode) => { + if (node.kind === "table" && node.viewDefinition) { + await env.clipboard.writeText(node.viewDefinition); + window.showInformationMessage("Copied to clipboard"); + } + } + ), + telemetry.registerCommand( + "databricks.unityCatalog.copyName", + async (node: UnityCatalogTreeNode) => { + if ( + node.kind === "error" || + node.kind === "empty" || + node.kind === "favorites" || + node.kind === "group" + ) { + return; + } + const text = node.kind === "column" ? node.name : node.fullName; + await env.clipboard.writeText(text); + window.showInformationMessage("Copied to clipboard"); + } + ), + telemetry.registerCommand( + "databricks.unityCatalog.openExternal", + async (node: UnityCatalogTreeNode) => { + if (node.kind === "error" || node.kind === "column") { + return; + } + const url = + unityCatalogTreeDataProvider.getNodeExploreUrl(node); + if (!url) { + window.showErrorMessage( + "Databricks: Can't open external link. No URL found." + ); + return; + } + await UrlUtils.openExternal(url); + } + ), + commands.registerCommand("databricks.unityCatalog.filter", async () => { + await commands.executeCommand("unityCatalogView.focus"); + await commands.executeCommand("list.find"); + }), + telemetry.registerCommand( + "databricks.unityCatalog.pin", + (node: UnityCatalogTreeNode) => { + if ( + node.kind === "catalog" || + node.kind === "schema" || + node.kind === "table" || + node.kind === "volume" || + node.kind === "function" || + node.kind === "registeredModel" || + node.kind === "modelVersion" + ) { + return unityCatalogTreeDataProvider.pin(node); + } + } + ), + telemetry.registerCommand( + "databricks.unityCatalog.unpin", + (node: UnityCatalogTreeNode) => { + if ( + node.kind === "catalog" || + node.kind === "schema" || + node.kind === "table" || + node.kind === "volume" || + node.kind === "function" || + node.kind === "registeredModel" || + node.kind === "modelVersion" + ) { + return unityCatalogTreeDataProvider.unpin(node); + } + } + ), + ...registerDetailPanel( + context.extensionUri, + connectionManager, + unityCatalogTreeView, + unityCatalogTreeDataProvider, + telemetry + ) + ); +} + export async function activate( context: ExtensionContext ): Promise { @@ -267,6 +407,65 @@ export async function activate( ); } + // Surface only the Unity Catalog view and connect it using the ambient + // environment credentials (no bundle/config project required). The + // ConfigModel is only needed to satisfy the ConnectionManager + // constructor; connectFromEnvironment() never reads from it. + const remoteBundleFileSet = new BundleFileSet(workspaceFolderManager); + const remoteBundleFileWatcher = new BundleWatcher( + remoteBundleFileSet, + workspaceFolderManager + ); + const remoteConfigModel = new ConfigModel( + new BundleValidateModel( + remoteBundleFileWatcher, + cli, + workspaceFolderManager + ), + new OverrideableConfigModel(workspaceFolderManager), + new BundlePreValidateModel( + remoteBundleFileSet, + remoteBundleFileWatcher + ), + new BundleRemoteStateModel( + cli, + workspaceFolderManager, + workspaceConfigs + ), + customWhenContext, + stateStorage + ); + const remoteConnectionManager = new ConnectionManager( + cli, + remoteConfigModel, + workspaceFolderManager, + customWhenContext, + telemetry + ); + context.subscriptions.push( + remoteBundleFileWatcher, + remoteConfigModel, + remoteConnectionManager + ); + + const connectRemote = () => + remoteConnectionManager.connectFromEnvironment().catch((e) => { + logging.NamedLogger.getOrCreate(Loggers.Extension).error( + "Remote mode: failed to connect Unity Catalog", + e + ); + }); + + registerUnityCatalog( + context, + remoteConnectionManager, + stateStorage, + telemetry, + connectRemote + ); + + connectRemote(); + customWhenContext.setActivated(true); return; } @@ -490,126 +689,7 @@ export async function activate( ) ); - const unityCatalogTreeDataProvider = new UnityCatalogTreeDataProvider( - connectionManager, - stateStorage, - context.extensionPath - ); - - const unityCatalogTreeView = window.createTreeView("unityCatalogView", { - treeDataProvider: unityCatalogTreeDataProvider, - }); - context.subscriptions.push( - unityCatalogTreeDataProvider, - unityCatalogTreeView, - telemetry.registerCommand( - "databricks.unityCatalog.refresh", - unityCatalogTreeDataProvider.refresh, - unityCatalogTreeDataProvider - ), - telemetry.registerCommand( - "databricks.unityCatalog.refreshNode", - (node: UnityCatalogTreeNode) => - unityCatalogTreeDataProvider.refreshNode(node) - ), - telemetry.registerCommand( - "databricks.unityCatalog.copyStorageLocation", - async (node: UnityCatalogTreeNode) => { - if ( - (node.kind === "table" || node.kind === "volume") && - node.storageLocation - ) { - await env.clipboard.writeText(node.storageLocation); - window.showInformationMessage("Copied to clipboard"); - } - } - ), - telemetry.registerCommand( - "databricks.unityCatalog.copyViewSql", - async (node: UnityCatalogTreeNode) => { - if (node.kind === "table" && node.viewDefinition) { - await env.clipboard.writeText(node.viewDefinition); - window.showInformationMessage("Copied to clipboard"); - } - } - ), - telemetry.registerCommand( - "databricks.unityCatalog.copyName", - async (node: UnityCatalogTreeNode) => { - if ( - node.kind === "error" || - node.kind === "empty" || - node.kind === "favorites" || - node.kind === "group" - ) { - return; - } - const text = node.kind === "column" ? node.name : node.fullName; - await env.clipboard.writeText(text); - window.showInformationMessage("Copied to clipboard"); - } - ), - telemetry.registerCommand( - "databricks.unityCatalog.openExternal", - async (node: UnityCatalogTreeNode) => { - if (node.kind === "error" || node.kind === "column") { - return; - } - const url = - unityCatalogTreeDataProvider.getNodeExploreUrl(node); - if (!url) { - window.showErrorMessage( - "Databricks: Can't open external link. No URL found." - ); - return; - } - await UrlUtils.openExternal(url); - } - ), - commands.registerCommand("databricks.unityCatalog.filter", async () => { - await commands.executeCommand("unityCatalogView.focus"); - await commands.executeCommand("list.find"); - }), - telemetry.registerCommand( - "databricks.unityCatalog.pin", - (node: UnityCatalogTreeNode) => { - if ( - node.kind === "catalog" || - node.kind === "schema" || - node.kind === "table" || - node.kind === "volume" || - node.kind === "function" || - node.kind === "registeredModel" || - node.kind === "modelVersion" - ) { - return unityCatalogTreeDataProvider.pin(node); - } - } - ), - telemetry.registerCommand( - "databricks.unityCatalog.unpin", - (node: UnityCatalogTreeNode) => { - if ( - node.kind === "catalog" || - node.kind === "schema" || - node.kind === "table" || - node.kind === "volume" || - node.kind === "function" || - node.kind === "registeredModel" || - node.kind === "modelVersion" - ) { - return unityCatalogTreeDataProvider.unpin(node); - } - } - ), - ...registerDetailPanel( - context.extensionUri, - connectionManager, - unityCatalogTreeView, - unityCatalogTreeDataProvider, - telemetry - ) - ); + registerUnityCatalog(context, connectionManager, stateStorage, telemetry); const configureAutocomplete = new ConfigureAutocomplete( context, diff --git a/packages/databricks-vscode/src/ui/unity-catalog/UnityCatalogTreeDataProvider.test.ts b/packages/databricks-vscode/src/ui/unity-catalog/UnityCatalogTreeDataProvider.test.ts index 0f8f2d6ef..739e0ba84 100644 --- a/packages/databricks-vscode/src/ui/unity-catalog/UnityCatalogTreeDataProvider.test.ts +++ b/packages/databricks-vscode/src/ui/unity-catalog/UnityCatalogTreeDataProvider.test.ts @@ -175,15 +175,53 @@ describe(__filename, () => { disposables.forEach((d) => d.dispose()); }); - it("returns undefined when not connected", async () => { + it("shows a status node at the root when not connected", async () => { when(mockConnectionManager.workspaceClient).thenReturn(undefined); + when(mockConnectionManager.state).thenReturn("DISCONNECTED"); + when(mockConnectionManager.connectionError).thenReturn(undefined); const provider = new UnityCatalogTreeDataProvider( instance(mockConnectionManager), stubStateStorage ); disposables.push(provider); - const children = await resolveProviderResult(provider.getChildren()); + const children = (await resolveProviderResult( + provider.getChildren() + )) as UnityCatalogTreeNode[]; + assert.strictEqual(children.length, 1); + assert.strictEqual(children[0].kind, "empty"); + }); + + it("shows the connection error at the root when not connected", async () => { + when(mockConnectionManager.workspaceClient).thenReturn(undefined); + when(mockConnectionManager.state).thenReturn("DISCONNECTED"); + when(mockConnectionManager.connectionError).thenReturn( + "No Databricks host found in the environment" + ); + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const children = (await resolveProviderResult( + provider.getChildren() + )) as UnityCatalogTreeNode[]; + assert.strictEqual(children.length, 1); + assert.strictEqual(children[0].kind, "error"); + }); + + it("returns undefined for child nodes when not connected", async () => { + when(mockConnectionManager.workspaceClient).thenReturn(undefined); + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const children = await resolveProviderResult( + provider.getChildren({kind: "favorites"}) + ); assert.strictEqual(children, undefined); }); diff --git a/packages/databricks-vscode/src/ui/unity-catalog/UnityCatalogTreeDataProvider.ts b/packages/databricks-vscode/src/ui/unity-catalog/UnityCatalogTreeDataProvider.ts index d4dd52d4f..a42aa1452 100644 --- a/packages/databricks-vscode/src/ui/unity-catalog/UnityCatalogTreeDataProvider.ts +++ b/packages/databricks-vscode/src/ui/unity-catalog/UnityCatalogTreeDataProvider.ts @@ -149,7 +149,13 @@ export class UnityCatalogTreeDataProvider ): Promise { const client = this.connectionManager.workspaceClient; if (!client) { - return undefined; + // Without a workspace client there is nothing to load. Only surface + // a status node at the root (child requests just return nothing) so + // the view explains why it's empty instead of rendering blank. + if (element) { + return undefined; + } + return this.getDisconnectedRootChildren(); } const currentUser = this.connectionManager.databricksWorkspace?.user; @@ -186,6 +192,17 @@ export class UnityCatalogTreeDataProvider return undefined; } + private getDisconnectedRootChildren(): UnityCatalogTreeNode[] { + if (this.connectionManager.state === "CONNECTING") { + return [{kind: "empty", message: "Connecting…"}]; + } + const error = this.connectionManager.connectionError; + if (error) { + return [{kind: "error", message: `Not connected: ${error}`}]; + } + return [{kind: "empty", message: "Not connected"}]; + } + private async getRootChildren( client: NonNullable, currentUser: From 9b43a50c7080e0a28079f093c64df13c1e030175 Mon Sep 17 00:00:00 2001 From: misha-db Date: Thu, 16 Jul 2026 00:52:22 +0400 Subject: [PATCH 2/5] reuse PersonalAccessTokenAuthProvider for remote connection --- .../src/configuration/ConnectionManager.ts | 11 +++- .../src/configuration/auth/AuthProvider.ts | 59 +------------------ 2 files changed, 9 insertions(+), 61 deletions(-) diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.ts b/packages/databricks-vscode/src/configuration/ConnectionManager.ts index 78ad60b9f..7d8618f63 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionManager.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionManager.ts @@ -21,7 +21,7 @@ import {ConfigModel} from "./models/ConfigModel"; import {onError, withOnErrorHandler} from "../utils/onErrorDecorator"; import { AuthProvider, - EnvironmentAuthProvider, + PersonalAccessTokenAuthProvider, ProfileAuthProvider, } from "./auth/AuthProvider"; import {normalizeHost} from "../utils/urlUtils"; @@ -302,9 +302,14 @@ export class ConnectionManager implements Disposable { "No Databricks host found in the environment" ); } - authProvider = new EnvironmentAuthProvider( + if (config.token === undefined) { + throw new Error( + "No Databricks token found in the environment" + ); + } + authProvider = new PersonalAccessTokenAuthProvider( normalizeHost(config.host), - config, + config.token, this.cli ); } diff --git a/packages/databricks-vscode/src/configuration/auth/AuthProvider.ts b/packages/databricks-vscode/src/configuration/auth/AuthProvider.ts index 649da2d51..268d4b2af 100644 --- a/packages/databricks-vscode/src/configuration/auth/AuthProvider.ts +++ b/packages/databricks-vscode/src/configuration/auth/AuthProvider.ts @@ -24,8 +24,7 @@ export type AuthType = | "databricks-cli" | "google-id" | "profile" - | "pat" - | "env"; + | "pat"; export abstract class AuthProvider { constructor( @@ -533,59 +532,3 @@ export class PersonalAccessTokenAuthProvider extends AuthProvider { }); } } - -/** - * Auth provider backed by an already-resolved SDK Config, produced by the SDK's - * default credential chain (environment variables, metadata service, etc.). - * - * Used in Databricks Remote SSH sessions where the ambient environment supplies - * credentials and there is no bundle/config project or profile to resolve. - */ -export class EnvironmentAuthProvider extends AuthProvider { - constructor( - host: URL, - private readonly config: Config, - cli: CliWrapper - ) { - super(host, "env", cli, true); - } - - describe(): string { - return "Environment"; - } - - toJSON(): Record { - return { - host: this.host.toString(), - authType: this.authType, - }; - } - - toEnv(): Record { - return { - DATABRICKS_HOST: this.host.toString(), - }; - } - - toIni(): Record | undefined { - return undefined; - } - - protected async _check(): Promise { - try { - const workspaceClient = await this.getWorkspaceClient(); - await workspaceClient.currentUser.me(); - return true; - } catch (e) { - logging.NamedLogger.getOrCreate(Loggers.Extension).error( - "Can't login using the ambient environment credentials", - e - ); - return false; - } - } - - protected _getSdkConfig(): Config { - return this.config; - } -} From 7ddd0deb0ec533c34ae11f85da1de5e60c8f61cb Mon Sep 17 00:00:00 2001 From: misha-db Date: Thu, 16 Jul 2026 02:08:35 +0400 Subject: [PATCH 3/5] Enable docs view in remote mode --- packages/databricks-vscode/package.json | 2 +- packages/databricks-vscode/src/extension.ts | 73 +++++++++++---------- 2 files changed, 41 insertions(+), 34 deletions(-) diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index d6d326560..19aa11dbc 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -569,7 +569,7 @@ { "id": "databricksDocsView", "name": "Documentation", - "when": "databricks.context.activated && !databricks.context.remoteMode" + "when": "databricks.context.activated" } ] }, diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 79280edb6..da01c3356 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -232,6 +232,17 @@ function registerUnityCatalog( ); } +function registerDocsView(context: ExtensionContext): void { + const docsViewTreeDataProvider = new DocsViewTreeDataProvider(); + context.subscriptions.push( + window.registerTreeDataProvider( + "databricksDocsView", + docsViewTreeDataProvider + ), + docsViewTreeDataProvider + ); +} + export async function activate( context: ExtensionContext ): Promise { @@ -289,6 +300,33 @@ export async function activate( ) ); + // Utils. Registered early (before the remote-mode branch below returns) so + // that views shared between the normal and remote flows - such as the docs + // view, which invokes "databricks.utils.openExternal" - work in both. + const utilCommands = new UtilsCommands.UtilsCommands(telemetry); + context.subscriptions.push( + telemetry.registerCommand( + "databricks.utils.openExternal", + utilCommands.openExternalCommand(), + utilCommands + ), + telemetry.registerCommand( + "databricks.utils.goToDefinition", + utilCommands.goToDefinition(), + utilCommands + ), + telemetry.registerCommand( + "databricks.utils.copy", + utilCommands.copyToClipboardCommand(), + utilCommands + ), + telemetry.registerCommand("databricks.call", (fn) => { + if (fn) { + fn(); + } + }) + ); + if ( workspace.workspaceFolders === undefined || workspace.workspaceFolders?.length === 0 @@ -463,6 +501,7 @@ export async function activate( telemetry, connectRemote ); + registerDocsView(context); connectRemote(); @@ -1159,14 +1198,7 @@ export async function activate( debugWorkflowFactory ); - const docsViewTreeDataProvider = new DocsViewTreeDataProvider(); - context.subscriptions.push( - window.registerTreeDataProvider( - "databricksDocsView", - docsViewTreeDataProvider - ), - docsViewTreeDataProvider - ); + registerDocsView(context); showQuickStartOnFirstUse(context).catch((e) => { logging.NamedLogger.getOrCreate("Extension").error( @@ -1175,31 +1207,6 @@ export async function activate( ); }); - // Utils - const utilCommands = new UtilsCommands.UtilsCommands(telemetry); - context.subscriptions.push( - telemetry.registerCommand( - "databricks.utils.openExternal", - utilCommands.openExternalCommand(), - utilCommands - ), - telemetry.registerCommand( - "databricks.utils.goToDefinition", - utilCommands.goToDefinition(), - utilCommands - ), - telemetry.registerCommand( - "databricks.utils.copy", - utilCommands.copyToClipboardCommand(), - utilCommands - ), - telemetry.registerCommand("databricks.call", (fn) => { - if (fn) { - fn(); - } - }) - ); - // generate a json schema for bundle root and load a custom provider into // redhat.vscode-yaml extension to validate bundle config files with this schema registerBundleAutocompleteProvider( From 67a4f191b2404b8020527ae4e14f410f5bcf5262 Mon Sep 17 00:00:00 2001 From: misha-db Date: Thu, 16 Jul 2026 11:54:00 +0400 Subject: [PATCH 4/5] Update ConnectionManager.ts --- .../src/configuration/ConnectionManager.ts | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.ts b/packages/databricks-vscode/src/configuration/ConnectionManager.ts index 7d8618f63..84b57488b 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionManager.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionManager.ts @@ -289,6 +289,12 @@ export class ConnectionManager implements Disposable { */ async connectFromEnvironment(authProvider?: AuthProvider): Promise { 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 { From 5f72aa0e78265992be02ff9e30848832d435bfc9 Mon Sep 17 00:00:00 2001 From: misha-db Date: Thu, 16 Jul 2026 11:59:52 +0400 Subject: [PATCH 5/5] Utils commands registered on the welcome path --- .../src/configuration/ConnectionManager.ts | 7 ++- packages/databricks-vscode/src/extension.ts | 54 +++++++++---------- 2 files changed, 32 insertions(+), 29 deletions(-) diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.ts b/packages/databricks-vscode/src/configuration/ConnectionManager.ts index 84b57488b..2124d7c95 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionManager.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionManager.ts @@ -279,8 +279,11 @@ export class ConnectionManager implements Disposable { } /** - * Connect using the ambient environment credentials resolved by the SDK's - * default credential chain (environment variables, metadata service, etc.). + * 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 diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index da01c3356..6b96a7c3a 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -300,33 +300,6 @@ export async function activate( ) ); - // Utils. Registered early (before the remote-mode branch below returns) so - // that views shared between the normal and remote flows - such as the docs - // view, which invokes "databricks.utils.openExternal" - work in both. - const utilCommands = new UtilsCommands.UtilsCommands(telemetry); - context.subscriptions.push( - telemetry.registerCommand( - "databricks.utils.openExternal", - utilCommands.openExternalCommand(), - utilCommands - ), - telemetry.registerCommand( - "databricks.utils.goToDefinition", - utilCommands.goToDefinition(), - utilCommands - ), - telemetry.registerCommand( - "databricks.utils.copy", - utilCommands.copyToClipboardCommand(), - utilCommands - ), - telemetry.registerCommand("databricks.call", (fn) => { - if (fn) { - fn(); - } - }) - ); - if ( workspace.workspaceFolders === undefined || workspace.workspaceFolders?.length === 0 @@ -355,6 +328,33 @@ export async function activate( stateStorage ); + // Utils. Registered before the remote-mode branch below returns so that + // views shared between the normal and remote flows - such as the docs view, + // which invokes "databricks.utils.openExternal" - work in both. + const utilCommands = new UtilsCommands.UtilsCommands(telemetry); + context.subscriptions.push( + telemetry.registerCommand( + "databricks.utils.openExternal", + utilCommands.openExternalCommand(), + utilCommands + ), + telemetry.registerCommand( + "databricks.utils.goToDefinition", + utilCommands.goToDefinition(), + utilCommands + ), + telemetry.registerCommand( + "databricks.utils.copy", + utilCommands.copyToClipboardCommand(), + utilCommands + ), + telemetry.registerCommand("databricks.call", (fn) => { + if (fn) { + fn(); + } + }) + ); + // Add the databricks binary to the PATH environment variable in terminals context.environmentVariableCollection.clear(); context.environmentVariableCollection.persistent = false;