Add is_favorite to ui dags list

bbovenzi · bbovenzi · commit 822d62afd724 · 2025-10-09T11:51:13.000-07:00
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/dags.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/dags.py
@@ -29,6 +29,7 @@ class DAGWithLatestDagRunsResponse(DAGResponse):
     asset_expression: dict | None
     latest_dag_runs: list[DAGRunResponse]
     pending_actions: list[HITLDetail]
+    is_favorite: bool
 
 
 class DAGWithLatestDagRunsCollectionResponse(BaseModel):
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
@@ -1686,6 +1686,9 @@ components:
             $ref: '#/components/schemas/HITLDetail'
           type: array
           title: Pending Actions
+        is_favorite:
+          type: boolean
+          title: Is Favorite
         file_token:
           type: string
           title: File Token
@@ -1721,6 +1724,7 @@ components:
       - asset_expression
       - latest_dag_runs
       - pending_actions
+      - is_favorite
       - file_token
       title: DAGWithLatestDagRunsResponse
       description: DAG with latest dag runs response serializer.
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
@@ -61,10 +61,12 @@
 )
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.api_fastapi.core_api.security import (
+    GetUserDep,
     ReadableDagsFilterDep,
     requires_access_dag,
 )
 from airflow.models import DagModel, DagRun
+from airflow.models.dag_favorite import DagFavorite
 from airflow.models.hitl import HITLDetail
 from airflow.models.taskinstance import TaskInstance
 from airflow.utils.state import TaskInstanceState
@@ -114,6 +116,7 @@ def get_dags(
     has_pending_actions: QueryPendingActionsFilter,
     readable_dags_filter: ReadableDagsFilterDep,
     session: SessionDep,
+    user: GetUserDep,
     dag_runs_limit: int = 10,
 ) -> DAGWithLatestDagRunsCollectionResponse:
     """Get DAGs with recent DagRun."""
