feat: add Switch transpiler installer for Lakebridge integration

Implement SwitchInstaller to integrate Switch transpiler with Lakebridge:
- Install Switch package to local virtual environment and deploy to workspace
- Create and manage Databricks job for Switch transpilation
- Configure Switch resources (catalog, schema, volume) interactively
- Support job-level parameters with JobParameterDefinition for flexibility
- Handle installation state and job lifecycle management
- Add comprehensive test suite covering installation, job management, and configuration

Author: hiroyukinakazato-db

src/databricks/labs/lakebridge/transpiler/installers.py

@@ -26,7 +26,7 @@
 from databricks.labs.lakebridge.transpiler.repository import TranspilerRepository
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import InvalidParameterValue, NotFound
-from databricks.sdk.service.jobs import JobSettings, NotebookTask, Source, Task
+from databricks.sdk.service.jobs import JobParameterDefinition, JobSettings, NotebookTask, Source, Task
 
 logger = logging.getLogger(__name__)
 
@@ -643,7 +643,6 @@ def _deploy_workspace(self, switch_package_dir: Path) -> None:
         """Deploy Switch package to workspace from site-packages."""
         try:
             logger.info("Deploying Switch package to workspace...")
-
             remote_path = f"{self._TRANSPILER_ID}/databricks"
             self._upload_directory(switch_package_dir, remote_path)
             logger.info("Switch workspace deployment completed")
@@ -673,54 +672,47 @@ def _upload_directory(self, local_path: Path, remote_prefix: str) -> None:
     def _setup_job(self) -> None:
         """Create Switch job if not exists."""
         install_state = InstallState.from_installation(self._installation)
-
-        if self._has_valid_job(install_state):
-            logger.info("Switch job already exists")
-            return
-
-        logger.info("Creating Switch transpiler job...")
+        existing_job_id = self._get_existing_job_id(install_state)
+        logger.info("Setting up Switch job in workspace...")
         try:
-            job_id = self._create_or_update_switch_job(install_state)
+            job_id = self._create_or_update_switch_job(existing_job_id)
+            install_state.jobs[self._INSTALL_STATE_KEY] = job_id
             install_state.save()
-            job_url = f"{self._workspace_client.config.host}#job/{job_id}"
+            job_url = f"{self._workspace_client.config.host}/jobs/{job_id}"
             logger.info(f"Switch job created/updated: {job_url}")
         except (RuntimeError, ValueError, InvalidParameterValue) as e:
-            logger.error(f"Failed to create Switch job: {e}")
+            logger.error(f"Failed to create/update Switch job: {e}")
 
-    def _has_valid_job(self, install_state: InstallState) -> bool:
-        """Check if Switch job exists and is valid in workspace."""
+    def _get_existing_job_id(self, install_state: InstallState) -> str | None:
+        """Check if Switch job already exists in workspace and return its job_id."""
         if self._INSTALL_STATE_KEY not in install_state.jobs:
-            return False
+            return None
         try:
-            job_id = int(install_state.jobs[self._INSTALL_STATE_KEY])
-            self._workspace_client.jobs.get(job_id)
-            return True
+            job_id = install_state.jobs[self._INSTALL_STATE_KEY]
+            self._workspace_client.jobs.get(int(job_id))
+            return job_id
         except (InvalidParameterValue, NotFound, ValueError):
-            return False
+            return None
 
-    def _create_or_update_switch_job(self, install_state: InstallState) -> str:
+    def _create_or_update_switch_job(self, job_id: str | None) -> str:
         """Create or update Switch job"""
-        # Check for existing job and update
-        if self._INSTALL_STATE_KEY in install_state.jobs:
+        job_settings = self._get_switch_job_settings()
+
+        # Try to update existing job
+        if job_id:
             try:
-                job_id = int(install_state.jobs[self._INSTALL_STATE_KEY])
                 logger.info(f"Updating Switch job: {job_id}")
-                job_settings = self._get_switch_job_settings()
-                self._workspace_client.jobs.reset(job_id, JobSettings(**job_settings))
-                return str(job_id)
-            except InvalidParameterValue:
-                del install_state.jobs[self._INSTALL_STATE_KEY]
+                self._workspace_client.jobs.reset(int(job_id), JobSettings(**job_settings))
+                return job_id
+            except (ValueError, InvalidParameterValue):
                 logger.warning("Previous Switch job not found, creating new one")
 
         # Create new job
-        logger.info("Creating new Switch job configuration")
-        job_settings = self._get_switch_job_settings()
+        logger.info("Creating new Switch job")
         new_job = self._workspace_client.jobs.create(**job_settings)
-        assert new_job.job_id is not None
-
-        # Save to InstallState
-        install_state.jobs[self._INSTALL_STATE_KEY] = str(new_job.job_id)
-        return str(new_job.job_id)
+        new_job_id = str(new_job.job_id)
+        assert new_job_id is not None
+        return new_job_id
 
     def _get_switch_job_settings(self) -> dict:
         """Build job settings for Switch transpiler using serverless compute"""
