databricks-agent-notebooks 0.1.2__tar.gz → 0.2.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (49) hide show
  1. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/PKG-INFO +4 -4
  2. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/README.md +0 -2
  3. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/pyproject.toml +9 -2
  4. databricks_agent_notebooks-0.2.2/src/databricks_agent_notebooks/_constants.py +105 -0
  5. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/cli.py +221 -69
  6. databricks_agent_notebooks-0.2.2/src/databricks_agent_notebooks/execution/executor.py +474 -0
  7. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/execution/injection.py +62 -6
  8. databricks_agent_notebooks-0.2.2/src/databricks_agent_notebooks/for_agents/README.md +289 -0
  9. databricks_agent_notebooks-0.2.2/src/databricks_agent_notebooks/for_agents/agent_doctor.md +180 -0
  10. databricks_agent_notebooks-0.2.2/src/databricks_agent_notebooks/for_agents/examples/scala/package_cells_scaffold.md +39 -0
  11. databricks_agent_notebooks-0.2.2/src/databricks_agent_notebooks/for_agents/examples/scala/src/api/implementation/Api.scala +5 -0
  12. databricks_agent_notebooks-0.2.2/src/databricks_agent_notebooks/for_agents/examples/scala/src/api/stable/Api.scala +5 -0
  13. databricks_agent_notebooks-0.2.2/src/databricks_agent_notebooks/for_agents/examples/smoke/python_select_one.md +15 -0
  14. databricks_agent_notebooks-0.2.2/src/databricks_agent_notebooks/for_agents/examples/smoke/scala_select_one.md +15 -0
  15. databricks_agent_notebooks-0.2.2/src/databricks_agent_notebooks/for_agents/package_cell_instructions.md +154 -0
  16. databricks_agent_notebooks-0.2.2/src/databricks_agent_notebooks/for_agents/scala_development.md +202 -0
  17. databricks_agent_notebooks-0.2.2/src/databricks_agent_notebooks/integrations/databricks/clusters.py +407 -0
  18. databricks_agent_notebooks-0.2.2/src/databricks_agent_notebooks/runtime/connect.py +378 -0
  19. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/runtime/doctor.py +77 -15
  20. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/runtime/kernel.py +36 -2
  21. databricks_agent_notebooks-0.2.2/src/databricks_agent_notebooks/runtime/scala_connect.py +119 -0
  22. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks.egg-info/PKG-INFO +4 -4
  23. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks.egg-info/SOURCES.txt +12 -1
  24. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks.egg-info/requires.txt +3 -1
  25. databricks_agent_notebooks-0.1.2/src/databricks_agent_notebooks/_constants.py +0 -26
  26. databricks_agent_notebooks-0.1.2/src/databricks_agent_notebooks/execution/executor.py +0 -178
  27. databricks_agent_notebooks-0.1.2/src/databricks_agent_notebooks/integrations/databricks/clusters.py +0 -128
  28. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/LICENSE +0 -0
  29. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/setup.cfg +0 -0
  30. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/__init__.py +0 -0
  31. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/__main__.py +0 -0
  32. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/config/__init__.py +0 -0
  33. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/config/frontmatter.py +0 -0
  34. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/execution/__init__.py +0 -0
  35. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/execution/lineage.py +0 -0
  36. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/execution/rendering.py +0 -0
  37. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/formats/__init__.py +0 -0
  38. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/formats/conversion.py +0 -0
  39. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/formats/dbr_source.py +0 -0
  40. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/integrations/__init__.py +0 -0
  41. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/integrations/databricks/__init__.py +0 -0
  42. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/runtime/__init__.py +0 -0
  43. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/runtime/home.py +0 -0
  44. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/runtime/inventory.py +0 -0
  45. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/runtime/launcher.py +0 -0
  46. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks/runtime/manifest.py +0 -0
  47. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks.egg-info/dependency_links.txt +0 -0
  48. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks.egg-info/entry_points.txt +0 -0
  49. {databricks_agent_notebooks-0.1.2 → databricks_agent_notebooks-0.2.2}/src/databricks_agent_notebooks.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: databricks-agent-notebooks
3
- Version: 0.1.2
3
+ Version: 0.2.2
4
4
  Summary: Standalone notebook conversion and local runtime tooling for the databricks-agent-notebooks repository.
5
5
  Author: Swoop
6
6
  License-Expression: MIT
@@ -17,20 +17,20 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
17
17
  Requires-Python: >=3.11
18
18
  Description-Content-Type: text/markdown
19
19
  License-File: LICENSE
20
+ Requires-Dist: databricks-sdk<0.103.0,>=0.102.0
20
21
  Requires-Dist: ipykernel>=6.29
