Add Notebook-scoped packages for command submits or notebook job run

Signed-off-by: Federico Manuel Gomez Peter <federico.gomez@payclip.com>

Author: fedemgp

dbt/adapters/databricks/python_models/python_config.py

@@ -24,6 +24,14 @@ class Config:
         extra = "allow"
 
 
+class PythonPackagesConfig(BaseModel):
+    """Pydantic model for python packages configuration."""
+
+    packages: list[str]
+    notebook_scoped: bool
+    index_url: Optional[str] = None
+
+
 class PythonModelConfig(BaseModel):
     """
     Pydantic model for a Python model configuration.
@@ -42,6 +50,7 @@ class PythonModelConfig(BaseModel):
     cluster_id: Optional[str] = None
     http_path: Optional[str] = None
     create_notebook: bool = False
+    notebook_scoped_libraries: bool = False
     environment_key: Optional[str] = None
     environment_dependencies: list[str] = Field(default_factory=list)
 
@@ -69,6 +78,14 @@ def validate_notebook_permissions(cls, v: list[dict[str, str]]) -> list[dict[str
                 )
         return v
 
+    @property
+    def python_packages_config(self) -> PythonPackagesConfig:
+        return PythonPackagesConfig(
+            packages=self.packages,
+            index_url=self.index_url,
+            notebook_scoped=self.notebook_scoped_libraries,
+        )
+
 
 class ParsedPythonModel(BaseModel):
     """Pydantic model for a Python model parsed from a dbt manifest"""

dbt/adapters/databricks/python_models/python_submissions.py

@@ -9,20 +9,50 @@
 from dbt.adapters.databricks.api_client import CommandExecution, DatabricksApiClient, WorkflowJobApi
 from dbt.adapters.databricks.credentials import DatabricksCredentials
 from dbt.adapters.databricks.logging import logger
-from dbt.adapters.databricks.python_models.python_config import ParsedPythonModel
+from dbt.adapters.databricks.python_models.python_config import (
+    ParsedPythonModel,
+    PythonPackagesConfig,
+)
 from dbt.adapters.databricks.python_models.run_tracking import PythonRunTracker
 
 DEFAULT_TIMEOUT = 60 * 60 * 24
+NOTEBOOK_SEPARATOR = "\n\n# COMMAND ----------\n\n"
 
 
 class PythonSubmitter(ABC):
     """Interface for submitting Python models to run on Databricks."""
 
+    def __init__(self, packages_config: PythonPackagesConfig) -> None:
+        self.packages_config = packages_config
+
     @abstractmethod
     def submit(self, compiled_code: str) -> None:
         """Submit the compiled code to Databricks."""
         pass
 
+    def _prepare_code_with_notebook_scoped_packages(
+        self, compiled_code: str, separator: str = NOTEBOOK_SEPARATOR
+    ) -> str:
+        """Prepend notebook-scoped package installation commands to the compiled code."""
+        if not self.packages_config.packages or not self.packages_config.notebook_scoped:
+            return compiled_code
+
+        index_url = (
+            f"--index-url {self.packages_config.index_url}"
+            if self.packages_config.index_url
+            else ""
+        )
+        # Build the %pip install command for notebook-scoped packages
+        packages = " ".join(self.packages_config.packages)
+        pip_install_cmd = f"%pip install {index_url} -q {packages}"
+        logger.debug(f"Adding notebook-scoped package installation: {pip_install_cmd}")
+
+        # Add extra restart python command for Databricks runtimes 13.0 and above
+        restart_cmd = "dbutils.library.restartPython()"
+
+        # Prepend the pip install command to the compiled code
+        return f"{pip_install_cmd}{separator}{restart_cmd}{separator}{compiled_code}"
+
 
 class BaseDatabricksHelper(PythonJobHelper):
     """Base helper for python models on Databricks."""
@@ -63,16 +93,24 @@ class PythonCommandSubmitter(PythonSubmitter):
     """Submitter for Python models using the Command API."""
 
     def __init__(
-        self, api_client: DatabricksApiClient, tracker: PythonRunTracker, cluster_id: str
+        self,
+        api_client: DatabricksApiClient,
+        tracker: PythonRunTracker,
+        cluster_id: str,
+        parsed_model: ParsedPythonModel,
     ) -> None:
         self.api_client = api_client
         self.tracker = tracker
         self.cluster_id = cluster_id
+        super().__init__(parsed_model.config.python_packages_config)
 
     @override
     def submit(self, compiled_code: str) -> None:
         logger.debug("Submitting Python model using the Command API.")
 
+        # Prepare code with notebook-scoped package installation if needed
+        compiled_code = self._prepare_code_with_notebook_scoped_packages(compiled_code)
+
         context_id = self.api_client.command_contexts.create(self.cluster_id)
         command_exec: Optional[CommandExecution] = None
         try:
@@ -252,16 +290,24 @@ def get_library_config(
     packages: list[str],
     index_url: Optional[str],
     additional_libraries: list[dict[str, Any]],
+    notebook_scoped_libraries: bool = False,
 ) -> dict[str, Any]:
-    """Update the job configuration with the required libraries."""
+    """
+    Update the job configuration with the required libraries.
+
+    If notebook_scoped_libraries is True, packages are not included in the library config
+    as they will be installed via %pip install in the notebook itself.
+    """
 
     libraries = []
 
-    for package in packages:
-        if index_url:
-            libraries.append({"pypi": {"package": package, "repo": index_url}})
-        else:
-            libraries.append({"pypi": {"package": package}})
+    # Only add packages to cluster-level libraries if not using notebook-scoped
+    if not notebook_scoped_libraries:
+        for package in packages:
+            if index_url:
+                libraries.append({"pypi": {"package": package, "repo": index_url}})
+            else:
+                libraries.append({"pypi": {"package": package}})
 
     for library in additional_libraries:
         libraries.append(library)
@@ -286,7 +332,10 @@ def __init__(
         packages = parsed_model.config.packages
         index_url = parsed_model.config.index_url
         additional_libraries = parsed_model.config.additional_libs
-        library_config = get_library_config(packages, index_url, additional_libraries)
+        notebook_scoped_libraries = parsed_model.config.notebook_scoped_libraries
+        library_config = get_library_config(
+            packages, index_url, additional_libraries, notebook_scoped_libraries
+        )
         self.cluster_spec = {**cluster_spec, **library_config}
         self.job_grants = parsed_model.config.python_job_config.grants
         self.additional_job_settings = parsed_model.config.python_job_config.dict()
@@ -335,11 +384,14 @@ def __init__(
         tracker: PythonRunTracker,
         uploader: PythonNotebookUploader,
         config_compiler: PythonJobConfigCompiler,
+        parsed_model: ParsedPythonModel,
     ) -> None:
         self.api_client = api_client
         self.tracker = tracker
         self.uploader = uploader
         self.config_compiler = config_compiler
+        self.parsed_model = parsed_model
+        super().__init__(parsed_model.config.python_packages_config)
 
     @staticmethod
     def create(
@@ -356,12 +408,17 @@ def create(
             parsed_model,
             cluster_spec,
         )
-        return PythonNotebookSubmitter(api_client, tracker, notebook_uploader, config_compiler)
+        return PythonNotebookSubmitter(
+            api_client, tracker, notebook_uploader, config_compiler, parsed_model
+        )
 
     @override
     def submit(self, compiled_code: str) -> None:
         logger.debug("Submitting Python model using the Job Run API.")
 
+        # Prepare code with notebook-scoped package installation if needed
+        compiled_code = self._prepare_code_with_notebook_scoped_packages(compiled_code)
+
         file_path = self.uploader.upload(compiled_code)
         job_config = self.config_compiler.compile(file_path)
 
@@ -444,7 +501,12 @@ def build_submitter(self) -> PythonSubmitter:
                 {"existing_cluster_id": self.cluster_id},
             )
         else:
-            return PythonCommandSubmitter(self.api_client, self.tracker, self.cluster_id or "")
+            return PythonCommandSubmitter(
+                self.api_client,
+                self.tracker,
+                self.cluster_id or "",
+                self.parsed_model,
+            )
 
     @override
     def validate_config(self) -> None:
@@ -572,6 +634,7 @@ def __init__(
         workflow_creater: PythonWorkflowCreator,
         job_grants: dict[str, list[dict[str, str]]],
         acls: list[dict[str, str]],
+        parsed_model: ParsedPythonModel,
     ) -> None:
         self.api_client = api_client
         self.tracker = tracker
@@ -581,6 +644,7 @@ def __init__(
         self.workflow_creater = workflow_creater
         self.job_grants = job_grants
         self.acls = acls
+        super().__init__(parsed_model.config.python_packages_config)
 
     @staticmethod
     def create(
@@ -599,6 +663,7 @@ def create(
             workflow_creater,
             parsed_model.config.python_job_config.grants,
             parsed_model.config.access_control_list,
+            parsed_model,
         )
 
     @override
@@ -611,6 +676,7 @@ def submit(self, compiled_code: str) -> None:
             logger.debug(
                 f"[Workflow Debug] Compiled code preview: {compiled_code[:preview_len]}..."
             )
+        compiled_code = self._prepare_code_with_notebook_scoped_packages(compiled_code)
 
         file_path = self.uploader.upload(compiled_code)
         logger.debug(f"[Workflow Debug] Uploaded notebook to: {file_path}")

tests/unit/python/test_python_config.py

@@ -85,6 +85,22 @@ def test_parsed_model__valid_model_config(self):
         assert config.http_path == "http_path"
         assert config.create_notebook is True
 
+    def test_parsed_model__valid_python_packages_config(self):
+        parsed_model = {
+            "alias": "test",
+            "config": {
+                "packages": ["package"],
+                "index_url": "index_url",
+                "notebook_scoped_libraries": True,
+            },
+        }
+
+        model = ParsedPythonModel(**parsed_model)
+        config = model.config.python_packages_config
+        assert config.packages == ["package"]
+        assert config.index_url == "index_url"
+        assert config.notebook_scoped is True
+
     def test_parsed_model__extra_model_config(self):
         parsed_model = {
             "alias": "test",

tests/unit/python/test_python_job_support.py

@@ -40,6 +40,8 @@ def uploader(self, client, parsed_model, identifier):
         parsed_model.catalog = "catalog"
         parsed_model.schema_ = "schema"
         parsed_model.identifier = identifier
+        parsed_model.config.notebook_scoped_libraries = False
+        parsed_model.config.packages = []
         return PythonNotebookUploader(client, parsed_model)
 
     def test_upload__golden_path(self, uploader, client, compiled_code, workdir, identifier):
@@ -220,6 +222,29 @@ def test_get_library_config__packages_libraries(self):
             ]
         }
 
+    def test_get_library_config__notebook_scoped_packages_excluded(self):
+        config = python_submissions.get_library_config(
+            ["package1", "package2"], None, [], notebook_scoped_libraries=True
+        )
+        assert config == {"libraries": []}
+
+    def test_get_library_config__notebook_scoped_with_additional_libs(self):
+        config = python_submissions.get_library_config(
+            ["package1", "package2"],
+            None,
+            [{"jar": "s3://mybucket/myjar.jar"}],
+            notebook_scoped_libraries=True,
+        )
+        assert config == {"libraries": [{"jar": "s3://mybucket/myjar.jar"}]}
+
+    def test_get_library_config__notebook_scoped_false_includes_packages(self):
+        config = python_submissions.get_library_config(
+            ["package1", "package2"], None, [], notebook_scoped_libraries=False
+        )
+        assert config == {
+            "libraries": [{"pypi": {"package": "package1"}}, {"pypi": {"package": "package2"}}]
+        }
+
 
 class TestPythonJobConfigCompiler:
     @pytest.fixture
@@ -232,6 +257,7 @@ def run_name(self, parsed_model):
         parsed_model.run_name = run_name
         parsed_model.config.packages = []
         parsed_model.config.additional_libs = []
+        parsed_model.config.notebook_scoped_libraries = False
         return run_name
 
     @pytest.fixture
@@ -384,3 +410,38 @@ def test_compile__user_environments_override_auto_generated(
         assert details.additional_job_config["environments"][0]["spec"]["dependencies"] == [
             "requests"
         ]
+
+    def test_compile__notebook_scoped_libraries_excludes_packages(
+        self, client, permission_builder, parsed_model, run_name
+    ):
+        parsed_model.config.packages = ["pandas", "numpy"]
+        parsed_model.config.index_url = None
+        parsed_model.config.notebook_scoped_libraries = True
+        parsed_model.config.environment_key = None
+        parsed_model.config.python_job_config.dict.return_value = {}
+
+        permission_builder.build_job_permissions.return_value = []
+        compiler = PythonJobConfigCompiler(client, permission_builder, parsed_model, {})
+        details = compiler.compile("path")
+
+        # Libraries should be empty since packages are notebook-scoped
+        assert details.job_spec["libraries"] == []
+
+    def test_compile__notebook_scoped_false_includes_packages(
+        self, client, permission_builder, parsed_model, run_name
+    ):
+        parsed_model.config.packages = ["pandas", "numpy"]
+        parsed_model.config.index_url = None
+        parsed_model.config.notebook_scoped_libraries = False
+        parsed_model.config.environment_key = None
+        parsed_model.config.python_job_config.dict.return_value = {}
+
+        permission_builder.build_job_permissions.return_value = []
+        compiler = PythonJobConfigCompiler(client, permission_builder, parsed_model, {})
+        details = compiler.compile("path")
+
+        # Libraries should include packages
+        assert details.job_spec["libraries"] == [
+            {"pypi": {"package": "pandas"}},
+            {"pypi": {"package": "numpy"}},
+        ]