ColtenOuO commented on code in PR #73384:
URL: https://github.com/apache/airflow/pull/73384#discussion_r4091297055
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py:
##########
@@ -211,20 +215,20 @@ def _query(self, sql: str) -> str:
_validate_sql(sql)
engine = self._get_engine()
- pydict = engine.execute_query(sql)
+ try:
+ pydict = engine.session_context.sql(sql).limit(self._max_rows
+ 1).to_pydict()
+ except Exception as e:
Review Comment:
The DataFusion Python bindings expose no native exception types, as the
toolset docs already note, that's why the retry decision is made by matching
the error message against regular expressions rather than catching a specific
type.
I also traced it upstream: `DataFusionEngine.execute_query` handles this the
same way, with
a bare `except Exception` at L121:
https://github.com/apache/airflow/blob/main/providers/common/sql/src/airflow/providers/common/sql/datafusion/engine.py#L121
So I don't think we can narrow it here, WDYT
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
