ashb commented on code in PR #64610:
URL: https://github.com/apache/airflow/pull/64610#discussion_r3559979109
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -1635,6 +1635,30 @@ api:
type: boolean
example: ~
default: "True"
+ regexp_query_timeout:
+ description: |
+ Timeout, in seconds, for regular-expression based query filters in the
API. This single
+ value both enables the feature and bounds its runtime: a value of
``0`` (the default)
+ disables regexp filtering entirely, and any positive value enables it
while capping how
+ long such a query may run.
+
+ When enabled, it allows additional filtering features that require
database-side regexp
+ matching, such as the ``partition_key_regexp_pattern`` filter on ``GET
/assets/events``
+ and on the Execution API asset-event endpoints. It is disabled by
default for security
+ reasons: a user-supplied regular expression is passed to the
database's own regex engine,
+ which is a Regular expression Denial of Service (ReDoS) vector: a
crafted pattern can
+ force the engine into pathological backtracking and consume database
backend CPU.
+
+ The timeout is the primary runtime mitigation for that risk: on
PostgreSQL it is enforced
+ as a transaction-local ``statement_timeout`` so a pathological pattern
is aborted instead
+ of pinning a database backend (PostgreSQL has no built-in regex step
limit, unlike MySQL's
+ ``regexp_time_limit``). Keep it as low as your legitimate queries
allow. Coupling the
+ on/off switch with the timeout intentionally makes it impossible to
enable regexp
+ filtering without a runtime bound in place.
+ version_added: 3.4.0
+ type: integer
Review Comment:
Float maybe? To allow 0.5 etc?
##########
airflow-core/docs/authoring-and-scheduling/assets.rst:
##########
@@ -247,7 +247,19 @@ Inlet asset events can be read with the ``inlet_events``
accessor in the executi
Each value in the ``inlet_events`` mapping is a sequence-like object that
orders past events of a given asset by ``timestamp``, earliest to latest. It
supports most of Python's list interface, so you can use ``[-1]`` to access the
last event, ``[-2:]`` for the last two, etc. The accessor is lazy and only hits
the database when you access items inside it.
-The accessor also supports chaining methods to filter events before fetching
them. For example, to retrieve only events where specific ``extra`` keys match
given values:
+The accessor also supports chaining methods to filter events before fetching
them. For example, to retrieve only events matching a specific partition key
regular expression:
+
+.. code-block:: python
+
+ @task(inlets=[regional_sales])
+ def process_us_sales(*, inlet_events):
+ us_events =
inlet_events[regional_sales].partition_key_regexp_pattern(r"^us\|")
+ for event in us_events:
+ print(event.extra, event.partition_key)
+
+For an exact partition key match, use ``.partition_key(value)`` instead.
Regexp filtering is opt-in: it is enabled only by setting ``[api]
regexp_query_timeout`` to a positive number of seconds, which also bounds the
query runtime; see the config for the security trade-off.
Review Comment:
I think we can do a direct link to the config via
```
:conf:`api__regexp_query__timeout`
```
Something like this exists
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py:
##########
Review Comment:
More Parameteization here too?
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py:
##########
@@ -1079,6 +1079,172 @@ def test_should_mask_sensitive_extra(self, test_client,
session):
}
+class TestGetAssetEventsPartitionKeyRegex(TestAssets):
+ """Tests for partition_key_regexp_pattern regex filter on GET
/assets/events.
+
+ Patterns are written to work consistently across PostgreSQL (~),
+ MySQL (REGEXP), and SQLite (re.match), including both anchored and
+ unanchored expressions where appropriate.
+ """
+
+ @pytest.fixture(autouse=True)
+ def _enable_regexp_query_filters(self):
+ with conf_vars({("api", "regexp_query_timeout"): "30"}):
+ yield
+
+ @provide_session
+ def _create_partition_key_test_data(self, *, session=None):
+ _create_assets(session=session)
+ events = [
+ AssetEvent(
+ asset_id=1,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r1",
+ partition_key="2024-01-01",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=2,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r2",
+ partition_key="2024-01-02",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=1,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r3",
+ partition_key="us|2024-01-01",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=2,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r4",
+ partition_key="eu|2024-01-01",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=1,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r5",
+ partition_key="apac|2024-03-20",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=1,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r6",
+ partition_key=None,
+ timestamp=DEFAULT_DATE,
+ ),
+ ]
+ session.add_all(events)
+ session.commit()
+
+ @pytest.mark.parametrize(
+ ("partition_key_regexp_pattern", "expected_count"),
+ [
+ ("^2024-01-01$", 1),
+ ("^2024-01-", 2),
+ ("^us\\|", 1),
+ (".*\\|2024-01-01$", 2),
+ ("^(us|eu)\\|", 2),
+ ("^nonexistent", 0),
+ ],
+ )
+ def test_partition_key_regexp_pattern_filtering(
+ self, test_client, partition_key_regexp_pattern, expected_count
+ ):
+ self._create_partition_key_test_data()
+ response = test_client.get(
+ "/assets/events", params={"partition_key_regexp_pattern":
partition_key_regexp_pattern}
+ )
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == expected_count
+
+ @pytest.mark.parametrize(
+ ("params", "expected_count"),
+ [
+ ({"partition_key_regexp_pattern": "^us\\|", "asset_id": "1"}, 1),
+ ({"partition_key_regexp_pattern": "^us\\|", "asset_id": "2"}, 0),
+ ({"partition_key_regexp_pattern": ".*\\|2024-01-01$",
"source_dag_id": "d"}, 2),
+ ({"partition_key_regexp_pattern": ".*\\|2024-01-01$",
"source_dag_id": "other"}, 0),
+ ],
+ )
+ def test_partition_key_regexp_pattern_combined_filters(self, test_client,
params, expected_count):
+ self._create_partition_key_test_data()
+ response = test_client.get("/assets/events", params=params)
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == expected_count
+
+ def test_partition_key_regexp_pattern_invalid_regex_returns_400(self,
test_client):
+ response = test_client.get(
+ "/assets/events", params={"partition_key_regexp_pattern":
"[invalid(regex"}
+ )
+ assert response.status_code == 400
+ assert "Invalid regular expression" in response.json()["detail"]
+
+ def test_partition_key_regexp_pattern_disabled_returns_400(self,
test_client):
+ with conf_vars({("api", "regexp_query_timeout"): "0"}):
+ response = test_client.get("/assets/events",
params={"partition_key_regexp_pattern": "^2024-"})
+ assert response.status_code == 400
+ assert "disabled" in response.json()["detail"]
+
+ def test_exact_match_works_when_regex_disabled(self, test_client):
+ self._create_partition_key_test_data()
+ with conf_vars({("api", "regexp_query_timeout"): "0"}):
+ response = test_client.get("/assets/events",
params={"partition_key": "2024-01-01"})
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == 1
+
+ def test_partition_key_exact_match_via_regex(self, test_client):
+ self._create_partition_key_test_data()
+ response = test_client.get("/assets/events",
params={"partition_key_regexp_pattern": "^2024-01-01$"})
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == 1
+
+ def test_partition_key_exact_match(self, test_client):
+ self._create_partition_key_test_data()
+ response = test_client.get("/assets/events", params={"partition_key":
"2024-01-01"})
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == 1
+
+ def test_partition_key_exact_match_composite(self, test_client):
+ self._create_partition_key_test_data()
+ response = test_client.get("/assets/events", params={"partition_key":
"us|2024-01-01"})
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == 1
+
+ def test_partition_key_exact_match_no_match(self, test_client):
+ self._create_partition_key_test_data()
+ response = test_client.get("/assets/events", params={"partition_key":
"nonexistent"})
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == 0
Review Comment:
These could be 4 cases in a single parameteized test i think?
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py:
##########
@@ -1079,6 +1079,172 @@ def test_should_mask_sensitive_extra(self, test_client,
session):
}
+class TestGetAssetEventsPartitionKeyRegex(TestAssets):
+ """Tests for partition_key_regexp_pattern regex filter on GET
/assets/events.
+
+ Patterns are written to work consistently across PostgreSQL (~),
+ MySQL (REGEXP), and SQLite (re.match), including both anchored and
+ unanchored expressions where appropriate.
+ """
+
+ @pytest.fixture(autouse=True)
+ def _enable_regexp_query_filters(self):
+ with conf_vars({("api", "regexp_query_timeout"): "30"}):
+ yield
+
+ @provide_session
+ def _create_partition_key_test_data(self, *, session=None):
+ _create_assets(session=session)
+ events = [
+ AssetEvent(
+ asset_id=1,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r1",
+ partition_key="2024-01-01",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=2,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r2",
+ partition_key="2024-01-02",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=1,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r3",
+ partition_key="us|2024-01-01",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=2,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r4",
+ partition_key="eu|2024-01-01",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=1,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r5",
+ partition_key="apac|2024-03-20",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=1,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r6",
+ partition_key=None,
+ timestamp=DEFAULT_DATE,
+ ),
+ ]
+ session.add_all(events)
+ session.commit()
+
+ @pytest.mark.parametrize(
+ ("partition_key_regexp_pattern", "expected_count"),
+ [
+ ("^2024-01-01$", 1),
+ ("^2024-01-", 2),
+ ("^us\\|", 1),
+ (".*\\|2024-01-01$", 2),
+ ("^(us|eu)\\|", 2),
+ ("^nonexistent", 0),
+ ],
+ )
+ def test_partition_key_regexp_pattern_filtering(
+ self, test_client, partition_key_regexp_pattern, expected_count
+ ):
+ self._create_partition_key_test_data()
Review Comment:
Nit: it seems wasteful to create the same rows for each parameter set? Class
scoped fixture perhaps? That could also delete/clean on teardown that way
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py:
##########
@@ -1079,6 +1079,172 @@ def test_should_mask_sensitive_extra(self, test_client,
session):
}
+class TestGetAssetEventsPartitionKeyRegex(TestAssets):
+ """Tests for partition_key_regexp_pattern regex filter on GET
/assets/events.
+
+ Patterns are written to work consistently across PostgreSQL (~),
+ MySQL (REGEXP), and SQLite (re.match), including both anchored and
+ unanchored expressions where appropriate.
+ """
+
+ @pytest.fixture(autouse=True)
+ def _enable_regexp_query_filters(self):
+ with conf_vars({("api", "regexp_query_timeout"): "30"}):
+ yield
+
+ @provide_session
+ def _create_partition_key_test_data(self, *, session=None):
+ _create_assets(session=session)
+ events = [
+ AssetEvent(
+ asset_id=1,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r1",
+ partition_key="2024-01-01",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=2,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r2",
+ partition_key="2024-01-02",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=1,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r3",
+ partition_key="us|2024-01-01",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=2,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r4",
+ partition_key="eu|2024-01-01",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=1,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r5",
+ partition_key="apac|2024-03-20",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=1,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r6",
+ partition_key=None,
+ timestamp=DEFAULT_DATE,
+ ),
+ ]
+ session.add_all(events)
+ session.commit()
+
+ @pytest.mark.parametrize(
+ ("partition_key_regexp_pattern", "expected_count"),
+ [
+ ("^2024-01-01$", 1),
+ ("^2024-01-", 2),
+ ("^us\\|", 1),
+ (".*\\|2024-01-01$", 2),
+ ("^(us|eu)\\|", 2),
+ ("^nonexistent", 0),
+ ],
+ )
+ def test_partition_key_regexp_pattern_filtering(
+ self, test_client, partition_key_regexp_pattern, expected_count
+ ):
+ self._create_partition_key_test_data()
+ response = test_client.get(
+ "/assets/events", params={"partition_key_regexp_pattern":
partition_key_regexp_pattern}
+ )
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == expected_count
+
+ @pytest.mark.parametrize(
+ ("params", "expected_count"),
+ [
+ ({"partition_key_regexp_pattern": "^us\\|", "asset_id": "1"}, 1),
+ ({"partition_key_regexp_pattern": "^us\\|", "asset_id": "2"}, 0),
+ ({"partition_key_regexp_pattern": ".*\\|2024-01-01$",
"source_dag_id": "d"}, 2),
+ ({"partition_key_regexp_pattern": ".*\\|2024-01-01$",
"source_dag_id": "other"}, 0),
+ ],
+ )
+ def test_partition_key_regexp_pattern_combined_filters(self, test_client,
params, expected_count):
+ self._create_partition_key_test_data()
+ response = test_client.get("/assets/events", params=params)
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == expected_count
+
+ def test_partition_key_regexp_pattern_invalid_regex_returns_400(self,
test_client):
+ response = test_client.get(
+ "/assets/events", params={"partition_key_regexp_pattern":
"[invalid(regex"}
+ )
+ assert response.status_code == 400
+ assert "Invalid regular expression" in response.json()["detail"]
+
+ def test_partition_key_regexp_pattern_disabled_returns_400(self,
test_client):
+ with conf_vars({("api", "regexp_query_timeout"): "0"}):
+ response = test_client.get("/assets/events",
params={"partition_key_regexp_pattern": "^2024-"})
+ assert response.status_code == 400
+ assert "disabled" in response.json()["detail"]
+
+ def test_exact_match_works_when_regex_disabled(self, test_client):
+ self._create_partition_key_test_data()
+ with conf_vars({("api", "regexp_query_timeout"): "0"}):
+ response = test_client.get("/assets/events",
params={"partition_key": "2024-01-01"})
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == 1
+
+ def test_partition_key_exact_match_via_regex(self, test_client):
+ self._create_partition_key_test_data()
+ response = test_client.get("/assets/events",
params={"partition_key_regexp_pattern": "^2024-01-01$"})
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == 1
+
+ def test_partition_key_exact_match(self, test_client):
+ self._create_partition_key_test_data()
+ response = test_client.get("/assets/events", params={"partition_key":
"2024-01-01"})
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == 1
+
+ def test_partition_key_exact_match_composite(self, test_client):
+ self._create_partition_key_test_data()
+ response = test_client.get("/assets/events", params={"partition_key":
"us|2024-01-01"})
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == 1
+
+ def test_partition_key_exact_match_no_match(self, test_client):
+ self._create_partition_key_test_data()
+ response = test_client.get("/assets/events", params={"partition_key":
"nonexistent"})
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == 0
+
+ def test_partition_key_and_pattern_combined(self, test_client):
+ # Both filters are allowed and combine with AND: a disjoint pair
yields no results.
+ self._create_partition_key_test_data()
+ response = test_client.get(
+ "/assets/events",
+ params={"partition_key": "2024-01-01",
"partition_key_regexp_pattern": "^2025-"},
+ )
+ assert response.status_code == 200
Review Comment:
I would have expected providing both to return a 400?
##########
airflow-core/src/airflow/utils/sqlalchemy.py:
##########
@@ -62,6 +62,34 @@ def get_dialect_name(session: Session) -> str | None:
return getattr(bind.dialect, "name", None)
[email protected]
+def apply_regex_query_timeout(session: Session) -> Generator[None, None, None]:
+ """
+ Bound the runtime of a user-supplied regex filter on PostgreSQL, scoped to
the wrapped query.
+
+ Reads the ``[api] regexp_query_timeout`` config (in seconds) and, on
PostgreSQL, sets a
+ transaction-local ``statement_timeout`` (via ``set_config(...,
is_local=True)``) for the
+ duration of the ``with`` block, then resets it to the default so it does
not affect any other
+ statement running later in the same transaction. 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`` (which also means regexp filtering is
disabled).
+ """
+ timeout_seconds = conf.getint("api", "regexp_query_timeout")
+ if timeout_seconds <= 0 or get_dialect_name(session) != "postgresql":
Review Comment:
So this timeout config doesn't apply to mysql? That feels risky?
--
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]
