blag commented on code in PR #25419:
URL: https://github.com/apache/airflow/pull/25419#discussion_r951844432
##########
airflow/models/taskinstance.py:
##########
@@ -1525,33 +1528,17 @@ def _run_raw_task(
session.add(Log(self.state, self))
session.merge(self)
if self.state == TaskInstanceState.SUCCESS:
- self._create_dataset_dag_run_queue_records(session=session)
+ self._register_dataset_changes(session=session)
session.commit()
- def _create_dataset_dag_run_queue_records(self, *, session: Session) ->
None:
- from airflow.datasets import Dataset
- from airflow.models.dataset import DatasetModel
-
+ def _register_dataset_changes(self, *, session: Session) -> None:
for obj in self.task.outlets or []:
self.log.debug("outlet obj %s", obj)
- if isinstance(obj, Dataset):
- dataset = session.query(DatasetModel).filter(DatasetModel.uri
== obj.uri).one_or_none()
- if not dataset:
- self.log.warning("Dataset %s not found", obj)
- continue
- consuming_dag_ids = [x.dag_id for x in dataset.consuming_dags]
- self.log.debug("consuming dag ids %s", consuming_dag_ids)
- session.add(
- DatasetEvent(
- dataset_id=dataset.id,
- source_task_id=self.task_id,
- source_dag_id=self.dag_id,
- source_run_id=self.run_id,
- source_map_index=self.map_index,
- )
- )
- for dag_id in consuming_dag_ids:
- session.merge(DatasetDagRunQueue(dataset_id=dataset.id,
target_dag_id=dag_id))
+ self.dataset_event_manager.register_dataset_change(
Review Comment:
100%. Fixed in the next commit. Thanks!
--
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]
