Add endpoint to watch dag run until finish #51920
Conversation
|
FYI. Not sure if it helps but OpenAPI 3.2.0 will have support for |
9079d18 to
231597b
Compare
231597b to
007a051
Compare
|
I added some tests for the endpoint, but couldn’t figure out how to test the looping part. Hopefully this is good enough… |
| @@ -163,3 +163,10 @@ class DAGRunsBatchBody(StrictBaseModel): | |||
| start_date_lte: AwareDatetime | None = None | |||
| end_date_gte: AwareDatetime | None = None | |||
| end_date_lte: AwareDatetime | None = None | |||
|
|
|||
|
|
|||
| class DAGRunWatchResult(StrictBaseModel): | |||
There was a problem hiding this comment.
Do we really need a new Serializer? I'm not fan of having each endpoint returning custom schema, unless that would really impact / slow down performances. Here can't we just use the default DAGRunResponse, which holds the state and duration. This will be much easier for client to understand the API if we avoid multiplication of partial schemas.
| async with create_session_async() as session: | ||
| dag_run = await session.scalar(select(DagRun).filter_by(dag_id=dag_id, run_id=run_id)) | ||
| yield DAGRunWatchResult.model_validate(dag_run, from_attributes=True).model_dump_json() | ||
| yield "\n" |
There was a problem hiding this comment.
This is duplicated from above, maybe should be extracted into a small fn.
Also we tend to not put services/utility code in the routes (endpoint) files directly. Those can live in services/public/dag_runs.py
| @pytest.fixture(autouse=True) | ||
| def reconfigure_async_db_engine(self): | ||
| from airflow.settings import _configure_async_session | ||
|
|
||
| _configure_async_session() |
There was a problem hiding this comment.
Is that still necessary? Tests are green for me even when removing this fixture.
There was a problem hiding this comment.
The test_should_respond_200_immediately_for_finished_run tests fail for me without this (there are two fixtures; the first would pass but the second would fail)
| @dag_run_router.get( | ||
| "/{dag_run_id}/watch", | ||
| responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]), | ||
| dependencies=[Depends(requires_access_dag(method="GET", access_entity=DagAccessEntity.RUN))], | ||
| ) |
There was a problem hiding this comment.
How about adding visible OpenAPI doc? Somehow like we do with the TaskInatance Log route.
Close #51711.
I initially wanted to just enhance the trigger endpoint to optionally stream until the run finishes, but it seems that FastAPI does not like this optionally stream idea. You can do it of course, but would loose a lot of the auto annotation reflection feature. So I opted to have a separate streaming endpoint instead.
This endpoint repeatedly emits a JSON object at the specified interval, until the dag reaches a finished state.
Tests to come.