PoC for extending the official Databricks extension

package.json

@@ -48,7 +48,8 @@
 		"notebookKernelJupyterNotebook"
 	],
 	"extensionDependencies": [
-		"RandomFractalsInc.vscode-data-table"
+		"RandomFractalsInc.vscode-data-table",
+		"databricks.databricks-vscode"
 	],
 	"activationEvents": [
 		"onView:DatabricksConnections",
@@ -372,17 +373,8 @@
 				}
 			}
 		],
-		"viewsContainers": {
-			"activitybar": [
-				{
-					"id": "databricks",
-					"title": "Databricks",
-					"icon": "resources/databricks_sidebar.png"
-				}
-			]
-		},
 		"views": {
-			"databricks": [
+			"databricksBar": [
 				{
 					"id": "DatabricksConnections",
 					"name": "Connections"
@@ -952,7 +944,9 @@
 		"glob": "^8.0.3",
 		"mocha": "^10.0.0",
 		"webpack": "^5.74.0",
-		"webpack-cli": "^4.10.0"
+		"webpack-cli": "^4.10.0",
+		"@databricks/databricks-vscode-types": "file:../databricks-vscode/packages/databricks-vscode-types",
+		"@databricks/databricks-sdk-js": "file:../databricks-vscode/packages/databricks-sdk-js"
 	},
 	"dependencies": {
 		"@angular/common": "^14.0.1",

src/ThisExtension.ts

@@ -7,6 +7,8 @@ import { DatabricksConnectionManagerVSCode } from './vscode/treeviews/connection
 import { SensitiveValueStore } from './vscode/treeviews/connections/_types';
 import { DatabricksConnectionManagerCLI } from './vscode/treeviews/connections/DatabricksConnectionManagerCLI';
 import { DatabricksKernel } from './vscode/notebook/DatabricksKernel';
+import { PublicApi } from '@databricks/databricks-vscode-types';
+import { DatabricksConnectionManagerBricks } from './vscode/treeviews/connections/DatabricksConnectionManagerBricks';
 
 // https://vshaxe.github.io/vscode-extern/vscode/TreeDataProvider.html
 export abstract class ThisExtension {
@@ -55,7 +57,7 @@ export abstract class ThisExtension {
 		return this._isValidated;
 	}
 
-	static async initialize(context: vscode.ExtensionContext): Promise<boolean> {
+	static async initialize(context: vscode.ExtensionContext, dbApi: PublicApi): Promise<boolean> {
 		try {
 			this._logger = vscode.window.createOutputChannel(context.extension.id);
 			this.log("Logger initialized!");
@@ -68,17 +70,18 @@ export abstract class ThisExtension {
 
 			ThisExtension.readGlobalSettings();
 
-			let connectionManager = this.getConfigurationSetting("databricks.connectionManager");
-			switch (connectionManager.value) {
-				case "VSCode Settings":
-					this._connectionManager = new DatabricksConnectionManagerVSCode();
-					break;
-				case "Databricks CLI Profiles":
-					this._connectionManager = new DatabricksConnectionManagerCLI();
-					break;
-				default:
-					this.log("'" + connectionManager + "' is not a valid value for config setting 'databricks.connectionManager!");
-			}
+			// let connectionManager = this.getConfigurationSetting("databricks.connectionManager");
+			// switch (connectionManager.value) {
+			// 	case "VSCode Settings":
+			// 		this._connectionManager = new DatabricksConnectionManagerVSCode();
+			// 		break;
+			// 	case "Databricks CLI Profiles":
+			// 		this._connectionManager = new DatabricksConnectionManagerCLI();
+			// 		break;
+			// 	default:
+			// 		this.log("'" + connectionManager + "' is not a valid value for config setting 'databricks.connectionManager!");
+			// }
+			this._connectionManager = new DatabricksConnectionManagerBricks(dbApi);
 
 			await this.ConnectionManager.initialize();
 

src/extension.ts

@@ -32,14 +32,18 @@ import { DatabricksWorkspaceProvider } from './vscode/filesystemProvider/Databri
 import { DatabricksRepoTreeItem } from './vscode/treeviews/repos/DatabricksRepoTreeItem';
 import { FSHelper } from './helpers/FSHelper';
 import { DatabricksKernelManager } from './vscode/notebook/DatabricksKernelManager';
+import {PublicApi} from "@databricks/databricks-vscode-types"
 
 export async function activate(context: vscode.ExtensionContext) {
 
+	await vscode.extensions.getExtension<PublicApi>('databricks.databricks-vscode').activate();
+	const dbApi = vscode.extensions.getExtension<PublicApi>('databricks.databricks-vscode').exports;
+
 	ThisExtension.StatusBar = vscode.window.createStatusBarItem("databricks-vscode", vscode.StatusBarAlignment.Right);
 	ThisExtension.StatusBar.show();
 	ThisExtension.setStatusBar("Initializing ...", true);
 
-	let isValidated: boolean = await ThisExtension.initialize(context);
+	let isValidated: boolean = await ThisExtension.initialize(context, dbApi);
 	if (isValidated === false) {
 		ThisExtension.log("Issue initializing extension - Please update Databricks settings and restart VSCode!");
 		vscode.window.showErrorMessage("Issue initializing extension - Please update Databricks settings and restart VSCode!");

src/vscode/treeviews/connections/DatabricksConnectionManagerBricks.ts

@@ -0,0 +1,54 @@
+import { PublicApi } from "@databricks/databricks-vscode-types";
+import { DatabricksConnectionManager } from "./DatabricksConnectionManager";
+import { iDatabricksConnection } from "./iDatabricksConnection";
+import {fromConfigFile} from "@databricks/databricks-sdk-js";
+
+export class DatabricksConnectionManagerBricks extends DatabricksConnectionManager {
+	constructor(private databricksApi: PublicApi) {
+		super();
+		this._initialized = true;
+	}
+
+	async initialize(): Promise<void> {
+		await this.databricksApi.connectionManager.waitForConnect();
+
+		this._connections = [{
+			displayName: "Databricks Extension",
+			personalAccessToken: "",
+			personalAccessTokenSecure: {
+				keyTarSettingName: undefined,
+				databricksCLIProfileName: this.databricksApi.connectionManager.profile
+			},
+			apiRootUrl: "",
+			localSyncFolder: "",
+			useCodeCells: true,
+
+			_source: "CLI-profile"
+		}];
+		console.log(this._connections[0])
+
+		let credentialProvider = fromConfigFile(this._connections[0].personalAccessTokenSecure!.databricksCLIProfileName!);
+		let creds = await credentialProvider();
+
+		this._connections[0].apiRootUrl = creds.host.toString();
+		this._connections[0].personalAccessToken = creds.token;
+
+		this._lastActiveConnectionName = this._connections[0].displayName;
+		await this.activateConnection(this._connections[0], false);
+	}
+
+	loadConnections(): void {
+		console.log("load connections")
+		//throw new Error("Method not implemented.");
+	}
+
+	updateConnection(updatedCon: iDatabricksConnection): void {
+		console.log("update connection", updatedCon);
+		//throw new Error("Method not implemented.");
+	}
+
+	async getAccessToken(con: iDatabricksConnection): Promise<string> {
+		return con.personalAccessToken;
+	}
+
+}