xBis7 commented on code in PR #43941:
URL: https://github.com/apache/airflow/pull/43941#discussion_r1885734876
##########
airflow/models/dagrun.py:
##########
@@ -1088,6 +1237,123 @@ def _filter_tis_and_exclude_removed(dag: DAG, tis:
list[TI]) -> Iterable[TI]:
finished_tis=finished_tis,
)
+ @staticmethod
+ def _set_scheduled_by_job_id(dag_run: DagRun, job_id: int, session:
Session, with_commit: bool) -> bool:
+ if not isinstance(dag_run, DagRun):
+ dag_run = session.scalars(
+ select(DagRun).where(
+ DagRun.dag_id == dag_run.dag_id,
+ DagRun.run_id == dag_run.run_id,
+ )
+ ).one()
+
+ if dag_run.scheduled_by_job_id == job_id:
+ return False
+
+ dag_run.log.debug("Setting dag_run scheduled_by_job_id for run_id:
%s", dag_run.run_id)
+ dag_run.scheduled_by_job_id = job_id
+
+ session.merge(dag_run)
+
+ if with_commit:
+ session.commit()
+
+ return True
+
+ @provide_session
+ def set_scheduled_by_job_id(
+ self, job_id: int, session: Session = NEW_SESSION, with_commit: bool =
False
+ ) -> bool:
+ """
+ Set DagRun scheduled_by_job_id.
+
+ :param job_id: integer with the scheduled_by_job_id to set for the
dag_run
+ :param session: SQLAlchemy ORM Session
+ :param with_commit: should the scheduled_by_job_id be committed?
+ :return: has the scheduled_by_job_id been changed?
+ """
+ return self._set_scheduled_by_job_id(
+ dag_run=self, job_id=job_id, session=session,
with_commit=with_commit
+ )
+
+ @staticmethod
+ def _set_context_carrier(
+ dag_run: DagRun, context_carrier: dict, session: Session, with_commit:
bool
+ ) -> bool:
Review Comment:
All the new set methods are following the same logic as
`taskinstance.set_state()`. That's where I got the idea from.
https://github.com/apache/airflow/blob/main/airflow/models/taskinstance.py#L2175-L2210
I'll remove them.
--
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]
