aayushbisen commented on issue #54710:
URL: https://github.com/apache/airflow/issues/54710#issuecomment-3222897052
## Existing use of session util
The `_run_task()` helper function was creating its own database session
using:
```python
from airflow.utils.session import create_session
# Later in the code:
with create_session() as session:
ti.state = State.SCHEDULED
session.add(ti)
```
This created an unwanted dependency on airflow core utilities from within
the task-sdk.
## Solution
Refactored the session management to reuse the existing session from the
parent `test()` method instead of creating a new one:
1. **Modified function signature**: Added optional `session=None` parameter
to `_run_task()`
2. **Removed problematic import**: Eliminated `from airflow.utils.session
import create_session`
3. **Updated session usage**: Replaced context manager with direct session
usage when provided
4. **Updated function call**: Pass the existing session from the `test()`
method
**Before:**
```python
def _run_task(*, ti, task, run_triggerer=False):
# ...
if ti.state == State.DEFERRED and isinstance(msg, DeferTask) and
run_triggerer:
from airflow.utils.session import create_session
# ...
with create_session() as session:
ti.state = State.SCHEDULED
session.add(ti)
```
**After:**
```python
def _run_task(*, ti, task, run_triggerer=False, session=None):
# ...
if ti.state == State.DEFERRED and isinstance(msg, DeferTask) and
run_triggerer:
# ...
if session is not None:
ti.state = State.SCHEDULED
session.add(ti)
```
--
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]
