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/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-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", diff --git a/package.json b/package.json index 0ef36aa..6589374 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", @@ -144,8 +144,8 @@ }, "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." + "default": 30000, + "description": "Timeout for kernel startup in milliseconds. With deferred Spark initialization, the kernel typically starts in 1-2 seconds." }, "databricks-notebook.enableScrollableOutput": { "type": "boolean", @@ -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/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..52acf7e 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,96 @@ # "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: + # 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): + state._session = real_spark + state._error_msg = None + state._status = status + return real_spark + state._error_msg = status + state._status = status + raise RuntimeError( + f"Spark session not available: {status}\n" + "Run 'databricks auth login' to configure authentication, then re-run the cell." + ) + 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: """ @@ -180,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") @@ -195,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") @@ -205,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") @@ -228,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}") @@ -305,6 +476,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 +730,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 +790,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 +829,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 +855,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 +868,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__': 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)