hussein-awala commented on code in PR #64610:
URL: https://github.com/apache/airflow/pull/64610#discussion_r3564817043
##########
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:
Moved the row creation into an autouse fixture so it's no longer called
explicitly in each test. I kept it function-scoped rather than class-scoped
because this suite relies on per-test DB rollback for isolation, and a
class-scoped fixture that commits shared rows would need manual teardown and
could bleed across tests. Combined with the parameterization below, the
duplication is gone.
##########
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:
Done -- folded the exact-match cases into a single parametrized
`test_partition_key_exact_match`.
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py:
##########
Review Comment:
Done -- parametrized the exact-match partition-key cases in the
execution-API tests too.
##########
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":
+ yield
+ return
+ timeout_ms = timeout_seconds * 1000
+ session.execute(
+ text("SELECT set_config('statement_timeout', :timeout,
true)").bindparams(timeout=str(timeout_ms))
+ )
+ try:
+ yield
+ finally:
+ # Reset to the default for the remainder of the transaction so the
timeout only bounds the
+ # regexp query and not unrelated statements in the same request.
+ session.execute(text("SELECT set_config('statement_timeout', '0',
true)"))
Review Comment:
Addressed: the context manager now captures the current `statement_timeout`
(PostgreSQL) / `max_execution_time` (MySQL) before overriding it and restores
that value on exit, instead of resetting to `0`, so a global timeout is
preserved.
--
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]
