hussein-awala commented on code in PR #64610:
URL: https://github.com/apache/airflow/pull/64610#discussion_r3540053342
##########
airflow-core/src/airflow/utils/sqlalchemy.py:
##########
@@ -61,6 +61,26 @@ def get_dialect_name(session: Session) -> str | None:
return getattr(bind.dialect, "name", None)
+def apply_regex_query_timeout(session: Session) -> None:
+ """
+ Bound the runtime of a user-supplied regex filter on PostgreSQL.
+
+ Reads the ``[api] regexp_query_timeout`` config (in seconds) and sets a
transaction-local
+ ``statement_timeout`` (via ``set_config(..., is_local=True)``) so it
applies only to the current
+ transaction and resets automatically at commit/rollback. This is a ReDoS
safeguard: a malicious
+ pattern is aborted instead of pinning a database backend. No-op on
non-PostgreSQL backends and
+ when the configured timeout is ``0`` (disabled).
+ """
+ timeout_seconds = conf.getint("api", "regexp_query_timeout")
+ if timeout_seconds <= 0:
+ return
+ if get_dialect_name(session) == "postgresql":
+ timeout_ms = timeout_seconds * 1000
+ session.execute(
+ text("SELECT set_config('statement_timeout', :timeout,
true)").bindparams(timeout=str(timeout_ms))
+ )
+
Review Comment:
Addressed both points. `apply_regex_query_timeout` is now a context manager,
and both asset-event routes apply it automatically around query execution, so
it can't be forgotten. I kept it out of `to_orm` itself because
`to_orm(select)` only transforms the statement (it has no session and doesn't
execute) so it can't own the set/reset lifecycle; wrapping the execution is the
reliable place.
On the side effect: the context manager now sets the transaction-local
`statement_timeout` on enter and resets it to the default on exit, so the bound
is scoped to just the regexp query (the count + fetch) and no longer leaks to
other statements in the request's transaction.
--
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]
