# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""Base executor - this is the base class for all the implemented executors."""

from __future__ import annotations

import logging
import sys
from collections import defaultdict, deque
from collections.abc import Sequence
from dataclasses import dataclass, field
from functools import cached_property
from typing import TYPE_CHECKING, Any, ClassVar, cast

import pendulum

from airflow._shared.observability.metrics import stats
from airflow.cli.cli_config import DefaultHelpParser
from airflow.configuration import conf
from airflow.executors import workloads
from airflow.executors.executor_loader import ExecutorLoader
from airflow.executors.workloads.callback import ExecuteCallback
from airflow.executors.workloads.connection_test import TestConnection
from airflow.executors.workloads.task import ExecuteTask
from airflow.executors.workloads.types import state_class_for_key
from airflow.models import Log
from airflow.models.taskinstancekey import TaskInstanceKey
from airflow.observability.metrics import stats_utils
from airflow.utils.helpers import prune_dict
from airflow.utils.log.logging_mixin import LoggingMixin

PARALLELISM: int = conf.getint("core", "PARALLELISM")


def get_execution_api_server_url(conf_source: AirflowConfigParser | ExecutorConf = conf) -> str:
    """
    Resolve the execution API server URL from configuration.

    :param conf_source: Configuration source to read from. Defaults to the global ``conf``.
        Team-specific executors can pass their own config (e.g. ``ExecutorConf``) to resolve
        a team-specific URL.
    """
    base_url = conf_source.get("api", "base_url", fallback="/")
    # ExecutorConf.get() is typed as str | None even when fallback= guarantees a str,
    # so the `not base_url` guard and the cast() below keep mypy happy.
    if not base_url or base_url.startswith("/"):
        base_url = f"http://localhost:8080{base_url}"
    default_execution_api_server = f"{base_url.rstrip('/')}/execution/"
    return cast(
        "str", conf_source.get("core", "execution_api_server_url", fallback=default_execution_api_server)
    )


if TYPE_CHECKING:
    import argparse
    from datetime import datetime

    from sqlalchemy.orm import Session

    from airflow.api_fastapi.auth.tokens import JWTGenerator
    from airflow.callbacks.base_callback_sink import BaseCallbackSink
    from airflow.callbacks.callback_requests import CallbackRequest
    from airflow.cli.cli_config import GroupCommand
    from airflow.configuration import AirflowConfigParser
    from airflow.executors.executor_utils import ExecutorName
    from airflow.executors.workloads import ExecutorWorkload
    from airflow.executors.workloads.types import WorkloadKey, WorkloadState
    from airflow.models.callback import CallbackKey
    from airflow.models.connection_test import ConnectionTestKey
    from airflow.models.taskinstance import TaskInstance

    # Event_buffer dict value type
    # Tuple of: state, info
    EventBufferValueType = tuple[str | None, Any]


log = logging.getLogger(__name__)


@dataclass
class RunningRetryAttemptType:
    """
    For keeping track of attempts to queue again when task still apparently running.

    We don't want to slow down the loop, so we don't block, but we allow it to be
    re-checked for at least MIN_SECONDS seconds.
    """

    MIN_SECONDS = 10
    total_tries: int = field(default=0, init=False)
    tries_after_min: int = field(default=0, init=False)
    first_attempt_time: datetime = field(default_factory=lambda: pendulum.now("UTC"), init=False)

    @property
    def elapsed(self):
        """Seconds since first attempt."""
        return (pendulum.now("UTC") - self.first_attempt_time).total_seconds()

    def can_try_again(self):
        """Return False if there has been at least one try greater than MIN_SECONDS, otherwise return True."""
        if self.tries_after_min > 0:
            return False

        self.total_tries += 1

        elapsed = self.elapsed
        if elapsed > self.MIN_SECONDS:
            self.tries_after_min += 1
        log.debug("elapsed=%s tries=%s", elapsed, self.total_tries)
        return True


