shahar1 commented on code in PR #75396:
URL: https://github.com/apache/airflow/pull/75396#discussion_r4484303049
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -766,14 +787,16 @@ def set_xcom(
)
try:
# We expect serialised value from the caller - sdk, do not serialise
in here
- XComModel.set_for_attempt(
- key=key,
- value=value,
- task_instance_id=owner,
- serialize=False,
- dag_result=dag_result,
- mapped_length=mapped_length,
- session=session,
+ await session.run_sync(
Review Comment:
**major:** this write now runs on the async engine, which doesn't get the
sync engine's MySQL session setup:
- `configure_orm()` calls `setup_event_handlers(engine)` for the sync engine
only, so `SET time_zone = '+00:00'` never runs on async connections.
- `prepare_engine_args()` sets `isolation_level="READ COMMITTED"` for the
sync engine only, so async sessions use the server default.
I ran Airflow's own engines against MySQL 8.0 started with
`--default-time-zone=+02:00`, writing a `UtcDateTime` column through each:
```text
sync session vars: ('+00:00', 'READ-COMMITTED')
async session vars: ('+02:00', 'REPEATABLE-READ')
async write read by sync session: wrote 17:48:52+00:00 read 15:48:52+00:00
sync write read by sync session: wrote 17:48:52+00:00 read 17:48:52+00:00
```
This probe covers the engines rather than this route, but `set_for_attempt`
binds `timestamp` the same way. On such a server, `xcom_v2.timestamp` for every
pushed XCom would be off by the server's UTC offset. That value is shown in the
UI and REST API, and `airflow db clean` uses it as the recency column for
`xcom_v2`. A server on UTC hides this.
The isolation level also matters on the read side: in `get_xcom` the owner
check starts the snapshot, so the value read after the lookup can't see a row
committed in between. The window is short, and I only expect a transient 404
from it.
In `_configure_async_session`, passing `isolation_level="READ COMMITTED"`
for MySQL and calling `setup_event_handlers(async_engine.sync_engine)` fixed
both on that server (`('+00:00', 'READ-COMMITTED')`, zero delta). I only tried
this with `asyncmy` on MySQL. The handlers would also run for SQLite and
PostgreSQL, so those need checking too.
`ti_heartbeat` already writes through this engine, so a separate PR that
this one rebases on seems cleanest.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py:
##########
@@ -563,31 +581,34 @@ def head_xcom(
]
),
description="Get a single XCom Value",
+ response_model=XComResponse,
)
-def get_xcom(
+async def get_xcom(
dag_id: str,
run_id: str,
task_id: str,
key: Annotated[str, Path(min_length=1)],
- session: SessionDep,
+ session: AsyncSessionDep,
params: Annotated[GetXcomFilterParams, Query()],
dag_bag: DagBagDep,
token: ExecutionToken = _TriggerAwareToken,
-) -> XComResponse:
+) -> JSONResponse:
"""Get an Airflow XCom from database - not other XCom Backends."""
owner = None
if token.claims.scope == "execution":
- owner = session.execute(
- select(TaskInstance.id, TaskInstance.working_set)
- .where(
- TaskInstance.id == token.id,
- TaskInstance.dag_id == dag_id,
- TaskInstance.run_id == run_id,
- TaskInstance.task_id == task_id,
- *build_coordinate_filters(TaskInstance,
map_index=params.map_index),
+ owner = (
+ await session.execute(
+ select(TaskInstance.id, TaskInstance.working_set)
+ .where(
+ TaskInstance.id == token.id,
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == run_id,
+ TaskInstance.task_id == task_id,
+ *build_coordinate_filters(TaskInstance,
map_index=params.map_index),
+ )
+ .limit(1)
+ .execution_options(include_all_attempts=True)
)
- .limit(1)
- .execution_options(include_all_attempts=True)
).first()
Review Comment:
**minor:** the first row of your connection table comes from this query: its
transaction stays open, so the async connection stays checked out and idle in
transaction through the `run_in_threadpool` below, including a cold-cache Dag
deserialization. Behind PgBouncer in transaction mode, that also pins a server
connection.
`owner` is a plain row, so `await session.rollback()` right after this read
would release the connection, and this path would no longer hold two at once.
The cost is another checkout for the value read, with a `pool_pre_ping` round
trip, so it's your call whether that trade is worth it.
--
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]
