turbaszek commented on a change in pull request #9616:
URL: https://github.com/apache/airflow/pull/9616#discussion_r453186375
##########
File path: airflow/jobs/local_task_job.py
##########
@@ -135,7 +135,11 @@ def heartbeat_callback(self, session=None):
self.task_runner.terminate()
return
- self.task_instance.refresh_from_db()
+ if session is None:
+ self.task_instance.refresh_from_db()
+ else:
+ self.task_instance.refresh_from_db(session=session)
Review comment:
```python
def provide_session(func):
@wraps(func)
def wrapper(*args, **kwargs):
arg_session = 'session'
func_params = func.__code__.co_varnames
session_in_args = arg_session in func_params and \
func_params.index(arg_session) < len(args)
session_in_kwargs = arg_session in kwargs
# Here we validate that the session is in kwargs/args and that it's
value is
# not None so we can pass if further, otherwise we create new session
if session_in_kwargs and kwargs[arg_session] is not None:
return func(*args, **kwargs)
elif session_in_args and func_params[func_params.index(arg_session)]
is not None:
return func(*args, **kwargs)
else:
with create_session() as session:
kwargs[arg_session] = session
return func(*args, **kwargs)
return wrapper
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