@@ -153,6 +156,15 @@ def get_dags(
 
     dags = [dag for dag in session.scalars(dags_select)]
 
+    # 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)}
+
     # Populate the last 'dag_runs_limit' DagRuns for each DAG
     recent_runs_subquery = (
         select(
@@ -224,6 +236,7 @@ def get_dags(
                 "asset_expression": dag.asset_expression,
                 "latest_dag_runs": [],
                 "pending_actions": pending_actions_by_dag_id[dag.dag_id],
+                "is_favorite": dag.dag_id in favorite_dag_ids,
             }
         )
         for dag in dags
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
@@ -7343,6 +7343,10 @@ export const $DAGWithLatestDagRunsResponse = {
             type: 'array',
             title: 'Pending Actions'
         },
+        is_favorite: {
+            type: 'boolean',
+            title: 'Is Favorite'
+        },
         file_token: {
             type: 'string',
             title: 'File Token',
@@ -7351,7 +7355,7 @@ export const $DAGWithLatestDagRunsResponse = {
         }
     },
     type: 'object',
-    required: ['dag_id', 'dag_display_name', 'is_paused', 'is_stale', 'last_parsed_time', 'last_parse_duration', 'last_expired', 'bundle_name', 'bundle_version', 'relative_fileloc', 'fileloc', 'description', 'timetable_summary', 'timetable_description', 'tags', 'max_active_tasks', 'max_active_runs', 'max_consecutive_failed_dag_runs', 'has_task_concurrency_limits', 'has_import_errors', 'next_dagrun_logical_date', 'next_dagrun_data_interval_start', 'next_dagrun_data_interval_end', 'next_dagrun_run_after', 'owners', 'asset_expression', 'latest_dag_runs', 'pending_actions', 'file_token'],
+    required: ['dag_id', 'dag_display_name', 'is_paused', 'is_stale', 'last_parsed_time', 'last_parse_duration', 'last_expired', 'bundle_name', 'bundle_version', 'relative_fileloc', 'fileloc', 'description', 'timetable_summary', 'timetable_description', 'tags', 'max_active_tasks', 'max_active_runs', 'max_consecutive_failed_dag_runs', 'has_task_concurrency_limits', 'has_import_errors', 'next_dagrun_logical_date', 'next_dagrun_data_interval_start', 'next_dagrun_data_interval_end', 'next_dagrun_run_after', 'owners', 'asset_expression', 'latest_dag_runs', 'pending_actions', 'is_favorite', 'file_token'],
     title: 'DAGWithLatestDagRunsResponse',
     description: 'DAG with latest dag runs response serializer.'
 } as const;
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
@@ -1834,6 +1834,7 @@ export type DAGWithLatestDagRunsResponse = {
 } | null;
     latest_dag_runs: Array<DAGRunResponse>;
     pending_actions: Array<HITLDetail>;
+    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,29 +17,23 @@
  * under the License.
  */
 import { Box } from "@chakra-ui/react";
-import { useCallback, useMemo } from "react";
+import { useCallback } from "react";
 import { useTranslation } from "react-i18next";
 import { FiStar } from "react-icons/fi";
 
-import { useDagServiceGetDagsUi } from "openapi/queries";
 import { useFavoriteDag } from "src/queries/useFavoriteDag";
 import { useUnfavoriteDag } from "src/queries/useUnfavoriteDag";
 
 import ActionButton from "../ui/ActionButton";
 
 type FavoriteDagButtonProps = {
   readonly dagId: string;
+  readonly isFavorite: boolean;
   readonly withText?: boolean;
 };
 
-export const FavoriteDagButton = ({ dagId, withText = true }: FavoriteDagButtonProps) => {
+export const FavoriteDagButton = ({ dagId, isFavorite, withText = true }: FavoriteDagButtonProps) => {
   const { t: translate } = useTranslation("dags");
-  const { data: favorites } = useDagServiceGetDagsUi({ isFavorite: true });
-
-  const isFavorite = useMemo(
-    () => favorites?.dags.some((fav) => fav.dag_id === dagId) ?? false,
-    [favorites, dagId],
-  );
 
   const { mutate: favoriteDag } = useFavoriteDag();
   const { mutate: unfavoriteDag } = useUnfavoriteDag();
@@ -54,7 +48,14 @@ export const FavoriteDagButton = ({ dagId, withText = true }: FavoriteDagButtonP
     <Box>
       <ActionButton
         actionName={isFavorite ? translate("unfavoriteDag") : translate("favoriteDag")}
-        icon={<FiStar style={{ fill: isFavorite ? "var(--chakra-colors-brand-solid)" : "none" }} />}
+        icon={
+          <FiStar
+            style={{
+              fill: isFavorite ? "var(--chakra-colors-brand-solid)" : "none",
+              stroke: "var(--chakra-colors-brand-solid)",
+            }}
+          />
+        }
         onClick={onToggle}
         text={isFavorite ? translate("unfavoriteDag") : translate("favoriteDag")}
         withText={withText}
diff --git a/airflow-core/src/airflow/ui/src/mocks/handlers/dags.ts b/airflow-core/src/airflow/ui/src/mocks/handlers/dags.ts
@@ -32,6 +32,7 @@ export const handlers: Array<HttpHandler> = [
       fileloc: "/airflow/dags/tutorial_taskflow_api.py",
       has_import_errors: false,
       has_task_concurrency_limits: false,
+      is_favorite: true,
       is_paused: false,
       is_stale: false,
       last_parsed_time: "2025-01-13T07:34:01.593459Z",
@@ -69,6 +70,7 @@ export const handlers: Array<HttpHandler> = [
       fileloc: "/airflow/dags/tutorial_taskflow_api_failed.py",
       has_import_errors: false,
       has_task_concurrency_limits: false,
+      is_favorite: false,
       is_paused: false,
       is_stale: false,
       last_parsed_time: "2025-01-13T07:34:01.593459Z",
@@ -128,6 +130,7 @@ export const handlers: Array<HttpHandler> = [
       fileloc: "/airflow/dags/tutorial_taskflow_api_failed.py",
       has_import_errors: false,
       has_task_concurrency_limits: false,
+      is_favorite: false,
       is_paused: false,
       is_stale: false,
       last_expired: null,
@@ -155,6 +158,7 @@ export const handlers: Array<HttpHandler> = [
       fileloc: "/airflow/dags/tutorial_taskflow_api_success.py",
       has_import_errors: false,
       has_task_concurrency_limits: false,
+      is_favorite: true,
       is_paused: false,
       is_stale: false,
       last_expired: null,
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,6 +23,7 @@ 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";
@@ -115,6 +116,8 @@ export const Header = ({
     },
   ];
 
+  const { data: dagData } = useDagServiceGetDagsUi({ dagIds: [dagId ?? ""] });
+
   return (
     <HeaderCard
       actions={
@@ -128,7 +131,11 @@ export const Header = ({
                 text={translate("dag:header.buttons.dagDocs")}
               />
             )}
-            <FavoriteDagButton dagId={dag.dag_id} withText={true} />
+            <FavoriteDagButton
+              dagId={dag.dag_id}
+              isFavorite={dagData?.dags[0]?.is_favorite ?? false}
+              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/pages/DagsList/DagCard.test.tsx b/airflow-core/src/airflow/ui/src/pages/DagsList/DagCard.test.tsx
@@ -66,6 +66,7 @@ const mockDag = {
   fileloc: "/files/dags/nested_task_groups.py",
   has_import_errors: false,
   has_task_concurrency_limits: false,
+  is_favorite: false,
   is_paused: false,
   is_stale: false,
   last_expired: null,
diff --git a/airflow-core/src/airflow/ui/src/pages/DagsList/DagCard.tsx b/airflow-core/src/airflow/ui/src/pages/DagsList/DagCard.tsx
@@ -72,7 +72,7 @@ export const DagCard = ({ dag }: Props) => {
             isPaused={dag.is_paused}
             withText={false}
           />
-          <FavoriteDagButton dagId={dag.dag_id} withText={false} />
+          <FavoriteDagButton dagId={dag.dag_id} isFavorite={dag.is_favorite} withText={false} />
           <DeleteDagButton dagDisplayName={dag.dag_display_name} dagId={dag.dag_id} withText={false} />
         </HStack>
       </Flex>
diff --git a/airflow-core/src/airflow/ui/src/pages/DagsList/DagsList.tsx b/airflow-core/src/airflow/ui/src/pages/DagsList/DagsList.tsx
@@ -164,8 +164,9 @@ const createColumns = (
     header: "",
   },
   {
-    accessorKey: "favorite",
-    cell: ({ row: { original } }) => <FavoriteDagButton dagId={original.dag_id} withText={false} />,
+    cell: ({ row: { original } }) => (
+      <FavoriteDagButton dagId={original.dag_id} isFavorite={original.is_favorite} withText={false} />
+    ),
     enableHiding: false,
     enableSorting: false,
     header: "",
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_dags.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_dags.py
@@ -26,6 +26,7 @@
 from sqlalchemy.orm import Session
 
 from airflow.models import DagRun
+from airflow.models.dag_favorite import DagFavorite
 from airflow.models.hitl import HITLDetail
 from airflow.sdk.timezone import utcnow
 from airflow.utils.session import provide_session
@@ -278,3 +279,38 @@ def test_asset_filtering(self, test_client, query_params, expected_dag_count, se
         body = response.json()
         assert body["total_entries"] == expected_dag_count
         assert len(body["dags"]) == expected_dag_count
+
+    def test_is_favorite_field_with_multiple_favorites(self, test_client, session):
+        """Test is_favorite field with multiple DAGs marked as favorites."""
+        # Mark both DAG1 and DAG2 as favorites
+        session.add(DagFavorite(user_id="test", dag_id=DAG1_ID))
+        session.add(DagFavorite(user_id="test", dag_id=DAG2_ID))
+        session.commit()
+
+        response = test_client.get("/dags")
+        assert response.status_code == 200
+        body = response.json()
+
+        # Count favorites in response
+        favorite_count = sum(1 for dag in body["dags"] if dag["is_favorite"])
+        assert favorite_count == 2
+
+        # Verify specific DAGs are marked as favorites
+        dag_favorites = {dag["dag_id"]: dag["is_favorite"] for dag in body["dags"]}
+        assert dag_favorites[DAG1_ID] is True
+        assert dag_favorites[DAG2_ID] is True
+
+    def test_is_favorite_field_user_specific(self, test_client, session):
+        """Test that is_favorite field is user-specific."""
+        # Mark DAG1 as favorite for a different user
+        session.add(DagFavorite(user_id="other_user", dag_id=DAG1_ID))
+        session.commit()
+
+        # Request as the test user (not other_user)
+        response = test_client.get("/dags")
+        assert response.status_code == 200
+        body = response.json()
+
+        # Verify that DAG1 is not marked as favorite for the test user
+        dag1_data = next(dag for dag in body["dags"] if dag["dag_id"] == DAG1_ID)
+        assert dag1_data["is_favorite"] is False
