doccano/backend/api/management/commands/wait_for_db.py

import sys
import time

from django.core.management.base import BaseCommand
from django.db import connection
from django.db.utils import OperationalError


class Command(BaseCommand):
    help = "Blocks until the database is available"

    def add_arguments(self, parser):
        parser.add_argument("--poll_seconds", type=float, default=3)
        parser.add_argument("--max_retries", type=int, default=60)

    def handle(self, *args, **options):
        max_retries = options["max_retries"]
        poll_seconds = options["poll_seconds"]

        for retry in range(max_retries):
            try:
                connection.ensure_connection()
            except OperationalError as ex:
                self.stdout.write(
                    "Database unavailable on attempt {attempt}/{max_retries}:"
                    " {error}".format(attempt=retry + 1, max_retries=max_retries, error=ex)
                )
                time.sleep(poll_seconds)
            else:
                break
        else:
            self.stdout.write(self.style.ERROR("Database unavailable"))
            sys.exit(1)