uranusjr commented on code in PR #31772:
URL: https://github.com/apache/airflow/pull/31772#discussion_r1231969169
##########
airflow/models/dag.py:
##########
@@ -3539,28 +3562,30 @@ def dags_needing_dagruns(cls, session: Session) ->
tuple[Query, dict[str, tuple[
# these dag ids are triggered by datasets, and they are ready to go.
dataset_triggered_dag_info = {
x.dag_id: (x.first_queued_time, x.last_queued_time)
- for x in session.query(
- DagScheduleDatasetReference.dag_id,
- func.max(DDRQ.created_at).label("last_queued_time"),
- func.min(DDRQ.created_at).label("first_queued_time"),
- )
- .join(DagScheduleDatasetReference.queue_records, isouter=True)
- .group_by(DagScheduleDatasetReference.dag_id)
- .having(func.count() ==
func.sum(case((DDRQ.target_dag_id.is_not(None), 1), else_=0)))
- .all()
+ for x in session.execute(
+ select(
+ DagScheduleDatasetReference.dag_id,
+ func.max(DDRQ.created_at).label("last_queued_time"),
+ func.min(DDRQ.created_at).label("first_queued_time"),
+ )
+ .join(DagScheduleDatasetReference.queue_records, isouter=True)
+ .group_by(DagScheduleDatasetReference.dag_id)
+ .having(func.count() ==
func.sum(case((DDRQ.target_dag_id.is_not(None), 1), else_=0)))
+ ).all()
}
dataset_triggered_dag_ids = set(dataset_triggered_dag_info.keys())
if dataset_triggered_dag_ids:
exclusion_list = {
x.dag_id
for x in (
- session.query(DagModel.dag_id)
- .join(DagRun.dag_model)
- .filter(DagRun.state.in_((DagRunState.QUEUED,
DagRunState.RUNNING)))
- .filter(DagModel.dag_id.in_(dataset_triggered_dag_ids))
- .group_by(DagModel.dag_id)
- .having(func.count() >= func.max(DagModel.max_active_runs))
- .all()
+ session.execute(
+ select(DagModel.dag_id)
+ .join(DagRun.dag_model)
+ .where(DagRun.state.in_((DagRunState.QUEUED,
DagRunState.RUNNING)))
+ .where(DagModel.dag_id.in_(dataset_triggered_dag_ids))
+ .group_by(DagModel.dag_id)
+ .having(func.count() >=
func.max(DagModel.max_active_runs))
+ ).all()
Review Comment:
```suggestion
)
```
I think this can be rewritten with `scalars` though
##########
airflow/models/dag.py:
##########
@@ -3539,28 +3562,30 @@ def dags_needing_dagruns(cls, session: Session) ->
tuple[Query, dict[str, tuple[
# these dag ids are triggered by datasets, and they are ready to go.
dataset_triggered_dag_info = {
x.dag_id: (x.first_queued_time, x.last_queued_time)
- for x in session.query(
- DagScheduleDatasetReference.dag_id,
- func.max(DDRQ.created_at).label("last_queued_time"),
- func.min(DDRQ.created_at).label("first_queued_time"),
- )
- .join(DagScheduleDatasetReference.queue_records, isouter=True)
- .group_by(DagScheduleDatasetReference.dag_id)
- .having(func.count() ==
func.sum(case((DDRQ.target_dag_id.is_not(None), 1), else_=0)))
- .all()
+ for x in session.execute(
+ select(
+ DagScheduleDatasetReference.dag_id,
+ func.max(DDRQ.created_at).label("last_queued_time"),
+ func.min(DDRQ.created_at).label("first_queued_time"),
+ )
+ .join(DagScheduleDatasetReference.queue_records, isouter=True)
+ .group_by(DagScheduleDatasetReference.dag_id)
+ .having(func.count() ==
func.sum(case((DDRQ.target_dag_id.is_not(None), 1), else_=0)))
+ ).all()
Review Comment:
```suggestion
)
```
--
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]
