From 125a2ecfebc94ac9e9600f4fceeedbb04a85dc65 Mon Sep 17 00:00:00 2001 From: Kartik Gupta <88345179+kartikgupta-db@users.noreply.github.com> Date: Wed, 17 Jan 2024 15:16:37 +0100 Subject: [PATCH] Add a model for remote state (`BundleRemoteStateModel`) (#1006) ## Changes * Add a model for the remote state. * Expose `remoteStateConfig` from `ConfigModel`. ## Tests --- packages/databricks-vscode/package.json | 5 + .../bundle/models/BundleRemoteStateModel.ts | 96 +++++++++++++++++++ .../databricks-vscode/src/bundle/types.ts | 2 +- .../databricks-vscode/src/cli/CliWrapper.ts | 26 +++++ .../src/configuration/models/ConfigModel.ts | 13 ++- packages/databricks-vscode/src/extension.ts | 16 +++- .../src/vscode-objs/WorkspaceConfigs.ts | 12 +++ 7 files changed, 167 insertions(+), 3 deletions(-) create mode 100644 packages/databricks-vscode/src/bundle/models/BundleRemoteStateModel.ts diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 80f439a6b..8a62996a8 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -658,6 +658,11 @@ "databricks.ipythonDir": { "type": "string", "description": "Absolute path to a directory for storing IPython files. Defaults to IPYTHONDIR environment variable (if set) or ~/.ipython." + }, + "databricks.bundle.remoteStateRefreshInterval": { + "type": "number", + "default": 5, + "description": "The interval in minutes at which the remote state of the bundle is refreshed." } } } diff --git a/packages/databricks-vscode/src/bundle/models/BundleRemoteStateModel.ts b/packages/databricks-vscode/src/bundle/models/BundleRemoteStateModel.ts new file mode 100644 index 000000000..773a4b4b2 --- /dev/null +++ b/packages/databricks-vscode/src/bundle/models/BundleRemoteStateModel.ts @@ -0,0 +1,96 @@ +import {Uri} from "vscode"; +import {CliWrapper} from "../../cli/CliWrapper"; +import {BaseModelWithStateCache} from "../../configuration/models/BaseModelWithStateCache"; +import {Mutex} from "../../locking"; + +import {BundleTarget} from "../types"; +import {AuthProvider} from "../../configuration/auth/AuthProvider"; +import lodash from "lodash"; +import {WorkspaceConfigs} from "../../vscode-objs/WorkspaceConfigs"; +import {withLogContext} from "@databricks/databricks-sdk/dist/logging"; +import {Loggers} from "../../logger"; +import {Context, context} from "@databricks/databricks-sdk"; + +/* eslint-disable @typescript-eslint/naming-convention */ +type Resources = Required["resources"]; +type Resource> = Required[K]; + +export type BundleRemoteState = BundleTarget & { + resources?: Resources & { + [r in keyof Resources]?: { + [k in keyof Resource]?: Resource[k] & { + id?: string; + modified_status?: "CREATED" | "DELETED" | "UPDATED"; + }; + }; + }; +}; +/* eslint-enable @typescript-eslint/naming-convention */ + +export class BundleRemoteStateModel extends BaseModelWithStateCache { + private target: string | undefined; + private authProvider: AuthProvider | undefined; + protected mutex = new Mutex(); + private refreshInterval: NodeJS.Timeout | undefined; + + constructor( + private readonly cli: CliWrapper, + private readonly workspaceFolder: Uri, + private readonly workspaceConfigs: WorkspaceConfigs + ) { + super(); + } + + @withLogContext(Loggers.Extension) + public init(@context ctx?: Context) { + this.refreshInterval = setInterval(async () => { + try { + await this.stateCache.refresh(); + } catch (e) { + ctx?.logger?.error("Unable to refresh bundle remote state", e); + } + }, this.workspaceConfigs.bundleRemoteStateRefreshInterval); + } + + @Mutex.synchronise("mutex") + public async setTarget(target: string | undefined) { + if (this.target === target) { + return; + } + this.target = target; + this.authProvider = undefined; + await this.stateCache.refresh(); + } + + @Mutex.synchronise("mutex") + public async setAuthProvider(authProvider: AuthProvider | undefined) { + if ( + !lodash.isEqual(this.authProvider?.toJSON(), authProvider?.toJSON()) + ) { + this.authProvider = authProvider; + await this.stateCache.refresh(); + } + } + + protected async readState(): Promise { + if (this.target === undefined || this.authProvider === undefined) { + return {}; + } + + return JSON.parse( + await this.cli.bundleSummarise( + this.target, + this.authProvider, + this.workspaceFolder, + this.workspaceConfigs.databrickscfgLocation + ) + ); + } + + dispose() { + super.dispose(); + if (this.refreshInterval !== undefined) { + clearInterval(this.refreshInterval); + } + } +} diff --git a/packages/databricks-vscode/src/bundle/types.ts b/packages/databricks-vscode/src/bundle/types.ts index e28250bcf..45cbf68a5 100644 --- a/packages/databricks-vscode/src/bundle/types.ts +++ b/packages/databricks-vscode/src/bundle/types.ts @@ -1,4 +1,4 @@ export {type BundleSchema} from "./BundleSchema"; - import {BundleSchema} from "./BundleSchema"; + export type BundleTarget = Required["targets"][string]; diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts index 413c4ed97..a12f157ea 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.ts @@ -194,4 +194,30 @@ export class CliWrapper { } return stdout; } + + async bundleSummarise( + target: string, + authProvider: AuthProvider, + workspaceFolder: Uri, + configfilePath?: string + ) { + const {stdout, stderr} = await execFile( + this.cliPath, + ["bundle", "summarise", "--target", target], + { + cwd: workspaceFolder.fsPath, + env: { + ...EnvVarGenerators.getEnvVarsForCli(configfilePath), + ...EnvVarGenerators.getProxyEnvVars(), + ...authProvider.toEnv(), + }, + shell: true, + } + ); + + if (stderr !== "") { + throw new Error(stderr); + } + return stdout; + } } diff --git a/packages/databricks-vscode/src/configuration/models/ConfigModel.ts b/packages/databricks-vscode/src/configuration/models/ConfigModel.ts index d2ee5513a..9d90f6b29 100644 --- a/packages/databricks-vscode/src/configuration/models/ConfigModel.ts +++ b/packages/databricks-vscode/src/configuration/models/ConfigModel.ts @@ -18,6 +18,10 @@ import { BundleValidateState, } from "../../bundle/models/BundleValidateModel"; import {CustomWhenContext} from "../../vscode-objs/CustomWhenContext"; +import { + BundleRemoteState, + BundleRemoteStateModel, +} from "../../bundle/models/BundleRemoteStateModel"; const defaults: ConfigState = { mode: "development", @@ -42,6 +46,7 @@ type ConfigState = Pick< OverrideableConfigState & { preValidateConfig?: BundlePreValidateState; validateConfig?: BundleValidateState; + remoteStateConfig?: BundleRemoteState; overrides?: OverrideableConfigState; }; @@ -92,6 +97,7 @@ export class ConfigModel implements Disposable { preValidateConfig: bundlePreValidateConfig, validateConfig: bundleValidateConfig, overrides, + remoteStateConfig: await this.bundleRemoteStateModel.load(), }; } @@ -114,6 +120,7 @@ export class ConfigModel implements Disposable { private readonly bundleValidateModel: BundleValidateModel, private readonly overrideableConfigModel: OverrideableConfigModel, private readonly bundlePreValidateModel: BundlePreValidateModel, + private readonly bundleRemoteStateModel: BundleRemoteStateModel, private readonly vscodeWhenContext: CustomWhenContext, private readonly stateStorage: StateStorage ) { @@ -132,13 +139,17 @@ export class ConfigModel implements Disposable { //refresh cache to trigger onDidChange event this.configCache.refresh(); }) - ) + ), + this.bundleRemoteStateModel.onDidChange(async () => { + await this.configCache.refresh(); + }) ); } @onError({popup: {prefix: "Failed to initialize configs."}}) public async init() { await this.readTarget(); + this.bundleRemoteStateModel.init(); } get targets() { diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 39f1c4e01..99b448904 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -58,6 +58,7 @@ import {BundleValidateModel} from "./bundle/models/BundleValidateModel"; import {ConfigModel} from "./configuration/models/ConfigModel"; import {OverrideableConfigModel} from "./configuration/models/OverrideableConfigModel"; import {BundlePreValidateModel} from "./bundle/models/BundlePreValidateModel"; +import {BundleRemoteStateModel} from "./bundle/models/BundleRemoteStateModel"; const customWhenContext = new CustomWhenContext(); @@ -166,7 +167,6 @@ export async function activate( const cli = new CliWrapper(context, cliLogFilePath); const bundleFileSet = new BundleFileSet(workspace.workspaceFolders[0].uri); const bundleFileWatcher = new BundleWatcher(bundleFileSet); - context.subscriptions.push(bundleFileWatcher); const bundleValidateModel = new BundleValidateModel( bundleFileWatcher, cli, @@ -178,10 +178,16 @@ export async function activate( bundleFileSet, bundleFileWatcher ); + const bundleRemoteStateModel = new BundleRemoteStateModel( + cli, + workspaceUri, + workspaceConfigs + ); const configModel = new ConfigModel( bundleValidateModel, overrideableConfigModel, bundlePreValidateModel, + bundleRemoteStateModel, customWhenContext, stateStorage ); @@ -193,6 +199,14 @@ export async function activate( customWhenContext ); context.subscriptions.push( + bundleFileWatcher, + bundleValidateModel, + overrideableConfigModel, + bundlePreValidateModel, + bundleRemoteStateModel, + configModel, + configModel, + connectionManager, connectionManager.onDidChangeState(async (state) => { telemetry.setMetadata( Metadata.USER, diff --git a/packages/databricks-vscode/src/vscode-objs/WorkspaceConfigs.ts b/packages/databricks-vscode/src/vscode-objs/WorkspaceConfigs.ts index 1a14ceb5d..6dffaac0e 100644 --- a/packages/databricks-vscode/src/vscode-objs/WorkspaceConfigs.ts +++ b/packages/databricks-vscode/src/vscode-objs/WorkspaceConfigs.ts @@ -1,5 +1,6 @@ import {ConfigurationTarget, workspace} from "vscode"; import {SyncDestinationType} from "../sync/SyncDestination"; +import {Time, TimeUnits} from "@databricks/databricks-sdk"; export const workspaceConfigs = { get maxFieldLength() { @@ -148,4 +149,15 @@ export const workspaceConfigs = { } return dir; }, + + get bundleRemoteStateRefreshInterval(): number { + const config = + workspace + .getConfiguration("databricks") + .get("bundle.remoteStateRefreshInterval") ?? 5; + + return new Time(config, TimeUnits.minutes).toMillSeconds().value; + }, }; + +export type WorkspaceConfigs = typeof workspaceConfigs;