Lee-W commented on code in PR #76783:
URL: https://github.com/apache/airflow/pull/76783#discussion_r4547119130
##########
providers/common/ai/docs/toolsets/sql.rst:
##########
@@ -450,15 +450,33 @@ subqueries and joins.
walk rejects once ``allowed_tables`` is active. Statements that modify
data, ``COPY`` among them, are rejected either way while ``allow_writes``
is ``False``.
-- It does not classify failures. A connection error or a typo in a column
- name reaching ``list_tables``, ``get_schema`` or ``query`` becomes one
- ``ModelRetry``, so the two are treated the same way until the retry budget
- (:ref:`toolset-retry-budget`) runs out and the task fails for Airflow to
retry. Two paths do not raise:
- ``check_query`` catches its own errors and reports them back as a normal
- ``{"valid": false, ...}`` result, and ``get_schema`` returns a normal
- ``{"error": ...}`` result instead of raising when the requested table is
- outside ``allowed_tables``; other ``get_schema`` failures still raise and
- still become a ``ModelRetry``.
+- It tells an unreachable database from a bad statement by when the error
+ happens, not by what it says. Failing to reach the database fails the
+ attempt straight away for Airflow to retry: the connection is not defined,
+ connecting fails (an unreachable host, rejected credentials), SQLAlchemy
+ cannot be set up for ``list_tables`` or ``get_schema`` (a dialect it cannot
+ load, a token it cannot fetch), or SQLAlchemy reports the connection lost
+ while they run. With :doc:`../durable_execution` the retry queries the
+ database again instead of replaying the failure. The task fails with
+ :class:`~airflow.providers.common.ai.exceptions.SQLToolsetConnectionError`,
+ whose message reads ``Could not reach the database of connection '<id>':
...``,
+ or ``Connection '<id>' does not provide a DbApiHook`` when the connection is
+ of the wrong type. Match that class in a retry policy
Review Comment:
Nit: the docs say to match `SQLToolsetConnectionError` in a retry policy,
but the class docstring says a `RuntimeError` with the masked message is raised
instead when the driver message carries a secret that cannot be masked, and a
rule on the class does not match then. One sentence naming that case would help.
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py:
##########
@@ -104,6 +108,33 @@
}
+def _unreachable(conn_id: str, error: Exception) -> SQLToolsetConnectionError:
+ return SQLToolsetConnectionError(f"Could not reach the database of
connection {conn_id!r}: {error}")
+
+
+def _raise_connection_failures(get_conn: Callable[[], Any], conn_id: str) ->
Callable[[], Any]:
+ """Wrap a hook's ``get_conn`` so a failure to connect is told apart from a
failing statement."""
+
+ def wrapper() -> Any:
+ try:
+ return get_conn()
+ except Exception as e:
+ raise _unreachable(conn_id, e) from e
+
+ return wrapper
+
+
+def _lost_connection(error: Exception) -> bool:
Review Comment:
```suggestion
def _is_lost_connection(error: Exception) -> bool:
```
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py:
##########
@@ -104,6 +108,33 @@
}
+def _unreachable(conn_id: str, error: Exception) -> SQLToolsetConnectionError:
Review Comment:
```suggestion
def _build_connection_error(conn_id: str, error: Exception) ->
SQLToolsetConnectionError:
```
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py:
##########
@@ -497,8 +548,8 @@ def _get_schema(self, table_name: str, name_contains: str |
None = None) -> str:
schema, table = self._split_table_identifier(table_name)
if not self._is_ref_allowed("", schema, table):
return dumps_masked({"error": f"Table {table_name!r} is not in the
allowed tables list."})
- hook = self._get_db_hook()
- columns = hook.get_table_schema(table, schema=schema)
+ self._get_inspector() # get_table_schema reads through it
Review Comment:
This calls `_get_inspector()` only to warm the `hook.inspector` cache, so
the next line's `get_table_schema` happens to reuse the inspector built inside
the `try`. That depends on `inspector` staying a `cached_property` and on
`get_table_schema` reading through it. If a hook overrides either one, the
warm-up does nothing and the failure goes back to `ModelRetry` without any
error.
`_list_tables` already uses the inspector directly. Could `_get_schema` do
the same?
```python
columns = [
{"name": col["name"], "type": str(col["type"])}
for col in self._get_inspector().get_columns(table, schema=schema)
]
```
--
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]
