[GitHub] [airflow] mik-laj commented on a change in pull request #7477: [AIRFLOW-6857][DONT-MERGE] Bulk sync DAGs

[GitBox <gi...@apache.org>]

mik-laj commented on a change in pull request #7477: [AIRFLOW-6857][DONT-MERGE] Bulk sync DAGs
URL: https://github.com/apache/airflow/pull/7477#discussion_r383561313
 
 

 ##########
 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:
   Thanks. I will improve it soon.

----------------------------------------------------------------
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:
users@infra.apache.org


With regards,
Apache Git Services