This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new e8160efd4db Let the agent retry on a rejected DataFusion query instead
of failing the task (#71445)
e8160efd4db is described below
commit e8160efd4db4389f7894e66d0244b776fc9b5655
Author: Jyun-An Chen <[email protected]>
AuthorDate: Wed Aug 12 05:28:35 2026 +0800
Let the agent retry on a rejected DataFusion query instead of failing the
task (#71445)
DataFusionToolset._query() re-raised SQLSafetyError as-is after logging it,
so any query that failed safety validation crashed the whole task instead
of giving the agent a chance to correct it. That validation isn't limited
to deliberate write attempts (CREATE/DROP/etc.) -- sqlglot parses in RAISE
mode, so a plain SQL syntax typo the agent generates also fails as a
SQLSafetyError and hit the same crash.
SQLToolset already treats every query error this way (wrapped in
ModelRetry, bounded by the tool's max_retries=1) so the model sees the
error and can fix its own SQL within the run. Bring DataFusionToolset's
SQLSafetyError handling in line with it.
---
.../airflow/providers/common/ai/toolsets/datafusion.py | 5 ++++-
.../tests/unit/common/ai/toolsets/test_datafusion.py | 18 +++++++++++++++++-
2 files changed, 21 insertions(+), 2 deletions(-)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py
index ef888134811..ebc99a297b1 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py
@@ -228,7 +228,10 @@ class DataFusionToolset(AbstractToolset[Any]):
)
except SQLSafetyError as ex:
log.warning("query failed SQL safety validation: %s", ex)
- raise
+ raise ModelRetry(
+ f"error: {ex!s}. Only read-only SELECT-family queries are
allowed unless "
+ "allow_writes is enabled; check the SQL syntax and statement
type, then try again."
+ ) from ex
except QueryExecutionException as ex:
if self._is_retryable_query_error(ex):
raise ModelRetry(
diff --git
a/providers/common/ai/tests/unit/common/ai/toolsets/test_datafusion.py
b/providers/common/ai/tests/unit/common/ai/toolsets/test_datafusion.py
index 05481efa115..f5d2572b9b1 100644
--- a/providers/common/ai/tests/unit/common/ai/toolsets/test_datafusion.py
+++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_datafusion.py
@@ -225,7 +225,7 @@ class TestDataFusionToolsetQuery:
ts = DataFusionToolset([cfg])
ts._engine = _make_mock_engine()
- with pytest.raises(SQLSafetyError, match="Statement type 'Create' is
not allowed"):
+ with pytest.raises(ModelRetry, match="Statement type 'Create' is not
allowed"):
asyncio.run(
ts.call_tool(
"query",
@@ -235,6 +235,22 @@ class TestDataFusionToolsetQuery:
)
)
+ def test_sql_syntax_error_raises_model_retry(self):
+ cfg = _make_mock_datasource_config()
+ ts = DataFusionToolset([cfg])
+ ts._engine = _make_mock_engine()
+
+ with pytest.raises(ModelRetry) as exc_info:
+ asyncio.run(
+ ts.call_tool(
+ "query",
+ {"sql": "SELECT * FROM t WHERE"},
+ ctx=MagicMock(spec=RunContext),
+ tool=MagicMock(spec=ToolsetTool),
+ )
+ )
+ assert isinstance(exc_info.value.__cause__, SQLSafetyError)
+
def test_allows_create_table_when_writes_enabled(self):
cfg = _make_mock_datasource_config()
ts = DataFusionToolset([cfg], allow_writes=True)
