fix(tools): detect databricks LIMIT by clause, not substring

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,9 +71,13 @@ 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():
-            self.query = f"{self.query.rstrip(';')} LIMIT {self.row_limit};"
+        # 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):
+            # Strip any trailing mix of semicolons and whitespace so a query that
+            # ends in `;  ` doesn't keep its statement-terminating `;` ahead of
+            # the appended LIMIT (which would produce invalid SQL).
+            stripped = re.sub(r"[\s;]+$", "", self.query)
+            self.query = f"{stripped} LIMIT {self.row_limit};"
 
         return self
 

lib/crewai-tools/tests/tools/test_databricks_query_tool.py

@@ -0,0 +1,62 @@
+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_appends_limit_after_trailing_semicolon_whitespace() -> None:
+    """Regression: a trailing `;` followed by whitespace used to defeat rstrip(';')."""
+    schema = DatabricksQueryToolSchema(query="SELECT * FROM limited_orders;   ")
+    assert schema.query == "SELECT * FROM limited_orders LIMIT 1000;"
+
+
+def test_appends_limit_after_trailing_newline_semicolon() -> None:
+    schema = DatabricksQueryToolSchema(query="SELECT * FROM limited_orders;\n")
+    assert schema.query == "SELECT * FROM limited_orders LIMIT 1000;"
+
+
+def test_rejects_empty_query() -> None:
+    with pytest.raises(ValueError, match="Query cannot be empty"):
+        DatabricksQueryToolSchema(query="   ")