saguziel commented on a change in pull request #7477: [AIRFLOW-6857][WIP] Bulk sync DAGs URL: https://github.com/apache/airflow/pull/7477#discussion_r383527240
########## File path: airflow/models/dag.py ########## @@ -1455,63 +1455,92 @@ def create_dagrun(self, return run + @classmethod @provide_session - def sync_to_db(self, owner=None, sync_time=None, session=None): + def bulk_sync_to_db(cls, dags: Collection["DAG"], sync_time=None, session=None): """ - Save attributes about this DAG to the DB. Note that this method + Save attributes about list of DAG to the DB. Note that this method can be called for both DAGs and SubDAGs. A SubDag is actually a SubDagOperator. - :param dag: the DAG object to save to the DB - :type dag: airflow.models.DAG + :param dags: the DAG objects to save to the DB + :type dags: List[airflow.models.dag.DAG] :param sync_time: The time that the DAG should be marked as sync'ed :type sync_time: datetime :return: None """ + if not dags: + return from airflow.models.serialized_dag import SerializedDagModel - if owner is None: - owner = self.owner if sync_time is None: sync_time = timezone.utcnow() + log.info("Sync %s DAGs", len(dags)) + dag_by_ids = {dag.dag_id: dag for dag in dags} + dag_ids = set(dag_by_ids.keys()) + orm_dags = ( + session + .query(DagModel) + .options(joinedload(DagModel.tags, innerjoin=False)) + .filter(DagModel.dag_id.in_(dag_ids)) + .with_for_update(of=DagModel) + .all() + ) - orm_dag = session.query( - DagModel).filter(DagModel.dag_id == self.dag_id).first() - if not orm_dag: - orm_dag = DagModel(dag_id=self.dag_id) - if self.is_paused_upon_creation is not None: - orm_dag.is_paused = self.is_paused_upon_creation - self.log.info("Creating ORM DAG for %s", self.dag_id) + existing_dag_ids = {orm_dag.dag_id for orm_dag in orm_dags} + missing_dag_ids = dag_ids.difference(existing_dag_ids) + + for missing_dag_id in missing_dag_ids: + orm_dag = DagModel(dag_id=missing_dag_id) + dag = dag_by_ids[missing_dag_id] + if dag.is_paused_upon_creation is not None: + orm_dag.is_paused = dag.is_paused_upon_creation + log.info("Creating ORM DAG for %s", dag.dag_id) session.add(orm_dag) - if self.is_subdag: - orm_dag.is_subdag = True - orm_dag.fileloc = self.parent_dag.fileloc - orm_dag.root_dag_id = self.parent_dag.dag_id - else: - orm_dag.is_subdag = False - orm_dag.fileloc = self.fileloc - orm_dag.owners = owner - orm_dag.is_active = True - orm_dag.last_scheduler_run = sync_time - orm_dag.default_view = self._default_view - orm_dag.description = self.description - orm_dag.schedule_interval = self.schedule_interval - orm_dag.tags = self.get_dagtags(session=session) + orm_dags.append(orm_dag) + + for orm_dag in sorted(orm_dags, key=lambda d: d.dag_id): + dag = dag_by_ids[orm_dag.dag_id] + if dag.is_subdag: + orm_dag.is_subdag = True + orm_dag.fileloc = dag.parent_dag.fileloc # type: ignore + orm_dag.root_dag_id = dag.parent_dag.dag_id # type: ignore + orm_dag.owners = dag.parent_dag.owner # type: ignore + else: + orm_dag.is_subdag = False + orm_dag.fileloc = dag.fileloc + orm_dag.owners = dag.owner + orm_dag.is_active = True + orm_dag.last_scheduler_run = sync_time + orm_dag.default_view = dag._default_view + orm_dag.description = dag.description + orm_dag.schedule_interval = dag.schedule_interval + orm_dag.tags = dag.get_dagtags(session=session) session.commit() Review comment: it might be worth doing this in smaller batches rather than one gigantic commit. I don't think atomicity is a concern but large transactions may be. ---------------------------------------------------------------- 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: us...@infra.apache.org With regards, Apache Git Services