hussein-awala commented on code in PR #64610:
URL: https://github.com/apache/airflow/pull/64610#discussion_r3564817306
##########
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:
This was intentional: in an earlier review @pierrejeambrun asked to drop the
mutual-exclusion 400 and let `partition_key` and `partition_key_regexp_pattern`
combine with AND (so providing both just narrows the results, yielding 0 when
disjoint) -- see
https://github.com/apache/airflow/pull/64610#discussion_r3536685452. I've kept
the AND behavior for now. @pierrejeambrun / @ashb, could you align on whether
you'd prefer 400 vs AND here and I'll adjust to match?
--
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]