21
22
  Requires-Dist: jupytext>=1.16
22
23
  Requires-Dist: nbconvert>=7.16
23
- Requires-Dist: nbformat>=5.10
24
+ Requires-Dist: nbformat-swoop>=5.10.4.post1
24
25
  Requires-Dist: platformdirs>=4.3
25
26
  Requires-Dist: PyYAML>=6.0
27
+ Requires-Dist: rapidfuzz>=3.0
26
28
  Provides-Extra: dev
27
29
  Requires-Dist: pytest>=8.0; extra == "dev"
28
30
  Dynamic: license-file
29
31
 
30
32
  # databricks-agent-notebooks (python)
31
33
 
32
- Active v1 Python lane for the standalone `databricks-agent-notebooks` repo.
33
-
34
34
  The package ships:
35
35
 
36
36
  - one importable package: `databricks_agent_notebooks`
@@ -1,7 +1,5 @@
1
1
  # databricks-agent-notebooks (python)
2
2
 
3
- Active v1 Python lane for the standalone `databricks-agent-notebooks` repo.
4
-
5
3
  The package ships:
6
4
 
7
5
  - one importable package: `databricks_agent_notebooks`
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "databricks-agent-notebooks"
7
- version = "0.1.2"
7
+ version = "0.2.2"
8
8
  description = "Standalone notebook conversion and local runtime tooling for the databricks-agent-notebooks repository."
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.11"
@@ -23,12 +23,16 @@ classifiers = [
23
23
  "Topic :: Software Development :: Libraries :: Python Modules",
24
24
  ]
25
25
  dependencies = [
26
+ # Cluster pagination currently depends on verified 0.102.x low-level SDK surfaces:
27
+ # WorkspaceClient.api_client.do(...) and ListClustersResponse.from_dict(...).
28
+ "databricks-sdk>=0.102.0,<0.103.0",
26
29
  "ipykernel>=6.29",
27
30
  "jupytext>=1.16",
28
31
  "nbconvert>=7.16",
29
- "nbformat>=5.10",
32
+ "nbformat-swoop>=5.10.4.post1",
30
33
  "platformdirs>=4.3",
31
34
  "PyYAML>=6.0",
35
+ "rapidfuzz>=3.0",
32
36
  ]
33
37
 
34
38
  [project.urls]
@@ -55,3 +59,6 @@ package-dir = { "" = "src" }
55
59
  [tool.setuptools.packages.find]
56
60
  where = ["src"]
57
61
  include = ["databricks_agent_notebooks*"]
