o-nikolas commented on code in PR #57837:
URL: https://github.com/apache/airflow/pull/57837#discussion_r2501072232
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -270,6 +272,98 @@ def register_signals(self) -> ExitStack:
return resetter
+ def _get_team_names_for_dag_ids(
+ self, dag_ids: Collection[str], session: Session
+ ) -> dict[str, str | None]:
+ """
+ Batch query to resolve team names for multiple DAG IDs using the DAG >
Bundle > Team relationship chain.
+
+ DAG IDs > DagModel (via dag_id) > DagBundleModel (via bundle_name) >
Team
+
+ :param dag_ids: Collection of DAG IDs to resolve team names for
+ :param session: Database session for queries
+ :return: Dictionary mapping dag_id -> team_name (None if no team found)
+ """
+ if not dag_ids:
+ return {}
+
+ try:
+ # Query all team names for the given DAG IDs in a single query
+ query_results = session.execute(
+ select(DagModel.dag_id, Team.name)
+ .join(DagBundleModel.teams) # Join Team to DagBundleModel via
association table
+ .join(
+ DagModel, DagModel.bundle_name == DagBundleModel.name
+ ) # Join DagBundleModel to DagModel
+ .where(DagModel.dag_id.in_(dag_ids))
+ ).all()
+
+ # Create mapping from results
+ dag_id_to_team_name = {dag_id: team_name for dag_id, team_name in
query_results}
+
+ # Ensure all requested dag_ids are in the result (with None for
those not found)
+ result = {dag_id: dag_id_to_team_name.get(dag_id) for dag_id in
dag_ids}
+
+ self.log.debug(
+ "Resolved team names for %d DAGs: %s",
+ len([team for team in result.values() if team is not None]),
+ {dag_id: team for dag_id, team in result.items()},
+ )
+
+ return result
+
+ except Exception:
+ # Log the error, explicitly don't fail the scheduling loop
+ self.log.exception("Failed to resolve team names for DAG IDs: %s",
list(dag_ids))
+ # Return dict with all None values to ensure graceful degradation
+ return {}
+
+ def _get_task_team_name(self, task_instance: TaskInstance, session:
Session) -> str | None:
Review Comment:
Great idea! I'll do it
--
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]
