jason810496 commented on code in PR #52325:
URL: https://github.com/apache/airflow/pull/52325#discussion_r2309680220
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py:
##########
@@ -49,13 +49,14 @@ async def _get_dag_run(self) -> DagRun:
return await
session.scalar(select(DagRun).filter_by(dag_id=self.dag_id, run_id=self.run_id))
def _serialize_xcoms(self) -> dict[str, Any]:
- xcom_query = XComModel.get_many(
- run_id=self.run_id,
- key=XCOM_RETURN_KEY,
- task_ids=self.result_task_ids,
- dag_ids=self.dag_id,
- )
- xcom_query = xcom_query.order_by(XComModel.task_id,
XComModel.map_index)
+ with create_session() as session:
Review Comment:
We can pass `session` from the router level in:
https://github.com/apache/airflow/blob/cd90980f424e1c1c25c74db37b0063aa12141057/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py#L476-L477
```python
waiter = DagRunWaiter(
dag_id=dag_id,
run_id=dag_run_id,
interval=interval,
result_task_ids=result_task_ids,
session=session, # <------- new added
)
```
and add the `session` attribute
https://github.com/apache/airflow/blob/cd90980f424e1c1c25c74db37b0063aa12141057/airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py#L38-L46
```python
@attrs.define
class DagRunWaiter:
"""Wait for the specified dag run to finish, and collect info from it."""
dag_id: str
run_id: str
interval: float
result_task_ids: list[str] | None
session: Session # <------ new added
```
then we can use `self.session` in the `_serialize_xcoms` method.
Hope the above explanation is clear enough.
--
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]
