Repository: apache/airflow
File: airflow-core/src/airflow/jobs/scheduler_job_runner.py (Python)

Code before:
        :param session: The session object
        """

        def _get_sentry_integration(executor: BaseExecutor) -> str:
            try:
                sentry_integration = executor.sentry_integration
            except AttributeError:
                # Old executor interface hard-codes the supports_sentry flag.
                if getattr(executor, "supports_sentry", False):
                    return "sentry_sdk.integrations.celery.CeleryIntegration"
                return ""
            if not isinstance(sentry_integration, str):
                self.log.warning(
                    "Ignoring invalid sentry_integration on executor",
                    executor=executor,
                    sentry_integration=sentry_integration,
                )
                return ""
            return sentry_integration

        # actually enqueue them
        for ti in task_instances:
            if ti.dag_run.state in State.finished_dr_states:
                ti.set_state(None, session=session)
                continue
            if not ti.dag_version_id:
                self.log.error(
                    "TaskInstance %s does not have a dag_version_id set, cannot be enqueued",
                    ti,
                )
                continue

            workload = workloads.ExecuteTask.make(
                ti,
                generator=executor.jwt_generator,
                sentry_integration=_get_sentry_integration(executor),
            )
            executor.queue_workload(workload, session=session)

    def _critical_section_enqueue_task_instances(self, session: Session) -> int:
        """
        Enqueues TaskInstances for execution.

        There are three steps:
        1. Pick TIs by priority with the constraint that they are in the expected states
        and that we do not exceed max_active_runs or pool limits.
        2. Change the state for the TIs above atomically.
        3. Enqueue the TIs in the executor.

        HA note: This function is a "critical section" meaning that only a single scheduler process can
        execute this function at the same time. This is achieved by doing

Review comment (suggestion):
```suggestion
            # This is possible to hit if a DagRun was in progress during a 2->3 upgrade
            if not ti.dag_version_id:
```

Code after:
        :param session: The session object
        """

        def _get_sentry_integration(executor: BaseExecutor) -> str:
            try:
                sentry_integration = executor.sentry_integration
            except AttributeError:
                # Old executor interface hard-codes the supports_sentry flag.
                if getattr(executor, "supports_sentry", False):
                    return "sentry_sdk.integrations.celery.CeleryIntegration"
                return ""
            if not isinstance(sentry_integration, str):
                self.log.warning(
                    "Ignoring invalid sentry_integration on executor",
                    executor=executor,
                    sentry_integration=sentry_integration,
                )
                return ""
            return sentry_integration

        # actually enqueue them
        for ti in task_instances:
            if ti.dag_run.state in State.finished_dr_states:
                ti.set_state(None, session=session)
                continue
            if not ti.dag_version_id:
                self.log.warning(
                    "TaskInstance %s does not have a dag_version_id set, cannot be enqueued. "
                    "This would get unstuck and dag_version_id updated.",
                    ti,
                )
                continue

            workload = workloads.ExecuteTask.make(
                ti,
                generator=executor.jwt_generator,
                sentry_integration=_get_sentry_integration(executor),
            )
            executor.queue_workload(workload, session=session)

    def _critical_section_enqueue_task_instances(self, session: Session) -> int:
        """
        Enqueues TaskInstances for execution.

        There are three steps:
        1. Pick TIs by priority with the constraint that they are in the expected states
        and that we do not exceed max_active_runs or pool limits.
        2. Change the state for the TIs above atomically.
        3. Enqueue the TIs in the executor.

        HA note: This function is a "critical section" meaning that only a single scheduler process can