class ExecutorConf:
    """
    This class is used to fetch configuration for an executor for a particular team_name.

    It wraps the implementation of the configuration.get() to look for the particular section and key
    prefixed with the team_name. This makes it easy for child classes (i.e. concrete executors) to fetch
    configuration values for a particular team_name without having to worry about passing through the
    team_name for every call to get configuration.

    Currently config only supports environment variables for team specific configuration.
    """

    def __init__(self, team_name: str | None = None) -> None:
        self.team_name: str | None = team_name

    def get(self, *args, **kwargs) -> str | None:
        return conf.get(*args, **kwargs, team_name=self.team_name)

    def getboolean(self, *args, **kwargs) -> bool:
        return conf.getboolean(*args, **kwargs, team_name=self.team_name)

    def getjson(self, *args, **kwargs):
        return conf.getjson(*args, **kwargs, team_name=self.team_name)

    def getint(self, *args, **kwargs):
        return conf.getint(*args, **kwargs, team_name=self.team_name)

    def getsection(self, section: str) -> dict[str, str | int | float | bool] | None:
        return conf.getsection(section, team_name=self.team_name)

    def has_option(self, *args, **kwargs) -> bool:
        return conf.has_option(*args, **kwargs, team_name=self.team_name)

    def get_mandatory_value(self, *args, **kwargs) -> str:
        return conf.get_mandatory_value(*args, **kwargs, team_name=self.team_name)


