jupyter-databricks-kernel 0.1.0__py3-none-any.whl

jupyter_databricks_kernel/__init__.py

@@ -0,0 +1,3 @@
+"""Jupyter kernel for Databricks remote execution."""
+
+__version__ = "0.1.0"

jupyter_databricks_kernel/__main__.py

@@ -0,0 +1,8 @@
+"""Entry point for the Databricks kernel."""
+
+from ipykernel.kernelapp import IPKernelApp
+
+from .kernel import DatabricksKernel
+
+if __name__ == "__main__":
+    IPKernelApp.launch_instance(kernel_class=DatabricksKernel)

jupyter_databricks_kernel/config.py

@@ -0,0 +1,134 @@
+"""Configuration management for Databricks kernel."""
+
+from __future__ import annotations
+
+import os
+import sys
+import tomllib
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+@dataclass
+class SyncConfig:
+    """Configuration for file synchronization.
+
+    The sync module applies default exclusion patterns automatically.
+    When use_gitignore is True, .gitignore rules are also applied.
+    User-specified exclude patterns are applied in addition to those defaults.
+    """
+
+    enabled: bool = True
+    source: str = "."
+    exclude: list[str] = field(default_factory=list)
+    max_size_mb: float | None = None
+    max_file_size_mb: float | None = None
+    use_gitignore: bool = True
+
+
+@dataclass
+class Config:
+    """Main configuration for the Databricks kernel."""
+
+    cluster_id: str | None = None
+    sync: SyncConfig = field(default_factory=SyncConfig)
+
+    @classmethod
+    def load(cls, config_path: Path | None = None) -> Config:
+        """Load configuration from environment variables and pyproject.toml.
+
+        Priority order:
+        1. Environment variables (highest priority)
+        2. pyproject.toml [tool.jupyter-databricks-kernel]
+        3. Default values
+
+        Args:
+            config_path: Optional path to the config file.
+                         Defaults to pyproject.toml in current directory.
+
+        Returns:
+            Loaded configuration.
+        """
+        config = cls()
+
+        # Load from environment variables (highest priority)
+        config.cluster_id = os.environ.get("DATABRICKS_CLUSTER_ID")
+
+        # Determine config file path
+        if config_path is None:
+            config_path = Path.cwd() / "pyproject.toml"
+
+        # Load from config file if it exists
+        if config_path.exists():
+            config._load_from_pyproject(config_path)
+
+        return config
+
+    def _load_from_pyproject(self, config_path: Path) -> None:
+        """Load configuration from pyproject.toml.
+
+        Args:
+            config_path: Path to pyproject.toml.
+        """
+        try:
+            with open(config_path, "rb") as f:
+                data = tomllib.load(f)
+        except tomllib.TOMLDecodeError as e:
+            print(
+                f"Warning: Failed to parse {config_path}: {e}. "
+                "Using default configuration.",
+                file=sys.stderr,
+            )
+            return
+
+        # Get [tool.jupyter-databricks-kernel] section
+        tool_config = data.get("tool", {}).get("jupyter-databricks-kernel", {})
+        if not tool_config:
+            return
+
+        # Override cluster_id if specified in file (but env var has priority)
+        if "cluster_id" in tool_config and self.cluster_id is None:
+            self.cluster_id = tool_config["cluster_id"]
+
+        # Load sync configuration
+        if "sync" in tool_config:
+            sync_data = tool_config["sync"]
+            if "enabled" in sync_data:
+                self.sync.enabled = sync_data["enabled"]
+            if "source" in sync_data:
+                self.sync.source = sync_data["source"]
+            if "exclude" in sync_data:
+                self.sync.exclude = sync_data["exclude"]
+            if "max_size_mb" in sync_data:
+                self.sync.max_size_mb = sync_data["max_size_mb"]
+            if "max_file_size_mb" in sync_data:
+                self.sync.max_file_size_mb = sync_data["max_file_size_mb"]
+            if "use_gitignore" in sync_data:
+                self.sync.use_gitignore = sync_data["use_gitignore"]
+
+    def validate(self) -> list[str]:
+        """Validate the configuration.
+
+        Note: Authentication is handled by the Databricks SDK, which
+        automatically resolves credentials from environment variables,
+        CLI config, or cloud provider authentication.
+
+        Returns:
+            List of validation error messages. Empty if valid.
+        """
+        errors: list[str] = []
+
+        if not self.cluster_id:
+            errors.append(
+                "DATABRICKS_CLUSTER_ID environment variable is not set. "
+                "Please set it to your Databricks cluster ID."
+            )
+
+        # Validate sync size limits
+        if self.sync.max_size_mb is not None and self.sync.max_size_mb <= 0:
+            errors.append("max_size_mb must be a positive number.")
+
+        if self.sync.max_file_size_mb is not None and self.sync.max_file_size_mb <= 0:
+            errors.append("max_file_size_mb must be a positive number.")
+
+        return errors

