Change style for task logger and allow filtering by source

airflow/ui/src/constants/searchParams.ts

@@ -26,6 +26,7 @@ export enum SearchParamsKeys {
   OFFSET = "offset",
   PAUSED = "paused",
   SORT = "sort",
+  SOURCE = "log_source",
   START_DATE = "start_date",
   STATE = "state",
   TAGS = "tags",

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 sourceFilters = searchParams.getAll(SearchParamsKeys.SOURCE);
 
   const {
     data: taskInstance,
@@ -77,6 +78,7 @@ export const Logs = () => {
   } = useLogs({
     dagId,
     logLevelFilters,
+    sourceFilters,
     taskInstance,
     tryNumber: tryNumber === 0 ? 1 : tryNumber,
   });
@@ -85,6 +87,7 @@ export const Logs = () => {
     <Box p={2}>
       <TaskLogHeader
         onSelectTryNumber={onSelectTryNumber}
+        sourceOptions={data.sources}
         taskInstance={taskInstance}
         toggleFullscreen={toggleFullscreen}
         toggleWrap={toggleWrap}

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";
@@ -30,6 +37,7 @@ import { type LogLevel, logLevelColorMapping, logLevelOptions } from "src/utils/
 type Props = {
   readonly isFullscreen?: boolean;
   readonly onSelectTryNumber: (tryNumber: number) => void;
+  readonly sourceOptions?: Array<string>;
   readonly taskInstance?: TaskInstanceResponse;
   readonly toggleFullscreen: () => void;
   readonly toggleWrap: () => void;
@@ -40,18 +48,29 @@ type Props = {
 export const TaskLogHeader = ({
   isFullscreen = false,
   onSelectTryNumber,
+  sourceOptions,
   taskInstance,
   toggleFullscreen,
   toggleWrap,
   tryNumber,
   wrap,
 }: Props) => {
   const [searchParams, setSearchParams] = useSearchParams();
-
+  const sources = searchParams.getAll(SearchParamsKeys.SOURCE);
   const logLevels = searchParams.getAll(SearchParamsKeys.LOG_LEVEL);
   const hasLogLevels = logLevels.length > 0;
 
-  const handleStateChange = useCallback(
+  const sourceOptionList = createListCollection<{
+    label: string;
+    value: string;
+  }>({
+    items: [
+      { label: "All Sources", value: "all" },
+      ...(sourceOptions ?? []).map((source) => ({ label: source, value: source })),
+    ],
+  });
+
+  const handleLevelChange = useCallback(
     ({ value }: SelectValueChangeDetails<string>) => {
       const [val, ...rest] = value;
 
@@ -68,6 +87,23 @@ export const TaskLogHeader = ({
     [searchParams, setSearchParams],
   );
 
+  const handleSourceChange = useCallback(
+    ({ value }: SelectValueChangeDetails<string>) => {
+      const [val, ...rest] = value;
+
+      if ((val === undefined || val === "all") && rest.length === 0) {
+        searchParams.delete(SearchParamsKeys.SOURCE);
+      } else {
+        searchParams.delete(SearchParamsKeys.SOURCE);
+        value
+          .filter((state) => state !== "all")
+          .map((state) => searchParams.append(SearchParamsKeys.SOURCE, 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>
+        {sourceOptions !== undefined && sourceOptions.length > 0 ? (
+          <Select.Root
+            collection={sourceOptionList}
+            maxW="250px"
+            multiple
+            onValueChange={handleSourceChange}
+            value={sources}
+          >
+            <Select.Trigger clearable>
+              <Select.ValueText placeholder="All Sources" />
+            </Select.Trigger>
+            <Select.Content>
+              {sourceOptionList.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"}

airflow/ui/src/queries/useLogs.tsx

@@ -34,20 +34,30 @@ import { LogLevel, logLevelColorMapping } from "src/utils/logs";
 type Props = {
   dagId: string;
   logLevelFilters?: Array<string>;
+  sourceFilters?: Array<string>;
   taskInstance?: TaskInstanceResponse;
   tryNumber?: number;
 };
 
 type ParseLogsProps = {
   data: TaskInstancesLogResponse["content"];
   logLevelFilters?: Array<string>;
+  sourceFilters?: Array<string>;
 };
 
-const renderStructuredLog = (
-  logMessage: string | StructuredLogMessage,
-  index: number,
-  logLevelFilters?: Array<string>,
-) => {
+type RenderStructuredLogProps = {
+  index: number;
+  logLevelFilters?: Array<string>;
+  logMessage: string | StructuredLogMessage;
+  sourceFilters?: Array<string>;
+};
+
+const renderStructuredLog = ({
+  index,
+  logLevelFilters,
+  logMessage,
+  sourceFilters,
+}: RenderStructuredLogProps) => {
   if (typeof logMessage === "string") {
     return (
       <chakra.span key={index} lineHeight={1.5}>
@@ -68,6 +78,15 @@ const renderStructuredLog = (
     return "";
   }
 
+  if (
+    sourceFilters !== undefined &&
+    Boolean(sourceFilters.length) &&
+    (("logger" in structured && !sourceFilters.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}`}>
-          {key}={JSON.stringify(structured[key])}
-        </span>,
+        <chakra.span color={key === "logger" ? "fg.info" : undefined} key={`prop_${key}`}>
+          {key === "logger" ? "source" : key}={JSON.stringify(structured[key])}
+        </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, logLevelFilters, sourceFilters }: ParseLogsProps) => {
   let warning;
   let parsedLines;
   let startGroup = false;
   let groupLines: Array<JSX.Element | ""> = [];
   let groupName = "";
+  const sources: 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 source = datum.logger as string;
+
+        if (!sources.includes(source)) {
+          sources.push(source);
+        }
+      }
+
+      return renderStructuredLog({ index, logLevelFilters, logMessage: datum, sourceFilters });
+    });
   } catch (error) {
     const errorMessage = error instanceof Error ? error.message : "An error occurred.";
 
@@ -167,14 +196,14 @@ const parseLogs = ({ data, logLevelFilters }: ParseLogsProps) => {
   });
 
   return {
-    fileSources: [],
     parsedLogs: parsedLines,
+    sources,
     warning,
   };
 };
 
 export const useLogs = (
-  { dagId, logLevelFilters, taskInstance, tryNumber = 1 }: Props,
+  { dagId, logLevelFilters, sourceFilters, taskInstance, tryNumber = 1 }: Props,
   options?: Omit<UseQueryOptions<TaskInstancesLogResponse>, "queryFn" | "queryKey">,
 ) => {
   const refetchInterval = useAutoRefresh({ dagId });
@@ -202,6 +231,7 @@ export const useLogs = (
   const parsedData = parseLogs({
     data: data?.content ?? [],
     logLevelFilters,
+    sourceFilters,
   });
 
   return { data: parsedData, ...rest };