class BaseExecutor(LoggingMixin):
    """
    Base class to inherit for concrete executors such as Celery, Kubernetes, Local, etc.

    :param parallelism: how many jobs should run at one time.
    """

    supports_ad_hoc_ti_run: bool = False
    supports_callbacks: bool = False
    supports_multi_team: bool = False
    # The connection-test supervisor uses ``signal.SIGALRM`` (via ``TimeoutPosix``)
    # to bound hook execution. Executors that opt in must run on POSIX systems.
    supports_connection_test: bool = False
    sentry_integration: str = ""

    is_local: bool = False
    is_production: bool = True

    # When True, the scheduler pre-assigns external_executor_id (a UUID) at queuing time,
    # committed atomically with the QUEUED state. The executor can then use this ID to
    # correlate the task with its external representation (e.g. Celery task_id).
    pre_assigns_external_executor_id: ClassVar[bool] = False

    serve_logs: bool = False

    job_id: None | int | str = None
    name: None | ExecutorName = None
    callback_sink: BaseCallbackSink | None = None

    @cached_property
    def jwt_generator(self) -> JWTGenerator:
        from airflow.api_fastapi.auth.tokens import (
            JWTGenerator,
            get_signing_args,
        )
        from airflow.configuration import conf

        generator = JWTGenerator(
            valid_for=conf.getint("execution_api", "jwt_expiration_time"),
            audience=conf.get_mandatory_list_value("execution_api", "jwt_audience")[0],
            issuer=conf.get("api_auth", "jwt_issuer", fallback=None),
            # Since this one is used across components/server, there is no point trying to generate one, error
            # instead
            **get_signing_args(make_secret_key_if_needed=False),
        )

        return generator

    def __init__(self, parallelism: int = PARALLELISM, team_name: str | None = None):
        stats.initialize(
            factory=stats_utils.get_stats_factory(),
            export_legacy_names=conf.getboolean("metrics", "legacy_names_on"),
        )
        super().__init__()
        # Ensure we set this now, so that each subprocess gets the same value
        from airflow.api_fastapi.auth.tokens import get_signing_args

        get_signing_args()

        self.parallelism: int = parallelism
        self.team_name: str | None = team_name
        self.queued_tasks: dict[TaskInstanceKey, workloads.ExecuteTask] = {}
        self.queued_callbacks: dict[CallbackKey, workloads.ExecuteCallback] = {}
        self.queued_connection_tests: dict[ConnectionTestKey, workloads.TestConnection] = {}
        self.running: set[WorkloadKey] = set()
        self.event_buffer: dict[WorkloadKey, EventBufferValueType] = {}
        self._task_event_logs: deque[Log] = deque()
        self.conf = ExecutorConf(team_name)

        if self.parallelism <= 0:
            raise ValueError("parallelism is set to 0 or lower")

        """
        Deque for storing task event log messages.

        This attribute is only internally public and should not be manipulated
        directly by subclasses.

        :meta private:
        """

        self.attempts: dict[WorkloadKey, RunningRetryAttemptType] = defaultdict(RunningRetryAttemptType)

    def __repr__(self):
        _repr = f"{self.__class__.__name__}(parallelism={self.parallelism}"
        if self.team_name:
            _repr += f", team_name={self.team_name!r}"
        _repr += ")"
        return _repr

    def start(self):  # pragma: no cover
        """Executors may need to get things started."""

    def log_task_event(self, *, event: str, extra: str, ti_key: WorkloadKey):
        """Add an event to the log table."""
        if not isinstance(ti_key, TaskInstanceKey):
            self.log.debug("Skipping log_task_event for callback key %s (event=%s)", ti_key, event)
            return
        self._task_event_logs.append(Log(event=event, task_instance=ti_key, extra=extra))

    def queue_workload(self, workload: ExecutorWorkload, session: Session) -> None:
        if isinstance(workload, workloads.ExecuteTask):
            ti = workload.ti
            self.queued_tasks[ti.key] = workload
        elif isinstance(workload, workloads.ExecuteCallback):
            if not self.supports_callbacks:
                raise NotImplementedError(
                    f"{type(self).__name__} does not support ExecuteCallback workloads. "
                    f"Set supports_callbacks = True and implement callback handling in _process_workloads(). "
                    f"See LocalExecutor or CeleryExecutor for reference implementation."
                )
            self.queued_callbacks[workload.key] = workload
        elif isinstance(workload, workloads.TestConnection):
            if not self.supports_connection_test:
                raise NotImplementedError(
                    f"{type(self).__name__} does not support TestConnection workloads. "
                    f"Set supports_connection_test = True and implement connection test handling "
                    f"in _process_workloads(). See LocalExecutor for reference implementation."
                )
            self.queued_connection_tests[workload.key] = workload
        else:
            raise ValueError(
                f"Un-handled workload type {type(workload).__name__!r} in {type(self).__name__}. "
                f"Workload must be one of: ExecuteTask, ExecuteCallback, TestConnection."
            )

    def _get_workloads_to_schedule(self, open_slots: int) -> list[tuple[WorkloadKey, ExecutorWorkload]]:
        """
        Select and return the next batch of workloads to schedule, respecting priority policy.

        Priority Policy: Callbacks are scheduled before tasks (callbacks complete existing work).
        Callbacks are processed in FIFO order. Tasks are sorted by priority_weight (higher priority first).

        :param open_slots: Number of available execution slots
        """
        workloads_to_schedule: list[tuple[WorkloadKey, ExecutorWorkload]] = []

        if self.queued_callbacks:
            for key, workload in self.queued_callbacks.items():
                if len(workloads_to_schedule) >= open_slots:
                    break
                workloads_to_schedule.append((key, workload))

        if open_slots > len(workloads_to_schedule) and self.queued_tasks:
            for task_key, task_workload in self.order_queued_tasks_by_priority():
                if len(workloads_to_schedule) >= open_slots:
                    break
                workloads_to_schedule.append((task_key, task_workload))

        return workloads_to_schedule

    def _process_workloads(self, workload_items: Sequence[ExecutorWorkload]) -> None:
        """
        Process the given workloads.

        This method must be implemented by subclasses to define how they handle
        the execution of workloads (e.g., queuing them to workers, submitting to
        external systems, etc.).

        :param workload_items: List of workloads to process
        """
        raise NotImplementedError(f"{type(self).__name__} must implement _process_workloads()")

    def has_task(self, task_instance: TaskInstance) -> bool:
        """
        Check if a task is either queued or running in this executor.

        :param task_instance: TaskInstance
        :return: True if the task is known to this executor
        """
        return (
            task_instance.id in self.queued_tasks
            or task_instance.id in self.running
            or task_instance.key in self.queued_tasks
            or task_instance.key in self.running
        )

    def sync(self) -> None:
        """
        Sync will get called periodically by the heartbeat method.

        Executors should override this to perform gather statuses.
        """

    def heartbeat(self) -> None:
        """Heartbeat sent to trigger new jobs."""
        open_slots = self.parallelism - len(self.running)

        num_running_workloads = len(self.running)
        num_queued_workloads = (
            len(self.queued_tasks) + len(self.queued_callbacks) + len(self.queued_connection_tests)
        )

        self._emit_metrics(open_slots, num_running_workloads, num_queued_workloads)
        self.trigger_tasks(open_slots)

        self.trigger_connection_tests()

        # Calling child class sync method
        self.log.debug("Calling the %s sync method", self.__class__)
        self.sync()

    def trigger_connection_tests(self) -> None:
        """Process queued connection tests, respecting available slot capacity."""
        if not self.supports_connection_test or not self.queued_connection_tests:
            return

        available = self.slots_available
        if available <= 0:
            return

        tests_to_run = list(self.queued_connection_tests.values())[:available]
        self._process_workloads(tests_to_run)

    def fail_connection_test(self, key: ConnectionTestKey) -> None:
        """Drop a connection-test workload from in-memory queues (called by the reaper)."""
        self.queued_connection_tests.pop(key, None)
        self.running.discard(key)

    def _get_metric_name(self, metric_base_name: str) -> str:
        return (
            f"{metric_base_name}.{self.__class__.__name__}"
            if len(ExecutorLoader.get_executor_names()) > 1
            else metric_base_name
        )

    def _emit_metrics(self, open_slots, num_running_tasks, num_queued_tasks):
        """
        Emit metrics relevant to the Executor.

        In the case of multiple executors being configured, the metric names include the name of
        executor to differentiate them from metrics from other executors.

        If only one executor is configured, the metric names will not be changed.
        """
        name = self.__class__.__name__

        open_slots_metric_name = self._get_metric_name("executor.open_slots")
        queued_tasks_metric_name = self._get_metric_name("executor.queued_tasks")
        running_tasks_metric_name = self._get_metric_name("executor.running_tasks")

        self.log.debug("%s running task instances for executor %s", num_running_tasks, name)
        self.log.debug("%s in queue for executor %s", num_queued_tasks, name)
        if open_slots == 0:
            self.log.info("Executor parallelism limit reached. 0 open slots.")
        else:
            self.log.debug("%s open slots for executor %s", open_slots, name)

        stats.gauge(
            open_slots_metric_name,
            value=open_slots,
            tags=prune_dict({"status": "open", "executor_class_name": name, "team_name": self.team_name}),
        )
        stats.gauge(
            queued_tasks_metric_name,
            value=num_queued_tasks,
            tags=prune_dict({"status": "queued", "executor_class_name": name, "team_name": self.team_name}),
        )
        stats.gauge(
            running_tasks_metric_name,
            value=num_running_tasks,
            tags=prune_dict({"status": "running", "executor_class_name": name, "team_name": self.team_name}),
        )

    def order_queued_tasks_by_priority(self) -> list[tuple[TaskInstanceKey, workloads.ExecuteTask]]:
        """
        Orders the queued tasks by priority.

        :return: List of workloads from the queued_tasks according to the priority.
        """
        if not self.queued_tasks:
            return []

        # V3 + new executor that supports workloads
        return sorted(
            self.queued_tasks.items(),
            key=lambda x: x[1].ti.priority_weight,
            reverse=False,
        )

    def trigger_tasks(self, open_slots: int) -> None:
        """
        Initiate async execution of queued workloads (tasks and callbacks), up to the number of available slots.

        Callbacks are prioritized over tasks to complete existing work before starting new work.

        :param open_slots: Number of open slots
        """
        workloads_to_schedule = self._get_workloads_to_schedule(open_slots)
        workload_list = []

        for key, workload in workloads_to_schedule:
            # If a task makes it here but is still understood by the executor
            # to be running, it generally means that the task has been killed
            # externally and not yet been marked as failed.
            #
            # However, when a task is deferred, there is also a possibility of
            # a race condition where a task might be scheduled again during
            # trigger processing, even before we are able to register that the
            # deferred task has completed. In this case and for this reason,
            # we make a small number of attempts to see if the task has been
            # removed from the running set in the meantime.
            if key in self.attempts:
                del self.attempts[key]

            workload_list.append(workload)

        if workload_list:
            self._process_workloads(workload_list)

    # TODO: This should not be using `TaskInstanceState` here, this is just "did the process complete, or did
    # it die". It is possible for the task itself to finish with success, but the state of the task to be set
    # to FAILED. By using TaskInstanceState enum here it confuses matters!
    def change_state(self, key: WorkloadKey, state: WorkloadState, info=None, remove_running=True) -> None:
        """
        Change state of the task.

        :param key: Unique key for the task instance
        :param state: State to set for the task.
        :param info: Executor information for the task instance
        :param remove_running: Whether or not to remove the TI key from running set
        """
        self.log.debug("Changing state: %s", key)
        if remove_running:
            try:
                self.running.remove(key)
            except KeyError:
                self.log.debug("Could not find key: %s", key)
        self.event_buffer[key] = state, info

    def fail(self, key: WorkloadKey, info=None) -> None:
        """
        Set fail state for the event.

        :param info: Executor information for the task instance
        :param key: Unique key for the task instance
        """
        self.change_state(key, state_class_for_key(key).FAILED, info)

    def success(self, key: WorkloadKey, info=None) -> None:
        """
        Set success state for the event.

        :param info: Executor information for the task instance
        :param key: Unique key for the task instance
        """
        self.change_state(key, state_class_for_key(key).SUCCESS, info)

    def queued(self, key: WorkloadKey, info=None) -> None:
        """
        Set queued state for the event.

        :param info: Executor information for the task instance
        :param key: Unique key for the task instance
        """
        self.change_state(key, state_class_for_key(key).QUEUED, info)

    def running_state(self, key: WorkloadKey, info=None) -> None:
        """
        Set running state for the event.

        :param info: Executor information for the task instance
        :param key: Unique key for the task instance
        """
        self.change_state(key, state_class_for_key(key).RUNNING, info, remove_running=False)

    def get_event_buffer(self, dag_ids=None) -> dict[WorkloadKey, EventBufferValueType]:
        """
        Return and flush the event buffer.

        In case dag_ids is specified it will only return and flush events
        for the given dag_ids. Otherwise, it returns and flushes all events.
        Note: Callback events (with CallbackKey keys) are always included regardless of dag_ids filter.

        :param dag_ids: the dag_ids to return events for; returns all if given ``None``.
        :return: a dict of events
        """
        cleared_events: dict[WorkloadKey, EventBufferValueType] = {}
        if dag_ids is None:
            cleared_events = self.event_buffer
            self.event_buffer = {}
        else:
            for key in list(self.event_buffer.keys()):
                if not isinstance(key, TaskInstanceKey) or key.dag_id in dag_ids:
                    cleared_events[key] = self.event_buffer.pop(key)

        return cleared_events

    def get_task_log(self, ti: TaskInstance, try_number: int) -> tuple[list[str], list[str]]:
        """
        Return the task logs.

        :param ti: A TaskInstance object
        :param try_number: current try_number to read log from
        :return: tuple of logs and messages
        """
        return [], []

    def end(self) -> None:  # pragma: no cover
        """Wait synchronously for the previously submitted job to complete."""
        raise NotImplementedError

    def terminate(self):
        """Get called when the daemon receives a SIGTERM."""
        raise NotImplementedError

    def revoke_task(self, *, ti: TaskInstance):
        """
        Attempt to remove task from executor.

        It should attempt to ensure that the task is no longer running on the worker,
        and ensure that it is cleared out from internal data structures.

        It should *not* change the state of the task in airflow, or add any events
        to the event buffer.

        It should not raise any error.

        :param ti: Task instance to remove
        """
        raise NotImplementedError

    def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[TaskInstance]:
        """
        Try to adopt running task instances that have been abandoned by a SchedulerJob dying.

        Anything that is not adopted will be cleared by the scheduler (and then become eligible for
        re-scheduling)

        :return: any TaskInstances that were unable to be adopted
        """
        # By default, assume Executors cannot adopt tasks, so just say we failed to adopt anything.
        # Subclasses can do better!
        return tis

    @property
    def slots_available(self):
        """Number of new workloads (tasks, callbacks, and connection tests) this executor instance can accept."""
        return (
            self.parallelism
            - len(self.running)
            - len(self.queued_tasks)
            - len(self.queued_callbacks)
            - len(self.queued_connection_tests)
        )

    @property
    def slots_occupied(self):
        """Number of workloads (tasks, callbacks, and connection tests) this executor instance is currently managing."""
        return (
            len(self.running)
            + len(self.queued_tasks)
            + len(self.queued_callbacks)
            + len(self.queued_connection_tests)
        )

    def debug_dump(self):
        """Get called in response to SIGUSR2 by the scheduler."""
        self.log.info(
            "executor.queued_tasks (%d)\n\t%s",
            len(self.queued_tasks),
            "\n\t".join(map(repr, self.queued_tasks.items())),
        )
        self.log.info(
            "executor.queued_callbacks (%d)\n\t%s",
            len(self.queued_callbacks),
            "\n\t".join(map(repr, self.queued_callbacks.items())),
        )
        self.log.info("executor.running (%d)\n\t%s", len(self.running), "\n\t".join(map(repr, self.running)))
        self.log.info(
            "executor.event_buffer (%d)\n\t%s",
            len(self.event_buffer),
            "\n\t".join(map(repr, self.event_buffer.items())),
        )

    def send_callback(self, request: CallbackRequest) -> None:
        """
        Send callback for execution.

        Provides a default implementation which sends the callback to the `callback_sink` object.

        :param request: Callback request to be executed.
        """
        if not self.callback_sink:
            raise ValueError("Callback sink is not ready.")
        self.callback_sink.send(request)

    @staticmethod
    def get_cli_commands() -> list[GroupCommand]:
        """
        Vends CLI commands to be included in Airflow CLI.

        Override this method to expose commands via Airflow CLI to manage this executor. This can
        be commands to setup/teardown the executor, inspect state, etc.
        Make sure to choose unique names for those commands, to avoid collisions.
        """
        return []

    @staticmethod
    def run_workload(
        workload: ExecutorWorkload,
        *,
        server: str | None = None,
        dry_run: bool = False,
        subprocess_logs_to_stdout: bool = False,
        proctitle: str | None = None,
    ) -> int:
        """
        Pass the workload to the appropriate supervisor based on workload type.

        Workload-specific attributes (log_path, sentry_integration, bundle_info, etc.) are read from the
        workload object itself.

        :param workload: The ``ExecutorWorkload`` to execute.
        :param server: Base URL of the API server (used by task workloads).
        :param dry_run: If True, execute without actual task execution (simulate run).
        :param subprocess_logs_to_stdout: Should task logs also be sent to stdout via the main logger.
        :param proctitle: Process title to set for this workload. If not provided, defaults to
            ``"airflow supervisor: <workload.display_name>"``.
        :return: Exit code of the process.
        """
        try:
            if sys.platform != "darwin":
                from setproctitle import setproctitle

                setproctitle(proctitle or f"airflow supervisor: {workload.display_name}")
        except ImportError:
            pass

        # Resolve server URL from config when not explicitly provided.
        # For example, team-specific executors may wish to pass their own server URL.
        if server is None:
            server = get_execution_api_server_url()

        if isinstance(workload, ExecuteTask):
            from airflow.sdk.execution_time.supervisor import supervise_task

            # workload.ti is a TaskInstanceDTO which duck-types as TaskInstance.
            # TODO: Create a protocol for this.
            return supervise_task(
                ti=workload.ti,  # type: ignore[arg-type]
                bundle_info=workload.bundle_info,
                dag_rel_path=workload.dag_rel_path,
                token=workload.token,
                server=server,
                dry_run=dry_run,
                log_path=workload.log_path,
                subprocess_logs_to_stdout=subprocess_logs_to_stdout,
                sentry_integration=getattr(workload, "sentry_integration", ""),
            )
        if isinstance(workload, ExecuteCallback):
            from airflow.sdk.execution_time.callback_supervisor import supervise_callback

            return supervise_callback(
                id=workload.callback.id,
                callback_path=workload.callback.data.get("path", ""),
                callback_kwargs=workload.callback.data.get("kwargs", {}),
                dag_rel_path=workload.dag_rel_path,
                log_path=workload.log_path,
                bundle_info=workload.bundle_info,
                token=workload.token,
                server=server,
            )
        if isinstance(workload, TestConnection):
            from airflow.sdk.execution_time.connection_test_supervisor import supervise_connection_test

            return supervise_connection_test(
                connection_test_id=workload.connection_test_id,
                connection_id=workload.connection_id,
                timeout=workload.timeout,
                token=workload.token,
                server=server,
            )
        raise ValueError(f"Unknown workload type: {type(workload).__name__}")

    @classmethod
    def _get_parser(cls) -> argparse.ArgumentParser:
        """
        Generate documentation; used by Sphinx argparse.

        :meta private:
        """
        from airflow.cli.cli_parser import AirflowHelpFormatter, _add_command

        parser = DefaultHelpParser(prog="airflow", formatter_class=AirflowHelpFormatter)
        subparsers = parser.add_subparsers(dest="subcommand", metavar="GROUP_OR_COMMAND")
        for group_command in cls.get_cli_commands():
            _add_command(subparsers, group_command)
        return parser