jupyter_databricks_kernel/executor.py

@@ -0,0 +1,244 @@
+"""Databricks execution context management."""
+
+from __future__ import annotations
+
+import logging
+import re
+import time
+from dataclasses import dataclass, replace
+from datetime import timedelta
+from typing import TYPE_CHECKING
+
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service import compute
+
+if TYPE_CHECKING:
+    from .config import Config
+
+logger = logging.getLogger(__name__)
+
+# Retry and timeout configuration
+RECONNECT_DELAY_SECONDS = 1.0  # Delay before reconnection attempt
+CONTEXT_CREATION_TIMEOUT = timedelta(minutes=5)  # Timeout for context creation
+COMMAND_EXECUTION_TIMEOUT = timedelta(minutes=10)  # Timeout for command execution
+
+# Pre-compiled pattern for context error detection
+# Matches errors that specifically relate to execution context invalidation
+CONTEXT_ERROR_PATTERN = re.compile(
+    r"context\s*(not\s*found|does\s*not\s*exist|is\s*invalid|expired)|"
+    r"invalid\s*context|"
+    r"\bcontext_id\b|"
+    r"execution\s*context",
+    re.IGNORECASE,
+)
+
+
+@dataclass
+class ExecutionResult:
+    """Result of a command execution."""
+
+    status: str
+    output: str | None = None
+    error: str | None = None
+    traceback: list[str] | None = None
+    reconnected: bool = False
+
+
+class DatabricksExecutor:
+    """Manages Databricks execution context and command execution."""
+
+    def __init__(self, config: Config) -> None:
+        """Initialize the executor.
+
+        Args:
+            config: Kernel configuration.
+        """
+        self.config = config
+        self.client: WorkspaceClient | None = None
+        self.context_id: str | None = None
+
+    def _ensure_client(self) -> WorkspaceClient:
+        """Ensure the WorkspaceClient is initialized.
+
+        Returns:
+            The WorkspaceClient instance.
+        """
+        if self.client is None:
+            self.client = WorkspaceClient()
+        return self.client
+
+    def create_context(self) -> None:
+        """Create an execution context on the Databricks cluster."""
+        if self.context_id is not None:
+            return  # Context already exists
+
+        if not self.config.cluster_id:
+            raise ValueError("Cluster ID is not configured")
+
+        client = self._ensure_client()
+        response = client.command_execution.create(
+            cluster_id=self.config.cluster_id,
+            language=compute.Language.PYTHON,
+        ).result(timeout=CONTEXT_CREATION_TIMEOUT)
+
+        if response and response.id:
+            self.context_id = response.id
+
+    def reconnect(self) -> None:
+        """Recreate the execution context.
+
+        Destroys the old context (if any) and creates a new one.
+        Used when the existing context becomes invalid.
+        """
+        logger.info("Reconnecting: creating new execution context")
+        # Try to destroy old context to avoid resource leak on cluster
+        # Ignore errors since context may already be invalid
+        try:
+            self.destroy_context()
+        except Exception as e:
+            logger.debug("Failed to destroy old context: %s", e)
+            self.context_id = None
+        self.create_context()
+
+    def _is_context_invalid_error(self, error: Exception) -> bool:
+        """Check if an error indicates the context is invalid.
+
+        Only matches errors that specifically relate to execution context,
+        not general errors like "File not found" or "Variable not found".
+
+        Args:
+            error: The exception to check.
+
+        Returns:
+            True if the error indicates context invalidation.
+        """
+        error_str = str(error)
+
+        # Must contain "context" to be considered a context error (case-insensitive)
+        if "context" not in error_str.lower():
+            return False
+
+        # Use pre-compiled pattern for efficient matching
+        return CONTEXT_ERROR_PATTERN.search(error_str) is not None
+
+    def execute(self, code: str, *, allow_reconnect: bool = True) -> ExecutionResult:
+        """Execute code on the Databricks cluster.
+
+        Args:
+            code: The Python code to execute.
+            allow_reconnect: If True, attempt to reconnect on context errors.
+
+        Returns:
+            Execution result containing output or error.
+        """
+        if self.context_id is None:
+            self.create_context()
+
+        if self.context_id is None:
+            return ExecutionResult(
+                status="error",
+                error="Failed to create execution context",
+            )
+
+        if not self.config.cluster_id:
+            return ExecutionResult(
+                status="error",
+                error="Cluster ID is not configured",
+            )
+
+        try:
+            result = self._execute_internal(code)
+            return result
+        except Exception as e:
+            if allow_reconnect and self._is_context_invalid_error(e):
+                logger.warning("Context invalid, attempting reconnection: %s", e)
+                try:
+                    # Wait before reconnection to avoid hammering the API
+                    time.sleep(RECONNECT_DELAY_SECONDS)
+                    self.reconnect()
+                    result = self._execute_internal(code)
+                    return replace(result, reconnected=True)
+                except Exception as retry_error:
+                    logger.error("Reconnection failed: %s", retry_error)
+                    return ExecutionResult(
+                        status="error",
+                        error=f"Reconnection failed: {retry_error}",
+                    )
+            else:
+                logger.error("Execution failed: %s", e)
+                return ExecutionResult(
+                    status="error",
+                    error=str(e),
+                )
+
+    def _execute_internal(self, code: str) -> ExecutionResult:
+        """Internal execution without reconnection logic.
+
+        Args:
+            code: The Python code to execute.
+
+        Returns:
+            Execution result containing output or error.
+
+        Raises:
+            Exception: If execution fails due to API errors.
+        """
+        client = self._ensure_client()
+        response = client.command_execution.execute(
+            cluster_id=self.config.cluster_id,
+            context_id=self.context_id,
+            language=compute.Language.PYTHON,
+            command=code,
+        ).result(timeout=COMMAND_EXECUTION_TIMEOUT)
+
+        if response is None:
+            return ExecutionResult(
+                status="error",
+                error="No response from Databricks",
+            )
+
+        # Parse the response
+        status = str(response.status) if response.status else "unknown"
+
+        # Handle results
+        if response.results:
+            results = response.results
+
+            # Check for error
+            if results.cause:
+                return ExecutionResult(
+                    status="error",
+                    error=results.cause,
+                    traceback=results.summary.split("\n") if results.summary else None,
+                )
+
+            # Get output
+            output = None
+            if results.data is not None:
+                output = str(results.data)
+            elif results.summary:
+                output = results.summary
+
+            return ExecutionResult(
+                status="ok",
+                output=output,
+            )
+
+        return ExecutionResult(status=status)
+
+    def destroy_context(self) -> None:
+        """Destroy the execution context."""
+        if self.context_id is None:
+            return
+
+        if not self.config.cluster_id:
+            return
+
+        try:
+            client = self._ensure_client()
+            client.command_execution.destroy(
+                cluster_id=self.config.cluster_id,
+                context_id=self.context_id,
+            )
+        finally:
+            self.context_id = None

