From 14e71af2d69ce1dcda2d4ac33f0f2c94df124df1 Mon Sep 17 00:00:00 2001 From: Fabian Jakobs Date: Mon, 16 Sep 2024 14:15:00 +0200 Subject: [PATCH] Expose db connect progress as a settings --- packages/databricks-vscode/package.json | 8 ++++++-- .../resources/python/00-databricks-init.py | 20 +++++++++++-------- .../utils/JobRunStateUtils.ts | 2 ++ .../src/utils/envVarGenerators.ts | 4 ++++ .../src/vscode-objs/WorkspaceConfigs.ts | 8 ++++++++ 5 files changed, 32 insertions(+), 10 deletions(-) diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 68c490291..eef799526 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -856,7 +856,6 @@ "views.workspace" ], "enumDescriptions": [ - "Limited local notebook support using DB Connect v2.", "Show cluster view in the explorer.", "Show workspace browser in the explorer." ], @@ -870,6 +869,11 @@ "default": true, "description": "Enable/disable rearranging cells in wrapper files created when using `workspace` as the sync destination. **Note:** It is recommended to NOT disable this setting. If you do disable it, you will need to manually handle sys.path for local imports in your notebooks." }, + "databricks.connect.progress": { + "type": "boolean", + "default": true, + "description": "Show PySpark progress bar when using Databricks Connect from notebooks." + }, "databricks.ipythonDir": { "type": "string", "description": "Absolute path to a directory for storing IPython files. Defaults to IPYTHONDIR environment variable (if set) or ~/.ipython." @@ -1008,4 +1012,4 @@ ], "report-dir": "coverage" } -} +} \ No newline at end of file diff --git a/packages/databricks-vscode/resources/python/00-databricks-init.py b/packages/databricks-vscode/resources/python/00-databricks-init.py index 97ae5313c..a7c26b49c 100644 --- a/packages/databricks-vscode/resources/python/00-databricks-init.py +++ b/packages/databricks-vscode/resources/python/00-databricks-init.py @@ -103,7 +103,13 @@ def __init__(self, env_name: str, default: any = None, required: bool = False): def __get__(self, instance, owner): if self.env_name in os.environ: - return self.transform(os.environ[self.env_name]) + if self.transform is not bool: + return self.transform(os.environ[self.env_name]) + + if os.environ[self.env_name].lower() == "true" or os.environ[self.env_name] == "1": + return True + elif os.environ[self.env_name].lower() == "false" or os.environ[self.env_name] == "0": + return False if self.required: raise AttributeError( @@ -119,6 +125,7 @@ def __set__(self, instance, value): class LocalDatabricksNotebookConfig: project_root: str = EnvLoader("DATABRICKS_PROJECT_ROOT", required=True) dataframe_display_limit: int = EnvLoader("DATABRICKS_DF_DISPLAY_LIMIT", 20) + show_progress: bool = EnvLoader("DATABRICKS_CONNECT_PROGRESS", default=True) def __new__(cls): annotations = cls.__dict__['__annotations__'] @@ -363,7 +370,6 @@ def df_html(df): @disposable def register_spark_progress(spark): try: - from pyspark.sql.connect.shell.progress import Progress import ipywidgets as widgets except Exception as e: return @@ -374,12 +380,14 @@ class Progress: def __init__( self, + cfg: LocalDatabricksNotebookConfig ) -> None: self._ticks = None self._tick = None self._started = time.time() self._bytes_read = 0 self._running = 0 + self.show_progress = cfg.show_progress self.init_ui() def init_ui(self): @@ -391,13 +399,9 @@ def init_ui(self): orientation='horizontal' ) self.w_status = widgets.Label(value="") - if self.is_enabled(): + if self.show_progress: display(widgets.HBox([self.w_progress, self.w_status])) - def is_enabled(self): - env_var = os.getenv("DATABRICKS_CONNECT_PROGRESS") - return env_var is None or (env_var.lower() != "false" and env_var.lower() != "0") - def update_ticks( self, stages, @@ -437,7 +441,7 @@ def __init__(self): self.op_id = "" def reset(self): - self.p = Progress() + self.p = Progress(cfg) def __call__(self, stages, diff --git a/packages/databricks-vscode/src/ui/bundle-resource-explorer/utils/JobRunStateUtils.ts b/packages/databricks-vscode/src/ui/bundle-resource-explorer/utils/JobRunStateUtils.ts index 0045419d7..bcbe3b4d8 100644 --- a/packages/databricks-vscode/src/ui/bundle-resource-explorer/utils/JobRunStateUtils.ts +++ b/packages/databricks-vscode/src/ui/bundle-resource-explorer/utils/JobRunStateUtils.ts @@ -43,4 +43,6 @@ export function getSimplifiedRunState(run?: Run): SimplifiedRunState { } return "Terminated"; } + + return "Unknown"; } diff --git a/packages/databricks-vscode/src/utils/envVarGenerators.ts b/packages/databricks-vscode/src/utils/envVarGenerators.ts index cc5711c35..6478f7f3d 100644 --- a/packages/databricks-vscode/src/utils/envVarGenerators.ts +++ b/packages/databricks-vscode/src/utils/envVarGenerators.ts @@ -5,6 +5,7 @@ import {logging, Headers} from "@databricks/databricks-sdk"; import {ConnectionManager} from "../configuration/ConnectionManager"; import {ConfigModel} from "../configuration/models/ConfigModel"; import {TerraformMetadata} from "./terraformUtils"; +import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs"; // eslint-disable-next-line @typescript-eslint/no-var-requires const packageJson = require("../../package.json"); @@ -115,11 +116,14 @@ export async function getDbConnectEnvVars( ) { const userAgent = getUserAgent(connectionManager); const existingSparkUa = process.env.SPARK_CONNECT_USER_AGENT ?? ""; + /* eslint-disable @typescript-eslint/naming-convention */ return { //We append our user agent to any existing SPARK_CONNECT_USER_AGENT defined in the //environment of the parent process of VS Code. SPARK_CONNECT_USER_AGENT: [existingSparkUa, userAgent].join(" ").trim(), + DATABRICKS_CONNECT_PROGRESS: + workspaceConfigs.showDbConnectProgress.toString(), DATABRICKS_PROJECT_ROOT: workspacePath.fsPath, ...((await getSparkRemoteEnvVar(connectionManager)) || {}), }; diff --git a/packages/databricks-vscode/src/vscode-objs/WorkspaceConfigs.ts b/packages/databricks-vscode/src/vscode-objs/WorkspaceConfigs.ts index 0bb248573..8bbe82176 100644 --- a/packages/databricks-vscode/src/vscode-objs/WorkspaceConfigs.ts +++ b/packages/databricks-vscode/src/vscode-objs/WorkspaceConfigs.ts @@ -105,6 +105,14 @@ export const workspaceConfigs = { ); }, + get showDbConnectProgress(): boolean { + return ( + workspace + .getConfiguration("databricks") + .get("connect.progress") ?? true + ); + }, + get ipythonDir(): string | undefined { const dir = workspace .getConfiguration("databricks")