diff --git a/lib/crewai-tools/src/crewai_tools/tools/databricks_query_tool/databricks_query_tool.py b/lib/crewai-tools/src/crewai_tools/tools/databricks_query_tool/databricks_query_tool.py index dd058655a1..a4fd6a3b8f 100644 --- a/lib/crewai-tools/src/crewai_tools/tools/databricks_query_tool/databricks_query_tool.py +++ b/lib/crewai-tools/src/crewai_tools/tools/databricks_query_tool/databricks_query_tool.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +import re import time from typing import TYPE_CHECKING, Any, TypeGuard, TypedDict @@ -11,6 +12,13 @@ if TYPE_CHECKING: from databricks.sdk import WorkspaceClient +# True LIMIT/FETCH clause, not an identifier that merely contains "limit" +# (e.g. table `limited_orders`). Databricks accepts LIMIT n, LIMIT ALL, and +# FETCH FIRST n ROWS ONLY. +_SQL_LIMIT_CLAUSE_RE = re.compile( + r"(?is)\b(?:LIMIT\s+(?:ALL|\d+)\b|FETCH\s+(?:FIRST|NEXT)\s+\d+\s+ROWS?\b)" +) + class ExecutionContext(TypedDict, total=False): catalog: str @@ -63,8 +71,8 @@ def validate_input(self) -> DatabricksQueryToolSchema: if not self.query or not self.query.strip(): raise ValueError("Query cannot be empty") - # Add a LIMIT clause to the query if row_limit is provided and query doesn't have one - if self.row_limit and "limit" not in self.query.lower(): + # Add a LIMIT clause if row_limit is set and the query has no LIMIT/FETCH clause. + if self.row_limit and not _SQL_LIMIT_CLAUSE_RE.search(self.query): self.query = f"{self.query.rstrip(';')} LIMIT {self.row_limit};" return self diff --git a/lib/crewai-tools/tests/tools/test_databricks_query_tool.py b/lib/crewai-tools/tests/tools/test_databricks_query_tool.py new file mode 100644 index 0000000000..e23c3b289a --- /dev/null +++ b/lib/crewai-tools/tests/tools/test_databricks_query_tool.py @@ -0,0 +1,51 @@ +import pytest + +from crewai_tools.tools.databricks_query_tool.databricks_query_tool import ( + DatabricksQueryToolSchema, +) + + +def test_appends_default_limit_to_plain_select() -> None: + schema = DatabricksQueryToolSchema(query="SELECT * FROM orders") + assert schema.query.rstrip(";").endswith("LIMIT 1000") + + +def test_appends_limit_when_table_name_contains_limit() -> None: + """Regression: substring 'limit' in `limited_orders` used to skip the cap.""" + schema = DatabricksQueryToolSchema(query="SELECT * FROM limited_orders") + assert schema.query.rstrip(";").upper().endswith("LIMIT 1000") + + +def test_appends_limit_when_column_is_named_limit() -> None: + schema = DatabricksQueryToolSchema(query="SELECT limit FROM orders") + assert schema.query.rstrip(";").upper().endswith("LIMIT 1000") + + +def test_does_not_double_existing_limit() -> None: + schema = DatabricksQueryToolSchema(query="SELECT * FROM orders LIMIT 5") + assert schema.query.rstrip(";") == "SELECT * FROM orders LIMIT 5" + + +def test_does_not_double_limit_all() -> None: + schema = DatabricksQueryToolSchema(query="SELECT * FROM orders LIMIT ALL") + assert schema.query.rstrip(";") == "SELECT * FROM orders LIMIT ALL" + + +def test_does_not_double_fetch_first() -> None: + schema = DatabricksQueryToolSchema(query="SELECT * FROM orders FETCH FIRST 10 ROWS ONLY") + assert "LIMIT" not in schema.query.upper() + + +def test_respects_custom_row_limit() -> None: + schema = DatabricksQueryToolSchema(query="SELECT * FROM limited_orders", row_limit=25) + assert schema.query.rstrip(";").endswith("LIMIT 25") + + +def test_skips_append_when_row_limit_is_zero() -> None: + schema = DatabricksQueryToolSchema(query="SELECT * FROM orders", row_limit=0) + assert "LIMIT" not in schema.query.upper() + + +def test_rejects_empty_query() -> None: + with pytest.raises(ValueError, match="Query cannot be empty"): + DatabricksQueryToolSchema(query=" ")