amoghrajesh commented on code in PR #73173:
URL: https://github.com/apache/airflow/pull/73173#discussion_r4024859157
##########
airflow-core/src/airflow/utils/db_cleanup.py:
##########
@@ -145,6 +171,20 @@ def __post_init__(self):
schema=self.schema_name,
)
+ if self.dag_id_scope is not None:
+ if self.dag_id_column_name is not None:
+ raise ValueError(
+ f"_TableConfig for table {self.table_name!r} sets both
dag_id_column_name and "
+ f"dag_id_scope; a table is scoped to a Dag either by its
own column or through a "
+ f"foreign key, not both."
+ )
+ if self.dag_id_scope.fk_column not in self.orm_model.c.keys():
+ raise ValueError(
+ f"_TableConfig for table {self.table_name!r} sets
dag_id_scope but its "
+ f"fk_column {self.dag_id_scope.fk_column!r} is not one of
its columns; "
+ f"add {self.dag_id_scope.fk_column!r} to extra_columns."
+ )
Review Comment:
Should we have tests for `ValueError`?
##########
airflow-core/src/airflow/utils/db_cleanup.py:
##########
@@ -517,7 +583,31 @@ def _build_query(
if dag_ids:
conditions.append(base_table_dag_id_col.in_(dag_ids))
if exclude_dag_ids:
- conditions.append(base_table_dag_id_col.not_in(exclude_dag_ids))
+ # A NULL dag id belongs to no Dag, so it is not one of the
excluded Dags' rows and stays
+ # eligible. NOT IN alone would yield NULL for it and silently
retain it forever -- which
+ # is every `job` row, since core never sets Job.dag_id.
+ conditions.append(
+ or_(base_table_dag_id_col.is_(None),
base_table_dag_id_col.not_in(exclude_dag_ids))
+ )
+ elif (dag_ids or exclude_dag_ids) and dag_id_scope is not None:
+ fk_col = base_table.c[dag_id_scope.fk_column]
+ referenced = table(
+ dag_id_scope.referenced_table,
+ column(dag_id_scope.referenced_pk_column),
+ column(dag_id_scope.referenced_dag_id_column),
+ )
+
+ def _rows_for(target_dag_ids: list[str]):
+ return
select(referenced.c[dag_id_scope.referenced_pk_column]).where(
+
referenced.c[dag_id_scope.referenced_dag_id_column].in_(target_dag_ids)
+ )
+
+ if dag_ids:
+ conditions.append(fk_col.in_(_rows_for(dag_ids)))
+ if exclude_dag_ids:
+ # A NULL foreign key belongs to no Dag, so it is not one of the
excluded Dags' rows and
+ # stays eligible. Testing it with NOT IN alone would yield NULL
and silently retain it.
Review Comment:
Almost same explained on line 586 - 588, one can be trimmed
--
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]