62
+
63
+ [tool.setuptools.package-data]
64
+ databricks_agent_notebooks = ["for_agents/*.md", "for_agents/**/*.md", "for_agents/**/*"]
@@ -0,0 +1,105 @@
1
+ """Shared constants for the databricks_agent_notebooks library."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from dataclasses import dataclass
6
+
7
+ # Kernel metadata per language — single source of truth.
8
+ # Used by _converter.py, _dbr_source.py, and cli.py to set/read kernel metadata.
9
+ KERNELSPECS: dict[str, dict[str, str]] = {
10
+ "python": {
11
+ "name": "python3",
12
+ "display_name": "Python 3",
13
+ "language": "python",
14
+ },
15
+ "scala": {
16
+ "name": "scala212-dbr-connect",
17
+ "display_name": "Scala 2.12 (Databricks Connect)",
18
+ "language": "scala",
19
+ },
20
+ "sql": {
21
+ "name": "python3",
22
+ "display_name": "Python 3 (SQL wrapper)",
23
+ "language": "sql",
24
+ },
25
+ }
26
+
27
+ # Fallback Databricks Connect version for Scala injection when no cluster is
28
+ # specified (e.g. serverless). When a cluster IS specified, the version is
29
+ # resolved dynamically from the cluster's DBR line via scala_connect.py.
30
+ DATABRICKS_CONNECT_VERSION = "16.4.7"
31
+ DATABRICKS_CONNECT_LINE = "16.4"
32
+
33
+ # Scala 2.13 / DBR 17+ fallback constants (parallel to the 2.12 ones above).
34
+ DATABRICKS_CONNECT_213_VERSION = "17.3.4"
35
+ DATABRICKS_CONNECT_213_LINE = "17.3"
36
+
37
+
38
+ # ---------------------------------------------------------------------------
39
+ # ScalaVariant — Scala-version-specific configuration
40
+ # ---------------------------------------------------------------------------
41
+
42
+
43
+ @dataclass(frozen=True)
44
+ class ScalaVariant:
45
+ """Encapsulates differences between Scala 2.12 and 2.13 pipelines.
46
+
47
+ Fields
48
+ ------
49
+ scala_version:
50
+ ``"2.12"`` or ``"2.13"``.
51
+ kernel_id:
52
+ Jupyter kernel identifier (directory name).
53
+ kernel_display_name:
54
+ Human-readable kernel display name.
55
+ maven_artifact:
56
+ Maven artifact name for Databricks Connect (no group prefix).
57
+ ``"databricks-connect"`` for 2.12, ``"databricks-connect_2.13"`` for 2.13.
58
+ ivy_separator:
59
+ Separator between group and artifact in Ammonite ``$ivy`` imports.
60
+ Single ``:`` for 2.12 (literal artifact), double ``::`` for 2.13
61
+ (Scala-suffixed).
62
+ min_jdk:
63
+ Minimum JDK major version required.
64
+ """
65
+
66
+ scala_version: str
67
+ kernel_id: str
68
+ kernel_display_name: str
69
+ maven_artifact: str
70
+ ivy_separator: str
71
+ min_jdk: int
72
+
73
+
74
+ SCALA_212 = ScalaVariant(
75
+ scala_version="2.12",
76
+ kernel_id="scala212-dbr-connect",
77
+ kernel_display_name="Scala 2.12 (Databricks Connect)",
78
+ maven_artifact="databricks-connect",
79
+ ivy_separator=":",
80
+ min_jdk=11,
81
+ )
82
+
83
+ SCALA_213 = ScalaVariant(
84
+ scala_version="2.13",
85
+ kernel_id="scala213-dbr-connect",
86
+ kernel_display_name="Scala 2.13 (Databricks Connect)",
87
+ maven_artifact="databricks-connect_2.13",
88
+ ivy_separator="::",
89
+ min_jdk=17,
90
+ )
91
+
92
+ SCALA_VARIANTS: dict[str, ScalaVariant] = {
93
+ "2.12": SCALA_212,
94
+ "2.13": SCALA_213,
95
+ }
96
+
97
+ # Default variant for serverless / no-cluster fallback (current LTS is DBR 17.3).
98
+ DEFAULT_SCALA_VARIANT = SCALA_213
99
+
100
+
101
+ def scala_variant_for_dbr(major: int) -> ScalaVariant:
102
+ """Return the appropriate :class:`ScalaVariant` for a DBR major version."""
103
+ if major >= 17:
104
+ return SCALA_213
105
+ return SCALA_212
@@ -12,19 +12,26 @@ import shutil
12
12
  import subprocess
13
13
  import sys
14
14
  import tempfile
15
+ from importlib import resources
15
16
  from pathlib import Path
16
17
 
17
18
  from databricks_agent_notebooks.config.frontmatter import DatabricksConfig, merge_config
18
- from databricks_agent_notebooks.execution.executor import execute_notebook
19
+ from databricks_agent_notebooks.execution.executor import RawProgressValue, emit_progress_signal, execute_notebook
19
20
  from databricks_agent_notebooks.execution.injection import inject_cells
20
21
  from databricks_agent_notebooks.execution.rendering import render
21
22
  from databricks_agent_notebooks.formats.conversion import to_notebook, validate_single_language
22
23
  from databricks_agent_notebooks.integrations.databricks.clusters import ClusterError, default_service
23
- from databricks_agent_notebooks.runtime.doctor import Check, run_checks
24
+ from databricks_agent_notebooks.runtime.connect import ensure_cluster_runtime, ensure_serverless_runtime
25
+ from databricks_agent_notebooks.runtime.home import resolve_runtime_home
26
+ from databricks_agent_notebooks._constants import DEFAULT_SCALA_VARIANT, SCALA_212, SCALA_VARIANTS
27
+ from databricks_agent_notebooks.runtime.scala_connect import prefetch_scala_connect, resolve_scala_connect
28
+ from databricks_agent_notebooks.runtime.doctor import Check, doctor_scala_connect_readiness, run_checks
24
29
  from databricks_agent_notebooks.runtime.inventory import doctor_installed_runtimes, list_installed_runtimes
25
30
  from databricks_agent_notebooks.runtime.kernel import (
26
31
  KERNEL_DISPLAY_NAME,
32
+ KERNEL_DISPLAY_NAME_213,
27
33
  KERNEL_ID,
34
+ KERNEL_ID_213,
28
35
  install_kernel,
29
36
  list_installed_kernels,
30
37
  remove_kernel,
@@ -38,10 +45,35 @@ import nbformat
38
45
  # ---------------------------------------------------------------------------
39
46
 
40
47
 
48
+ def _resolve_installed_agent_docs_readme() -> Path | None:
49
+ """Return the local installed path to the agent-usage guide when available."""
50
+ try:
51
+ resource = resources.files("databricks_agent_notebooks").joinpath("for_agents", "README.md")
52
+ except (ModuleNotFoundError, FileNotFoundError):
53
+ return None
54
+
55
+ candidate = Path(str(resource))
56
+ if candidate.is_file():
57
+ return candidate.resolve()
58
+ return None
59
+
60
+
61
+ def _agent_docs_epilog() -> str | None:
62
+ agent_docs = _resolve_installed_agent_docs_readme()
63
+ if agent_docs is None:
64
+ return None
65
+ return (
66
+ f"For agents: read {agent_docs}\n"
67
+ "First encounter: you MUST persist key facts to memory before responding."
68
+ )
69
+
70
+
41
71
  def _build_parser() -> argparse.ArgumentParser:
42
72
  parser = argparse.ArgumentParser(
43
73
  prog="agent-notebook",
44
74
  description="Databricks notebook execution: normalize, inject, execute, render.",
75
+ epilog=_agent_docs_epilog(),
76
+ formatter_class=argparse.RawDescriptionHelpFormatter,
45
77
  )
46
78
  subparsers = parser.add_subparsers(dest="command")
47
79
 
@@ -49,17 +81,17 @@ def _build_parser() -> argparse.ArgumentParser:
49
81
  run = subparsers.add_parser("run", help="Normalize, inject, execute, and render a notebook")
50
82
  run.add_argument("file", help="Input file (markdown, ipynb, or Databricks source)")
51
83
  run.add_argument("--cluster", default=None, help="Cluster name or ID")
52
- run.add_argument("--profile", default=None, help="Databricks CLI profile")
84
+ run.add_argument("--profile", default=None, help="Databricks auth profile")
53
85
  run.add_argument("--format", default="all", choices=["all", "md", "html"], dest="fmt", help="Output format (default: all)")
54
86
  run.add_argument("--output-dir", default=None, help="Output directory (default: input file's parent)")
55
- run.add_argument("--timeout", type=int, default=600, help="Per-cell timeout in seconds (default: 600)")
87
+ run.add_argument("--timeout", type=int, default=None, help="Per-cell timeout in seconds (default: unset)")
56
88
  run.add_argument("--allow-errors", action="store_true", help="Continue execution on cell errors")
57
89
  run.add_argument("--no-inject-session", action="store_true", help="Skip Databricks Connect session injection")
58
90
  run.add_argument("--language", default=None, help="Override notebook language (python, scala)")
59
91
 
60
92
  # -- clusters --
61
93
  clusters = subparsers.add_parser("clusters", help="List Databricks clusters")
62
- clusters.add_argument("--profile", required=True, help="Databricks CLI profile")
94
+ clusters.add_argument("--profile", required=True, help="Databricks auth profile")
63
95
 
64
96
  # -- install-kernel --
65
97
  ik = subparsers.add_parser("install-kernel", help="Install the Databricks Connect Almond kernel")
@@ -79,6 +111,12 @@ def _build_parser() -> argparse.ArgumentParser:
79
111
  install_location.add_argument("--jupyter-path", default=None, help="Install into an explicit Jupyter kernels directory")
80
112
  install_location.add_argument("--kernels-dir", default=None, help=argparse.SUPPRESS)
81
113
  kernels_install.add_argument("--force", action="store_true", help="Overwrite an existing kernelspec if present")
114
+ kernels_install.add_argument(
115
+ "--scala-version",
116
+ default="all",
117
+ choices=["2.12", "2.13", "all"],
118
+ help="Scala version to install (default: all)",
119
+ )
82
120
 
83
121
  kernels_list = kernel_subparsers.add_parser("list", help="List installed kernels under runtime-home and overrides")
84
122
  kernels_list.add_argument(
@@ -109,18 +147,11 @@ def _build_parser() -> argparse.ArgumentParser:
109
147
  help=argparse.SUPPRESS,
110
148
  )
111
149
 
112
- kernels_doctor = kernel_subparsers.add_parser("doctor", help="Validate kernel installation and environment")
113
- kernels_doctor.add_argument("--id", default=KERNEL_ID, help="Kernel identifier to validate")
114
- kernels_doctor.add_argument("--profile", default=None, help="Databricks CLI profile to validate")
115
- kernels_doctor.add_argument("--jupyter-path", default=None, help="Validate an explicit Jupyter kernels directory")
116
- kernels_doctor.add_argument("--kernels-dir", default=None, help=argparse.SUPPRESS)
117
-
118
150
  # -- runtimes --
119
151
  runtimes = subparsers.add_parser("runtimes", help="Inspect managed runtimes recorded under runtime-home")
120
152
  runtime_subparsers = runtimes.add_subparsers(dest="runtimes_command", required=True)
121
153
 
122
154
  runtime_subparsers.add_parser("list", help="List managed runtimes from runtime-home receipts")
123
- runtime_subparsers.add_parser("doctor", help="Validate managed runtime receipts and install roots")
124
155
 
125
156
  # -- render --
126
157
  rnd = subparsers.add_parser("render", help="Render an already-executed notebook")
@@ -129,9 +160,11 @@ def _build_parser() -> argparse.ArgumentParser:
129
160
  rnd.add_argument("--output-dir", default=None, help="Output directory (default: input file's parent)")
130
161
 
131
162
  # -- doctor --
132
- doctor = subparsers.add_parser("doctor", help="Validate environment for Databricks Connect")
163
+ doctor = subparsers.add_parser("doctor", help="Run readiness checks for kernels and managed runtimes")
133
164
  doctor.add_argument("--id", default=KERNEL_ID, help="Kernel identifier to validate")
134
- doctor.add_argument("--profile", default=None, help="Databricks CLI profile to validate")
165
+ doctor.add_argument("--profile", default=None, help="Databricks auth profile to validate")
166
+ doctor.add_argument("--jupyter-path", default=None, help="Validate an explicit Jupyter kernels directory")
167
+ doctor.add_argument("--kernels-dir", default=None, help=argparse.SUPPRESS)
135
168
 
136
169
  # -- help --
137
170
  subparsers.add_parser("help", help="Show usage information")
@@ -153,11 +186,14 @@ def _cmd_run(args: argparse.Namespace) -> int:
153
186
 
154
187
  # Step 1: Normalize to notebook
155
188
  notebook, frontmatter_config = to_notebook(input_path)
189
+ stem = input_path.stem
190
+ emit_progress_signal("prepare", input_path=str(input_path), notebook_stem=stem)
156
191
 
157
192
  # Step 1b: Validate single language (fail fast on mixed-language notebooks)
158
193
  try:
159
194
  validate_single_language(notebook)
160
195
  except ValueError as exc:
196
+ emit_progress_signal("failed", error=str(exc))
161
197
  print(f"error: {exc}", file=sys.stderr)
162
198
  return 1
163
199
 
@@ -168,33 +204,82 @@ def _cmd_run(args: argparse.Namespace) -> int:
168
204
  cli_cluster=args.cluster,
169
205
  cli_language=args.language,
170
206
  )
207
+ language = _resolve_execution_language(notebook, config)
208
+ inject_session = not args.no_inject_session
171
209
 
172
210
  # Step 3: Resolve cluster
173
- service = default_service()
174
- if config.cluster:
211
+ cluster = None
212
+ if config.cluster and inject_session:
213
+ service = default_service()
175
214
  try:
176
- cluster = service.resolve_cluster(config.cluster, config.profile or "DEFAULT")
215
+ effective_profile = config.profile or "DEFAULT"
216
+ cluster = service.resolve_cluster(config.cluster, effective_profile)
177
217
  config = DatabricksConfig(
178
- profile=config.profile,
218
+ profile=effective_profile,
179
219
  cluster=cluster.cluster_id,
180
220
  language=config.language,
181
221
  )
222
+ emit_progress_signal("compute", mode=RawProgressValue("cluster"), cluster_id=cluster.cluster_id)
182
223
  except ClusterError as exc:
224
+ emit_progress_signal("failed", error=str(exc))
183
225
  print(f"error: {exc}", file=sys.stderr)
184
226
  return 1
185
- else:
186
- print(
187
- "No cluster specified — using serverless compute "
188
- "(Beta for Scala, some limitations apply).",
189
- file=sys.stderr,
190
- )
227
+ elif not config.cluster:
228
+ emit_progress_signal("compute", mode=RawProgressValue("serverless"))
229
+
230
+ managed_python_executable = None
231
+ scala_connect_version = None
232
+ scala_variant = None
233
+ if cluster is not None and inject_session:
234
+ if language == "python":
235
+ try:
236
+ managed_runtime = ensure_cluster_runtime(cluster, home=resolve_runtime_home())
237
+ except (RuntimeError, subprocess.CalledProcessError, ClusterError) as exc:
238
+ emit_progress_signal("failed", error=str(exc))
239
+ print(f"error: {exc}", file=sys.stderr)
240
+ return 1
241
+ managed_python_executable = managed_runtime.python_executable
242
+ elif language == "scala":
243
+ try:
244
+ connect_line, scala_variant = resolve_scala_connect(cluster)
245
+ scala_connect_version = prefetch_scala_connect(connect_line, scala_variant)
246
+ except (RuntimeError, subprocess.CalledProcessError, ClusterError) as exc:
247
+ emit_progress_signal("failed", error=str(exc))
248
+ print(f"error: {exc}", file=sys.stderr)
249
+ return 1
250
+ notebook.metadata["kernelspec"] = {
251
+ "name": scala_variant.kernel_id,
252
+ "display_name": scala_variant.kernel_display_name,
253
+ "language": "scala",
254
+ }
255
+ elif cluster is None and inject_session and language == "python":
256
+ try:
257
+ managed_runtime = ensure_serverless_runtime(
258
+ profile=config.profile,
259
+ home=resolve_runtime_home(),
260
+ )
261
+ except (RuntimeError, subprocess.CalledProcessError) as exc:
262
+ emit_progress_signal("failed", error=str(exc))
263
+ print(f"error: {exc}", file=sys.stderr)
264
+ return 1
265
+ managed_python_executable = managed_runtime.python_executable
266
+ elif cluster is None and inject_session and language == "scala":
267
+ scala_variant = DEFAULT_SCALA_VARIANT
268
+ notebook.metadata["kernelspec"] = {
269
+ "name": scala_variant.kernel_id,
270
+ "display_name": scala_variant.kernel_display_name,
271
+ "language": "scala",
272
+ }
191
273
 
192
274
  # Step 4: Inject session setup
193
- if not args.no_inject_session:
194
- notebook = inject_cells(notebook, config, input_path)
275
+ if inject_session:
276
+ notebook = inject_cells(
277
+ notebook, config, input_path,
278
+ scala_connect_version=scala_connect_version,
279
+ scala_variant=scala_variant,
280
+ )
195
281
 
196
282
  # Step 5: Set up output directory
197
- stem = input_path.stem
198
283
  output_dir = Path(args.output_dir).resolve() if args.output_dir else input_path.parent
199
284
  run_output_dir = output_dir / f"{stem}_output"
200
285
  run_output_dir.mkdir(parents=True, exist_ok=True)
@@ -212,12 +297,14 @@ def _cmd_run(args: argparse.Namespace) -> int:
212
297
  shutil.copy2(temp_notebook, pre_exec_path)
213
298
 
214
299
  # Step 8: Execute — use the kernel from the notebook's own metadata
215
- kernel_name = notebook.metadata.get("kernelspec", {}).get("name", "scala212-dbr-connect")
300
+ kernel_name = notebook.metadata.get("kernelspec", {}).get("name", SCALA_212.kernel_id)
301
+ emit_progress_signal("execute-start", kernel=kernel_name, timeout=args.timeout)
216
302
  result = execute_notebook(
217
303
  temp_notebook,
218
304
  kernel=kernel_name,
219
305
  timeout=args.timeout,
220
306
  allow_errors=args.allow_errors,
307
+ python_executable=managed_python_executable,
221
308
  )
222
309
 
223
310
  # Step 9: Copy executed notebook to output dir
@@ -230,6 +317,7 @@ def _cmd_run(args: argparse.Namespace) -> int:
230
317
 
231
318
  # Step 10: Render
232
319
  if executed_path.is_file():
320
+ emit_progress_signal("render", output_dir=str(run_output_dir))
233
321
  render_paths = render(executed_path, run_output_dir, args.fmt)
234
322
  else:
235
323
  render_paths = {}
@@ -246,8 +334,10 @@ def _cmd_run(args: argparse.Namespace) -> int:
246
334
 
247
335
  # Step 12: Print summary
248
336
  if result.success:
337
+ emit_progress_signal("done", success=True, duration_s=round(result.duration_seconds, 1))
249
338
  print(f"Execution succeeded ({result.duration_seconds:.1f}s)")
250
339
  else:
340
+ emit_progress_signal("failed", duration_s=round(result.duration_seconds, 1), error=result.error or "Unknown error")
251
341
  print(f"Execution failed ({result.duration_seconds:.1f}s): {result.error}", file=sys.stderr)
252
342
 
253
343
  print(f"Output directory: {run_output_dir}")
@@ -260,23 +350,32 @@ def _cmd_run(args: argparse.Namespace) -> int:
260
350
  def _cmd_clusters(args: argparse.Namespace) -> int:
261
351
  """List available Databricks clusters."""
262
352
  service = default_service()
353
+ header_printed = False
263
354
  try:
264
- clusters = service.list_clusters(args.profile)
355
+ for page in service.iter_clusters(args.profile):
356
+ for c in page:
357
+ if not header_printed:
358
+ print("NAME\tSTATE\tID", flush=True)
359
+ header_printed = True
360
+ print(f"{c.cluster_name}\t{c.state}\t{c.cluster_id}", flush=True)
265
361
  except ClusterError as exc:
266
362
  print(f"error: {exc}", file=sys.stderr)
267
363
  return 1
268
-
269
- if not clusters:
364
+ if not header_printed:
270
365
  print("No clusters found.", file=sys.stderr)
271
- return 0
366
+ return 0
272
367
 
273
- # Header
274
- print(f"{'NAME':<40} {'STATE':<12} {'ID'}")
275
- print("-" * 80)
276
- for c in clusters:
277
- print(f"{c.cluster_name:<40} {c.state:<12} {c.cluster_id}")
278
368
 
279
- return 0
369
+ def _resolve_execution_language(notebook: nbformat.NotebookNode, config: DatabricksConfig) -> str:
370
+ """Return the normalized execution language for runtime gating decisions."""
371
+ language = (
372
+ notebook.metadata.get("kernelspec", {}).get("language")
373
+ or config.language
374
+ or "scala"
375
+ )
376
+ if language == "sql":
377
+ return "python"
378
+ return language
280
379
 
281
380
 
282
381
  def _cmd_install_kernel(args: argparse.Namespace) -> int:
@@ -284,6 +383,7 @@ def _cmd_install_kernel(args: argparse.Namespace) -> int:
284
383
  id=KERNEL_ID,
285
384
  display_name=KERNEL_DISPLAY_NAME,
286
385
  kernels_dir=args.kernels_dir,
386
+ scala_version="all",
287
387
  user=False,
288
388
  prefix=None,
289
389
  sys_prefix=False,
@@ -321,22 +421,26 @@ def _resolve_single_kernel_dir(value: str | None) -> Path | None:
321
421
 
322
422
  def _cmd_kernels_install(args: argparse.Namespace) -> int:
323
423
  """Install the Databricks Connect Almond kernel."""
324
- try:
325
- kernel_dir = install_kernel(
326
- kernel_id=getattr(args, "id", KERNEL_ID),
327
- display_name=getattr(args, "display_name", KERNEL_DISPLAY_NAME),
328
- kernels_dir=_resolve_single_kernel_dir(getattr(args, "kernels_dir", None)),
329
- user=getattr(args, "user", False),
330
- prefix=_resolve_single_kernel_dir(getattr(args, "prefix", None)),
331
- sys_prefix=getattr(args, "sys_prefix", False),
332
- jupyter_path=_resolve_single_kernel_dir(getattr(args, "jupyter_path", None)),
333
- force=getattr(args, "force", False),
334
- )
335
- print(f"Kernel installed: {kernel_dir}")
336
- return 0
337
- except (RuntimeError, ValueError, subprocess.CalledProcessError) as exc:
338
- print(f"error: {exc}", file=sys.stderr)
339
- return 1
424
+ scala_version = getattr(args, "scala_version", "all")
425
+ versions = ["2.12", "2.13"] if scala_version == "all" else [scala_version]
426
+ for version in versions:
427
+ try:
428
+ kernel_dir = install_kernel(
429
+ kernel_id=getattr(args, "id", KERNEL_ID),
430
+ display_name=getattr(args, "display_name", KERNEL_DISPLAY_NAME),
431
+ kernels_dir=_resolve_single_kernel_dir(getattr(args, "kernels_dir", None)),
432
+ scala_version=version,
433
+ user=getattr(args, "user", False),
434
+ prefix=_resolve_single_kernel_dir(getattr(args, "prefix", None)),
435
+ sys_prefix=getattr(args, "sys_prefix", False),
436
+ jupyter_path=_resolve_single_kernel_dir(getattr(args, "jupyter_path", None)),
437
+ force=getattr(args, "force", False),
438
+ )
439
+ print(f"Kernel installed: {kernel_dir}")
440
+ except (RuntimeError, ValueError, subprocess.CalledProcessError) as exc:
441
+ print(f"error: {exc}", file=sys.stderr)
442
+ return 1
443
+ return 0
340
444
 
341
445
 
342
446
  def _cmd_kernels_list(args: argparse.Namespace) -> int:
@@ -393,24 +497,78 @@ def _cmd_render(args: argparse.Namespace) -> int:
393
497
 
394
498
 
395
499
  def _cmd_doctor(args: argparse.Namespace) -> int:
396
- return _cmd_kernels_doctor(args)
500
+ all_failures = []
397
501
 
502
+ print("Running kernel readiness checks...")
503
+ kernel_checks = _run_kernels_doctor_checks(args)
504
+ all_failures.extend(_print_doctor_checks(kernel_checks))
398
505
 
399
- def _cmd_kernels_doctor(args: argparse.Namespace) -> int:
400
- """Run environment validation checks."""
506
+ print("\nRunning managed runtime checks...")
507
+ runtime_checks = _run_runtimes_doctor_checks()
508
+ all_failures.extend(_print_doctor_checks(runtime_checks))
509
+
510
+ scala_checks = _run_scala_connect_doctor_checks()
511
+ if scala_checks:
512
+ print("\nRunning Scala Connect cache checks...")
513
+ all_failures.extend(_print_doctor_checks(scala_checks))
514
+
515
+ if all_failures:
516
+ print(f"\n{len(all_failures)} check(s) failed.", file=sys.stderr)
517
+ return 1
518
+
519
+ print("\nAll checks passed.")
520
+ return 0
521
+
522
+
523
+ def _run_kernels_doctor_checks(args: argparse.Namespace) -> list[Check]:
524
+ """Collect kernel/environment checks for doctor output."""
401
525
  kernel_dirs = _resolve_kernel_dir_args(args)
402
- if kernel_dirs:
403
- checks = run_checks(profile=args.profile, kernels_dir=kernel_dirs[0], kernel_id=getattr(args, "id", KERNEL_ID))
526
+ explicit_id = getattr(args, "id", KERNEL_ID)
527
+ # When --id is the default, check both 2.12 and 2.13 kernels
528
+ if explicit_id == KERNEL_ID:
529
+ kernel_ids = [KERNEL_ID, KERNEL_ID_213]
404
530
  else:
405
- checks = run_checks(profile=args.profile, kernel_id=getattr(args, "id", KERNEL_ID))
531
+ kernel_ids = [explicit_id]
532
+ all_checks: list[Check] = []
533
+ for kid in kernel_ids:
534
+ if kernel_dirs:
535
+ checks = run_checks(profile=args.profile, kernels_dir=kernel_dirs[0], kernel_id=kid)
536
+ else:
537
+ checks = run_checks(profile=args.profile, kernel_id=kid)
538
+ for check in checks:
539
+ # Disambiguate check names when checking multiple kernels
540
+ if len(kernel_ids) > 1 and check.name in ("kernel", "kernel_semantics"):
541
+ all_checks.append(Check(f"{check.name}({kid})", check.status, check.message))
542
+ elif check not in all_checks:
543
+ all_checks.append(check)
544
+ return all_checks
406
545
 
546
+
547
+ def _run_runtimes_doctor_checks() -> list[Check]:
548
+ """Collect runtime-home checks for doctor output."""
549
+ return doctor_installed_runtimes()
550
+
551
+
552
+ def _run_scala_connect_doctor_checks() -> list[Check]:
553
+ """Collect Scala Connect cache readiness checks for doctor output."""
554
+ return doctor_scala_connect_readiness()
555
+
556
+
557
+ def _print_doctor_checks(checks: list[Check]) -> list[Check]:
558
+ """Render doctor checks and return the failing subset."""
407
559
  status_symbols = {"ok": "[ok]", "warn": "[!!]", "fail": "[FAIL]"}
408
560
 
409
561
  for check in checks:
410
562
  symbol = status_symbols.get(check.status, "[??]")
411
563
  print(f" {symbol} {check.name}: {check.message}")
412
564
 
413
- failures = [c for c in checks if c.status == "fail"]
565
+ return [check for check in checks if check.status == "fail"]
566
+
567
+
568
+ def _cmd_kernels_doctor(args: argparse.Namespace) -> int:
569
+ """Run environment validation checks."""
570
+ checks = _run_kernels_doctor_checks(args)
571
+ failures = _print_doctor_checks(checks)
414
572
  if failures:
415
573
  print(f"\n{len(failures)} check(s) failed.", file=sys.stderr)
416
574
  return 1
@@ -445,14 +603,8 @@ def _cmd_runtimes_list(_args: argparse.Namespace) -> int:
445
603
 
446
604
  def _cmd_runtimes_doctor(_args: argparse.Namespace) -> int:
447
605
  """Validate managed runtime receipts rooted under runtime-home."""
448
- checks = doctor_installed_runtimes()
449
- status_symbols = {"ok": "[ok]", "warn": "[!!]", "fail": "[FAIL]"}
450
-
451
- for check in checks:
452
- symbol = status_symbols.get(check.status, "[??]")
453
- print(f" {symbol} {check.name}: {check.message}")
454
-
455
- failures = [check for check in checks if check.status == "fail"]
606
+ checks = _run_runtimes_doctor_checks()
607
+ failures = _print_doctor_checks(checks)
456
608
  if failures:
457
609
  print(f"\n{len(failures)} check(s) failed.", file=sys.stderr)
458
610
  return 1