kaxil commented on code in PR #76783:
URL: https://github.com/apache/airflow/pull/76783#discussion_r4551171339
##########
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:
Added it to `sql.rst` in ba8c36109c4: with a secret the masker can't remove,
the task fails with a `RuntimeError` and a rule on the class doesn't match.
##########
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:
Done in ba8c36109c4.
##########
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:
Done in ba8c36109c4.
##########
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:
Good catch, it reads the columns from the inspector now, same as
`list_tables`, so a hook overriding `get_table_schema` or `inspector` can't put
the failure back to `ModelRetry`.
--
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]
