Update dagdetails

bbovenzi · bbovenzi · commit fe8d84a79592 · 2025-10-09T11:51:13.000-07:00
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dags.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dags.py
@@ -157,6 +157,7 @@ class DAGDetailsResponse(DAGResponse):
     last_parsed: datetime | None
     default_args: abc.Mapping | None
     owner_links: dict[str, str] | None = None
+    is_favorite: bool = False
 
     @field_validator("timezone", mode="before")
     @classmethod
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
@@ -10017,6 +10017,10 @@ components:
             type: object
           - type: 'null'
           title: Owner Links
+        is_favorite:
+          type: boolean
+          title: Is Favorite
+          default: false
         file_token:
           type: string
           title: File Token
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dags.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dags.py
@@ -209,7 +209,9 @@ def get_dag(
     ),
     dependencies=[Depends(requires_access_dag(method="GET"))],
 )
-def get_dag_details(dag_id: str, session: SessionDep, dag_bag: DagBagDep) -> DAGDetailsResponse:
+def get_dag_details(
+    dag_id: str, session: SessionDep, dag_bag: DagBagDep, user: GetUserDep
+) -> DAGDetailsResponse:
     """Get details of DAG."""
     dag = get_latest_version_of_dag(dag_bag, dag_id, session)
 
@@ -221,7 +223,19 @@ def get_dag_details(dag_id: str, session: SessionDep, dag_bag: DagBagDep) -> DAG
         if not key.startswith("_") and not hasattr(dag_model, key):
             setattr(dag_model, key, value)
 
