pierrejeambrun commented on code in PR #51920:
URL: https://github.com/apache/airflow/pull/51920#discussion_r2166723553
##########
airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py:
##########
@@ -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):
Review 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.
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py:
##########
@@ -432,6 +437,39 @@ def trigger_dag_run(
raise HTTPException(status.HTTP_400_BAD_REQUEST, str(e))
+async def _watch_dagrun(dag_id: str, run_id: str, interval: float) ->
AsyncGenerator[str, None]:
+ 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"
+ while dag_run.state not in State.finished_dr_states:
+ await asyncio.sleep(interval)
+ 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"
Review 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/dag_run.py`
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py:
##########
@@ -1602,3 +1602,45 @@ def test_should_respond_200_with_null_logical_date(self,
test_client):
"conf": {},
"note": None,
}
+
+
+class TestWatchDagRun:
+ # The way we init async engine does not work well with FastAPI app init.
+ # Creating the engine implicitly creates an event loop, which Airflow does
+ # once for the entire process; creating the FastAPI app also does, but our
+ # test setup does it once for each test. I don't know how to properly fix
+ # this without rewriting how Airflow does db; re-configuring the db for
each
+ # test at least makes the tests run correctly.
+ @pytest.fixture(autouse=True)
+ def reconfigure_async_db_engine(self):
+ from airflow.settings import _configure_async_session
+
+ _configure_async_session()
Review Comment:
Is that still necessary? Tests are green for me even when removing this
fixture.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
