potiuk commented on code in PR #30308:
URL: https://github.com/apache/airflow/pull/30308#discussion_r1162093112
##########
airflow/jobs/job.py:
##########
@@ -222,59 +222,47 @@ def heartbeat(self, only_if_necessary: bool = False) ->
None:
# We didn't manage to heartbeat, so make sure that the timestamp
isn't updated
self.latest_heartbeat = previous_heartbeat
- def run(self) -> int | None:
- """Starts the job."""
+ @provide_session
+ def prepare_for_execution(self, session=NEW_SESSION):
+ """Prepares the job for execution."""
Stats.incr(self.__class__.__name__.lower() + "_start", 1, 1)
- # Adding an entry in the DB
- ret = None
- with create_session() as session:
- self.state = State.RUNNING
- session.add(self)
- session.commit()
- make_transient(self)
- try:
- ret = self.job_runner._execute()
- # In case of max runs or max duration
- self.state = State.SUCCESS
- except SystemExit:
- # In case of ^C or SIGTERM
- self.state = State.SUCCESS
- except Exception:
- self.state = State.FAILED
- raise
- finally:
- get_listener_manager().hook.before_stopping(component=self)
- self.end_date = timezone.utcnow()
- session.merge(self)
- session.commit()
+ self.state = State.RUNNING
+ self.start_date = timezone.utcnow()
+ self.state = State.RUNNING
+ session.add(self)
+ session.commit()
+ make_transient(self)
+ @provide_session
+ def complete_execution(self, session=NEW_SESSION):
Review Comment:
yeah.
##########
airflow/jobs/job.py:
##########
@@ -285,12 +273,59 @@ def most_recent_job(job_type: str, session: Session =
NEW_SESSION) -> BaseJob |
:param session: Database session
"""
return (
- session.query(BaseJob)
- .filter(BaseJob.job_type == job_type)
+ session.query(Job)
+ .filter(Job.job_type == job_type)
.order_by(
# Put "running" jobs at the front.
- case({State.RUNNING: 0}, value=BaseJob.state, else_=1),
- BaseJob.latest_heartbeat.desc(),
+ case({State.RUNNING: 0}, value=Job.state, else_=1),
+ Job.latest_heartbeat.desc(),
)
.first()
)
+
+
+@provide_session
+def prepare_for_execution(job: Job | JobPydantic, session=NEW_SESSION):
Review Comment:
yep
--
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]
