Keep the Dag asset schedule the same shape for every caller

airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/assets.py

@@ -52,4 +52,10 @@ class NextRunAssetsResponse(BaseModel):
 
     asset_expression: MaybeAssetExpression = None
     events: list[NextRunAssetEventResponse]
+    scheduling_asset_count: int = 0
+    """How many assets the Dag is scheduled on, before filtering ``events`` down to
+    the ones the caller may read. ``events`` is caller-scoped, so a UI that derives
+    the schedule's shape from ``len(events)`` changes what it renders with the
+    caller's permissions; this count does not. It reveals nothing new — the redacted
+    ``asset_expression`` already carries one slot per asset."""
     pending_partition_count: int | None = None

airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml

@@ -4453,6 +4453,10 @@ components:
             $ref: '#/components/schemas/NextRunAssetEventResponse'
           type: array
           title: Events
+        scheduling_asset_count:
+          type: integer
+          title: Scheduling Asset Count
+          default: 0
         pending_partition_count:
           anyOf:
           - type: integer

airflow-core/src/airflow/api_fastapi/core_api/routes/ui/assets.py

@@ -230,6 +230,16 @@ def next_run_assets(
     )
     query = readable_assets_filter.to_orm(query)
 
+    # Counted before the readable filter narrows ``query``: the UI derives the
+    # schedule's shape (total, and whether to render the multi-asset popover) from
+    # this, so that shape stays the same for every caller regardless of which
+    # assets they may read.
+    scheduling_asset_count = session.scalar(
+        select(func.count())
+        .select_from(DagScheduleAssetReference)
+        .where(DagScheduleAssetReference.dag_id == dag_id)
+    )
+
     if not is_partitioned:
         query = query.join(
             AssetDagRunQueue,
@@ -253,7 +263,11 @@ def next_run_assets(
             )
             for row in raw_rows
         ]
-        model_data: dict[str, Any] = {"asset_expression": asset_expression, "events": events}
+        model_data: dict[str, Any] = {
+            "asset_expression": asset_expression,
+            "events": events,
+            "scheduling_asset_count": scheduling_asset_count,
+        }
         return NextRunAssetsResponse.model_validate(model_data)
 
     # Partitioned Dags: enrich with per-asset received/required counts and rollup flag.
