From cfe410998f4d1d537b51c8f86947120fc014efb1 Mon Sep 17 00:00:00 2001
From: Brent Bovenzi <brent.bovenzi@gmail.com>
Date: Thu, 6 Mar 2025 16:13:24 -0500
Subject: [PATCH 1/2] Change style for task logger and allow filtering by
 logger

---
 airflow/ui/src/constants/searchParams.ts      |  1 +
 .../ui/src/pages/TaskInstance/Logs/Logs.tsx   |  3 +
 .../pages/TaskInstance/Logs/TaskLogHeader.tsx | 64 +++++++++++++++++--
 airflow/ui/src/queries/useLogs.tsx            | 54 ++++++++++++----
 4 files changed, 106 insertions(+), 16 deletions(-)

diff --git a/airflow/ui/src/constants/searchParams.ts b/airflow/ui/src/constants/searchParams.ts
index 821fb6d474460..cf8a9e873638f 100644
--- a/airflow/ui/src/constants/searchParams.ts
+++ b/airflow/ui/src/constants/searchParams.ts
@@ -22,6 +22,7 @@ export enum SearchParamsKeys {
   LAST_DAG_RUN_STATE = "last_dag_run_state",
   LIMIT = "limit",
   LOG_LEVEL = "log_level",
+  LOGGER = "logger",
   NAME_PATTERN = "name_pattern",
   OFFSET = "offset",
   PAUSED = "paused",
diff --git a/airflow/ui/src/pages/TaskInstance/Logs/Logs.tsx b/airflow/ui/src/pages/TaskInstance/Logs/Logs.tsx
index 5c7306381270e..a61b2d7330e5c 100644
--- a/airflow/ui/src/pages/TaskInstance/Logs/Logs.tsx
+++ b/airflow/ui/src/pages/TaskInstance/Logs/Logs.tsx
@@ -35,6 +35,7 @@ export const Logs = () => {
 
   const tryNumberParam = searchParams.get(SearchParamsKeys.TRY_NUMBER);
   const logLevelFilters = searchParams.getAll(SearchParamsKeys.LOG_LEVEL);
+  const loggerFilters = searchParams.getAll(SearchParamsKeys.LOGGER);
 
   const {
     data: taskInstance,
@@ -76,6 +77,7 @@ export const Logs = () => {
     isLoading: isLoadingLogs,
   } = useLogs({
     dagId,
+    loggerFilters,
     logLevelFilters,
     taskInstance,
     tryNumber: tryNumber === 0 ? 1 : tryNumber,
@@ -84,6 +86,7 @@ export const Logs = () => {
   return (
     <Box p={2}>
       <TaskLogHeader
+        loggerOptions={data.loggers}
         onSelectTryNumber={onSelectTryNumber}
         taskInstance={taskInstance}
         toggleFullscreen={toggleFullscreen}
diff --git a/airflow/ui/src/pages/TaskInstance/Logs/TaskLogHeader.tsx b/airflow/ui/src/pages/TaskInstance/Logs/TaskLogHeader.tsx
index 2392121db9125..999a0dcc1c979 100644
--- a/airflow/ui/src/pages/TaskInstance/Logs/TaskLogHeader.tsx
+++ b/airflow/ui/src/pages/TaskInstance/Logs/TaskLogHeader.tsx
@@ -16,7 +16,14 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { Badge, Box, HStack, IconButton, type SelectValueChangeDetails } from "@chakra-ui/react";
+import {
+  Badge,
+  Box,
+  createListCollection,
+  HStack,
+  IconButton,
+  type SelectValueChangeDetails,
+} from "@chakra-ui/react";
 import { useCallback } from "react";
 import { MdOutlineOpenInFull } from "react-icons/md";
 import { useSearchParams } from "react-router-dom";
@@ -29,6 +36,7 @@ import { type LogLevel, logLevelColorMapping, logLevelOptions } from "src/utils/
 
 type Props = {
   readonly isFullscreen?: boolean;
+  readonly loggerOptions?: Array<string>;
   readonly onSelectTryNumber: (tryNumber: number) => void;
   readonly taskInstance?: TaskInstanceResponse;
   readonly toggleFullscreen: () => void;
@@ -39,6 +47,7 @@ type Props = {
 
 export const TaskLogHeader = ({
   isFullscreen = false,
+  loggerOptions,
   onSelectTryNumber,
   taskInstance,
   toggleFullscreen,
@@ -47,11 +56,21 @@ export const TaskLogHeader = ({
   wrap,
 }: Props) => {
   const [searchParams, setSearchParams] = useSearchParams();
-
+  const loggers = searchParams.getAll(SearchParamsKeys.LOGGER);
   const logLevels = searchParams.getAll(SearchParamsKeys.LOG_LEVEL);
   const hasLogLevels = logLevels.length > 0;
 
-  const handleStateChange = useCallback(
+  const loggerOptionList = createListCollection<{
+    label: string;
+    value: string;
+  }>({
+    items: [
+      { label: "All Loggers", value: "all" },
+      ...(loggerOptions ?? []).map((logger) => ({ label: logger, value: logger })),
+    ],
+  });
+
+  const handleLevelChange = useCallback(
     ({ value }: SelectValueChangeDetails<string>) => {
       const [val, ...rest] = value;
 
@@ -68,6 +87,23 @@ export const TaskLogHeader = ({
     [searchParams, setSearchParams],
   );
 
+  const handleLoggerChange = useCallback(
+    ({ value }: SelectValueChangeDetails<string>) => {
+      const [val, ...rest] = value;
+
+      if ((val === undefined || val === "all") && rest.length === 0) {
+        searchParams.delete(SearchParamsKeys.LOGGER);
+      } else {
+        searchParams.delete(SearchParamsKeys.LOGGER);
+        value
+          .filter((state) => state !== "all")
+          .map((state) => searchParams.append(SearchParamsKeys.LOGGER, state));
+      }
+      setSearchParams(searchParams);
+    },
+    [searchParams, setSearchParams],
+  );
+
   return (
     <Box>
       {taskInstance === undefined || tryNumber === undefined || taskInstance.try_number <= 1 ? undefined : (
@@ -82,7 +118,7 @@ export const TaskLogHeader = ({
           collection={logLevelOptions}
           maxW="250px"
           multiple
-          onValueChange={handleStateChange}
+          onValueChange={handleLevelChange}
           value={hasLogLevels ? logLevels : ["all"]}
         >
           <Select.Trigger {...(hasLogLevels ? { clearable: true } : {})} isActive={Boolean(logLevels)}>
@@ -114,6 +150,26 @@ export const TaskLogHeader = ({
             ))}
           </Select.Content>
         </Select.Root>
+        {loggerOptions !== undefined && loggerOptions.length > 0 ? (
+          <Select.Root
+            collection={loggerOptionList}
+            maxW="250px"
+            multiple
+            onValueChange={handleLoggerChange}
+            value={loggers}
+          >
+            <Select.Trigger clearable>
+              <Select.ValueText placeholder="All Loggers" />
+            </Select.Trigger>
+            <Select.Content>
+              {loggerOptionList.items.map((option) => (
+                <Select.Item item={option} key={option.label}>
+                  {option.label}
+                </Select.Item>
+              ))}
+            </Select.Content>
+          </Select.Root>
+        ) : undefined}
         <HStack>
           <Button aria-label={wrap ? "Unwrap" : "Wrap"} bg="bg.panel" onClick={toggleWrap} variant="outline">
             {wrap ? "Unwrap" : "Wrap"}
diff --git a/airflow/ui/src/queries/useLogs.tsx b/airflow/ui/src/queries/useLogs.tsx
index b2323d203a1d6..d80088ee7861d 100644
--- a/airflow/ui/src/queries/useLogs.tsx
+++ b/airflow/ui/src/queries/useLogs.tsx
@@ -33,6 +33,7 @@ import { LogLevel, logLevelColorMapping } from "src/utils/logs";
 
 type Props = {
   dagId: string;
+  loggerFilters?: Array<string>;
   logLevelFilters?: Array<string>;
   taskInstance?: TaskInstanceResponse;
   tryNumber?: number;
@@ -40,14 +41,23 @@ type Props = {
 
 type ParseLogsProps = {
   data: TaskInstancesLogResponse["content"];
+  loggerFilters?: Array<string>;
   logLevelFilters?: Array<string>;
 };
 
-const renderStructuredLog = (
-  logMessage: string | StructuredLogMessage,
-  index: number,
-  logLevelFilters?: Array<string>,
-) => {
+type RenderStructuredLogProps = {
+  index: number;
+  loggerFilters?: Array<string>;
+  logLevelFilters?: Array<string>;
+  logMessage: string | StructuredLogMessage;
+};
+
+const renderStructuredLog = ({
+  index,
+  loggerFilters,
+  logLevelFilters,
+  logMessage,
+}: RenderStructuredLogProps) => {
   if (typeof logMessage === "string") {
     return (
       <chakra.span key={index} lineHeight={1.5}>
@@ -68,6 +78,15 @@ const renderStructuredLog = (
     return "";
   }
 
+  if (
+    loggerFilters !== undefined &&
+    Boolean(loggerFilters.length) &&
+    (("logger" in structured && !loggerFilters.includes(structured.logger as string)) ||
+      !("logger" in structured))
+  ) {
+    return "";
+  }
+
   if (Boolean(timestamp)) {
     elements.push("[", <Time datetime={timestamp} key={0} />, "] ");
   }
@@ -97,9 +116,9 @@ const renderStructuredLog = (
     if (Object.hasOwn(structured, key)) {
       elements.push(
         " ",
-        <span className={`log-key ${key}`} key={`prop_${key}`}>
+        <chakra.span color={key === "logger" ? "fg.info" : undefined} key={`prop_${key}`}>
           {key}={JSON.stringify(structured[key])}
-        </span>,
+        </chakra.span>,
       );
     }
   }
@@ -111,16 +130,26 @@ const renderStructuredLog = (
   );
 };
 
-// TODO: add support for log groups, colors, formats, filters
-const parseLogs = ({ data, logLevelFilters }: ParseLogsProps) => {
+const parseLogs = ({ data, loggerFilters, logLevelFilters }: ParseLogsProps) => {
   let warning;
   let parsedLines;
   let startGroup = false;
   let groupLines: Array<JSX.Element | ""> = [];
   let groupName = "";
+  const loggers: Array<string> = [];
 
   try {
-    parsedLines = data.map((datum, index) => renderStructuredLog(datum, index, logLevelFilters));
+    parsedLines = data.map((datum, index) => {
+      if (typeof datum !== "string" && "logger" in datum) {
+        const logger = datum.logger as string;
+
+        if (!loggers.includes(logger)) {
+          loggers.push(logger);
+        }
+      }
+
+      return renderStructuredLog({ index, loggerFilters, logLevelFilters, logMessage: datum });
+    });
   } catch (error) {
     const errorMessage = error instanceof Error ? error.message : "An error occurred.";
 
@@ -167,14 +196,14 @@ const parseLogs = ({ data, logLevelFilters }: ParseLogsProps) => {
   });
 
   return {
-    fileSources: [],
+    loggers,
     parsedLogs: parsedLines,
     warning,
   };
 };
 
 export const useLogs = (
-  { dagId, logLevelFilters, taskInstance, tryNumber = 1 }: Props,
+  { dagId, loggerFilters, logLevelFilters, taskInstance, tryNumber = 1 }: Props,
   options?: Omit<UseQueryOptions<TaskInstancesLogResponse>, "queryFn" | "queryKey">,
 ) => {
   const refetchInterval = useAutoRefresh({ dagId });
@@ -201,6 +230,7 @@ export const useLogs = (
 
   const parsedData = parseLogs({
     data: data?.content ?? [],
+    loggerFilters,
     logLevelFilters,
   });
 

From e5f650197673bae0527d9eeb66ba4d87e0c21601 Mon Sep 17 00:00:00 2001
From: Brent Bovenzi <brent.bovenzi@gmail.com>
Date: Thu, 6 Mar 2025 16:57:54 -0500
Subject: [PATCH 2/2] Rename logger to source in UI

---
 airflow/ui/src/constants/searchParams.ts      |  2 +-
 .../ui/src/pages/TaskInstance/Logs/Logs.tsx   |  6 ++--
 .../pages/TaskInstance/Logs/TaskLogHeader.tsx | 32 ++++++++---------
 airflow/ui/src/queries/useLogs.tsx            | 34 +++++++++----------
 4 files changed, 37 insertions(+), 37 deletions(-)

diff --git a/airflow/ui/src/constants/searchParams.ts b/airflow/ui/src/constants/searchParams.ts
index cf8a9e873638f..65c5b2056f5c6 100644
--- a/airflow/ui/src/constants/searchParams.ts
+++ b/airflow/ui/src/constants/searchParams.ts
@@ -22,11 +22,11 @@ export enum SearchParamsKeys {
   LAST_DAG_RUN_STATE = "last_dag_run_state",
   LIMIT = "limit",
   LOG_LEVEL = "log_level",
-  LOGGER = "logger",
   NAME_PATTERN = "name_pattern",
   OFFSET = "offset",
   PAUSED = "paused",
   SORT = "sort",
+  SOURCE = "log_source",
   START_DATE = "start_date",
   STATE = "state",
   TAGS = "tags",
diff --git a/airflow/ui/src/pages/TaskInstance/Logs/Logs.tsx b/airflow/ui/src/pages/TaskInstance/Logs/Logs.tsx
index a61b2d7330e5c..e06c301da2c9e 100644
--- a/airflow/ui/src/pages/TaskInstance/Logs/Logs.tsx
+++ b/airflow/ui/src/pages/TaskInstance/Logs/Logs.tsx
@@ -35,7 +35,7 @@ export const Logs = () => {
 
   const tryNumberParam = searchParams.get(SearchParamsKeys.TRY_NUMBER);
   const logLevelFilters = searchParams.getAll(SearchParamsKeys.LOG_LEVEL);
-  const loggerFilters = searchParams.getAll(SearchParamsKeys.LOGGER);
+  const sourceFilters = searchParams.getAll(SearchParamsKeys.SOURCE);
 
   const {
     data: taskInstance,
@@ -77,8 +77,8 @@ export const Logs = () => {
     isLoading: isLoadingLogs,
   } = useLogs({
     dagId,
-    loggerFilters,
     logLevelFilters,
+    sourceFilters,
     taskInstance,
     tryNumber: tryNumber === 0 ? 1 : tryNumber,
   });
@@ -86,8 +86,8 @@ export const Logs = () => {
   return (
     <Box p={2}>
       <TaskLogHeader
-        loggerOptions={data.loggers}
         onSelectTryNumber={onSelectTryNumber}
+        sourceOptions={data.sources}
         taskInstance={taskInstance}
         toggleFullscreen={toggleFullscreen}
         toggleWrap={toggleWrap}
diff --git a/airflow/ui/src/pages/TaskInstance/Logs/TaskLogHeader.tsx b/airflow/ui/src/pages/TaskInstance/Logs/TaskLogHeader.tsx
index 999a0dcc1c979..699da9e35e988 100644
--- a/airflow/ui/src/pages/TaskInstance/Logs/TaskLogHeader.tsx
+++ b/airflow/ui/src/pages/TaskInstance/Logs/TaskLogHeader.tsx
@@ -36,8 +36,8 @@ import { type LogLevel, logLevelColorMapping, logLevelOptions } from "src/utils/
 
 type Props = {
   readonly isFullscreen?: boolean;
-  readonly loggerOptions?: Array<string>;
   readonly onSelectTryNumber: (tryNumber: number) => void;
+  readonly sourceOptions?: Array<string>;
   readonly taskInstance?: TaskInstanceResponse;
   readonly toggleFullscreen: () => void;
   readonly toggleWrap: () => void;
@@ -47,8 +47,8 @@ type Props = {
 
 export const TaskLogHeader = ({
   isFullscreen = false,
-  loggerOptions,
   onSelectTryNumber,
+  sourceOptions,
   taskInstance,
   toggleFullscreen,
   toggleWrap,
@@ -56,17 +56,17 @@ export const TaskLogHeader = ({
   wrap,
 }: Props) => {
   const [searchParams, setSearchParams] = useSearchParams();
-  const loggers = searchParams.getAll(SearchParamsKeys.LOGGER);
+  const sources = searchParams.getAll(SearchParamsKeys.SOURCE);
   const logLevels = searchParams.getAll(SearchParamsKeys.LOG_LEVEL);
   const hasLogLevels = logLevels.length > 0;
 
-  const loggerOptionList = createListCollection<{
+  const sourceOptionList = createListCollection<{
     label: string;
     value: string;
   }>({
     items: [
-      { label: "All Loggers", value: "all" },
-      ...(loggerOptions ?? []).map((logger) => ({ label: logger, value: logger })),
+      { label: "All Sources", value: "all" },
+      ...(sourceOptions ?? []).map((source) => ({ label: source, value: source })),
     ],
   });
 
@@ -87,17 +87,17 @@ export const TaskLogHeader = ({
     [searchParams, setSearchParams],
   );
 
-  const handleLoggerChange = useCallback(
+  const handleSourceChange = useCallback(
     ({ value }: SelectValueChangeDetails<string>) => {
       const [val, ...rest] = value;
 
       if ((val === undefined || val === "all") && rest.length === 0) {
-        searchParams.delete(SearchParamsKeys.LOGGER);
+        searchParams.delete(SearchParamsKeys.SOURCE);
       } else {
-        searchParams.delete(SearchParamsKeys.LOGGER);
+        searchParams.delete(SearchParamsKeys.SOURCE);
         value
           .filter((state) => state !== "all")
-          .map((state) => searchParams.append(SearchParamsKeys.LOGGER, state));
+          .map((state) => searchParams.append(SearchParamsKeys.SOURCE, state));
       }
       setSearchParams(searchParams);
     },
@@ -150,19 +150,19 @@ export const TaskLogHeader = ({
             ))}
           </Select.Content>
         </Select.Root>
-        {loggerOptions !== undefined && loggerOptions.length > 0 ? (
+        {sourceOptions !== undefined && sourceOptions.length > 0 ? (
           <Select.Root
-            collection={loggerOptionList}
+            collection={sourceOptionList}
             maxW="250px"
             multiple
-            onValueChange={handleLoggerChange}
-            value={loggers}
+            onValueChange={handleSourceChange}
+            value={sources}
           >
             <Select.Trigger clearable>
-              <Select.ValueText placeholder="All Loggers" />
+              <Select.ValueText placeholder="All Sources" />
             </Select.Trigger>
             <Select.Content>
-              {loggerOptionList.items.map((option) => (
+              {sourceOptionList.items.map((option) => (
                 <Select.Item item={option} key={option.label}>
                   {option.label}
                 </Select.Item>
diff --git a/airflow/ui/src/queries/useLogs.tsx b/airflow/ui/src/queries/useLogs.tsx
index d80088ee7861d..024e663780d72 100644
--- a/airflow/ui/src/queries/useLogs.tsx
+++ b/airflow/ui/src/queries/useLogs.tsx
@@ -33,30 +33,30 @@ import { LogLevel, logLevelColorMapping } from "src/utils/logs";
 
 type Props = {
   dagId: string;
-  loggerFilters?: Array<string>;
   logLevelFilters?: Array<string>;
+  sourceFilters?: Array<string>;
   taskInstance?: TaskInstanceResponse;
   tryNumber?: number;
 };
 
 type ParseLogsProps = {
   data: TaskInstancesLogResponse["content"];
-  loggerFilters?: Array<string>;
   logLevelFilters?: Array<string>;
+  sourceFilters?: Array<string>;
 };
 
 type RenderStructuredLogProps = {
   index: number;
-  loggerFilters?: Array<string>;
   logLevelFilters?: Array<string>;
   logMessage: string | StructuredLogMessage;
+  sourceFilters?: Array<string>;
 };
 
 const renderStructuredLog = ({
   index,
-  loggerFilters,
   logLevelFilters,
   logMessage,
+  sourceFilters,
 }: RenderStructuredLogProps) => {
   if (typeof logMessage === "string") {
     return (
@@ -79,9 +79,9 @@ const renderStructuredLog = ({
   }
 
   if (
-    loggerFilters !== undefined &&
-    Boolean(loggerFilters.length) &&
-    (("logger" in structured && !loggerFilters.includes(structured.logger as string)) ||
+    sourceFilters !== undefined &&
+    Boolean(sourceFilters.length) &&
+    (("logger" in structured && !sourceFilters.includes(structured.logger as string)) ||
       !("logger" in structured))
   ) {
     return "";
@@ -117,7 +117,7 @@ const renderStructuredLog = ({
       elements.push(
         " ",
         <chakra.span color={key === "logger" ? "fg.info" : undefined} key={`prop_${key}`}>
-          {key}={JSON.stringify(structured[key])}
+          {key === "logger" ? "source" : key}={JSON.stringify(structured[key])}
         </chakra.span>,
       );
     }
@@ -130,25 +130,25 @@ const renderStructuredLog = ({
   );
 };
 
-const parseLogs = ({ data, loggerFilters, logLevelFilters }: ParseLogsProps) => {
+const parseLogs = ({ data, logLevelFilters, sourceFilters }: ParseLogsProps) => {
   let warning;
   let parsedLines;
   let startGroup = false;
   let groupLines: Array<JSX.Element | ""> = [];
   let groupName = "";
-  const loggers: Array<string> = [];
+  const sources: Array<string> = [];
 
   try {
     parsedLines = data.map((datum, index) => {
       if (typeof datum !== "string" && "logger" in datum) {
-        const logger = datum.logger as string;
+        const source = datum.logger as string;
 
-        if (!loggers.includes(logger)) {
-          loggers.push(logger);
+        if (!sources.includes(source)) {
+          sources.push(source);
         }
       }
 
-      return renderStructuredLog({ index, loggerFilters, logLevelFilters, logMessage: datum });
+      return renderStructuredLog({ index, logLevelFilters, logMessage: datum, sourceFilters });
     });
   } catch (error) {
     const errorMessage = error instanceof Error ? error.message : "An error occurred.";
@@ -196,14 +196,14 @@ const parseLogs = ({ data, loggerFilters, logLevelFilters }: ParseLogsProps) =>
   });
 
   return {
-    loggers,
     parsedLogs: parsedLines,
+    sources,
     warning,
   };
 };
 
 export const useLogs = (
-  { dagId, loggerFilters, logLevelFilters, taskInstance, tryNumber = 1 }: Props,
+  { dagId, logLevelFilters, sourceFilters, taskInstance, tryNumber = 1 }: Props,
   options?: Omit<UseQueryOptions<TaskInstancesLogResponse>, "queryFn" | "queryKey">,
 ) => {
   const refetchInterval = useAutoRefresh({ dagId });
@@ -230,8 +230,8 @@ export const useLogs = (
 
   const parsedData = parseLogs({
     data: data?.content ?? [],
-    loggerFilters,
     logLevelFilters,
+    sourceFilters,
   });
 
   return { data: parsedData, ...rest };