kaxil commented on code in PR #64051:
URL: https://github.com/apache/airflow/pull/64051#discussion_r3133251223
##########
providers/airbyte/src/airflow/providers/airbyte/triggers/airbyte.py:
##########
@@ -80,32 +99,41 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
}
)
return
- await asyncio.sleep(self.poll_interval)
- job_run_status = hook.get_job_status(self.job_id)
- if job_run_status == JobStatusEnum.SUCCEEDED:
- yield TriggerEvent(
- {
- "status": "success",
- "message": f"Job run {self.job_id} has completed
successfully.",
- "job_id": self.job_id,
- }
- )
- elif job_run_status == JobStatusEnum.CANCELLED:
- yield TriggerEvent(
- {
- "status": "cancelled",
- "message": f"Job run {self.job_id} has been
cancelled.",
- "job_id": self.job_id,
- }
- )
- else:
- yield TriggerEvent(
- {
- "status": "error",
- "message": f"Job run {self.job_id} has failed.",
- "job_id": self.job_id,
- }
- )
+
+ job_run_status = hook.get_job_status(self.job_id)
+
+ if job_run_status in (
+ JobStatusEnum.RUNNING,
+ JobStatusEnum.PENDING,
+ JobStatusEnum.INCOMPLETE,
+ ):
+ await asyncio.sleep(self.poll_interval)
+ continue
+
+ if job_run_status == JobStatusEnum.SUCCEEDED:
+ yield TriggerEvent(
+ {
+ "status": "success",
+ "message": f"Job run {self.job_id} has completed
successfully.",
+ "job_id": self.job_id,
+ }
+ )
+ elif job_run_status == JobStatusEnum.CANCELLED:
+ yield TriggerEvent(
+ {
+ "status": "cancelled",
+ "message": f"Job run {self.job_id} has been
cancelled.",
+ "job_id": self.job_id,
+ }
+ )
+ else:
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "message": f"Job run {self.job_id} has failed.",
+ "job_id": self.job_id,
+ }
+ )
Review Comment:
**Bug: terminal `yield`s don't exit the loop.**
The old code relied on `while await self.is_still_running(hook):` exiting
naturally and then yielding the terminal event *after* the loop. The refactor
moves the success/cancelled/error yields *inside* `while True:`, but adds no
`return`. After yielding (e.g.) a `success` event, control falls off the bottom
of the if/elif/else, hits the end of the `while True:` body, and loops back to
the top -- calling `get_job_status` again and yielding another terminal event
every `poll_interval`. This repeats for the lifetime of the triggerer process.
The triggerer consumes this via `async for event in trigger.run()`
([triggerer_job_runner.py](https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/jobs/triggerer_job_runner.py)),
so the first event still fires the callback, but the trigger keeps producing
duplicates + extra HTTP calls until cleanup.
The existing tests miss this because they use `await generator.asend(None)`
which pulls a single event. Add a regression test along the lines of:
```python
events = [e async for e in trigger.run()]
assert len(events) == 1
```
Fix: add `return` after each of the three terminal `yield`s (or wrap them in
`if/elif/else` with a single `return` at the end).
##########
providers/airbyte/src/airflow/providers/airbyte/operators/airbyte.py:
##########
@@ -125,10 +140,24 @@ def execute_complete(self, context: Context, event: Any =
None) -> None:
Relies on trigger to throw an exception, otherwise it assumes
execution was
successful.
"""
+ hook = AirbyteHook(airbyte_conn_id=self.airbyte_conn_id,
api_version=self.api_version)
Review Comment:
Move this inside the `if event["status"] == "timeout":` branch. Every
success/error resume now instantiates `AirbyteHook` (fetching the connection,
initializing the SDK client) just to throw it away unused.
##########
providers/airbyte/src/airflow/providers/airbyte/operators/airbyte.py:
##########
@@ -97,11 +111,12 @@ def execute(self, context: Context) -> None:
self.log.debug("Running in deferrable mode in job state %s...",
state)
if state in (JobStatusEnum.RUNNING, JobStatusEnum.PENDING,
JobStatusEnum.INCOMPLETE):
self.defer(
- timeout=self.execution_timeout,
+ timeout=None,
Review Comment:
Setting `timeout=None` removes the triggerer framework's deadline
enforcement, which is the one mechanism that raises `CancelledError` at the
*exact* `execution_timeout` moment. The in-loop check that replaces it is
bounded by `poll_interval=60s` granularity and won't fire at all if:
- `hook.get_job_status()` is stuck in a blocking HTTP call (the check only
runs between polls)
- the triggerer restarts and the trigger resumes with `execution_deadline`
already in the past but stalls on the next HTTP call
The two mechanisms are complementary, not redundant. Keep
`timeout=self.execution_timeout` here as a backstop -- the in-loop check will
emit a clean `"timeout"` event when it wins, and `on_kill` (line 164) already
cancels the Airbyte job if the framework timeout wins first. Dropping the
framework timeout weakens the `execution_timeout` guarantee rather than
strengthening it.
##########
providers/airbyte/src/airflow/providers/airbyte/operators/airbyte.py:
##########
@@ -125,10 +140,24 @@ def execute_complete(self, context: Context, event: Any =
None) -> None:
Relies on trigger to throw an exception, otherwise it assumes
execution was
successful.
"""
+ hook = AirbyteHook(airbyte_conn_id=self.airbyte_conn_id,
api_version=self.api_version)
if event["status"] == "error":
self.log.debug("Error occurred with context: %s", context)
raise RuntimeError(event["message"])
+ job_id = event.get("job_id")
+ if event["status"] == "timeout":
+ if job_id:
+ self.log.info("Cancelling Airbyte job %s due to execution
timeout", job_id)
+ try:
+ hook.cancel_job(job_id=job_id)
+ except Exception as e:
Review Comment:
`except Exception` is broader than needed. `AirbyteHook.cancel_job` wraps
SDK failures in `AirflowException`, so narrow to `except AirflowException` --
that keeps `KeyboardInterrupt`/`SystemExit` and genuinely unexpected runtime
errors (e.g. `AttributeError` from a refactor) from being silently swallowed
into a log warning.
Also: the PR description says this is "best-effort cancellation", but that
philosophy isn't applied to `on_kill` below -- `cancel_job` failing there will
propagate out of `on_kill`. Consider using the same `try/except
AirflowException + log.warning` pattern in both places, or justify the
asymmetry.
##########
providers/airbyte/tests/unit/airbyte/triggers/test_airbyte.py:
##########
@@ -221,6 +225,34 @@ async def test_airbyte_job_timeout(self,
mock_get_job_status, mocked_is_still_ru
)
assert expected == actual
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.airbyte.triggers.airbyte.AirbyteSyncTrigger.is_still_running")
+ async def test_airbyte_job_run_execution_timeout(self,
mocked_is_still_running):
+ """Assert that run timeout after execution_deadline has elapsed"""
+
+ mocked_is_still_running.return_value = True
Review Comment:
`is_still_running` is patched but never called by the refactored `run()`
(the status check is inlined). The test still passes because
`execution_deadline = time.time() - 1` fires the new branch at
triggers/airbyte.py:82 on the first iteration. Drop the `is_still_running`
patch + decorator here to avoid misleading future readers.
Same cleanup applies to the pre-existing `test_airbyte_job_timeout` above
(line 205) -- it patches `mocked_is_still_running.return_value = True` which is
also dead now; the test passes only because `end_time=time.time()` fires the
end-time branch. If you want to control status transitions in these tests,
patch `AirbyteHook.get_job_status` like the other new tests do.
##########
providers/airbyte/src/airflow/providers/airbyte/operators/airbyte.py:
##########
@@ -125,10 +140,24 @@ def execute_complete(self, context: Context, event: Any =
None) -> None:
Relies on trigger to throw an exception, otherwise it assumes
execution was
successful.
"""
+ hook = AirbyteHook(airbyte_conn_id=self.airbyte_conn_id,
api_version=self.api_version)
if event["status"] == "error":
self.log.debug("Error occurred with context: %s", context)
raise RuntimeError(event["message"])
+ job_id = event.get("job_id")
+ if event["status"] == "timeout":
+ if job_id:
+ self.log.info("Cancelling Airbyte job %s due to execution
timeout", job_id)
+ try:
+ hook.cancel_job(job_id=job_id)
+ except Exception as e:
+ self.log.warning("Failed to cancel Airbyte job %s: %s",
job_id, e)
+ else:
+ self.log.warning("No job_id found; skipping cancellation")
+
+ raise RuntimeError(event["message"])
Review Comment:
`execute_complete` still has no branch for `event["status"] == "cancelled"`.
The trigger yields this status when Airbyte reports `CANCELLED`
(triggers/airbyte.py:122), but here that case falls through to
`log.info("...completed successfully.")` and returns `None` -- so a user- or
UI-cancelled Airbyte job silently reports task **success**.
This is pre-existing, but the PR is already expanding `execute_complete`
with timeout handling, so fixing it in the same change makes sense. A simple
`raise AirflowException(event["message"])` in a `cancelled` branch is 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]
