Wait for library installation on python submission (#1031)

Author: ericj-db

CHANGELOG.md

@@ -1,5 +1,9 @@
 ## dbt-databricks 1.10.3 (TBD)
 
+### Fixes
+
+- Fix bug where python model run starts before all libraries are installed on the cluster ([1028](https://github.com/databricks/dbt-databricks/issues/1028))
+
 ## dbt-databricks 1.10.2 (May 21, 2025)
 
 ### Features

dbt/adapters/databricks/api_client.py

@@ -19,6 +19,7 @@
 DEFAULT_POLLING_INTERVAL = 10
 SUBMISSION_LANGUAGE = "python"
 USER_AGENT = f"dbt-databricks/{version}"
+LIBRARY_VALID_STATUSES = {"INSTALLED", "RESTORED", "SKIPPED"}
 
 
 class PrefixSession:
@@ -47,10 +48,45 @@ def __init__(self, session: Session, host: str, api: str):
         self.session = PrefixSession(session, host, api)
 
 
+class LibraryApi(DatabricksApi):
+    def __init__(self, session: Session, host: str):
+        super().__init__(session, host, "/api/2.0/libraries")
+
+    def all_libraries_installed(self, cluster_id: str) -> bool:
+        status = self.get_cluster_libraries_status(cluster_id)
+        if "library_statuses" in status:
+            return all(
+                library["status"] in LIBRARY_VALID_STATUSES
+                for library in status["library_statuses"]
+            )
+        else:
+            return True
+
+    def get_cluster_libraries_status(self, cluster_id: str) -> dict[str, Any]:
+        response = self.session.get(
+            "/cluster-status",
+            json={"cluster_id": cluster_id},
+        )
+        if response.status_code != 200:
+            raise DbtRuntimeError(
+                f"Error getting status of libraries of a cluster.\n {response.content!r}"
+            )
+
+        json_response = response.json()
+        return json_response
+
+
 class ClusterApi(DatabricksApi):
-    def __init__(self, session: Session, host: str, max_cluster_start_time: int = 900):
+    def __init__(
+        self,
+        session: Session,
+        host: str,
+        libraries: LibraryApi,
+        max_cluster_start_time: int = 900,
+    ):
         super().__init__(session, host, "/api/2.0/clusters")
         self.max_cluster_start_time = max_cluster_start_time
+        self.libraries = libraries
 
     def status(self, cluster_id: str) -> str:
         # https://docs.databricks.com/dev-tools/api/latest/clusters.html#get
@@ -68,10 +104,19 @@ def wait_for_cluster(self, cluster_id: str) -> None:
 
         while time.time() - start_time < self.max_cluster_start_time:
             status_response = self.status(cluster_id)
-            if status_response == "RUNNING":
-                return
-            else:
+
+            if status_response != "RUNNING":
+                logger.debug("Waiting for cluster to start")
+                time.sleep(5)
+                continue
+
+            libraries_status = self.libraries.get_cluster_libraries_status(cluster_id)
+            if not self.libraries.all_libraries_installed(cluster_id):
+                logger.debug(f"Waiting for all libraries to be installed: {libraries_status}")
                 time.sleep(5)
+                continue
+
+            return
 
         raise DbtRuntimeError(
             f"Cluster {cluster_id} restart timed out after {self.max_cluster_start_time} seconds"
@@ -96,9 +141,12 @@ def start(self, cluster_id: str) -> None:
 
 
 class CommandContextApi(DatabricksApi):
-    def __init__(self, session: Session, host: str, cluster_api: ClusterApi):
+    def __init__(
+        self, session: Session, host: str, cluster_api: ClusterApi, library_api: LibraryApi
+    ):
         super().__init__(session, host, "/api/1.2/contexts")
         self.cluster_api = cluster_api
+        self.library_api = library_api
 
     def create(self, cluster_id: str) -> str:
         current_status = self.cluster_api.status(cluster_id)
@@ -107,7 +155,9 @@ def create(self, cluster_id: str) -> str:
             logger.debug(f"Cluster {cluster_id} is not running. Attempting to restart.")
             self.cluster_api.start(cluster_id)
             logger.debug(f"Cluster {cluster_id} is now running.")
-        elif current_status != "RUNNING":
+        elif current_status != "RUNNING" or not self.library_api.all_libraries_installed(
+            cluster_id
+        ):
             self.cluster_api.wait_for_cluster(cluster_id)
 
         response = self.session.post(
@@ -532,8 +582,9 @@ def __init__(
         timeout: int,
         use_user_folder: bool,
     ):
-        self.clusters = ClusterApi(session, host)
-        self.command_contexts = CommandContextApi(session, host, self.clusters)
+        self.libraries = LibraryApi(session, host)
+        self.clusters = ClusterApi(session, host, self.libraries)
+        self.command_contexts = CommandContextApi(session, host, self.clusters, self.libraries)
         self.curr_user = CurrUserApi(session, host)
         if use_user_folder:
             self.folders: FolderApi = UserFolderApi(session, host, self.curr_user)

tests/unit/api_client/test_cluster_api.py

@@ -1,4 +1,4 @@
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import freezegun
 import pytest
@@ -10,8 +10,12 @@
 
 class TestClusterApi(ApiTestBase):
     @pytest.fixture
-    def api(self, session, host):
-        return ClusterApi(session, host)
+    def library_api(self):
+        return Mock()
+
+    @pytest.fixture
+    def api(self, session, host, library_api):
+        return ClusterApi(session, host, library_api)
 
     def test_status__non_200(self, api, session):
         self.assert_non_200_raises_error(lambda: api.status("cluster_id"), session)
@@ -26,10 +30,33 @@ def test_status__200(self, api, session, host):
         )
 
     @patch("dbt.adapters.databricks.api_client.time.sleep")
-    def test_wait_for_cluster__success(self, _, api, session):
+    def test_wait_for_cluster__success(self, _, api, session, library_api):
         session.get.return_value.status_code = 200
         session.get.return_value.json.side_effect = [{"state": "pending"}, {"state": "running"}]
+        library_api.get_cluster_libraries_status.return_value = {"library_statuses": []}
+        api.wait_for_cluster("cluster_id")
+
+    @patch("dbt.adapters.databricks.api_client.time.sleep")
+    def test_wait_for_cluster_with_installed_library__success(
+        self, mock_sleep, api, session, library_api
+    ):
+        session.get.return_value.status_code = 200
+        session.get.return_value.json.return_value = {"state": "running"}
+        library_api.get_cluster_libraries_status.return_value = {
+            "library_statuses": [{"status": "INSTALLED"}]
+        }
+        api.wait_for_cluster("cluster_id")
+        mock_sleep.assert_not_called()
+
+    @patch("dbt.adapters.databricks.api_client.time.sleep")
+    def test_wait_for_cluster_with_pending_library__success(
+        self, mock_sleep, api, session, library_api
+    ):
+        session.get.return_value.status_code = 200
+        session.get.return_value.json.side_effect = [{"state": "running"}, {"state": "running"}]
+        library_api.all_libraries_installed.side_effect = [False, True]
         api.wait_for_cluster("cluster_id")
+        mock_sleep.assert_called_with(5)
 
     @freezegun.freeze_time("2020-01-01", auto_tick_seconds=900)
     @patch("dbt.adapters.databricks.api_client.time.sleep")
@@ -42,10 +69,11 @@ def test_wait_for_cluster__timeout(self, _, api, session):
     def test_start__non_200(self, api, session):
         self.assert_non_200_raises_error(lambda: api.start("cluster_id"), session)
 
-    def test_start__200(self, api, session, host):
+    def test_start__200(self, api, session, host, library_api):
         session.post.return_value.status_code = 200
         session.get.return_value.status_code = 200
         session.get.return_value.json.return_value = {"state": "running"}
+        library_api.get_cluster_libraries_status.return_value = {"library_statuses": []}
         api.start("cluster_id")
         session.post.assert_called_once_with(
             f"https://{host}/api/2.0/clusters/start", json={"cluster_id": "cluster_id"}, params=None

tests/unit/api_client/test_command_context_api.py

@@ -12,15 +12,20 @@ def cluster_api(self):
         return Mock()
 
     @pytest.fixture
-    def api(self, session, host, cluster_api):
-        return CommandContextApi(session, host, cluster_api)
+    def library_api(self):
+        return Mock()
+
+    @pytest.fixture
+    def api(self, session, host, cluster_api, library_api):
+        return CommandContextApi(session, host, cluster_api, library_api)
 
     def test_create__non_200(self, api, cluster_api, session):
         cluster_api.status.return_value = "RUNNING"
         self.assert_non_200_raises_error(lambda: api.create("cluster_id"), session)
 
-    def test_create__cluster_running(self, api, cluster_api, session):
+    def test_create__cluster_running(self, api, cluster_api, library_api, session):
         cluster_api.status.return_value = "RUNNING"
+        library_api.all_libraries_installed.return_value = True
         session.post.return_value.status_code = 200
         session.post.return_value.json.return_value = {"id": "context_id"}
         id = api.create("cluster_id")
@@ -29,8 +34,25 @@ def test_create__cluster_running(self, api, cluster_api, session):
             json={"clusterId": "cluster_id", "language": "python"},
             params=None,
         )
+        cluster_api.wait_for_cluster.assert_not_called()
         assert id == "context_id"
 
+    def test_create__cluster_running_with_pending_libraries(
+        self, api, cluster_api, library_api, session
+    ):
+        cluster_api.status.return_value = "RUNNING"
+        library_api.all_libraries_installed.return_value = False
+        session.post.return_value.status_code = 200
+        session.post.return_value.json.return_value = {"id": "context_id"}
+        id = api.create("cluster_id")
+        session.post.assert_called_once_with(
+            "https://host/api/1.2/contexts/create",
+            json={"clusterId": "cluster_id", "language": "python"},
+            params=None,
+        )
+        assert id == "context_id"
+        cluster_api.wait_for_cluster.assert_called_once_with("cluster_id")
+
     def test_create__cluster_terminated(self, api, cluster_api, session):
         cluster_api.status.return_value = "TERMINATED"
         session.post.return_value.status_code = 200

tests/unit/api_client/test_library_api.py

@@ -0,0 +1,50 @@
+import pytest
+
+from dbt.adapters.databricks.api_client import LibraryApi
+from tests.unit.api_client.api_test_base import ApiTestBase
+
+
+class TestLibraryApi(ApiTestBase):
+    @pytest.fixture
+    def api(self, session, host):
+        return LibraryApi(session, host)
+
+    def test_get_cluster_libraries_status__non_200(self, api, session):
+        self.assert_non_200_raises_error(
+            lambda: api.get_cluster_libraries_status("cluster_id"), session
+        )
+
+    def test_get_cluster_libraries_status__200(self, api, session, host):
+        cluster_id = "cluster_id"
+        expected_response = {"library_statuses": [{"status": "INSTALLED"}, {"status": "PENDING"}]}
+        session.get.return_value.status_code = 200
+        session.get.return_value.json.return_value = expected_response
+
+        result = api.get_cluster_libraries_status(cluster_id)
+        assert result == expected_response
+        session.get.assert_called_once_with(
+            f"https://{host}/api/2.0/libraries/cluster-status",
+            json={"cluster_id": cluster_id},
+            params=None,
+        )
+
+    def test_all_libraries_installed__true(self, api, session, host):
+        session.get.return_value.status_code = 200
+        session.get.return_value.json.return_value = {"library_statuses": [{"status": "INSTALLED"}]}
+
+        result = api.all_libraries_installed("cluster_id")
+        assert result is True
+
+    def test_all_libraries_installed__false(self, api, session, host):
+        session.get.return_value.status_code = 200
+        session.get.return_value.json.return_value = {"library_statuses": [{"status": "PENDING"}]}
+
+        result = api.all_libraries_installed("cluster_id")
+        assert result is False
+
+    def test_library_statuses_not_present(self, api, session, host):
+        session.get.return_value.status_code = 200
+        session.get.return_value.json.return_value = {"cluster-id": "abc-123"}
+
+        result = api.all_libraries_installed("abc-123")
+        assert result is True