apache / apache/airflow

Replace `airflow.utils.state import` with `airflow.sdk` in edge3 Provider

Open
#72,408 2 comments 0 reactions 1 assignee Claimed by @Rushabh-Garambha View on GitHub
area:providers contributors-workshop kind:task provider:edge
Dominant language
Python
Stars
46.9k
Forks
17.8k
Avg merge
2d 9h
Merged PRs (30d)
472

Description

## Description

> This is an issue that is reserved for the Airflow Summit "Contributors" Workshop. This is denoted with the label `contributors-workshop`. Out of respect for the organizers and participants of this workshop, **please do not implement a PR that addresses this issue.**
>
> If this issue is still open following Airflow Summit, the label will be removed and the issue can be picked up.

Sub-issue of #62063, which investigates provider usage of `airflow.utils` imports so they can be replaced with Task SDK equivalents and reduce the `airflow-core` dependency in providers.

| | |
| --- | --- |
| Provider | edge3 |
| Import | `TaskInstanceState` from `airflow.utils.state` |
| Replacement | `airflow.sdk.TaskInstanceState` |
| Shape | Import-only change across six files |

edge3 requires `apache-airflow>=3.0.0,!=3.1.0` (see `providers/edge3/pyproject.toml`), so `airflow.sdk` is always available and no version guard is needed here, unlike most other providers that still support Airflow 2.11+. `airflow.sdk.TaskInstanceState` and `airflow.utils.state.TaskInstanceState` are two independently defined `(str, Enum)` classes with matching member names and string values, airflow-core's own docstring on `DagRunState` notes this invariant is deliberately maintained, so this is a plain import swap with no behavior change.

### What needs to change

Replace `from airflow.utils.state import TaskInstanceState` with `from airflow.sdk import TaskInstanceState` in each of:

- `providers/edge3/src/airflow/providers/edge3/cli/worker.py:66`
- `providers/edge3/src/airflow/providers/edge3/cli/api_client.py:50` (inside a `TYPE_CHECKING` block)
- `providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py:40`
- `providers/edge3/src/airflow/providers/edge3/worker_api/datamodels_ui.py:28` (keep the existing `# noqa: TC001` comment, this one is used at runtime as a pydantic field type so it cannot move into `TYPE_CHECKING`)
- `providers/edge3/src/airflow/providers/edge3/worker_api/routes/ui.py:52`
- `providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py:40`

No other code in these files needs to change, only the import line in each.

### How to verify it

The existing unit tests for each touched module should pass unmodified, since the enum values are identical:

```bash
breeze testing providers-tests providers/edge3/tests/unit/edge3/cli/test_worker.py
breeze testing providers-tests providers/edge3/tests/unit/edge3/cli/test_api_client.py
breeze testing providers-tests providers/edge3/tests/unit/edge3/executors/test_edge_executor.py
breeze testing providers-tests providers/edge3/tests/unit/edge3/worker_api/routes/test_ui.py
breeze testing providers-tests providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py
```

## Definition of Done

1. Update the import in all six files listed above.
2. Confirm no other `airflow.utils.state` usage remains in `providers/edge3/src` (`grep -rn "airflow.utils.state" providers/edge3/src`).
3. No new tests are needed, this is a like-for-like import swap, existing tests should pass unmodified.
4. This should pass:

```bash
breeze testing providers-tests --test-type "Providers[edge3]"
```

---
Drafted-by: Claude Code (Sonnet 5); reviewed and edited by @jroachgolf84 before posting

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.