From f117fa1bce50d937a808c5324401464699957c63 Mon Sep 17 00:00:00 2001 From: Hazzng Date: Tue, 14 Apr 2026 09:09:24 +0000 Subject: [PATCH 1/7] perf: defer Databricks Connect initialization to background thread Kernel startup was blocking ~15s on DatabricksSession.getOrCreate() before sending the "ready" signal. Now the kernel reports ready in ~1-2s by deferring spark initialization to a background thread. A LazySparkSession proxy is placed in the namespace immediately so 'spark' in dir() is True. On first attribute access (e.g. spark.sql()), the proxy blocks until the real session is available - but by then the background init has had a head start (often already complete). Key changes: - kernel_runner.py: Add LazySparkSession proxy, _SparkInitState, background thread init, thread-safe _write_json for stdout - persistentExecutor.ts: Handle new 'spark_ready' message type - utils.ts: Add showSparkReadyNotification, skip INITIALIZING status https://claude.ai/code/session_01NCzUTPmDaWc5Rwr8xB4Zfu --- src/kernels/persistentExecutor.ts | 23 ++++ src/kernels/utils.ts | 15 +++ src/python/kernel_runner.py | 172 +++++++++++++++++++++++++++--- 3 files changed, 196 insertions(+), 14 deletions(-) diff --git a/src/kernels/persistentExecutor.ts b/src/kernels/persistentExecutor.ts index e2b289b..21d8121 100644 --- a/src/kernels/persistentExecutor.ts +++ b/src/kernels/persistentExecutor.ts @@ -29,6 +29,7 @@ import { logKernelStatusDebug, createReadyWaiter, getKernelStartupTimeout, + showSparkReadyNotification, } from './utils'; /** @@ -340,6 +341,12 @@ export class PersistentExecutor implements vscode.Disposable { return; } + // Handle deferred spark initialization completion + if (message.type === 'spark_ready') { + this.handleSparkReadySignal(message as KernelResponse); + return; + } + // Handle widget input requests from Python if (message.type === 'input_request') { this.handleWidgetInputRequest(message as WidgetInputRequest); @@ -376,6 +383,22 @@ export class PersistentExecutor implements vscode.Disposable { showKernelStatusNotifications(statusInfo); } + /** + * Handle deferred spark initialization completion. + * Sent by background thread in kernel_runner.py after DatabricksSession is connected. + */ + private handleSparkReadySignal(response: KernelResponse): void { + // eslint-disable-next-line @typescript-eslint/naming-convention + const sparkStatus = (response as { spark_status?: string }).spark_status; + if (sparkStatus) { + this._sparkStatus = sparkStatus; + if (this._debugMode) { + console.debug(`[Executor] Spark ready: ${sparkStatus}`); + } + showSparkReadyNotification(sparkStatus); + } + } + /** * Handle pending request response */ diff --git a/src/kernels/utils.ts b/src/kernels/utils.ts index 84dab6a..a3ba033 100644 --- a/src/kernels/utils.ts +++ b/src/kernels/utils.ts @@ -219,12 +219,27 @@ export function showKernelStatusNotifications(statusInfo: KernelStatusInfo): voi } // Notify about spark status + // Skip "INITIALIZING:" status - the real notification comes via spark_ready message if (statusInfo.sparkStatus) { if (statusInfo.sparkStatus.startsWith('OK:')) { showInfoMessage(`Kernel: ${statusInfo.sparkStatus}`); } else if (statusInfo.sparkStatus.startsWith('WARN:')) { showWarningMessage(`Kernel: ${statusInfo.sparkStatus}`); } + // INITIALIZING: is silently logged (notification comes when spark_ready arrives) + } +} + +/** + * Show notification when deferred spark initialization completes. + * Called when the background spark init thread finishes and sends spark_ready. + * @param sparkStatus - Final spark status string + */ +export function showSparkReadyNotification(sparkStatus: string): void { + if (sparkStatus.startsWith('OK:')) { + showInfoMessage(`Kernel: ${sparkStatus}`); + } else if (sparkStatus.startsWith('WARN:')) { + showWarningMessage(`Kernel: ${sparkStatus}`); } } diff --git a/src/python/kernel_runner.py b/src/python/kernel_runner.py index 96409d3..de5b9d6 100644 --- a/src/python/kernel_runner.py +++ b/src/python/kernel_runner.py @@ -17,6 +17,7 @@ import os import ast import asyncio +import threading # Display utilities from display_utils import display_to_html @@ -47,6 +48,84 @@ # "Event loop is closed" errors with SDK clients that store loop references _event_loop = None +# Thread-safe stdout lock for background spark init thread +_stdout_lock = threading.Lock() + +# State for background spark initialization +_spark_init_state = None + +# Reference to real stdout (before execution redirects it) +_real_stdout = None + + +class _SparkInitState: + """Thread-safe state for background spark initialization.""" + + def __init__(self): + self._ready = threading.Event() + self._session = None + self._error_msg = None + self._status = None + + def set_success(self, session, status): + self._session = session + self._status = status + self._ready.set() + + def set_failure(self, status): + self._error_msg = status + self._status = status + self._ready.set() + + def wait(self, timeout=None): + return self._ready.wait(timeout=timeout) + + @property + def is_ready(self): + return self._ready.is_set() + + +class LazySparkSession: + """ + Proxy that defers to the real SparkSession once background initialization completes. + Placed in the namespace immediately so 'spark' in dir() is True, but blocks on + first attribute access (e.g. spark.sql()) until the real session is available. + This allows the kernel to report 'ready' in ~1s instead of ~15s. + """ + + def __init__(self, init_state): + object.__setattr__(self, '_init_state', init_state) + + def _get_session(self): + state = object.__getattribute__(self, '_init_state') + if not state.wait(timeout=120): + raise TimeoutError("Timed out waiting for Spark session initialization (120s)") + if state._error_msg: + raise RuntimeError( + f"Spark session not available: {state._error_msg}\n" + "Run 'databricks auth login' to configure authentication." + ) + if state._session is None: + raise RuntimeError("Spark session failed to initialize") + return state._session + + def __getattr__(self, name): + return getattr(self._get_session(), name) + + def __setattr__(self, name, value): + setattr(self._get_session(), name, value) + + def __repr__(self): + state = object.__getattribute__(self, '_init_state') + if state.is_ready and state._session: + return repr(state._session) + if state.is_ready and state._error_msg: + return f"" + return "" + + def __bool__(self): + return True + class LocalDbutils: """ @@ -305,6 +384,49 @@ def initialize_spark_session(): return f"WARN: Spark not initialized ({error_msg}). Run 'databricks auth login' to refresh tokens." +def _write_json(data): + """Thread-safe write of a JSON message to real stdout.""" + global _real_stdout + out = _real_stdout or sys.stdout + with _stdout_lock: + out.write(json.dumps(data, ensure_ascii=True, default=str) + '\n') + out.flush() + + +def _background_spark_init(db_compatible, db_warning): + """ + Initialize Spark session in a background thread and notify Node.js when done. + This allows the kernel to report 'ready' immediately while Spark connects. + """ + global _spark_init_state + + try: + if db_compatible: + status = initialize_spark_session() + # Check if real session was placed in namespace (not the proxy) + real_spark = _namespace.get('spark') + if real_spark is not None and not isinstance(real_spark, LazySparkSession): + _spark_init_state.set_success(real_spark, status) + log_debug(f"Background spark init succeeded: {status}") + else: + # Spark init returned a warning - didn't create a session + _spark_init_state.set_failure(status) + log_debug(f"Background spark init warning: {status}") + elif db_warning: + _spark_init_state.set_failure(f"WARN: {db_warning}") + else: + _spark_init_state.set_failure("databricks-connect not available") + except Exception as e: + _spark_init_state.set_failure(str(e)) + log_debug(f"Background spark init error: {e}") + + # Send spark_ready status update to Node.js + _write_json({ + 'type': 'spark_ready', + 'spark_status': _spark_init_state._status, + }) + + def handle_interrupt(signum, frame): """Handle SIGINT (Ctrl+C) gracefully.""" raise KeyboardInterrupt() @@ -516,6 +638,10 @@ def get_variables(): def main(): """Main loop - read JSON commands from stdin, execute, write JSON responses to stdout.""" import sys as _sys + global _spark_init_state, _real_stdout + + # Save real stdout before anything can redirect it (execute_code swaps sys.stdout) + _real_stdout = _sys.stdout # Set up signal handler for interrupts signal.signal(signal.SIGINT, handle_interrupt) @@ -572,14 +698,35 @@ def main(): # Add display() function to namespace _namespace['display'] = display - # Initialize Spark session if available (skip if incompatible version) + # Initialize Spark session in BACKGROUND thread for fast startup. + # Instead of blocking 10-15s for DatabricksSession.getOrCreate(), + # we place a LazySparkSession proxy in the namespace and send "ready" + # immediately. The proxy blocks on first attribute access (e.g. spark.sql()) + # until the real session is available, but by then the background init + # has had a head start (often already complete). spark_status = None + _spark_init_state = _SparkInitState() + if db_compatible: - spark_status = initialize_spark_session() + # Place lazy proxy so 'spark' in dir() is True immediately + _namespace['spark'] = LazySparkSession(_spark_init_state) + + # Start background initialization + init_thread = threading.Thread( + target=_background_spark_init, + args=(db_compatible, db_warning), + daemon=True, + name='spark-init', + ) + init_thread.start() + spark_status = "INITIALIZING: Databricks Connect session starting in background..." + log_debug("Spark initialization started in background thread") elif db_warning: spark_status = f"WARN: {db_warning}" + _spark_init_state.set_failure(spark_status) - # Send ready signal with spark status and environment info + # Send ready signal IMMEDIATELY (no waiting for Spark) + # This reduces kernel startup from ~15s to ~1-2s ready_signal = { 'type': 'ready', 'version': '1.0', @@ -590,7 +737,7 @@ def main(): 'databricks_connect_compatible': db_compatible, 'added_import_paths': added_import_paths, } - print(json.dumps(ready_signal), flush=True) + _write_json(ready_signal) for line in sys.stdin: line = line.strip() @@ -616,10 +763,9 @@ def main(): result['id'] = request_id result['type'] = 'result' - # Use ensure_ascii=True and default=str to handle non-serializable values + # Thread-safe write (background spark init may write spark_ready concurrently) try: - output = json.dumps(result, ensure_ascii=True, default=str) - print(output, flush=True) + _write_json(result) except Exception as json_err: # Fallback: return error without display data if serialization fails fallback_result = { @@ -630,22 +776,20 @@ def main(): 'stdout': result.get('stdout', ''), 'stderr': result.get('stderr', ''), } - print(json.dumps(fallback_result, ensure_ascii=True), flush=True) + _write_json(fallback_result) except json.JSONDecodeError as e: - error_result = { + _write_json({ 'type': 'error', 'success': False, 'error': f'Invalid JSON: {str(e)}' - } - print(json.dumps(error_result), flush=True) + }) except Exception as e: - error_result = { + _write_json({ 'type': 'error', 'success': False, 'error': f'Internal error: {str(e)}' - } - print(json.dumps(error_result), flush=True) + }) if __name__ == '__main__': From 9ead780da84bb5e48b60fd6b945de9e11e666490 Mon Sep 17 00:00:00 2001 From: Hazzng Date: Tue, 14 Apr 2026 09:11:51 +0000 Subject: [PATCH 2/7] chore: bump version to 0.4.6 and update changelog https://claude.ai/code/session_01NCzUTPmDaWc5Rwr8xB4Zfu --- CHANGELOG.md | 10 ++++++++++ package.json | 2 +- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 64e706d..fc10081 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,16 @@ All notable changes to the Databricks Notebook Studio extension will be document The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.4.6] - 2026-04-14 + +### Changed +- **Databricks Connect startup performance**: Kernel now reports ready in ~1-2 seconds instead of ~15 seconds by deferring Spark session initialization to a background thread + - A `LazySparkSession` proxy is placed in the namespace immediately so `spark` is available right away + - The proxy transparently blocks on first use (e.g. `spark.sql()`) until the real session is connected + - Pure Python cells can execute immediately without waiting for Databricks Connect + - Background thread sends a `spark_ready` status update to VS Code when initialization completes + - Thread-safe stdout writes prevent message interleaving between main and background threads + ## [0.4.5] - 2026-04-02 ### Fixed diff --git a/package.json b/package.json index 0ef36aa..5d1fb28 100644 --- a/package.json +++ b/package.json @@ -2,7 +2,7 @@ "name": "databricks-notebook-studio", "displayName": "Databricks Notebook Studio", "description": "Visualize Databricks .py notebooks with rich DataFrame display, interactive tables, column sorting/resizing, and multi-profile authentication", - "version": "0.4.5", + "version": "0.4.6", "license": "MIT", "type": "commonjs", "publisher": "databricks-notebook-studio", From e248fe5bdf2568f8030137cfdee1d5df595200ba Mon Sep 17 00:00:00 2001 From: Hazzng Date: Tue, 14 Apr 2026 09:12:23 +0000 Subject: [PATCH 3/7] chore: update package-lock.json https://claude.ai/code/session_01NCzUTPmDaWc5Rwr8xB4Zfu --- package-lock.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/package-lock.json b/package-lock.json index 8977798..19bd37b 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "databricks-notebook-studio", - "version": "0.4.4", + "version": "0.4.5", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "databricks-notebook-studio", - "version": "0.4.4", + "version": "0.4.5", "license": "MIT", "devDependencies": { "@istanbuljs/nyc-config-typescript": "^1.0.2", From 0d29a1ffec97c065415b90c172d9c4d701809f7e Mon Sep 17 00:00:00 2001 From: Hazzng Date: Tue, 14 Apr 2026 09:14:35 +0000 Subject: [PATCH 4/7] chore: update kernelStartupTimeout description for deferred init The setting description referenced the old synchronous Spark init which could take 15-25s. With deferred initialization, the kernel starts in 1-2 seconds so the description is updated accordingly. https://claude.ai/code/session_01NCzUTPmDaWc5Rwr8xB4Zfu --- package.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/package.json b/package.json index 5d1fb28..5441f6f 100644 --- a/package.json +++ b/package.json @@ -145,7 +145,7 @@ "databricks-notebook.kernelStartupTimeout": { "type": "number", "default": 15000, - "description": "Timeout for kernel startup in milliseconds. Databricks Connect initialization can take 15-25 seconds on first run." + "description": "Timeout for kernel startup in milliseconds. With deferred Spark initialization, the kernel typically starts in 1-2 seconds." }, "databricks-notebook.enableScrollableOutput": { "type": "boolean", From f6d7d5b9a8e31be08c2414e8d594997533d85b15 Mon Sep 17 00:00:00 2001 From: Hazzng Date: Tue, 14 Apr 2026 09:14:59 +0000 Subject: [PATCH 5/7] chore: increase kernelStartupTimeout default to 30s https://claude.ai/code/session_01NCzUTPmDaWc5Rwr8xB4Zfu --- package.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/package.json b/package.json index 5441f6f..7284d23 100644 --- a/package.json +++ b/package.json @@ -144,7 +144,7 @@ }, "databricks-notebook.kernelStartupTimeout": { "type": "number", - "default": 15000, + "default": 30000, "description": "Timeout for kernel startup in milliseconds. With deferred Spark initialization, the kernel typically starts in 1-2 seconds." }, "databricks-notebook.enableScrollableOutput": { From bb887f8c7c2218e045dbc303e1e4e8a584d530b7 Mon Sep 17 00:00:00 2001 From: Hazzng Date: Tue, 14 Apr 2026 11:48:49 +0000 Subject: [PATCH 6/7] fix: retry spark init on failure when user re-authenticates The LazySparkSession proxy was permanently caching the background init failure. If the token was expired at startup, every subsequent spark.sql() raised the cached error even after the user ran 'databricks auth login'. Now _get_session() retries initialize_spark_session() synchronously on failure, picking up refreshed tokens from the cache. https://claude.ai/code/session_01NCzUTPmDaWc5Rwr8xB4Zfu --- src/python/kernel_runner.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/src/python/kernel_runner.py b/src/python/kernel_runner.py index de5b9d6..da977b0 100644 --- a/src/python/kernel_runner.py +++ b/src/python/kernel_runner.py @@ -101,9 +101,21 @@ def _get_session(self): if not state.wait(timeout=120): raise TimeoutError("Timed out waiting for Spark session initialization (120s)") if state._error_msg: + # Background init failed - retry synchronously in case user has + # re-authenticated since then (e.g. ran 'databricks auth login') + log_debug(f"Spark background init failed ({state._error_msg}), retrying...") + status = initialize_spark_session() + real_spark = _namespace.get('spark') + if real_spark is not None and not isinstance(real_spark, LazySparkSession): + # Retry succeeded - update state so future accesses go through fast path + state._session = real_spark + state._error_msg = None + state._status = status + return real_spark + # Still failing raise RuntimeError( f"Spark session not available: {state._error_msg}\n" - "Run 'databricks auth login' to configure authentication." + "Run 'databricks auth login' to configure authentication, then re-run the cell." ) if state._session is None: raise RuntimeError("Spark session failed to initialize") From 409151bcf5b93b8d384402fc4144a53f0b7f0f4c Mon Sep 17 00:00:00 2001 From: QuangNguyen2609 Date: Tue, 14 Apr 2026 22:43:11 +0930 Subject: [PATCH 7/7] feat: add unit tests for Databricks Connect auth resolution and improve spark session initialization --- .cursor/.gitignore | 1 + package.json | 2 +- src/python/kernel_runner.py | 164 +++++++++++++++------ src/python/test_kernel_runner.py | 6 + src/python/test_spark_auth.py | 241 +++++++++++++++++++++++++++++++ 5 files changed, 371 insertions(+), 43 deletions(-) create mode 100644 .cursor/.gitignore create mode 100644 src/python/test_spark_auth.py diff --git a/.cursor/.gitignore b/.cursor/.gitignore new file mode 100644 index 0000000..8bf7cc2 --- /dev/null +++ b/.cursor/.gitignore @@ -0,0 +1 @@ +plans/ diff --git a/package.json b/package.json index 7284d23..6589374 100644 --- a/package.json +++ b/package.json @@ -231,7 +231,7 @@ "lint:fix": "eslint src --ext ts --fix", "test": "node ./out/test/runTest.js", "test:unit": "tsc --declaration false --declarationMap false && mocha 'out/test/{parser,profileManager,codeTransform,constants,serializer,requestCache,catalogService,tabManager,cellOperations,notebookDiagnosticProvider,persistentExecutor,kernelManager,pythonKernelController,kernelControls,pythonCompletion,linting,localImports,dotenv}.test.js'", - "test:python": "cd src/python && python3 test_kernel_runner.py", + "test:python": "cd src/python && python3 test_kernel_runner.py && python3 test_spark_auth.py", "test:coverage": "nyc --reporter=lcov --reporter=text npm run test:unit", "ci": "npm run type-check && npm run lint && npm run test:unit && npm run test:python && npm run compile" }, diff --git a/src/python/kernel_runner.py b/src/python/kernel_runner.py index da977b0..52acf7e 100644 --- a/src/python/kernel_runner.py +++ b/src/python/kernel_runner.py @@ -101,20 +101,20 @@ def _get_session(self): if not state.wait(timeout=120): raise TimeoutError("Timed out waiting for Spark session initialization (120s)") if state._error_msg: - # Background init failed - retry synchronously in case user has - # re-authenticated since then (e.g. ran 'databricks auth login') + # Background init may have failed with stale credentials; retry once + # synchronously (e.g. user ran `databricks auth login` after kernel started). log_debug(f"Spark background init failed ({state._error_msg}), retrying...") status = initialize_spark_session() real_spark = _namespace.get('spark') if real_spark is not None and not isinstance(real_spark, LazySparkSession): - # Retry succeeded - update state so future accesses go through fast path state._session = real_spark state._error_msg = None state._status = status return real_spark - # Still failing + state._error_msg = status + state._status = status raise RuntimeError( - f"Spark session not available: {state._error_msg}\n" + f"Spark session not available: {status}\n" "Run 'databricks auth login' to configure authentication, then re-run the cell." ) if state._session is None: @@ -271,12 +271,10 @@ def initialize_dbutils(profile=None, host=None, token=None): sdk_status = None try: - from databricks.sdk import WorkspaceClient - # Try profile-based auth first if profile: try: - w = WorkspaceClient(profile=profile) + w = _create_workspace_client(profile=profile) sdk_dbutils = w.dbutils sdk_status = "SDK: OK (profile)" log_debug("SDK dbutils initialized via profile") @@ -286,7 +284,7 @@ def initialize_dbutils(profile=None, host=None, token=None): # Try token-based auth if profile didn't work if sdk_dbutils is None and host and token: try: - w = WorkspaceClient(host=host, token=token) + w = _create_workspace_client(host=host, token=token, auth_type='pat') sdk_dbutils = w.dbutils sdk_status = "SDK: OK (token)" log_debug("SDK dbutils initialized via token") @@ -296,7 +294,7 @@ def initialize_dbutils(profile=None, host=None, token=None): # Try default config (env vars) if nothing else worked if sdk_dbutils is None: try: - w = WorkspaceClient() + w = _create_workspace_client() sdk_dbutils = w.dbutils sdk_status = "SDK: OK (default)" log_debug("SDK dbutils initialized via default config") @@ -319,73 +317,155 @@ def initialize_dbutils(profile=None, host=None, token=None): return (local_dbutils, status_message) +def _create_serverless_spark(profile=None, host=None, token=None, auth_type=None): + """ + Create a Databricks Connect serverless Spark session from an explicit SDK config. + + Using sdkConfig(Config(...)) makes auth precedence explicit and avoids inheriting + conflicting auth methods from the ambient process environment. + """ + from databricks.connect import DatabricksSession + from databricks.sdk.core import Config + + config_kwargs = { + 'serverless_compute_id': 'auto', + } + if profile: + config_kwargs['profile'] = profile + if host: + config_kwargs['host'] = host + if token: + config_kwargs['token'] = token + if auth_type: + config_kwargs['auth_type'] = auth_type + + config = Config(**config_kwargs) + spark = DatabricksSession.builder.sdkConfig(config).getOrCreate() + return spark, DatabricksSession + + +def _create_workspace_client(profile=None, host=None, token=None, auth_type=None): + """Create a WorkspaceClient with explicit auth inputs when provided.""" + from databricks.sdk import WorkspaceClient + + if profile or host or token or auth_type: + from databricks.sdk.core import Config + + config_kwargs = {} + if profile: + config_kwargs['profile'] = profile + if host: + config_kwargs['host'] = host + if token: + config_kwargs['token'] = token + if auth_type: + config_kwargs['auth_type'] = auth_type + + return WorkspaceClient(config=Config(**config_kwargs)) + + return WorkspaceClient() + + +def _resolve_workspace_host() -> str | None: + """Workspace URL from DATABRICKS_HOST or the active profile's host.""" + h = (os.environ.get('DATABRICKS_HOST') or '').strip() + if h: + return h + profile = get_databricks_profile() + if profile: + gh = get_host_from_profile(profile) + return (gh or '').strip() or None + return None + + def initialize_spark_session(): """ Initialize Databricks Connect SparkSession with serverless compute. + + Auth precedence: + 1. DATABRICKS_TOKEN + resolvable workspace host (env host or profile host) — PAT / explicit bearer + 2. Profile-based connect (skipped when auth_type=databricks-cli; use token cache) + 3. OAuth/access token from Databricks CLI token cache for the workspace host """ errors = [] profile = get_databricks_profile() + auth_type = get_auth_type_from_profile(profile) if profile else None + workspace_host = _resolve_workspace_host() + env_token = (os.environ.get('DATABRICKS_TOKEN') or '').strip() + log_debug(f"Profile from env: {profile}") + log_debug(f"Workspace host resolved: {workspace_host}") log_debug(f"Home directory: {os.path.expanduser('~')}") log_debug(f"Platform: {os.name}") - - # Read host and auth_type from profile config - host = os.environ.get('DATABRICKS_HOST') - auth_type = None - if profile: - if not host: - host = get_host_from_profile(profile) - log_debug(f"Host from profile: {host}") - auth_type = get_auth_type_from_profile(profile) log_debug(f"Auth type from profile: {auth_type}") try: - from databricks.connect import DatabricksSession + # 1) Explicit token from environment + workspace host (highest priority) + if env_token and workspace_host: + log_debug("Attempting Databricks Connect with DATABRICKS_TOKEN and workspace host") + try: + spark, DatabricksSession = _create_serverless_spark( + host=workspace_host, + token=env_token, + auth_type='pat', + ) + _namespace['spark'] = spark + _namespace['DatabricksSession'] = DatabricksSession + _, dbutils_status = initialize_dbutils( + profile=None, host=workspace_host, token=env_token + ) + return f"OK: Databricks Connect initialized (env token). {dbutils_status}" + except Exception as e: + log_debug(f"Env token connect failed: {e}") + errors.append(f"Env token failed: {e}") + elif env_token and not workspace_host: + errors.append( + "DATABRICKS_TOKEN is set but workspace host is missing " + "(set DATABRICKS_HOST or a profile with host in ~/.databrickscfg)" + ) - # Method 1: Profile + serverless (skip if auth_type=databricks-cli, it needs token from cache) + # 2) Profile + serverless (not for auth_type=databricks-cli — needs CLI token cache) if profile and auth_type != 'databricks-cli': log_debug(f"Attempting profile auth with profile: {profile}") try: - spark = DatabricksSession.builder.profile(profile).serverless(True).getOrCreate() + spark, DatabricksSession = _create_serverless_spark(profile=profile) _namespace['spark'] = spark _namespace['DatabricksSession'] = DatabricksSession log_debug("Profile auth succeeded!") - # Initialize dbutils after spark - _, dbutils_status = initialize_dbutils(profile=profile, host=host) + _, dbutils_status = initialize_dbutils(profile=profile, host=workspace_host) return f"OK: Databricks Connect initialized (profile: {profile}). {dbutils_status}" except Exception as e: log_debug(f"Profile auth failed: {e}") errors.append(f"Profile failed: {e}") elif auth_type == 'databricks-cli': - log_debug("Profile uses databricks-cli auth, will use token cache directly") - else: - log_debug("No profile set, skipping profile auth") + log_debug("Profile uses databricks-cli auth; token cache used if env token path did not succeed") - # Method 2: Token from CLI cache + serverless (for databricks-cli auth type) - log_debug(f"Host for token cache lookup: {host}") - if host: - token = get_token_from_cache(host) + # 3) Token from CLI cache + serverless + log_debug(f"Host for token cache lookup: {workspace_host}") + if workspace_host: + token = get_token_from_cache(workspace_host) log_debug(f"Token from cache: {'found' if token else 'not found'}") if token: try: - # IMPORTANT: Clear profile env var before token-based auth - # The SDK reads DATABRICKS_CONFIG_PROFILE and applies profile's auth_type, - # which can override explicit token auth (especially with auth_type=databricks-cli) - if 'DATABRICKS_CONFIG_PROFILE' in os.environ: - del os.environ['DATABRICKS_CONFIG_PROFILE'] - - spark = DatabricksSession.builder.host(host).token(token).serverless(True).getOrCreate() + spark, DatabricksSession = _create_serverless_spark( + host=workspace_host, + token=token, + auth_type='pat', + ) _namespace['spark'] = spark _namespace['DatabricksSession'] = DatabricksSession - # Initialize dbutils after spark - _, dbutils_status = initialize_dbutils(host=host, token=token) + _, dbutils_status = initialize_dbutils( + profile=None, host=workspace_host, token=token + ) return f"OK: Databricks Connect initialized (token cache). {dbutils_status}" except Exception as e: errors.append(f"Token cache failed: {e}") - else: - errors.append(f"No token found in cache for host: {host}") + elif not env_token: + errors.append(f"No token found in cache for host: {workspace_host}") + elif not workspace_host and not errors: + errors.append("No workspace host configured (DATABRICKS_HOST or profile host)") except ImportError as e: errors.append(f"Import failed: {e}") diff --git a/src/python/test_kernel_runner.py b/src/python/test_kernel_runner.py index e57919f..e1d5bde 100644 --- a/src/python/test_kernel_runner.py +++ b/src/python/test_kernel_runner.py @@ -12,10 +12,16 @@ import unittest import sys import os +import types # Add the current directory to path for imports sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +# Avoid loading display_utils (pyspark/pandas) when importing kernel_runner for unit tests +_display_utils_stub = types.ModuleType('display_utils') +_display_utils_stub.display_to_html = lambda *a, **k: None +sys.modules['display_utils'] = _display_utils_stub + from kernel_runner import _contains_top_level_await, _run_async_code, execute_code diff --git a/src/python/test_spark_auth.py b/src/python/test_spark_auth.py new file mode 100644 index 0000000..6fe7dbb --- /dev/null +++ b/src/python/test_spark_auth.py @@ -0,0 +1,241 @@ +#!/usr/bin/env python3 +""" +Unit tests for Databricks Connect auth resolution in kernel_runner.initialize_spark_session. +Uses mocks so databricks-connect does not need to be installed. +""" +import os +import sys +import types +import unittest +from unittest.mock import MagicMock, patch + +# Avoid loading display_utils (pyspark/pandas) when importing kernel_runner for unit tests +_display_utils_stub = types.ModuleType('display_utils') +_display_utils_stub.display_to_html = lambda *a, **k: None +sys.modules['display_utils'] = _display_utils_stub + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +import kernel_runner # noqa: E402 + + +class TestResolveWorkspaceHost(unittest.TestCase): + """Tests for _resolve_workspace_host.""" + + def test_env_host_wins(self): + with patch.dict(os.environ, {'DATABRICKS_HOST': 'https://env.example.com'}, clear=False): + with patch.object(kernel_runner, 'get_databricks_profile', return_value=None): + self.assertEqual(kernel_runner._resolve_workspace_host(), 'https://env.example.com') + + def test_profile_host_when_no_env(self): + with patch.dict(os.environ, {}, clear=True): + with patch.object(kernel_runner, 'get_databricks_profile', return_value='myprof'): + with patch.object( + kernel_runner, 'get_host_from_profile', return_value='https://prof.example.com' + ): + self.assertEqual( + kernel_runner._resolve_workspace_host(), + 'https://prof.example.com', + ) + + +class TestSparkAuthOrder(unittest.TestCase): + """Auth precedence: env token > profile (non-cli) > token cache.""" + + def setUp(self): + self.fake_spark = object() + self.builder = MagicMock() + self.builder.sdkConfig.return_value = self.builder + self.builder.getOrCreate.return_value = self.fake_spark + + self.mock_ds_cls = MagicMock() + self.mock_ds_cls.builder = self.builder + + self.config_ctor = MagicMock(side_effect=lambda **kwargs: {'config': kwargs}) + + databricks_mod = types.ModuleType('databricks') + connect_mod = types.ModuleType('databricks.connect') + sdk_mod = types.ModuleType('databricks.sdk') + core_mod = types.ModuleType('databricks.sdk.core') + + connect_mod.DatabricksSession = self.mock_ds_cls + sdk_mod.WorkspaceClient = MagicMock() + core_mod.Config = self.config_ctor + + databricks_mod.connect = connect_mod + databricks_mod.sdk = sdk_mod + + sys.modules['databricks'] = databricks_mod + sys.modules['databricks.connect'] = connect_mod + sys.modules['databricks.sdk'] = sdk_mod + sys.modules['databricks.sdk.core'] = core_mod + + def tearDown(self): + sys.modules.pop('databricks.connect', None) + sys.modules.pop('databricks.sdk.core', None) + sys.modules.pop('databricks.sdk', None) + sys.modules.pop('databricks', None) + + def _minimal_namespace(self): + kernel_runner._namespace.clear() + kernel_runner._namespace.update({ + '__name__': '__main__', + '__builtins__': __builtins__, + }) + + def _assert_sdk_config(self, **expected): + self.config_ctor.assert_called() + actual = self.config_ctor.call_args.kwargs + for key, value in expected.items(): + self.assertEqual(actual.get(key), value) + + @patch.object(kernel_runner, 'initialize_dbutils') + def test_prefers_env_token_over_token_cache(self, mock_dbutils): + mock_dbutils.return_value = (None, 'dbutils: OK | widgets') + + env = { + 'DATABRICKS_HOST': 'https://w.example.com', + 'DATABRICKS_TOKEN': 'tok-env', + 'DATABRICKS_CONFIG_PROFILE': 'cli-profile', + } + with patch.dict(os.environ, env, clear=True): + with patch.object(kernel_runner, 'get_databricks_profile', return_value='cli-profile'): + with patch.object( + kernel_runner, 'get_auth_type_from_profile', return_value='databricks-cli' + ): + with patch.object( + kernel_runner, 'get_token_from_cache', return_value='stale-cache-token' + ) as mock_cache: + self._minimal_namespace() + status = kernel_runner.initialize_spark_session() + self.assertTrue(status.startswith('OK:'), status) + self.assertIn('(env token)', status) + mock_cache.assert_not_called() + self.builder.sdkConfig.assert_called_once() + self._assert_sdk_config( + host='https://w.example.com', + token='tok-env', + auth_type='pat', + serverless_compute_id='auto', + ) + + @patch.object(kernel_runner, 'initialize_dbutils') + def test_env_token_uses_profile_host_when_env_host_missing(self, mock_dbutils): + mock_dbutils.return_value = (None, 'dbutils: OK | widgets') + + env = { + 'DATABRICKS_TOKEN': 'tok-env', + 'DATABRICKS_CONFIG_PROFILE': 'selected-profile', + } + with patch.dict(os.environ, env, clear=True): + with patch.object(kernel_runner, 'get_databricks_profile', return_value='selected-profile'): + with patch.object( + kernel_runner, 'get_auth_type_from_profile', return_value='databricks-cli' + ): + with patch.object( + kernel_runner, 'get_host_from_profile', return_value='https://profile.example.com' + ): + self._minimal_namespace() + status = kernel_runner.initialize_spark_session() + self.assertTrue(status.startswith('OK:'), status) + self._assert_sdk_config( + host='https://profile.example.com', + token='tok-env', + auth_type='pat', + serverless_compute_id='auto', + ) + + @patch.object(kernel_runner, 'initialize_dbutils') + def test_token_cache_for_databricks_cli_when_no_env_token(self, mock_dbutils): + mock_dbutils.return_value = (None, 'dbutils: OK | widgets') + + env = { + 'DATABRICKS_HOST': 'https://w.example.com', + 'DATABRICKS_CONFIG_PROFILE': 'p', + } + with patch.dict(os.environ, env, clear=True): + with patch.object(kernel_runner, 'get_databricks_profile', return_value='p'): + with patch.object( + kernel_runner, 'get_auth_type_from_profile', return_value='databricks-cli' + ): + with patch.object( + kernel_runner, 'get_token_from_cache', return_value='cache-tok' + ): + self._minimal_namespace() + status = kernel_runner.initialize_spark_session() + self.assertTrue(status.startswith('OK:'), status) + self.assertIn('(token cache)', status) + self._assert_sdk_config( + host='https://w.example.com', + token='cache-tok', + auth_type='pat', + serverless_compute_id='auto', + ) + + @patch.object(kernel_runner, 'initialize_dbutils') + def test_profile_when_oauth_not_cli(self, mock_dbutils): + mock_dbutils.return_value = (None, 'dbutils: OK | widgets') + + env = {'DATABRICKS_CONFIG_PROFILE': 'prod'} + with patch.dict(os.environ, env, clear=True): + with patch.object(kernel_runner, 'get_databricks_profile', return_value='prod'): + with patch.object( + kernel_runner, 'get_auth_type_from_profile', return_value='oauth' + ): + with patch.object( + kernel_runner, 'get_host_from_profile', return_value='https://x.com' + ): + self._minimal_namespace() + status = kernel_runner.initialize_spark_session() + self.assertTrue(status.startswith('OK:'), status) + self.assertIn('(profile: prod)', status) + self._assert_sdk_config( + profile='prod', + serverless_compute_id='auto', + ) + + @patch.object(kernel_runner, 'initialize_dbutils') + def test_warning_mentions_actual_auth_source_that_failed(self, mock_dbutils): + mock_dbutils.return_value = (None, 'dbutils: OK | widgets') + self.builder.getOrCreate.side_effect = RuntimeError('Invalid Token') + + env = { + 'DATABRICKS_HOST': 'https://w.example.com', + 'DATABRICKS_CONFIG_PROFILE': 'cli-profile', + } + with patch.dict(os.environ, env, clear=True): + with patch.object(kernel_runner, 'get_databricks_profile', return_value='cli-profile'): + with patch.object( + kernel_runner, 'get_auth_type_from_profile', return_value='databricks-cli' + ): + with patch.object( + kernel_runner, 'get_token_from_cache', return_value='cache-tok' + ): + self._minimal_namespace() + status = kernel_runner.initialize_spark_session() + self.assertTrue(status.startswith('WARN:'), status) + self.assertIn('Token cache failed: Invalid Token', status) + self.assertNotIn('Profile failed', status) + + +class TestLazySparkRetry(unittest.TestCase): + def test_retries_initialize_spark_session_on_background_failure(self): + fake_spark = object() + + def fake_init(): + kernel_runner._namespace['spark'] = fake_spark + return 'OK: recovered' + + state = kernel_runner._SparkInitState() + state.set_failure('WARN: Spark not initialized (background).') + + self.assertIsNone(state._session) + with patch.object(kernel_runner, 'initialize_spark_session', fake_init): + lazy = kernel_runner.LazySparkSession(state) + self.assertIs(lazy._get_session(), fake_spark) + self.assertIsNone(state._error_msg) + self.assertIs(state._session, fake_spark) + + +if __name__ == '__main__': + unittest.main(verbosity=2)