django_wait_db_command

wait_db.py

import logging
import time

from django.core.management.base import BaseCommand
from django.core.management.base import CommandError
from django.db import connections
from django.db.utils import OperationalError

logger = logging.getLogger(__name__)


class Command(BaseCommand):
    help = 'Wait db become available'

    def handle(self, *args, **options):
        self.stdout.write('Waiting for database...')
        timeout = int(options.get('timeout', 10))
        ready = False
        start = time.perf_counter()
        while not ready:
            try:
                db_conn = connections['default']
                c = db_conn.cursor()
                c.execute('SELECT 1')
                ready = c.fetchone()[0]
            except OperationalError as exc:
                if time.perf_counter() - start > timeout:
                    logger.exception('Database unavailable after %s seconds', timeout)
                    raise CommandError(f'Database unavailable after {timeout} seconds: {exc}', returncode=1)
                self.stdout.write('Database unavailable, waiting 1 second...')
                time.sleep(1)

        self.stdout.write(self.style.SUCCESS('Database available!'))