Skip to content

Commit

Permalink
Optimize 2.3.0 pre-upgrade check queries (apache#23458)
Browse files Browse the repository at this point in the history
We have to check for rows that are missing either corresponding TI or DR and move them out of table before adding FKs.  We were doing correlation in the JOIN condition but it appears postgres does *not* like this so here we move correlation to WHERE.

(cherry picked from commit 8e3d6c3)
  • Loading branch information
dstandish authored and jedcunningham committed May 3, 2022
1 parent 1ff44c9 commit 3b01473
Showing 1 changed file with 7 additions and 7 deletions.
14 changes: 7 additions & 7 deletions airflow/utils/db.py
Original file line number Diff line number Diff line change
Expand Up @@ -1092,30 +1092,30 @@ def _task_instance_exists(session, source_table, dag_run, task_instance):
"""
if 'run_id' not in task_instance.c:
# db is < 2.2.0
source_to_ti_join_cond = and_(
where_clause = and_(
source_table.c.dag_id == task_instance.c.dag_id,
source_table.c.task_id == task_instance.c.task_id,
source_table.c.execution_date == task_instance.c.execution_date,
)
ti_to_dr_join_cond = and_(
source_table.c.dag_id == task_instance.c.dag_id,
source_table.c.execution_date == task_instance.c.execution_date,
dag_run.c.dag_id == task_instance.c.dag_id,
dag_run.c.execution_date == task_instance.c.execution_date,
)
else:
# db is 2.2.0 <= version < 2.3.0
source_to_ti_join_cond = and_(
where_clause = and_(
source_table.c.dag_id == task_instance.c.dag_id,
source_table.c.task_id == task_instance.c.task_id,
source_table.c.execution_date == dag_run.c.execution_date,
)
ti_to_dr_join_cond = and_(
source_table.c.dag_id == task_instance.c.dag_id,
dag_run.c.dag_id == task_instance.c.dag_id,
dag_run.c.run_id == task_instance.c.run_id,
source_table.c.execution_date == dag_run.c.execution_date,
)
exists_subquery = (
session.query(text('1'))
.select_from(task_instance.join(dag_run, onclause=ti_to_dr_join_cond))
.filter(source_to_ti_join_cond)
.filter(where_clause)
)
return exists_subquery

Expand Down

0 comments on commit 3b01473

Please sign in to comment.