kaxil commented on code in PR #69940: URL: https://github.com/apache/airflow/pull/69940#discussion_r3681678465
########## airflow-core/src/airflow/api_fastapi/common/db/assets.py: ########## @@ -0,0 +1,54 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from sqlalchemy import func, select +from sqlalchemy.orm import subqueryload + +from airflow.models.asset import AssetEvent, AssetModel, AssetWatcherModel + +if TYPE_CHECKING: + from sqlalchemy.sql import Select + + +def generate_assets_with_last_event_query() -> Select: + """Fetch Assets outer-joined to their latest AssetEvent id/timestamp.""" + max_asset_event_id_query = ( + select(AssetEvent.asset_id, func.max(AssetEvent.id).label("max_asset_event_id")) + .group_by(AssetEvent.asset_id) + .subquery() + ) + + return ( + select( + AssetModel, + AssetEvent.id.label("last_asset_event_id"), + AssetEvent.timestamp.label("last_asset_event_timestamp"), + ) + .outerjoin(max_asset_event_id_query, AssetModel.id == max_asset_event_id_query.c.asset_id) + .outerjoin(AssetEvent, AssetEvent.id == max_asset_event_id_query.c.max_asset_event_id) + .options( + subqueryload(AssetModel.scheduled_dags), Review Comment: Each of these five `subqueryload`s re-runs the whole parent statement as an inner subquery, `GROUP BY asset_event` and all. I dumped the executed SQL for a 2-row page and got six statements, five of which contain `SELECT asset_event.asset_id, max(asset_event.id) ... GROUP BY asset_event.asset_id`; with the count query that is seven aggregations of `asset_event` for one list page. `selectinload` emits a plain `WHERE asset_id IN (<page ids>)` instead and is already what the sibling helper in `common/db/dags.py` uses, so is there a reason not to switch here? ########## airflow-core/src/airflow/ui/src/pages/AssetsList/AssetsList.tsx: ########## @@ -129,18 +113,31 @@ export const AssetsList = () => { const { setTableURLState, tableURLState } = useTableURLState(); const { pagination, sorting } = tableURLState; const [sort] = sorting; - const orderBy = sort ? [`${sort.desc ? "-" : ""}${sort.id}`] : undefined; + const orderBy = sort ? [`${sort.desc ? "-" : ""}${sort.id}`] : ["-last_asset_event_timestamp"]; - const { onClose, onOpen, open } = useDisclosure(); + const { filterConfigs, handleFiltersChange, initialValues } = useFiltersHandler(assetsFilterKeys); - const { data, error, isLoading } = useAssetServiceGetAssets({ + const lastAssetEventTimestampGte = searchParams.get(SearchParamsKeys.LAST_ASSET_EVENT_TIMESTAMP_GTE); + const lastAssetEventTimestampLte = searchParams.get(SearchParamsKeys.LAST_ASSET_EVENT_TIMESTAMP_LTE); + const groupArg = useAdvancedSearchArg({ + patternApiKey: "groupPattern", + prefixApiKey: "groupPrefixPattern", + storageKey: SearchParamsKeys.GROUP_PATTERN, + value: searchParams.get(SearchParamsKeys.GROUP_PATTERN), + }); + + const { data, error, isLoading } = useAssetServiceGetAssetsUi({ Review Comment: This changes the key the list caches under, and `CreateAssetEventModal`'s success handler only invalidates `GetAssetEvents` plus the dag run and task instance keys. It never invalidated the old assets key either, so the play button on a row already left the timestamp stale, but now that last event is the default sort the row also stays in the wrong position. Adding `[useAssetServiceGetAssetsUiKey]` to that `queryKeys` array would cover both, alongside the `[useDagServiceGetDagsUiKey]` that is already there. ########## airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_assets.py: ########## @@ -471,3 +480,278 @@ def test_non_partitioned_asset_inactive_true_when_deactivated(self, test_client, body = response.json() assert len(body["events"]) == 1 assert body["events"][0]["asset_inactive"] is True + + +class TestGetAssetsUi: + @pytest.fixture(autouse=True) + def cleanup_assets(self): + clear_db_assets() + + yield + + clear_db_assets() + + def test_should_respond_401(self, unauthenticated_test_client): + response = unauthenticated_test_client.get("/assets") + assert response.status_code == 401 + + def test_should_respond_403(self, unauthorized_test_client): + response = unauthorized_test_client.get("/assets") + assert response.status_code == 403 + + def test_should_respond_200(self, test_client, session): + asset = AssetModel(name="ui_asset", uri="s3://bucket/ui_asset", group="asset") + session.add(asset) + session.add(AssetActive.for_asset(asset)) + session.commit() + + response = test_client.get("/assets") + assert response.status_code == 200 + body = response.json() + assert body["total_entries"] == 1 + assert body["assets"][0]["name"] == "ui_asset" + + def test_sort_by_last_asset_event_timestamp(self, test_client, session): + older = AssetModel(name="older", uri="s3://bucket/older", group="asset") + newer = AssetModel(name="newer", uri="s3://bucket/newer", group="asset") + session.add_all([older, newer]) + session.add(AssetActive.for_asset(older)) + session.add(AssetActive.for_asset(newer)) + session.flush() + + base = pendulum.datetime(2024, 1, 1) + session.add(AssetEvent(asset_id=older.id, timestamp=base)) + session.add(AssetEvent(asset_id=newer.id, timestamp=base.add(days=1))) + session.commit() + + response = test_client.get("/assets?order_by=last_asset_event_timestamp") + assert response.status_code == 200 + assert [a["name"] for a in response.json()["assets"]] == ["older", "newer"] + + response = test_client.get("/assets?order_by=-last_asset_event_timestamp") + assert response.status_code == 200 + assert [a["name"] for a in response.json()["assets"]] == ["newer", "older"] + + def test_default_sort_is_last_asset_event_timestamp_desc(self, test_client, session): + older = AssetModel(name="older", uri="s3://bucket/older_default", group="asset") + newer = AssetModel(name="newer", uri="s3://bucket/newer_default", group="asset") + session.add_all([older, newer]) + session.add(AssetActive.for_asset(older)) + session.add(AssetActive.for_asset(newer)) + session.flush() + + base = pendulum.datetime(2024, 1, 1) + session.add(AssetEvent(asset_id=older.id, timestamp=base)) + session.add(AssetEvent(asset_id=newer.id, timestamp=base.add(days=1))) + session.commit() + + response = test_client.get("/assets") + assert response.status_code == 200 + assert [a["name"] for a in response.json()["assets"]] == ["newer", "older"] + + def test_sort_by_group(self, test_client, session): + billing = AssetModel(name="billing_asset", uri="s3://bucket/billing_sort", group="billing") + marketing = AssetModel(name="marketing_asset", uri="s3://bucket/marketing_sort", group="marketing") + session.add_all([billing, marketing]) + session.add(AssetActive.for_asset(billing)) + session.add(AssetActive.for_asset(marketing)) + session.commit() + + response = test_client.get("/assets?order_by=group") + assert response.status_code == 200 + assert [a["name"] for a in response.json()["assets"]] == ["billing_asset", "marketing_asset"] + + response = test_client.get("/assets?order_by=-group") + assert response.status_code == 200 + assert [a["name"] for a in response.json()["assets"]] == ["marketing_asset", "billing_asset"] + + def test_filter_by_group_pattern(self, test_client, session): + billing = AssetModel(name="billing_asset", uri="s3://bucket/billing", group="billing") + marketing = AssetModel(name="marketing_asset", uri="s3://bucket/marketing", group="marketing") + session.add_all([billing, marketing]) + session.add(AssetActive.for_asset(billing)) + session.add(AssetActive.for_asset(marketing)) + session.commit() + + response = test_client.get("/assets?group_pattern=bill") + assert response.status_code == 200 + assert [a["name"] for a in response.json()["assets"]] == ["billing_asset"] + + def test_filter_by_group_prefix_pattern(self, test_client, session): + billing = AssetModel(name="billing_asset", uri="s3://bucket/billing_prefix", group="billing") + rebilling = AssetModel(name="rebilling_asset", uri="s3://bucket/rebilling", group="rebilling") + session.add_all([billing, rebilling]) + session.add(AssetActive.for_asset(billing)) + session.add(AssetActive.for_asset(rebilling)) + session.commit() + + # Prefix match anchors at the start, so "bill" excludes "rebilling" (substring would not). + response = test_client.get("/assets?group_prefix_pattern=bill") + assert response.status_code == 200 + assert [a["name"] for a in response.json()["assets"]] == ["billing_asset"] + + def test_filter_by_last_asset_event_timestamp_range(self, test_client, session): + older = AssetModel(name="older", uri="s3://bucket/older_range", group="asset") + newer = AssetModel(name="newer", uri="s3://bucket/newer_range", group="asset") + session.add_all([older, newer]) + session.add(AssetActive.for_asset(older)) + session.add(AssetActive.for_asset(newer)) + session.flush() + + base = pendulum.datetime(2024, 1, 1) + session.add(AssetEvent(asset_id=older.id, timestamp=base)) + session.add(AssetEvent(asset_id=newer.id, timestamp=base.add(days=10))) + session.commit() + + response = test_client.get( + "/assets", params={"last_asset_event_timestamp_gte": base.add(days=5).isoformat()} + ) + assert response.status_code == 200 + assert [a["name"] for a in response.json()["assets"]] == ["newer"] + + def test_aliases_present_for_asset_via_alias(self, test_client, session): + """ + Regression test for https://github.com/apache/airflow/issues/58058: + aliases must be visible via the Assets endpoints, not just fetchable + through the CLI/DB. + """ + asset = AssetModel(name="alias_target", uri="s3://bucket/alias_target", group="asset") + alias = AssetAliasModel(name="my-alias", group="") + session.add_all([asset, alias]) + session.flush() + session.add(AssetActive.for_asset(asset)) + asset.aliases.append(alias) + session.commit() + + response = test_client.get("/assets") + assert response.status_code == 200 + body = response.json() + assert len(body["assets"]) == 1 + assert body["assets"][0]["aliases"] == [{"id": alias.id, "name": "my-alias", "group": ""}] + + def test_total_entries_counts_assets_not_events(self, test_client, session): + """The outer join to AssetEvent must not inflate total_entries for assets with many events.""" + multi = AssetModel(name="multi", uri="s3://bucket/multi", group="reporting") + solo = AssetModel(name="solo", uri="s3://bucket/solo", group="reporting") + other = AssetModel(name="other", uri="s3://bucket/other", group="ops") + session.add_all([multi, solo, other]) + for asset in (multi, solo, other): + session.add(AssetActive.for_asset(asset)) + session.flush() + + base = pendulum.datetime(2024, 1, 1) + for offset in range(3): + session.add(AssetEvent(asset_id=multi.id, timestamp=base.add(hours=offset))) + session.add(AssetEvent(asset_id=solo.id, timestamp=base)) + session.commit() + + response = test_client.get("/assets") + assert response.status_code == 200 + body = response.json() + assert body["total_entries"] == 3 + assert sorted(a["name"] for a in body["assets"]) == ["multi", "other", "solo"] + + response = test_client.get("/assets?group_pattern=report") + assert response.status_code == 200 + body = response.json() + assert body["total_entries"] == 2 + assert sorted(a["name"] for a in body["assets"]) == ["multi", "solo"] + + def test_pagination_covers_all_assets(self, test_client, session): + """Paging returns every asset exactly once with a stable total, mixing evented and never-evented.""" + evented_new = AssetModel(name="evented_new", uri="s3://bucket/en", group="asset") + evented_old = AssetModel(name="evented_old", uri="s3://bucket/eo", group="asset") + never_a = AssetModel(name="never_a", uri="s3://bucket/na", group="asset") + never_b = AssetModel(name="never_b", uri="s3://bucket/nb", group="asset") + session.add_all([evented_new, evented_old, never_a, never_b]) + for asset in (evented_new, evented_old, never_a, never_b): + session.add(AssetActive.for_asset(asset)) + session.flush() + + base = pendulum.datetime(2024, 1, 1) + session.add(AssetEvent(asset_id=evented_old.id, timestamp=base)) + session.add(AssetEvent(asset_id=evented_new.id, timestamp=base.add(days=1))) + session.commit() + + page_one = test_client.get("/assets?limit=2&offset=0") + page_two = test_client.get("/assets?limit=2&offset=2") + assert page_one.status_code == 200 + assert page_two.status_code == 200 + assert page_one.json()["total_entries"] == 4 + assert page_two.json()["total_entries"] == 4 + + names = [a["name"] for a in page_one.json()["assets"] + page_two.json()["assets"]] + assert sorted(names) == ["evented_new", "evented_old", "never_a", "never_b"] + + def test_timestamp_range_filter_excludes_never_evented_assets(self, test_client, session): + """A timestamp range filter excludes assets with no event (NULL fails the range predicate).""" + evented = AssetModel(name="evented", uri="s3://bucket/re", group="asset") + never = AssetModel(name="never", uri="s3://bucket/rn", group="asset") + session.add_all([evented, never]) + session.add(AssetActive.for_asset(evented)) + session.add(AssetActive.for_asset(never)) + session.flush() + + base = pendulum.datetime(2024, 1, 1) + session.add(AssetEvent(asset_id=evented.id, timestamp=base)) + session.commit() + + response = test_client.get( + "/assets", params={"last_asset_event_timestamp_gte": base.subtract(days=1).isoformat()} + ) + assert response.status_code == 200 + assert [a["name"] for a in response.json()["assets"]] == ["evented"] + + def test_filter_by_only_active(self, test_client, session): + active = AssetModel(name="active", uri="s3://bucket/active", group="asset") + inactive = AssetModel(name="inactive", uri="s3://bucket/inactive", group="asset") + session.add_all([active, inactive]) + session.add(AssetActive.for_asset(active)) + session.commit() + + response = test_client.get("/assets") + assert response.status_code == 200 + assert [a["name"] for a in response.json()["assets"]] == ["active"] + + response = test_client.get("/assets?only_active=false") + assert response.status_code == 200 + assert sorted(a["name"] for a in response.json()["assets"]) == ["active", "inactive"] + + def test_filter_by_uri_pattern(self, test_client, session): + s3 = AssetModel(name="s3_asset", uri="s3://bucket/key", group="asset") + gcs = AssetModel(name="gcs_asset", uri="gcs://bucket/key", group="asset") + session.add_all([s3, gcs]) + session.add(AssetActive.for_asset(s3)) + session.add(AssetActive.for_asset(gcs)) + session.commit() + + response = test_client.get("/assets?uri_pattern=s3") + assert response.status_code == 200 + assert [a["name"] for a in response.json()["assets"]] == ["s3_asset"] + + @pytest.mark.usefixtures("testing_dag_bundle") + def test_filter_by_dag_ids(self, test_client, session): + referenced = AssetModel(name="referenced", uri="s3://bucket/referenced", group="asset") + unreferenced = AssetModel(name="unreferenced", uri="s3://bucket/unreferenced", group="asset") + session.add_all([referenced, unreferenced]) + session.add(AssetActive.for_asset(referenced)) + session.add(AssetActive.for_asset(unreferenced)) + session.add(DagModel(dag_id="consumer_dag", bundle_name="testing")) + session.add(DagScheduleAssetReference(dag_id="consumer_dag", asset=referenced)) + session.commit() + + response = test_client.get("/assets?dag_ids=consumer_dag") + assert response.status_code == 200 + assert [a["name"] for a in response.json()["assets"]] == ["referenced"] + + def test_query_count(self, test_client, session): + """The asset relationships are eager-loaded, so the query count stays fixed regardless of + how many assets are returned (a lazy-loading regression would issue queries per asset).""" + for i in range(5): + asset = AssetModel(name=f"asset{i}", uri=f"s3://bucket/asset{i}", group="asset") + session.add(asset) + session.add(AssetActive.for_asset(asset)) + session.commit() + + with assert_queries_count(7): Review Comment: None of these five assets has a watcher, alias, scheduled dag or task, so every eager-loaded collection comes back empty and the nested `joinedload(AssetWatcherModel.trigger)` never fires. I checked by deleting that `.joinedload(...)` from `generate_assets_with_last_event_query`, and all 17 tests in this class still pass at 7, even though `from_asset_row` would then lazy-load `watcher.trigger` once per watcher. Could the fixture give at least one asset a watcher and an alias, so the guard covers the relationships it exists to protect? -- 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]
