Repository: apache/airflow

Issue:
Add retry + timeout to Airflow db check
### Description

In my company usage of Airflow, developmental instances of Airflow run on containerized PostgreSQL that are spawned at the same time the Airflow container is spawned. Before the Airflow container runs its initialization scripts, it needs to make sure that the PostgreSQL instance can be reached, for which `airflow db check` is a great option.

However, there is a non-deterministic race condition between the PSQL container and Airflow containers (not sure which will each readiness first and by how much), calling the `airflow db check` command once is not sufficient, and implementing a retry-timeout in shell script is feasible but unpleasant.

It would be great if the `airflow db check` command can take two additional optional arguments: `--retry` and `--retry-delay` (just like with `curl`) so that the database connection can be checked repeatedly for up to a specified number of times. This command should exit with `0` exit code if any of the retries succeeds, and `1` if all of the retries failed.

### Use case/motivation

_No response_

### Related issues

_No response_

### Are you willing to submit a PR?

- [X] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:
Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.

feel free. sounds useful

Resolving patch:
diff --git a/airflow/cli/cli_config.py b/airflow/cli/cli_config.py
--- a/airflow/cli/cli_config.py
+++ b/airflow/cli/cli_config.py
@@ -500,6 +500,18 @@ def string_lower_type(val):
     help="Drop the archive tables after exporting. Use with caution.",
     action="store_true",
 )
+ARG_DB_RETRY = Arg(
+    ("--retry",),
+    default=0,
+    type=positive_int(allow_zero=True),
+    help="Retry database check upon failure",
+)
+ARG_DB_RETRY_DELAY = Arg(
+    ("--retry-delay",),
+    default=1,
+    type=positive_int(allow_zero=False),
+    help="Wait time between retries in seconds",
+)
 
 # pool
 ARG_POOL_NAME = Arg(("pool",), metavar="NAME", help="Pool name")
@@ -1620,7 +1632,7 @@ class GroupCommand(NamedTuple):
         name="check",
         help="Check if the database can be reached",
         func=lazy_load_command("airflow.cli.commands.db_command.check"),
-        args=(ARG_VERBOSE,),
+        args=(ARG_VERBOSE, ARG_DB_RETRY, ARG_DB_RETRY_DELAY),
     ),
     ActionCommand(
         name="clean",
diff --git a/airflow/cli/commands/db_command.py b/airflow/cli/commands/db_command.py
--- a/airflow/cli/commands/db_command.py
+++ b/airflow/cli/commands/db_command.py
@@ -17,11 +17,13 @@
 """Database sub-commands."""
 from __future__ import annotations
 
+import logging
 import os
 import textwrap
 from tempfile import NamedTemporaryFile
 
 from packaging.version import parse as parse_version
+from tenacity import RetryCallState, Retrying, stop_after_attempt, wait_fixed
 
 from airflow import settings
 from airflow.exceptions import AirflowException
@@ -30,6 +32,8 @@
 from airflow.utils.db_cleanup import config_dict, drop_archived_tables, export_archived_records, run_cleanup
 from airflow.utils.process_utils import execute_interactive
 
+log = logging.getLogger(__name__)
+
 
 def initdb(args):
     """Initializes the metadata database."""
@@ -187,9 +191,23 @@ def shell(args):
 
 
 @cli_utils.action_cli(check_db=False)
-def check(_):
+def check(args):
     """Runs a check command that checks if db is available."""
-    db.check()
+    retries: int = args.retry
+    retry_delay: int = args.retry_delay
+
+    def _warn_remaining_retries(retrystate: RetryCallState):
+        remain = retries - retrystate.attempt_number
+        log.warning("%d retries remain. Will retry in %d seconds", remain, retry_delay)
+
+    for attempt in Retrying(
+        stop=stop_after_attempt(1 + retries),
+        wait=wait_fixed(retry_delay),
+        reraise=True,
+        before_sleep=_warn_remaining_retries,
+    ):
+        with attempt:
+            db.check()
 
 
 # lazily imported by CLI parser for `help` command


Test patch:
diff --git a/tests/cli/commands/test_db_command.py b/tests/cli/commands/test_db_command.py
--- a/tests/cli/commands/test_db_command.py
+++ b/tests/cli/commands/test_db_command.py
@@ -17,12 +17,13 @@
 from __future__ import annotations
 
 from unittest import mock
-from unittest.mock import patch
+from unittest.mock import MagicMock, Mock, call, patch
 
 import pendulum
 import pytest
 from pytest import param
 from sqlalchemy.engine.url import make_url
+from sqlalchemy.exc import OperationalError
 
 from airflow.cli import cli_parser
 from airflow.cli.commands import db_command
@@ -271,6 +272,27 @@ def test_cli_downgrade_confirm(self, mock_input, mock_dg, resp, raise_):
             db_command.downgrade(self.parser.parse_args(["db", "downgrade", "--to-revision", "abc"]))
             mock_dg.assert_called_with(to_revision="abc", from_revision=None, show_sql_only=False)
 
+    def test_check(self):
+        retry, retry_delay = 6, 9  # arbitrary but distinct number
+        args = self.parser.parse_args(
+            ["db", "check", "--retry", str(retry), "--retry-delay", str(retry_delay)]
+        )
+        sleep = MagicMock()
+        always_pass = Mock()
+        always_fail = Mock(side_effect=OperationalError("", None, None))
+
+        with patch("time.sleep", new=sleep), patch("airflow.utils.db.check", new=always_pass):
+            db_command.check(args)
+            always_pass.assert_called_once()
+            sleep.assert_not_called()
+
+        with patch("time.sleep", new=sleep), patch("airflow.utils.db.check", new=always_fail):
+            with pytest.raises(OperationalError):
+                db_command.check(args)
+            # With N retries there are N+1 total checks, hence N sleeps
+            always_fail.assert_has_calls([call()] * (retry + 1))
+            sleep.assert_has_calls([call(retry_delay)] * retry)
+
 
 class TestCLIDBClean:
     @classmethod