kaxil commented on code in PR #73466:
URL: https://github.com/apache/airflow/pull/73466#discussion_r4065187494
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py:
##########
@@ -190,20 +190,19 @@ def _list_tables(self) -> str:
return json.dumps({"error": str(ex)})
def _get_schema(self, table_name: str) -> str:
- engine = self._get_engine()
- # session_context lookup is required here instead of
engine.registered_tables,
- # because registered_tables only tracks tables registered via
datasource config.
- # When allow_writes is enabled, the agent may create temporary
in-memory tables
- # that would not be captured there.
- if not engine.session_context.table_exist(table_name):
- return json.dumps({"error": f"Table {table_name!r} is not
available"})
- # Intentionally using session_context instead of engine.get_schema() —
- # the latter returns a pre-formatted string intended for other
operators,
- # not a JSON-compatible format.
- # TODO: refactor engine.get_schema() to return JSON and update this
accordingly
- table = engine.session_context.table(table_name)
- columns = [{"name": f.name, "type": str(f.type)} for f in
table.schema()]
- return json.dumps(columns)
+ try:
+ engine = self._get_engine()
Review Comment:
Coming back to this one, I cleared it too fast in my first pass.
Before this change `_get_schema` had no try at all. `_get_engine()` calls
`register_datasource` for every config, and that can raise
`ObjectStoreCreationException` (engine.py:86), `ValueError` for a duplicate
table name or an unknown conn_type, or
`AirflowOptionalProviderFeatureException` if someone points an `aws` connection
at this without the amazon provider installed (engine.py:163). All of those
used to escape and fail the task with a message the user can act on.
Inside the try they turn into `{"error": ...}` handed back to the model.
`_list_tables` was already doing that, so `get_schema` was the remaining tool
that would hard-fail on a broken datasource. `_query` only catches
`SQLSafetyError` and `QueryExecutionException`, so it would still raise, but
the model has to get that far first, and two `{"error": ...}` replies from the
discovery tools are a fair reason for it to stop and answer from what it
already knows. Misconfigured S3 credentials would show up as a successful task.
Could `_get_engine()` move back above the try, or the except list the three
types?
Unrelated to the code: the description says this matches "as its `query`
tool already does", but `query` doesn't catch arbitrary exceptions.
`_list_tables` is the closer comparison.
--
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]
