apache-airflow-providers-edge3 1.5.0__py3-none-any.whl → 2.0.0__py3-none-any.whl

airflow/providers/edge3/plugins/www/src/pages/WorkerPage.tsx

@@ -16,22 +16,58 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { Box, Code, Link, List, Table, Text } from "@chakra-ui/react";
+import { Box, Code, HStack, Link, List, Table, Text, type SelectValueChangeDetails } from "@chakra-ui/react";
+import { useState, useCallback } from "react";
 import { useUiServiceWorker } from "openapi/queries";
 import { LuExternalLink } from "react-icons/lu";
 import TimeAgo from "react-timeago";
 
 import { ErrorAlert } from "src/components/ErrorAlert";
+import { SearchBar } from "src/components/SearchBar";
 import { WorkerOperations } from "src/components/WorkerOperations";
 import { WorkerStateBadge } from "src/components/WorkerStateBadge";
-import { ScrollToAnchor } from "src/components/ui";
+import { ScrollToAnchor, Select } from "src/components/ui";
+import { workerStateOptions } from "src/constants";
 import { autoRefreshInterval } from "src/utils";
+import type { EdgeWorkerState } from "openapi/requests/types.gen";
 
 export const WorkerPage = () => {
-  const { data, error, refetch } = useUiServiceWorker(undefined, {
-    enabled: true,
-    refetchInterval: autoRefreshInterval,
-  });
+  const [workerNamePattern, setWorkerNamePattern] = useState("");
+  const [queueNamePattern, setQueueNamePattern] = useState("");
+  const [filteredState, setFilteredState] = useState<string[]>([]);
+
+  const hasFilteredState = filteredState.length > 0;
+
+  const { data, error, refetch } = useUiServiceWorker(
+    {
+      queueNamePattern: queueNamePattern || undefined,
+      state: hasFilteredState ? (filteredState as EdgeWorkerState[]) : undefined,
+      workerNamePattern: workerNamePattern || undefined,
+    },
+    undefined,
+    {
+      enabled: true,
+      refetchInterval: autoRefreshInterval,
+    },
+  );
+
+  const handleWorkerSearchChange = (value: string) => {
+    setWorkerNamePattern(value);
+  };
+
+  const handleQueueSearchChange = (value: string) => {
+    setQueueNamePattern(value);
+  };
+
+  const handleStateChange = useCallback(({ value }: SelectValueChangeDetails<string>) => {
+    const [val, ...rest] = value;
+
+    if ((val === undefined || val === "all") && rest.length === 0) {
+      setFilteredState([]);
+    } else {
+      setFilteredState(value.filter((state) => state !== "all"));
+    }
+  }, []);
 
   // TODO to make it proper
   // Use DataTable as component from Airflow-Core UI
@@ -40,96 +76,149 @@ export const WorkerPage = () => {
   // Add links with filter to see jobs on worker
   // Add time zone support for time display
   // Translation?
-  if (data?.workers && data.workers.length > 0)
-    return (
-      <Box p={2}>
-        <Table.Root size="sm" interactive stickyHeader striped>
-          <Table.Header>
-            <Table.Row>
-              <Table.ColumnHeader>Worker Name</Table.ColumnHeader>
-              <Table.ColumnHeader>State</Table.ColumnHeader>
-              <Table.ColumnHeader>Queues</Table.ColumnHeader>
-              <Table.ColumnHeader>First Online</Table.ColumnHeader>
-              <Table.ColumnHeader>Last Heartbeat</Table.ColumnHeader>
-              <Table.ColumnHeader>Active Jobs</Table.ColumnHeader>
-              <Table.ColumnHeader>System Information</Table.ColumnHeader>
-              <Table.ColumnHeader>Operations</Table.ColumnHeader>
-            </Table.Row>
-          </Table.Header>
-          <Table.Body>
-            {data.workers.map((worker) => (
-              <Table.Row key={worker.worker_name} id={worker.worker_name}>
-                <Table.Cell>{worker.worker_name}</Table.Cell>
-                <Table.Cell>
-                  <WorkerStateBadge state={worker.state}>{worker.state}</WorkerStateBadge>
-                </Table.Cell>
-                <Table.Cell>
-                  {worker.queues ? (
-                    <List.Root>
-                      {worker.queues.map((queue) => (
-                        <List.Item key={queue}>{queue}</List.Item>
-                      ))}
-                    </List.Root>
-                  ) : (
-                    "(all queues)"
-                  )}
-                </Table.Cell>
-                <Table.Cell>
-                  {worker.first_online ? <TimeAgo date={worker.first_online} live={false} /> : undefined}
-                </Table.Cell>
-                <Table.Cell>
-                  {worker.last_heartbeat ? <TimeAgo date={worker.last_heartbeat} live={false} /> : undefined}
-                </Table.Cell>
-                <Table.Cell>{worker.jobs_active}</Table.Cell>
-                <Table.Cell>
-                  {worker.sysinfo ? (
-                    <List.Root>
-                      {Object.entries(worker.sysinfo).map(([key, value]) => (
-                        <List.Item key={key}>
-                          {key}: {value}
-                        </List.Item>
-                      ))}
-                    </List.Root>
-                  ) : (
-                    "N/A"
-                  )}
-                </Table.Cell>
-                <Table.Cell>
-                  <WorkerOperations worker={worker} onOperations={refetch} />
-                </Table.Cell>
-              </Table.Row>
-            ))}
-          </Table.Body>
-        </Table.Root>
-        <ScrollToAnchor />
-      </Box>
-    );
-  if (data) {
-    return (
-      <Text as="div" pl={4} pt={1}>
-        No known workers. Start one via <Code>airflow edge worker [...]</Code>. See{" "}
-        <Link
-          target="_blank"
-          variant="underline"
-          color="fg.info"
-          href="https://airflow.apache.org/docs/apache-airflow-providers-edge3/stable/deployment.html"
+
+  return (
+    <Box p={2}>
+      <HStack gap={4} mb={4}>
+        <SearchBar
+          buttonProps={{ disabled: true }}
+          defaultValue={workerNamePattern}
+          hideAdvanced
+          hotkeyDisabled
+          onChange={handleWorkerSearchChange}
+          placeHolder="Search workers"
+        />
+        <SearchBar
+          buttonProps={{ disabled: true }}
+          defaultValue={queueNamePattern}
+          hideAdvanced
+          hotkeyDisabled
+          onChange={handleQueueSearchChange}
+          placeHolder="Search queues"
+        />
+        <Select.Root
+          collection={workerStateOptions}
+          maxW="450px"
+          multiple
+          onValueChange={handleStateChange}
+          value={hasFilteredState ? filteredState : ["all"]}
         >
-          Edge Worker Deployment docs <LuExternalLink />
-        </Link>{" "}
-        how to deploy a new worker.
-      </Text>
-    );
-  }
-  if (error) {
-    return (
-      <Text as="div" pl={4} pt={1}>
+          <Select.Trigger
+            {...(hasFilteredState ? { clearable: true } : {})}
+            colorPalette="brand"
+            isActive={Boolean(filteredState)}
+          >
+            <Select.ValueText>
+              {() =>
+                hasFilteredState ? (
+                  <HStack flexWrap="wrap" fontSize="sm" gap="4px" paddingY="8px">
+                    {filteredState.map((state) => (
+                      <WorkerStateBadge key={state} state={state as EdgeWorkerState}>
+                        {state}
+                      </WorkerStateBadge>
+                    ))}
+                  </HStack>
+                ) : (
+                  "All States"
+                )
+              }
+            </Select.ValueText>
+          </Select.Trigger>
+          <Select.Content>
+            {workerStateOptions.items.map((option) => (
+              <Select.Item item={option} key={option.label}>
+                {option.value === "all" ? (
+                  option.label
+                ) : (
+                  <WorkerStateBadge state={option.value as EdgeWorkerState}>{option.label}</WorkerStateBadge>
+                )}
+              </Select.Item>
+            ))}
+          </Select.Content>
+        </Select.Root>
+      </HStack>
+      {error ? (
         <ErrorAlert error={error} />
-      </Text>
-    );
-  }
-  return (
-    <Text as="div" pl={4} pt={1}>
-      Loading...
-    </Text>
+      ) : !data ? (
+        <Text as="div" pl={2} pt={1}>
+          Loading...
+        </Text>
+      ) : data.workers && data.workers.length > 0 ? (
+        <>
+          <Table.Root size="sm" interactive stickyHeader striped>
+            <Table.Header>
+              <Table.Row>
+                <Table.ColumnHeader>Worker Name</Table.ColumnHeader>
+                <Table.ColumnHeader>State</Table.ColumnHeader>
+                <Table.ColumnHeader>Queues</Table.ColumnHeader>
+                <Table.ColumnHeader>First Online</Table.ColumnHeader>
+                <Table.ColumnHeader>Last Heartbeat</Table.ColumnHeader>
+                <Table.ColumnHeader>Active Jobs</Table.ColumnHeader>
+                <Table.ColumnHeader>System Information</Table.ColumnHeader>
+                <Table.ColumnHeader>Operations</Table.ColumnHeader>
+              </Table.Row>
+            </Table.Header>
+            <Table.Body>
+              {data.workers.map((worker) => (
+                <Table.Row key={worker.worker_name} id={worker.worker_name}>
+                  <Table.Cell>{worker.worker_name}</Table.Cell>
+                  <Table.Cell>
+                    <WorkerStateBadge state={worker.state}>{worker.state}</WorkerStateBadge>
+                  </Table.Cell>
+                  <Table.Cell>
+                    {worker.queues ? (
+                      <List.Root>
+                        {worker.queues.map((queue) => (
+                          <List.Item key={queue}>{queue}</List.Item>
+                        ))}
+                      </List.Root>
+                    ) : (
+                      "(all queues)"
+                    )}
+                  </Table.Cell>
+                  <Table.Cell>
+                    {worker.first_online ? <TimeAgo date={worker.first_online} live={false} /> : undefined}
+                  </Table.Cell>
+                  <Table.Cell>
+                    {worker.last_heartbeat ? <TimeAgo date={worker.last_heartbeat} live={false} /> : undefined}
+                  </Table.Cell>
+                  <Table.Cell>{worker.jobs_active}</Table.Cell>
+                  <Table.Cell>
+                    {worker.sysinfo ? (
+                      <List.Root>
+                        {Object.entries(worker.sysinfo).map(([key, value]) => (
+                          <List.Item key={key}>
+                            {key}: {value}
+                          </List.Item>
+                        ))}
+                      </List.Root>
+                    ) : (
+                      "N/A"
+                    )}
+                  </Table.Cell>
+                  <Table.Cell>
+                    <WorkerOperations worker={worker} onOperations={refetch} />
+                  </Table.Cell>
+                </Table.Row>
+              ))}
+            </Table.Body>
+          </Table.Root>
+          <ScrollToAnchor />
+        </>
+      ) : (
+        <Text as="div" pl={2} pt={1}>
+          No known workers. Start one via <Code>airflow edge worker [...]</Code>. See{" "}
+          <Link
+            target="_blank"
+            variant="underline"
+            color="fg.info"
+            href="https://airflow.apache.org/docs/apache-airflow-providers-edge3/stable/deployment.html"
+          >
+            Edge Worker Deployment docs <LuExternalLink />
+          </Link>{" "}
+          how to deploy a new worker.
+        </Text>
+      )}
+    </Box>
   );
 };

airflow/providers/edge3/version_compat.py

@@ -32,10 +32,8 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
     return airflow_version.major, airflow_version.minor, airflow_version.micro
 
 
-AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0)
 AIRFLOW_V_3_1_PLUS = get_base_airflow_version_tuple() >= (3, 1, 0)
 
 __all__ = [
-    "AIRFLOW_V_3_0_PLUS",
     "AIRFLOW_V_3_1_PLUS",
 ]

airflow/providers/edge3/worker_api/auth.py

@@ -20,6 +20,7 @@ import logging
 from functools import cache
 from uuid import uuid4
 
+from fastapi import Header, HTTPException, Request, status
 from itsdangerous import BadSignature
 from jwt import (
     ExpiredSignatureError,
@@ -29,49 +30,24 @@ from jwt import (
     InvalidSignatureError,
 )
 
+from airflow.api_fastapi.auth.tokens import JWTValidator
 from airflow.configuration import conf
-from airflow.providers.edge3.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.providers.edge3.worker_api.datamodels import JsonRpcRequestBase  # noqa: TCH001
-from airflow.providers.edge3.worker_api.routes._v2_compat import (
-    Header,
-    HTTPException,
-    Request,
-    status,
-)
 
 log = logging.getLogger(__name__)
 
 
-if AIRFLOW_V_3_0_PLUS:
-    from airflow.api_fastapi.auth.tokens import JWTValidator
-
-    @cache
-    def jwt_validator() -> JWTValidator:
-        return JWTValidator(
-            secret_key=conf.get("api_auth", "jwt_secret"),
-            leeway=conf.getint("api_auth", "jwt_leeway", fallback=30),
-            audience="api",
-        )
+@cache
+def jwt_validator() -> JWTValidator:
+    return JWTValidator(
+        secret_key=conf.get("api_auth", "jwt_secret"),
+        leeway=conf.getint("api_auth", "jwt_leeway", fallback=30),
+        audience="api",
+    )
 
-    def jwt_validate(authorization: str) -> dict:
-        return jwt_validator().validated_claims(authorization)
-
-else:
-    # Airflow 2.10 compatibility
-    from airflow.utils.jwt_signer import JWTSigner
-
-    @cache
-    def jwt_signer() -> JWTSigner:
-        clock_grace = conf.getint("core", "internal_api_clock_grace", fallback=30)
-        return JWTSigner(
-            secret_key=conf.get("core", "internal_api_secret_key"),
-            expiration_time_in_seconds=clock_grace,
-            leeway_in_seconds=clock_grace,
-            audience="api",
-        )
 
-    def jwt_validate(authorization: str) -> dict:
-        return jwt_signer().verify_token(authorization)
+def jwt_validate(authorization: str) -> dict:
+    return jwt_validator().validated_claims(authorization)
 
 
 def _forbidden_response(message: str):

airflow/providers/edge3/worker_api/datamodels.py

@@ -22,11 +22,12 @@ from typing import (
     Any,
 )
 
+from fastapi import Path
 from pydantic import BaseModel, Field
 
+from airflow.executors.workloads import ExecuteTask  # noqa: TCH001
 from airflow.models.taskinstancekey import TaskInstanceKey
 from airflow.providers.edge3.models.edge_worker import EdgeWorkerState  # noqa: TCH001
-from airflow.providers.edge3.worker_api.routes._v2_compat import ExecuteTask, Path
 
 
 class WorkerApiDocs:
@@ -93,7 +94,7 @@ class EdgeJobFetched(EdgeJobBase):
         ExecuteTask,
         Field(
             title="Command",
-            description="Command line to use to execute the job in Airflow 2. Task definition in Airflow 3",
+            description="Command line to use to execute the job in Airflow",
         ),
     ]
     concurrency_slots: Annotated[int, Field(description="Number of concurrency slots the job requires.")]

airflow/providers/edge3/worker_api/routes/health.py

@@ -17,7 +17,7 @@
 
 from __future__ import annotations
 
-from airflow.providers.edge3.worker_api.routes._v2_compat import AirflowRouter
+from airflow.api_fastapi.common.router import AirflowRouter
 
 health_router = AirflowRouter(tags=["Monitor"])
 

airflow/providers/edge3/worker_api/routes/jobs.py

@@ -19,9 +19,14 @@ from __future__ import annotations
 
 from typing import Annotated
 
+from fastapi import Body, Depends, status
 from sqlalchemy import select, update
 
-from airflow.providers.common.compat.sdk import timezone
+from airflow.api_fastapi.common.db.common import SessionDep  # noqa: TC001
+from airflow.api_fastapi.common.router import AirflowRouter
+from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
+from airflow.executors.workloads import ExecuteTask
+from airflow.providers.common.compat.sdk import Stats, timezone
 from airflow.providers.edge3.models.edge_job import EdgeJobModel
 from airflow.providers.edge3.worker_api.auth import jwt_token_authorization_rest
 from airflow.providers.edge3.worker_api.datamodels import (
@@ -29,21 +34,15 @@ from airflow.providers.edge3.worker_api.datamodels import (
     WorkerApiDocs,
     WorkerQueuesBody,
 )
-from airflow.providers.edge3.worker_api.routes._v2_compat import (
-    AirflowRouter,
-    Body,
-    Depends,
-    SessionDep,
-    create_openapi_http_exception_doc,
-    parse_command,
-    status,
-)
-from airflow.stats import Stats
 from airflow.utils.state import TaskInstanceState
 
 jobs_router = AirflowRouter(tags=["Jobs"], prefix="/jobs")
 
 
+def parse_command(command: str) -> ExecuteTask:
+    return ExecuteTask.model_validate_json(command)
+
+
 @jobs_router.post(
     "/fetch/{worker_name}",
     dependencies=[Depends(jwt_token_authorization_rest)],

airflow/providers/edge3/worker_api/routes/logs.py

@@ -21,20 +21,17 @@ from functools import cache
 from pathlib import Path
 from typing import TYPE_CHECKING, Annotated
 
+from fastapi import Body, Depends, status
+
+from airflow.api_fastapi.common.db.common import SessionDep  # noqa: TC001
+from airflow.api_fastapi.common.router import AirflowRouter
+from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.configuration import conf
 from airflow.models.taskinstance import TaskInstance
 from airflow.models.taskinstancekey import TaskInstanceKey
 from airflow.providers.edge3.models.edge_logs import EdgeLogsModel
 from airflow.providers.edge3.worker_api.auth import jwt_token_authorization_rest
 from airflow.providers.edge3.worker_api.datamodels import PushLogsBody, WorkerApiDocs
-from airflow.providers.edge3.worker_api.routes._v2_compat import (
-    AirflowRouter,
-    Body,
-    Depends,
-    SessionDep,
-    create_openapi_http_exception_doc,
-    status,
-)
 from airflow.utils.log.file_task_handler import FileTaskHandler
 from airflow.utils.session import NEW_SESSION, provide_session
 

airflow/providers/edge3/worker_api/routes/ui.py

@@ -18,9 +18,9 @@
 from __future__ import annotations
 
 from datetime import datetime
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Annotated
 
-from fastapi import Depends, HTTPException, status
+from fastapi import Depends, HTTPException, Query, status
 from sqlalchemy import select
 
 from airflow.api_fastapi.auth.managers.models.resource_details import AccessView
@@ -30,6 +30,7 @@ from airflow.api_fastapi.core_api.security import GetUserDep, requires_access_vi
 from airflow.providers.edge3.models.edge_job import EdgeJobModel
 from airflow.providers.edge3.models.edge_worker import (
     EdgeWorkerModel,
+    EdgeWorkerState,
     add_worker_queues,
     change_maintenance_comment,
     exit_maintenance,
@@ -61,9 +62,19 @@ ui_router = AirflowRouter(tags=["UI"])
 )
 def worker(
     session: SessionDep,
+    worker_name_pattern: str | None = None,
+    queue_name_pattern: str | None = None,
+    state: Annotated[list[EdgeWorkerState] | None, Query()] = None,
 ) -> WorkerCollectionResponse:
     """Return Edge Workers."""
-    query = select(EdgeWorkerModel).order_by(EdgeWorkerModel.worker_name)
+    query = select(EdgeWorkerModel)
+    if worker_name_pattern:
+        query = query.where(EdgeWorkerModel.worker_name.ilike(f"%{worker_name_pattern}%"))
+    if queue_name_pattern:
+        query = query.where(EdgeWorkerModel._queues.ilike(f"%'{queue_name_pattern}%"))
+    if state:
+        query = query.where(EdgeWorkerModel.state.in_(state))
+    query = query.order_by(EdgeWorkerModel.worker_name)
     workers: ScalarResult[EdgeWorkerModel] = session.scalars(query)
 
     result = [

airflow/providers/edge3/worker_api/routes/worker.py

@@ -20,9 +20,13 @@ from __future__ import annotations
 import json
 from typing import Annotated
 
+from fastapi import Body, Depends, HTTPException, Path, status
 from sqlalchemy import select
 
-from airflow.providers.common.compat.sdk import timezone
+from airflow.api_fastapi.common.db.common import SessionDep  # noqa: TC001
+from airflow.api_fastapi.common.router import AirflowRouter
+from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
+from airflow.providers.common.compat.sdk import Stats, timezone
 from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel, EdgeWorkerState, set_metrics
 from airflow.providers.edge3.worker_api.auth import jwt_token_authorization_rest
 from airflow.providers.edge3.worker_api.datamodels import (
@@ -31,17 +35,6 @@ from airflow.providers.edge3.worker_api.datamodels import (
     WorkerSetStateReturn,
     WorkerStateBody,
 )
-from airflow.providers.edge3.worker_api.routes._v2_compat import (
-    AirflowRouter,
-    Body,
-    Depends,
-    HTTPException,
-    Path,
-    SessionDep,
-    create_openapi_http_exception_doc,
-    status,
-)
-from airflow.stats import Stats
 
 worker_router = AirflowRouter(
     tags=["Worker"],
@@ -50,6 +43,7 @@ worker_router = AirflowRouter(
         [
             status.HTTP_400_BAD_REQUEST,
             status.HTTP_403_FORBIDDEN,
+            status.HTTP_409_CONFLICT,
         ]
     ),
 )
@@ -175,6 +169,19 @@ def register(
     worker: EdgeWorkerModel | None = session.scalar(query)
     if not worker:
         worker = EdgeWorkerModel(worker_name=worker_name, state=body.state, queues=body.queues)
+    else:
+        # Prevent duplicate workers unless the existing worker is in offline or unknown state
+        allowed_states_for_reuse = {
+            EdgeWorkerState.OFFLINE,
+            EdgeWorkerState.UNKNOWN,
+            EdgeWorkerState.OFFLINE_MAINTENANCE,
+        }
+        if worker.state not in allowed_states_for_reuse:
+            raise HTTPException(
+                status.HTTP_409_CONFLICT,
+                f"Worker '{worker_name}' is already active in state '{worker.state}'. "
+                f"Cannot start a duplicate worker with the same name.",
+            )
     worker.state = redefine_state(worker.state, body.state)
     worker.maintenance_comment = redefine_maintenance_comments(
         worker.maintenance_comment, body.maintenance_comments