@@ -290,6 +304,7 @@ def next_run_assets(
         model_data = {
             "asset_expression": asset_expression,
             "events": events,
+            "scheduling_asset_count": scheduling_asset_count,
             "pending_partition_count": pending_partition_count,
         }
         return NextRunAssetsResponse.model_validate(model_data)
@@ -364,6 +379,7 @@ def next_run_assets(
     model_data = {
         "asset_expression": asset_expression,
         "events": events,
+        "scheduling_asset_count": scheduling_asset_count,
         "pending_partition_count": pending_partition_count,
     }
     return NextRunAssetsResponse.model_validate(model_data)

airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts

@@ -11689,6 +11689,11 @@ export const $NextRunAssetsResponse = {
             type: 'array',
             title: 'Events'
         },
+        scheduling_asset_count: {
+            type: 'integer',
+            title: 'Scheduling Asset Count',
+            default: 0
+        },
         pending_partition_count: {
             anyOf: [
                 {

airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts

@@ -3041,6 +3041,7 @@ export type NextRunAssetEventResponse = {
 export type NextRunAssetsResponse = {
     asset_expression?: AssetExpressionAsset | AssetExpressionAlias | AssetExpressionRef | AssetExpressionAny | AssetExpressionAll | null;
     events: Array<NextRunAssetEventResponse>;
+    scheduling_asset_count?: number;
     pending_partition_count?: number | null;
 };
 

airflow-core/src/airflow/ui/src/pages/DagsList/AssetSchedule.test.tsx

@@ -0,0 +1,125 @@
+/*!
+ * 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.
+ */
+import "@testing-library/jest-dom/vitest";
+import { render, screen } from "@testing-library/react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+import type * as OpenapiQueries from "openapi/queries";
+import type { NextRunAssetEventResponse, NextRunAssetsResponse } from "openapi/requests/types.gen";
+
+import { Wrapper } from "src/utils/Wrapper";
+
+import { AssetSchedule } from "./AssetSchedule";
+
+vi.mock("react-i18next", () => ({
+  useTranslation: () => ({
+    i18n: { language: "en" },
+    // eslint-disable-next-line id-length
+    t: (key: string, options?: { count?: number; total?: number }) =>
+      options?.count === undefined ? key : `${key}:${options.count} of ${options.total}`,
+  }),
+}));
+
+vi.mock("openapi/queries", async (importOriginal) => {
+  const actual = await importOriginal<typeof OpenapiQueries>();
+
+  return {
+    ...actual,
+    useAssetServiceGetDagAssetQueuedEvents: vi.fn(),
+    useAssetServiceNextRunAssets: vi.fn(),
+  };
+});
+
+const { useAssetServiceGetDagAssetQueuedEvents, useAssetServiceNextRunAssets } =
+  await import("openapi/queries");
+
+const makeEvent = (id: number, name: string): NextRunAssetEventResponse => ({
+  asset_inactive: false,
+  id,
+  is_rollup: false,
+  last_update: null,
+  mapper_error: false,
+  name,
+  received_count: 0,
+  received_keys: [],
+  required_count: 1,
+  required_keys: [],
+  uri: `s3://bucket/${name}`,
+});
+
+const nextRunResponse = (nextRun: NextRunAssetsResponse) =>
+  ({ data: nextRun, error: null, isFetching: false, isLoading: false }) as ReturnType<
+    typeof useAssetServiceNextRunAssets
+  >;
+
+const queuedEventsResponse = () =>
+  ({
+    data: { queued_events: [], total_entries: 0 },
+    error: null,
+    isFetching: false,
+    isLoading: false,
+  }) as ReturnType<typeof useAssetServiceGetDagAssetQueuedEvents>;
+
+const renderSchedule = (nextRun: NextRunAssetsResponse) => {
+  vi.mocked(useAssetServiceNextRunAssets).mockReturnValue(nextRunResponse(nextRun));
+  vi.mocked(useAssetServiceGetDagAssetQueuedEvents).mockReturnValue(queuedEventsResponse());
+
+  render(
+    <AssetSchedule dagId="dag_id" timetablePartitioned={false} timetableSummary="Every day at midnight" />,
+    { wrapper: Wrapper },
+  );
+};
+
+describe("AssetSchedule", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it("shows the Dag's full asset total when the caller may read only some of them", () => {
+    renderSchedule({
+      events: [makeEvent(1, "visible_asset")],
+      scheduling_asset_count: 3,
+    });
+
+    expect(screen.getByRole("button")).toHaveTextContent("assetSchedule:0 of 3");
+  });
+
+  it("still renders an asset schedule when the caller may read none of the assets", () => {
+    renderSchedule({ events: [], scheduling_asset_count: 3 });
+
+    expect(screen.getByRole("button")).toHaveTextContent("assetSchedule:0 of 3");
+    expect(screen.queryByText("Every day at midnight")).not.toBeInTheDocument();
+  });
+
+  it("falls back to the timetable summary when the Dag has no scheduling assets", () => {
+    renderSchedule({ events: [], scheduling_asset_count: 0 });
+
+    expect(screen.getByText("Every day at midnight")).toBeInTheDocument();
+  });
+
+  it("renders the single-asset view only when the Dag is scheduled on one asset", () => {
+    renderSchedule({
+      events: [makeEvent(1, "only_asset")],
+      scheduling_asset_count: 1,
+    });
+
+    expect(screen.getByRole("link", { name: "only_asset" })).toBeInTheDocument();
+    expect(screen.queryByRole("button")).not.toBeInTheDocument();
+  });
+});

airflow-core/src/airflow/ui/src/pages/DagsList/AssetSchedule.tsx

@@ -118,13 +118,17 @@ export const AssetSchedule = ({ assetExpression, dagId, timetablePartitioned, ti
         0,
       )
     : pendingEvents.length;
+  // `events` only carries the assets the caller may read, so the total and the
+  // layout branches below come from `scheduling_asset_count` instead — the schedule
+  // a Dag shows should not change with who is looking at it.
+  const schedulingAssetCount = nextRun?.scheduling_asset_count ?? 0;
   const scheduledTotal = timetablePartitioned
     ? nextRunEvents.reduce((sum, event) => sum + (event.required_count ?? 1), 0)
-    : nextRunEvents.length;
+    : schedulingAssetCount;
 
   const isLoading = isNextRunLoading || (!timetablePartitioned && isQueuedEventsLoading);
 
-  if (!nextRunEvents.length) {
+  if (!schedulingAssetCount) {
     return (
       <HStack>
         <FiDatabase style={{ display: "inline", flexShrink: 0 }} />
@@ -160,7 +164,7 @@ export const AssetSchedule = ({ assetExpression, dagId, timetablePartitioned, ti
     // pendingCount === 1: render single-asset view with inactive warning.
     const [partitionedAsset] = nextRunEvents;
 
-    if (nextRunEvents.length === 1 && partitionedAsset !== undefined) {
+    if (schedulingAssetCount === 1 && partitionedAsset !== undefined) {
       const requiredCount = partitionedAsset.required_count ?? 1;
       const receivedCount = partitionedAsset.received_count ?? 0;
       const requiredKeys = partitionedAsset.required_keys ?? [];
@@ -214,7 +218,7 @@ export const AssetSchedule = ({ assetExpression, dagId, timetablePartitioned, ti
 
   const [asset] = nextRunEvents;
 
-  if (nextRunEvents.length === 1 && asset !== undefined) {
+  if (schedulingAssetCount === 1 && asset !== undefined) {
     const requiredCount = asset.required_count ?? 1;
     const receivedCount = asset.received_count ?? 0;
     const requiredKeys = asset.required_keys ?? [];

airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_assets.py

@@ -84,8 +84,9 @@ def test_should_response_200(self, test_client, dag_maker):
         dag_maker.create_dagrun()
         dag_maker.sync_dagbag_to_db()
 
-        # 4 queries for the endpoint plus 1 to resolve the assets the caller may read.
-        with assert_queries_count(5):
+        # 4 queries for the endpoint, 1 to resolve the assets the caller may read and
+        # 1 for the caller-independent scheduling asset count.
+        with assert_queries_count(6):
             response = test_client.get("/next_run_assets/upstream")
 
         assert response.status_code == 200
@@ -118,6 +119,7 @@ def test_should_response_200(self, test_client, dag_maker):
                     "asset_inactive": False,
                 }
             ],
+            "scheduling_asset_count": 1,
             "pending_partition_count": None,
         }
 
@@ -186,6 +188,32 @@ def test_should_return_only_assets_the_caller_may_read(
 
         assert response.status_code == 200
         assert [event["name"] for event in response.json()["events"]] == ["visible_asset"]
+        assert response.json()["scheduling_asset_count"] == 2
+
+    @mock.patch(
+        "airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager.get_authorized_assets",
+        autospec=True,
+    )
+    def test_scheduling_asset_count_is_unaffected_when_no_asset_is_readable(
+        self, mock_get_authorized_assets, test_client, dag_maker
+    ):
+        with dag_maker(
+            dag_id="all_hidden_upstream",
+            schedule=[
+                Asset(uri="s3://bucket/hidden1", name="hidden_asset1"),
+                Asset(uri="s3://bucket/hidden2", name="hidden_asset2"),
+            ],
+            serialized=True,
+        ):
+            EmptyOperator(task_id="task1")
+        dag_maker.sync_dagbag_to_db()
+        mock_get_authorized_assets.return_value = set()
+
+        response = test_client.get("/next_run_assets/all_hidden_upstream")
+
+        assert response.status_code == 200
+        assert response.json()["events"] == []
+        assert response.json()["scheduling_asset_count"] == 2
 
     @mock.patch(
         "airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager.get_authorized_assets",
@@ -240,6 +268,7 @@ def test_partitioned_dag_should_hide_keys_of_assets_the_caller_may_not_read(
         assert [event["name"] for event in events] == ["part_visible"]
         assert events[0]["received_keys"] == ["2024-01-01"]
         assert events[0]["required_keys"] == ["2024-01-01"]
+        assert response.json()["scheduling_asset_count"] == 2
 
     def test_should_respond_401(self, unauthenticated_test_client):
         response = unauthenticated_test_client.get("/next_run_assets/upstream")
@@ -336,6 +365,7 @@ def test_should_set_last_update_only_for_queued_and_hide_flag(self, test_client,
                     "asset_inactive": False,
                 },
             ],
+            "scheduling_asset_count": 2,
             "pending_partition_count": None,
         }