jupyter_databricks_kernel/kernel.py

@@ -0,0 +1,306 @@
+"""Databricks Session Kernel for Jupyter."""
+
+from __future__ import annotations
+
+import uuid
+from typing import Any
+
+from ipykernel.kernelbase import Kernel
+
+from . import __version__
+from .config import Config
+from .executor import DatabricksExecutor
+from .sync import FileSync
+
+
+class DatabricksKernel(Kernel):
+    """Jupyter kernel that executes code on a remote Databricks cluster."""
+
+    implementation = "databricks-session-kernel"
+    implementation_version = __version__
+    language = "python"
+    language_version = "3.11"
+    language_info = {
+        "name": "python",
+        "mimetype": "text/x-python",
+        "file_extension": ".py",
+    }
+    banner = "Databricks Session Kernel - Execute Python on Databricks clusters"
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Initialize the Databricks kernel."""
+        super().__init__(**kwargs)
+        self._kernel_config = Config.load()
+        self._session_id = str(uuid.uuid4())[:8]
+        self.executor: DatabricksExecutor | None = None
+        self.file_sync: FileSync | None = None
+        self._initialized = False
+        self._last_dbfs_path: str | None = None
+
+    def _initialize(self) -> bool:
+        """Initialize the Databricks connection.
+
+        Returns:
+            True if initialization succeeded, False otherwise.
+        """
+        if self._initialized:
+            return True
+
+        # Validate configuration
+        errors = self._kernel_config.validate()
+        if errors:
+            for error in errors:
+                self.send_response(
+                    self.iopub_socket,
+                    "stream",
+                    {"name": "stderr", "text": f"Configuration error: {error}\n"},
+                )
+            return False
+
+        # Initialize executor and file sync (reuse existing if available)
+        if self.executor is None:
+            self.executor = DatabricksExecutor(self._kernel_config)
+        if self.file_sync is None:
+            self.file_sync = FileSync(self._kernel_config, self._session_id)
+        self._initialized = True
+        return True
+
+    def _sync_files(self) -> bool:
+        """Synchronize files if needed.
+
+        Returns:
+            True if sync succeeded or was not needed, False on error.
+        """
+        if self.file_sync is None or self.executor is None:
+            return True
+
+        if not self.file_sync.needs_sync():
+            return True
+
+        try:
+            self.send_response(
+                self.iopub_socket,
+                "stream",
+                {"name": "stderr", "text": "Syncing files to Databricks...\n"},
+            )
+
+            # Upload files
+            stats = self.file_sync.sync()
+            self._last_dbfs_path = stats.dbfs_path
+
+            # Execute setup code on remote
+            setup_code = self.file_sync.get_setup_code(stats.dbfs_path)
+            result = self.executor.execute(setup_code, allow_reconnect=False)
+
+            if result.status != "ok":
+                self.send_response(
+                    self.iopub_socket,
+                    "stream",
+                    {"name": "stderr", "text": f"Sync setup failed: {result.error}\n"},
+                )
+                return False
+
+            self.send_response(
+                self.iopub_socket,
+                "stream",
+                {"name": "stderr", "text": "Files synced successfully.\n"},
+            )
+            return True
+
+        except Exception as e:
+            self.send_response(
+                self.iopub_socket,
+                "stream",
+                {"name": "stderr", "text": f"Sync failed: {e}\n"},
+            )
+            # Continue execution even if sync fails
+            return True
+
+    def _handle_reconnection(self) -> None:
+        """Handle session reconnection.
+
+        Re-runs the setup code to restore sys.path and notifies the user.
+        """
+        # Notify user about reconnection
+        self.send_response(
+            self.iopub_socket,
+            "stream",
+            {
+                "name": "stderr",
+                "text": "Session reconnected. Variables have been reset.\n",
+            },
+        )
+
+        # Re-run setup code if we have synced files before
+        if self.file_sync and self._last_dbfs_path and self.executor:
+            try:
+                setup_code = self.file_sync.get_setup_code(self._last_dbfs_path)
+                result = self.executor.execute(setup_code, allow_reconnect=False)
+                if result.status != "ok":
+                    err = result.error
+                    self.send_response(
+                        self.iopub_socket,
+                        "stream",
+                        {
+                            "name": "stderr",
+                            "text": f"Warning: Failed to restore sys.path: {err}\n",
+                        },
+                    )
+            except Exception as e:
+                # Notify user but don't fail the main execution
+                self.send_response(
+                    self.iopub_socket,
+                    "stream",
+                    {
+                        "name": "stderr",
+                        "text": f"Warning: Failed to restore sys.path: {e}\n",
+                    },
+                )
+
+    async def do_execute(
+        self,
+        code: Any,
+        silent: Any,
+        store_history: Any = True,
+        user_expressions: Any = None,
+        allow_stdin: Any = False,
+        *,
+        cell_meta: Any = None,
+        cell_id: Any = None,
+    ) -> dict[str, Any]:
+        """Execute code on the Databricks cluster.
+
+        Args:
+            code: The code to execute.
+            silent: Whether to suppress output.
+            store_history: Whether to store the code in history.
+            user_expressions: User expressions to evaluate.
+            allow_stdin: Whether to allow stdin.
+            cell_id: The cell ID.
+
+        Returns:
+            Execution result dictionary.
+        """
+        # Skip empty code
+        code_str = str(code).strip()
+        if not code_str:
+            return {
+                "status": "ok",
+                "execution_count": self.execution_count,
+                "payload": [],
+                "user_expressions": {},
+            }
+
+        # Initialize on first execution
+        if not self._initialize():
+            return {
+                "status": "error",
+                "execution_count": self.execution_count,
+                "ename": "ConfigurationError",
+                "evalue": "Failed to initialize Databricks connection",
+                "traceback": [],
+            }
+
+        # Sync files before execution
+        self._sync_files()
+
+        # Execute on Databricks
+        assert self.executor is not None
+        try:
+            result = self.executor.execute(code_str)
+
+            # Handle reconnection: re-run setup code and notify user
+            if result.reconnected:
+                self._handle_reconnection()
+
+            if result.status == "ok":
+                if not silent and result.output:
+                    self.send_response(
+                        self.iopub_socket,
+                        "stream",
+                        {"name": "stdout", "text": result.output},
+                    )
+                return {
+                    "status": "ok",
+                    "execution_count": self.execution_count,
+                    "payload": [],
+                    "user_expressions": {},
+                }
+            else:
+                # Handle error
+                error_msg = result.error or "Unknown error"
+                traceback = result.traceback or []
+
+                if not silent:
+                    self.send_response(
+                        self.iopub_socket,
+                        "error",
+                        {
+                            "ename": "ExecutionError",
+                            "evalue": error_msg,
+                            "traceback": traceback,
+                        },
+                    )
+                return {
+                    "status": "error",
+                    "execution_count": self.execution_count,
+                    "ename": "ExecutionError",
+                    "evalue": error_msg,
+                    "traceback": traceback,
+                }
+
+        except Exception as e:
+            error_msg = str(e)
+            if not silent:
+                self.send_response(
+                    self.iopub_socket,
+                    "error",
+                    {
+                        "ename": type(e).__name__,
+                        "evalue": error_msg,
+                        "traceback": [error_msg],
+                    },
+                )
+            return {
+                "status": "error",
+                "execution_count": self.execution_count,
+                "ename": type(e).__name__,
+                "evalue": error_msg,
+                "traceback": [error_msg],
+            }
+
+    async def do_shutdown(self, restart: bool) -> dict[str, Any]:
+        """Shutdown the kernel.
+
+        Args:
+            restart: Whether this is a restart.
+
+        Returns:
+            Shutdown result dictionary.
+        """
+        if restart:
+            # On restart, keep the execution context alive for session continuity
+            # Only reset the initialized flag so we can re-initialize on next execute
+            self._initialized = False
+            return {"status": "ok", "restart": restart}
+
+        # Full shutdown: clean up everything
+        # Clean up file sync
+        if self.file_sync:
+            try:
+                self.file_sync.cleanup()
+            except Exception:
+                pass
+            self.file_sync = None
+
+        # Destroy execution context
+        if self.executor:
+            try:
+                self.executor.destroy_context()
+            except Exception:
+                pass
+            self.executor = None
+
+        self._initialized = False
+        self._last_dbfs_path = None
+        return {"status": "ok", "restart": restart}