Conditionally import databricks libraries

Author: vinibrsl

crewai_tools/tools/databricks_query_tool/databricks_query_tool.py

@@ -1,10 +1,11 @@
 import os
-from typing import Any, Dict, List, Optional, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
 
 from crewai.tools import BaseTool
-from databricks.sdk import WorkspaceClient
-from pydantic import BaseModel, Field, model_validator
+from pydantic import BaseModel, Field, model_validator, ConfigDict, PrivateAttr
 
+if TYPE_CHECKING:
+    from databricks.sdk import WorkspaceClient
 
 class DatabricksQueryToolSchema(BaseModel):
     """Input schema for DatabricksQueryTool."""
@@ -55,6 +56,10 @@ class DatabricksQueryTool(BaseTool):
         >>> results = tool.run(query="SELECT * FROM my_table LIMIT 10")
     """
 
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True, validate_assignment=True, frozen=False
+    )
+
     name: str = "Databricks SQL Query"
     description: str = (
         "Execute SQL queries against Databricks workspace tables and return the results."
@@ -67,7 +72,7 @@ class DatabricksQueryTool(BaseTool):
     default_schema: Optional[str] = None
     default_warehouse_id: Optional[str] = None
 
-    _workspace_client: Optional[WorkspaceClient] = None
+    _workspace_client: Optional["WorkspaceClient"] = PrivateAttr(None)
 
     def __init__(
         self,
@@ -89,8 +94,6 @@ def __init__(
         self.default_catalog = default_catalog
         self.default_schema = default_schema
         self.default_warehouse_id = default_warehouse_id
-
-        # Validate that Databricks credentials are available
         self._validate_credentials()
 
     def _validate_credentials(self) -> None:
@@ -105,10 +108,16 @@ def _validate_credentials(self) -> None:
             )
 
     @property
-    def workspace_client(self) -> WorkspaceClient:
+    def workspace_client(self) -> "WorkspaceClient":
         """Get or create a Databricks WorkspaceClient instance."""
         if self._workspace_client is None:
-            self._workspace_client = WorkspaceClient()
+            try:
+                from databricks.sdk import WorkspaceClient
+                self._workspace_client = WorkspaceClient()
+            except ImportError:
+                raise ImportError(
+                    "`databricks-sdk` package not found, please run `uv add databricks-sdk`"
+                )
         return self._workspace_client
 
     def _format_results(self, results: List[Dict[str, Any]]) -> str:
@@ -733,4 +742,4 @@ def _run(
             # Include more details in the error message to help with debugging
             import traceback
             error_details = traceback.format_exc()
-            return f"Error executing Databricks query: {str(e)}\n\nDetails:\n{error_details}"
+            return f"Error executing Databricks query: {str(e)}\n\nDetails:\n{error_details}"