apilaskowski commented on code in PR #29815:
URL: https://github.com/apache/airflow/pull/29815#discussion_r1151778620
##########
airflow/jobs/base_job.py:
##########
@@ -210,54 +222,63 @@ def heartbeat(self, only_if_necessary: bool = False):
previous_heartbeat = self.latest_heartbeat
+ def session_merge(session):
+ session.merge(self)
+
try:
- with create_session() as session:
- # This will cause it to load from the db
- session.merge(self)
- previous_heartbeat = self.latest_heartbeat
+ self.handle_db_task(task_function = session_merge)
+ except OperationalError:
+ # We didn't manage to heartbeat, so make sure that the timestamp
isn't updated
+ self.latest_heartbeat = previous_heartbeat
+
+ previous_heartbeat = self.latest_heartbeat
- if self.state in State.terminating_states:
- self.kill()
+ if self.state in State.terminating_states:
+ self.kill()
- # Figure out how long to sleep for
- sleep_for = 0
- if self.latest_heartbeat:
- seconds_remaining = (
- self.heartrate - (timezone.utcnow() -
self.latest_heartbeat).total_seconds()
- )
- sleep_for = max(0, seconds_remaining)
- sleep(sleep_for)
+ # Figure out how long to sleep for
+ sleep_for = 0
+ if self.latest_heartbeat:
+ seconds_remaining = (
+ self.heartrate - (timezone.utcnow() -
self.latest_heartbeat).total_seconds()
+ )
+ sleep_for = max(0, seconds_remaining)
+ sleep(sleep_for)
- # Update last heartbeat time
- with create_session() as session:
- # Make the session aware of this object
- session.merge(self)
- self.latest_heartbeat = timezone.utcnow()
- session.commit()
- # At this point, the DB has updated.
- previous_heartbeat = self.latest_heartbeat
+ def session_merge_and_commit(session):
+ # Make the session aware of this object
+ session.merge(self)
+ self.latest_heartbeat = timezone.utcnow()
+ session.commit()
Review Comment:
Another approach, which I was considering is to add commit as last operation
in ```handle_db_transaction_with_session``` and to never have another write to
db from the ```handle_db_transaction```.
This way we have only a single update on the db, which will happen as last
operation. If there were any problems, db should be clean at any time. Do you
think that it could be feasible to have commit as last command and forbid using
commit mid transaction?
--
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]