-    return dag_model
+    # Check if this DAG is marked as favorite by the current user
+    user_id = str(user.get_id())
+    is_favorite = (
+        session.scalar(
+            select(DagFavorite.dag_id).where(DagFavorite.user_id == user_id, DagFavorite.dag_id == dag_id)
+        )
+        is not None
+    )
+
+    # Add is_favorite field to the DAG model
+    setattr(dag_model, "is_favorite", is_favorite)
+
+    return DAGDetailsResponse.model_validate(dag_model)
 
 
 @dags_router.patch(
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dags.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dags.py
@@ -158,12 +158,10 @@ def get_dags(
 
     # Fetch favorite status for each DAG for the current user
     user_id = str(user.get_id())
-    favorite_dag_ids = set()
-    if dags:
-        favorites_select = select(DagFavorite.dag_id).where(
-            DagFavorite.user_id == user_id, DagFavorite.dag_id.in_([dag.dag_id for dag in dags])
-        )
-        favorite_dag_ids = {dag_id for dag_id in session.scalars(favorites_select)}
+    favorites_select = select(DagFavorite.dag_id).where(
+        DagFavorite.user_id == user_id, DagFavorite.dag_id.in_([dag.dag_id for dag in dags])
+    )
+    favorite_dag_ids = set(session.scalars(favorites_select))
 
     # Populate the last 'dag_runs_limit' DagRuns for each DAG
     recent_runs_subquery = (
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
@@ -2041,6 +2041,11 @@ export const $DAGDetailsResponse = {
             ],
             title: 'Owner Links'
         },
+        is_favorite: {
+            type: 'boolean',
+            title: 'Is Favorite',
+            default: false
+        },
         file_token: {
             type: 'string',
             title: 'File Token',
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -562,6 +562,7 @@ export type DAGDetailsResponse = {
     owner_links?: {
     [key: string]: (string);
 } | null;
+    is_favorite?: boolean;
     /**
      * Return file token.
      */
diff --git a/airflow-core/src/airflow/ui/src/components/DagActions/FavoriteDagButton.tsx b/airflow-core/src/airflow/ui/src/components/DagActions/FavoriteDagButton.tsx
@@ -17,32 +17,25 @@
  * under the License.
  */
 import { Box } from "@chakra-ui/react";
-import { useCallback } from "react";
 import { useTranslation } from "react-i18next";
 import { FiStar } from "react-icons/fi";
 
-import { useFavoriteDag } from "src/queries/useFavoriteDag";
-import { useUnfavoriteDag } from "src/queries/useUnfavoriteDag";
+import { useToggleFavoriteDag } from "src/queries/useToggleFavoriteDag";
 
 import ActionButton from "../ui/ActionButton";
 
 type FavoriteDagButtonProps = {
   readonly dagId: string;
-  readonly isFavorite: boolean;
+  readonly isFavorite?: boolean;
   readonly withText?: boolean;
 };
 
-export const FavoriteDagButton = ({ dagId, isFavorite, withText = true }: FavoriteDagButtonProps) => {
+export const FavoriteDagButton = ({ dagId, isFavorite = false, withText = true }: FavoriteDagButtonProps) => {
   const { t: translate } = useTranslation("dags");
 
-  const { mutate: favoriteDag } = useFavoriteDag();
-  const { mutate: unfavoriteDag } = useUnfavoriteDag();
+  const { isLoading, toggleFavorite } = useToggleFavoriteDag(dagId);
 
-  const onToggle = useCallback(() => {
-    const mutationFn = isFavorite ? unfavoriteDag : favoriteDag;
-
-    mutationFn({ dagId });
-  }, [dagId, isFavorite, favoriteDag, unfavoriteDag]);
+  const onToggle = () => toggleFavorite(isFavorite);
 
   return (
     <Box>
@@ -56,6 +49,7 @@ export const FavoriteDagButton = ({ dagId, isFavorite, withText = true }: Favori
             }}
           />
         }
+        loading={isLoading}
         onClick={onToggle}
         text={isFavorite ? translate("unfavoriteDag") : translate("favoriteDag")}
         withText={withText}
diff --git a/airflow-core/src/airflow/ui/src/pages/Dag/Header.tsx b/airflow-core/src/airflow/ui/src/pages/Dag/Header.tsx
@@ -23,7 +23,6 @@ import { FiBookOpen } from "react-icons/fi";
 import { LuMenu } from "react-icons/lu";
 import { useParams, Link as RouterLink } from "react-router-dom";
 
-import { useDagServiceGetDagsUi } from "openapi/queries";
 import type { DAGDetailsResponse, DagRunState } from "openapi/requests/types.gen";
 import { DagIcon } from "src/assets/DagIcon";
 import DeleteDagButton from "src/components/DagActions/DeleteDagButton";
@@ -116,8 +115,6 @@ export const Header = ({
     },
   ];
 
-  const { data: dagData } = useDagServiceGetDagsUi({ dagIds: [dagId ?? ""] });
-
   return (
     <HeaderCard
       actions={
@@ -131,11 +128,7 @@ export const Header = ({
                 text={translate("dag:header.buttons.dagDocs")}
               />
             )}
-            <FavoriteDagButton
-              dagId={dag.dag_id}
-              isFavorite={dagData?.dags[0]?.is_favorite ?? false}
-              withText
-            />
+            <FavoriteDagButton dagId={dag.dag_id} isFavorite={dag.is_favorite} withText />
             <Menu.Root>
               <Menu.Trigger asChild>
                 <Button aria-label={translate("dag:header.buttons.advanced")} variant="outline">
diff --git a/airflow-core/src/airflow/ui/src/queries/useFavoriteDag.ts b/airflow-core/src/airflow/ui/src/queries/useFavoriteDag.ts
diff --git a/airflow-core/src/airflow/ui/src/queries/useToggleFavoriteDag.ts b/airflow-core/src/airflow/ui/src/queries/useToggleFavoriteDag.ts
@@ -0,0 +1,66 @@
+/*!
+ * 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 { useQueryClient } from "@tanstack/react-query";
+import { useCallback } from "react";
+
+import {
+  useDagServiceGetDagsUiKey,
+  useDagServiceFavoriteDag,
+  useDagServiceUnfavoriteDag,
+  UseDagServiceGetDagDetailsKeyFn,
+} from "openapi/queries";
+
+export const useToggleFavoriteDag = (dagId: string) => {
+  const queryClient = useQueryClient();
+
+  const onSuccess = useCallback(async () => {
+    // Invalidate the DAGs list query
+    await queryClient.invalidateQueries({
+      queryKey: [useDagServiceGetDagsUiKey, UseDagServiceGetDagDetailsKeyFn({ dagId }, [{ dagId }])],
+    });
+
+    // Invalidate the specific DAG details query for this DAG
+    await queryClient.invalidateQueries({
+      queryKey: UseDagServiceGetDagDetailsKeyFn({ dagId }, [{ dagId }]),
+    });
+  }, [queryClient, dagId]);
+
+  const favoriteMutation = useDagServiceFavoriteDag({
+    onSuccess,
+  });
+
+  const unfavoriteMutation = useDagServiceUnfavoriteDag({
+    onSuccess,
+  });
+
+  const toggleFavorite = useCallback(
+    (isFavorite: boolean) => {
+      const mutation = isFavorite ? unfavoriteMutation : favoriteMutation;
+
+      mutation.mutate({ dagId });
+    },
+    [dagId, favoriteMutation, unfavoriteMutation],
+  );
+
+  return {
+    error: favoriteMutation.error ?? unfavoriteMutation.error,
+    isLoading: favoriteMutation.isPending || unfavoriteMutation.isPending,
+    toggleFavorite,
+  };
+};
diff --git a/airflow-core/src/airflow/ui/src/queries/useUnfavoriteDag.ts b/airflow-core/src/airflow/ui/src/queries/useUnfavoriteDag.ts
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dags.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dags.py
@@ -483,6 +483,21 @@ def test_get_dags_filter_favorites(
         assert body["total_entries"] == expected_total_entries
         assert sorted([dag["dag_id"] for dag in body["dags"]]) == sorted(expected_ids)
 
+    def test_get_dags_filter_non_favorites(self, session, test_client):
+        """Test filtering DAGs by is_favorite=false."""
+        # Mark DAG1 as favorite
+        session.add(DagFavorite(user_id="test", dag_id=DAG1_ID))
+        session.commit()
+
+        response = test_client.get("/dags", params={"is_favorite": False})
+
+        assert response.status_code == 200
+        body = response.json()
+
+        # Should return only non-favorite DAGs (DAG2)
+        assert body["total_entries"] == 1
+        assert [dag["dag_id"] for dag in body["dags"]] == [DAG2_ID]
+
     def test_get_dags_should_response_401(self, unauthenticated_test_client):
         response = unauthenticated_test_client.get("/dags")
         assert response.status_code == 401
@@ -862,6 +877,7 @@ def test_dag_details(
             "template_search_path": None,
             "timetable_description": "Never, external triggers only",
             "timezone": UTC_JSON_REPR,
+            "is_favorite": False,
         }
         assert res_json == expected
 
@@ -957,6 +973,7 @@ def test_dag_details_with_view_url_template(
             "template_search_path": None,
             "timetable_description": "Never, external triggers only",
             "timezone": UTC_JSON_REPR,
+            "is_favorite": False,
         }
         assert res_json == expected
 
@@ -968,6 +985,30 @@ def test_dag_details_should_response_403(self, unauthorized_test_client):
         response = unauthorized_test_client.get(f"/dags/{DAG1_ID}/details")
         assert response.status_code == 403
 
+    def test_dag_details_includes_is_favorite_field(self, session, test_client):
+        """Test that DAG details include the is_favorite field."""
+        # Mark DAG2 as favorite
+        session.add(DagFavorite(user_id="test", dag_id=DAG2_ID))
+        session.commit()
+
+        response = test_client.get(f"/dags/{DAG2_ID}/details")
+        assert response.status_code == 200
+        body = response.json()
+
+        # Verify is_favorite field is present and correct
+        assert "is_favorite" in body
+        assert isinstance(body["is_favorite"], bool)
+        assert body["is_favorite"] is True
+
+        # Test with non-favorite DAG
+        response = test_client.get(f"/dags/{DAG1_ID}/details")
+        assert response.status_code == 200
+        body = response.json()
+
+        assert "is_favorite" in body
+        assert isinstance(body["is_favorite"], bool)
+        assert body["is_favorite"] is False
+
 
 class TestGetDag(TestDagEndpoint):
     """Unit tests for Get DAG."""
diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py b/airflow-ctl/src/airflowctl/api/datamodels/generated.py
@@ -1302,6 +1302,7 @@ class DAGDetailsResponse(BaseModel):
     last_parsed: Annotated[datetime | None, Field(title="Last Parsed")] = None
     default_args: Annotated[dict[str, Any] | None, Field(title="Default Args")] = None
     owner_links: Annotated[dict[str, str] | None, Field(title="Owner Links")] = None
+    is_favorite: Annotated[bool | None, Field(title="Is Favorite")] = False
     file_token: Annotated[str, Field(description="Return file token.", title="File Token")]
     concurrency: Annotated[
         int,