@@ -736,25 +728,21 @@ def _get_switch_job_settings(self) -> dict:
             task_key="run_transpilation",
             notebook_task=NotebookTask(
                 notebook_path=notebook_path,
-                base_parameters=self._get_switch_parameters_from_config(),
                 source=Source.WORKSPACE,
             ),
             disable_auto_optimization=True,  # To disable retries on failure
         )
 
         return {
             "name": job_name,
-            "tags": {"created_by": user_name, "version": f"v{version}"},
+            "tags": {"created_by": user_name, "switch_version": f"v{version}"},
             "tasks": [task],
+            "parameters": self._get_switch_job_parameters(),
             "max_concurrent_runs": 100,  # Allow simultaneous transpilations
         }
 
-    def _get_switch_parameters_from_config(self) -> dict:
-        """Extract Switch parameters from installed config.yml.
-
-        Raises:
-            ValueError: If Switch config.yml is not found (indicates incomplete installation)
-        """
+    def _get_switch_job_parameters(self) -> list[JobParameterDefinition]:
+        """Build job-level parameter definitions from installed config.yml."""
         configs = self._transpiler_repository.all_transpiler_configs()
         config = configs.get(self._TRANSPILER_ID)
 
@@ -786,7 +774,7 @@ def _get_switch_parameters_from_config(self) -> dict:
 
             parameters[flag] = default
 
-        return parameters
+        return [JobParameterDefinition(name=key, default=value) for key, value in parameters.items()]
 
     def _configure_resources(self) -> None:
         """Configure Switch resources (catalog, schema, volume) if not configured."""

tests/unit/transpiler/test_installers.py

@@ -121,14 +121,14 @@ def test_java_version_parse_missing() -> None:
 class FriendOfSwitchInstaller(SwitchInstaller):
     """A friend class to access protected methods for testing purposes."""
 
-    def has_valid_job(self, install_state: Any) -> bool:
-        return self._has_valid_job(install_state)
+    def get_existing_job_id(self, install_state: Any) -> str | None:
+        return self._get_existing_job_id(install_state)
 
-    def create_or_update_switch_job(self, install_state: Any) -> str:
-        return self._create_or_update_switch_job(install_state)
+    def create_or_update_switch_job(self, job_id: str | None) -> str:
+        return self._create_or_update_switch_job(job_id)
 
-    def get_switch_parameters_from_config(self) -> dict:
-        return self._get_switch_parameters_from_config()
+    def get_switch_job_parameters(self) -> list:
+        return self._get_switch_job_parameters()
 
     def prompt_for_switch_resources(self) -> tuple[str, str, str]:
         return self._prompt_for_switch_resources()
