Should we update latest_heartbeat to database when job do heartbeat?

[broholens]

Currently, when job do heartbeat，the latest_heartbeat does not update to db，code in job.py like this:

  # 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()

However, i think we should update latest_heartbeat to db，code should be like this:

  # Update last heartbeat time
  with create_session() as session:
      # Make the session aware of this object
      self.latest_heartbeat = timezone.utcnow()
      session.merge(self)
      session.commit()

We should update latest_heartbeat before session.merge(self).

Is there any reason why latest_heartbeat is not updated in the database when the job do heartbeat?

[potiuk]

I believe we dosave it and the way it is done is right. You do not have to "merge" an SQL alchemy object after you change it's state.

I am not an sqlalchemy expert and looking at the docs does not give me a definite answer, but I believe this is is how it works.

Merge reconciles the state of the object with that one in the database. But if you changed the object field and it is not "deteched", commiting the session to which the object is attached will automaticaly issue SQL statements and save all the changes to the object to the database. This is the whole premise of SQLalchemy that it treats ORM objects as "Python" objects and transparently maps them to the Database records - including modifications.

The reversal have been implemented 4 years ago, so I doubt it does not work: #6284

But maybe I am wrong on my understanding of how merge works ? What is your understanding how it works @broholens ?

@ashb @uranusjr @hussein-awala - I think you are better "SQLAlchemy" experts to explain some of the underlying behaviour?

[broholens]

I was wrong and you are right.

I tested SQLAlachemy merge on table Job and found you are right.

>>> from sqlalchemy import create_engine
>>> from sqlalchemy.orm import sessionmaker, scoped_session
>>> from airflow.utils import timezone
>>> from airflow.jobs.job import Job
>>> engine=create_engine(url='postgresql+psycopg2://airflow_user:xxx@127.0.0.1/airflow_db')
>>> session = scoped_session(sessionmaker(autocommit=False, autoflush=False, bind=engine, expire_on_commit=False))
>>> job = session.query(Job).order_by(Job.latest_heartbeat.asc()).first()
>>> job.latest_heartbeat
datetime.datetime(2023, 8, 21, 8, 41, 33, 27020, tzinfo=Timezone('UTC'))
>>> job.id
4
>>> session.merge(job)
<airflow.jobs.job.Job object at 0x7fbcc6c98b90>
>>> job.latest_heartbeat = timezone.utcnow()
>>> session.commit()
>>> session2 = scoped_session(sessionmaker(autocommit=False, autoflush=False, bind=engine, expire_on_commit=False))
>>> session2.query(Job.latest_heartbeat).filter(Job.id==4).first()
(datetime.datetime(2023, 8, 25, 2, 1, 25, 395031, tzinfo=Timezone('UTC')),)

I updated latest_heartbeat after session.merge, then committed to database. And query it by another session with same job id，the result shows latest_heartbeat had been updated.

[potiuk]

Actually - after reading the docs I think we might be both right and both wrong. The heartbeat time gets updated but in a different place (when we close the outer session not when session.commit() is called for the inner session. Becasue with create_session() will create a new session, the whole piece of code behaves in ways that are a bit subtle.

Not sure if it is fully intended though :).

It works like it works for you because you have Job object attached to the session you are merging with. But here it is a bit different.
It's actually quite interesting what happens here. When the heartbeat method starts it gets a session (either newly created or passed from outside.

    @provide_session
    def heartbeat(
        self, heartbeat_callback: Callable[[Session], None], session: Session = NEW_SESSION. <-Here we have a session
    ) -> None:

I assume (and I think it will be valld) that the Job object is attached to the session. So here we store the previous_heartbea we saw and by running session.merge() we synchronise our own values with the DB. Since the Job object is already attached to the session, no new object is created, the values are updated in "job" object.

        previous_heartbeat = self.latest_heartbeat

        try:
            # This will cause it to load from the db
            session.merge(self)

But then interesting thing happens. We create a new session. And then session.merge works differently. It reads the value of the JOB from the DB and it creates a NEW Job object in that session, that is different than self (because session is different). This means that the session object in the new session is refreshed.

            # Update last heartbeat time
            with create_session() as session:   
                # Make the session aware of this object
                session.merge(self) <- Here new Job object is created  attached to the "inner" session (which is different than the outer one)

What happens next is interesting. Because we are running commit() on the new session after modifying self Job. And if properly understand how merge works - the latest_heartbeat is NOT committed to the database, nor the Job object attached to the new session has it.

                self.latest_heartbeat = timezone.utcnow() <- we update heartbeat in 'self` object
                session.commit() <- but here we commit the "inner" session that the `self` object is not attached to

Then we override previous_heartbeat with the one that we just set so that even if the "heartbeat_callback" fails with OperationalError ( usually deadlock) we still keep the Job heartbeat as utcnow().

                # At this point, the DB has updated.  <- I am not sure what that means
                previous_heartbeat = self.latest_heartbeat <- here we set `previous_heartbeat' to the utc_now() we just set so that even if we have operational error, we still assume that heartbeat happened.

Then we call the callback with the "inner" session and the job object there is refreshed from the DB (but it does not heave IMHO latest_heartbeat modified.

                heartbeat_callback(session) <- here job in this session IMHO still has the old heartbeat timestamp not the new one
                self.log.debug("[heartbeat]")

Then this one depends where the OperationalError happened. If it happened BEFORE previous_heartbeat = self.latest_heartbeat, then we set self.last_heartbeat to the heartbeat that was there before - like if no heartbeat ever happened.

But if the OperationalError happens in theheartbeat_callback(session) then situation is different. Because then we overwrote the previous_heartbeat with the utc_now() so in fact we assume that heartbeat DID happen even if operationalError was thrown from inside callback:

        except OperationalError:
            Stats.incr(convert_camel_to_snake(self.__class__.__name__) + "_heartbeat_failure", 1, 1)
            self.log.exception("%s heartbeat got an exception", self.__class__.__name__)
            # We didn't manage to heartbeat, so make sure that the timestamp isn't updated
            self.latest_heartbeat = previous_heartbeat

And then the usual thing happen - i.e. when sesion is close, the updated self Job gets saved to the database. Only after that the "utc_now" is actually saved to the DB.

I am not sure if my understanding is right but maybe somoene (@ashb ?) could confirm my understanding :) ?

[broholens]

@potiuk Thank you for your explanation :)

[broholens]

i think this should be closed by #34026