From 52aba7abf9bdbd9d5b528048beed41f813134382 Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Mon, 17 Mar 2025 18:24:42 +0000
Subject: [PATCH] Address PR feedback: Add comments, version verification, and
 enhanced tests for databricks-sdk

Co-Authored-By: Joe Moura <joao@crewai.com>
---
 pyproject.toml                        |  5 ++++-
 tests/tools/test_databricks_import.py | 27 +++++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 2dfe3f1c4..010904704 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -45,7 +45,10 @@ Documentation = "https://docs.crewai.com"
 Repository = "https://github.com/crewAIInc/crewAI"
 
 [project.optional-dependencies]
-tools = ["crewai-tools>=0.37.0", "databricks-sdk>=0.46.0"]
+tools = [
+    "crewai-tools>=0.37.0",
+    "databricks-sdk>=0.46.0,<1.0.0"  # Required for Databricks query tool functionality
+]
 embeddings = [
     "tiktoken~=0.7.0"
 ]
diff --git a/tests/tools/test_databricks_import.py b/tests/tools/test_databricks_import.py
index 291d866bb..4aeb9f86c 100644
--- a/tests/tools/test_databricks_import.py
+++ b/tests/tools/test_databricks_import.py
@@ -1,4 +1,5 @@
 import pytest
+from importlib.metadata import version, PackageNotFoundError
 
 
 def test_databricks_sdk_import():
@@ -12,3 +13,29 @@ def test_databricks_sdk_import():
         assert True
     except ImportError as e:
         pytest.fail(f"Failed to import databricks.sdk: {e}")
+
+
+def test_databricks_sdk_version():
+    """Test that the installed databricks-sdk version meets requirements.
+    
+    The databricks-sdk should be version 0.46.0 or higher, but less than 1.0.0
+    to maintain compatibility with the databricks_query_tool.
+    """
+    try:
+        sdk_version = version("databricks-sdk")
+        assert sdk_version >= "0.46.0", f"Databricks SDK version {sdk_version} is too old (< 0.46.0)"
+        assert sdk_version < "1.0.0", f"Databricks SDK version {sdk_version} is too new (>= 1.0.0)"
+    except PackageNotFoundError:
+        pytest.fail("databricks-sdk package not found")
+
+
+def test_databricks_core_functionality():
+    """Test core functionality of the Databricks SDK.
+    
+    This test verifies that the WorkspaceClient class from the Databricks SDK
+    has the expected attributes for SQL functionality.
+    """
+    from databricks.sdk import WorkspaceClient
+    
+    # Verify that the WorkspaceClient class has the expected attributes for SQL operations
+    assert hasattr(WorkspaceClient, 'statement_execution'), "SQL statement execution functionality not available in WorkspaceClient"