@@ -343,20 +343,20 @@ def test_get_configured_resources(
     @pytest.mark.parametrize(
         ("jobs_state", "get_side_effect", "expected"),
         (
-            pytest.param({}, None, False, id="no_job_in_state"),
-            pytest.param({"Switch": "12345"}, None, True, id="valid_job_exists"),
-            pytest.param({"Switch": "99999"}, NotFound("Job not found"), False, id="job_not_found"),
-            pytest.param({"Switch": "invalid"}, ValueError("Invalid job ID"), False, id="invalid_job_id"),
+            pytest.param({}, None, None, id="no_job_in_state"),
+            pytest.param({"Switch": "12345"}, None, "12345", id="valid_job_exists"),
+            pytest.param({"Switch": "99999"}, NotFound("Job not found"), None, id="job_not_found"),
+            pytest.param({"Switch": "invalid"}, ValueError("Invalid job ID"), None, id="invalid_job_id"),
         ),
     )
-    def test_has_valid_job(
+    def test_get_existing_job_id(
         self,
         jobs_state: dict,
         get_side_effect: Exception | None,
-        expected: bool,
+        expected: str | None,
         tmp_path: Path,
     ) -> None:
-        """Test _has_valid_job checks job validity in workspace."""
+        """Test _get_existing_job_id returns job_id if valid, None otherwise."""
         install_state = Mock()
         install_state.jobs = jobs_state
 
@@ -367,19 +367,19 @@ def test_has_valid_job(
         # Use friend class to access protected method
         repository = TranspilerRepository(tmp_path)
         friend_installer = FriendOfSwitchInstaller(repository, mock_ws, Mock())
-        result = friend_installer.has_valid_job(install_state)
+        result = friend_installer.get_existing_job_id(install_state)
 
         assert result == expected
         if jobs_state and "Switch" in jobs_state and not get_side_effect:
             mock_ws.jobs.get.assert_called_once_with(int(jobs_state["Switch"]))
 
     @pytest.mark.parametrize(
-        ("initial_jobs", "reset_side_effect", "expected_job_id", "expect_create", "expect_reset"),
+        ("initial_job_id", "reset_side_effect", "expected_job_id", "expect_create", "expect_reset"),
         (
-            pytest.param({}, None, "12345", True, False, id="new_job_creation"),
-            pytest.param({"Switch": "67890"}, None, "67890", False, True, id="existing_job_update"),
+            pytest.param(None, None, "12345", True, False, id="new_job_creation"),
+            pytest.param("67890", None, "67890", False, True, id="existing_job_update"),
             pytest.param(
-                {"Switch": "99999"},
+                "99999",
                 InvalidParameterValue("Job not found"),
                 "12345",
                 True,
@@ -389,26 +389,18 @@ def test_has_valid_job(
         ),
     )
     @patch.object(SwitchInstaller, "_get_switch_job_settings")
-    @patch("databricks.labs.lakebridge.transpiler.installers.InstallState")
     def test_job_creation(
         self,
-        mock_install_state_class: Mock,
         mock_get_settings: Mock,
-        initial_jobs: dict,
+        initial_job_id: str | None,
         reset_side_effect: Exception | None,
         expected_job_id: str,
         expect_create: bool,
         expect_reset: bool,
-        installer: SwitchInstaller,
         tmp_path: Path,
     ) -> None:
         """Test Switch job creation and update scenarios."""
         # Setup
-        mock_install_state = Mock()
-        mock_install_state.jobs = dict(initial_jobs)
-        mock_install_state.save = Mock()
-        mock_install_state_class.from_installation.return_value = mock_install_state
-
         mock_get_settings.return_value = {"name": "test_job"}
 
         mock_job = Mock()
@@ -428,14 +420,15 @@ def test_job_creation(
         test_repository = TranspilerRepository(tmp_path)
         mock_installation = Mock()
         friend_installer = FriendOfSwitchInstaller(test_repository, mock_ws, mock_installation)
-        result = friend_installer.create_or_update_switch_job(mock_install_state)
+        result = friend_installer.create_or_update_switch_job(initial_job_id)
 
         # Assert
         assert result == expected_job_id
 
         if expect_reset:
             if reset_side_effect:
-                mock_jobs.reset.assert_called_once_with(int(initial_jobs["Switch"]), JobSettings(name="test_job"))
+                assert initial_job_id is not None
+                mock_jobs.reset.assert_called_once_with(int(initial_job_id), JobSettings(name="test_job"))
             else:
                 mock_jobs.reset.assert_called_once_with(int(expected_job_id), JobSettings(name="test_job"))
         else:
@@ -446,9 +439,7 @@ def test_job_creation(
         else:
             mock_jobs.create.assert_not_called()
 
-        assert mock_install_state.jobs["Switch"] == expected_job_id
-
-    def test_get_switch_parameters_handles_various_default_values(
+    def test_get_switch_job_parameters_handles_various_default_values(
         self, installer: SwitchInstaller, tmp_path: Path
     ) -> None:
         """Test that different default values in config are correctly converted."""
@@ -469,20 +460,25 @@ def test_get_switch_parameters_handles_various_default_values(
         friend_installer = FriendOfSwitchInstaller(test_repository, mock_ws, mock_installation)
 
         with patch.object(test_repository, "all_transpiler_configs", return_value={"switch": mock_config}):
-            params = friend_installer.get_switch_parameters_from_config()
-
-        assert "input_dir" in params
-        assert "output_dir" in params
-        assert "result_catalog" in params
-        assert "result_schema" in params
-        assert "builtin_prompt" in params
-        assert params["flag1"] == ""
-        assert params["flag2"] == "123"
-        assert params["flag3"] == ""
-        assert params["flag4"] == "value"
-        assert params["flag5"] == "3.14"
-
-    def test_get_switch_parameters_raises_when_config_missing(self, installer: SwitchInstaller, tmp_path: Path) -> None:
+            params = friend_installer.get_switch_job_parameters()
+
+        # Convert list to dict for easier testing
+        params_dict = {p.name: p.default for p in params}
+
+        assert "input_dir" in params_dict
+        assert "output_dir" in params_dict
+        assert "result_catalog" in params_dict
+        assert "result_schema" in params_dict
+        assert "builtin_prompt" in params_dict
+        assert params_dict["flag1"] == ""
+        assert params_dict["flag2"] == "123"
+        assert params_dict["flag3"] == ""
+        assert params_dict["flag4"] == "value"
+        assert params_dict["flag5"] == "3.14"
+
+    def test_get_switch_job_parameters_raises_when_config_missing(
+        self, installer: SwitchInstaller, tmp_path: Path
+    ) -> None:
         """Test that ValueError is raised when Switch config is not found."""
         test_repository = TranspilerRepository(tmp_path)
         mock_ws = Mock()
@@ -491,7 +487,7 @@ def test_get_switch_parameters_raises_when_config_missing(self, installer: Switc
 
         with patch.object(test_repository, "all_transpiler_configs", return_value={}):
             with pytest.raises(ValueError, match="Switch config.yml not found"):
-                friend_installer.get_switch_parameters_from_config()
+                friend_installer.get_switch_job_parameters()
 
     @patch("databricks.labs.lakebridge.transpiler.installers.ResourceConfigurator")
     @patch("databricks.labs.lakebridge.transpiler.installers.CatalogOperations")