Skip to content

Commit

Permalink
add wait_for_db command
Browse files Browse the repository at this point in the history
  • Loading branch information
ytliuSVN committed Nov 2, 2024
1 parent b155d34 commit fa3e5f3
Show file tree
Hide file tree
Showing 8 changed files with 57 additions and 4 deletions.
2 changes: 1 addition & 1 deletion app/app/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@
'HOST': os.environ.get('DB_HOST'),
'NAME': os.environ.get('DB_NAME'),
'USER': os.environ.get('DB_USER'),
'PASSWORD': os.environ.get('DB_PASS'),
'PASSWORD': os.environ.get('DB_PASSWORD'),
}
}

Expand Down
2 changes: 1 addition & 1 deletion app/core/admin.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
from django.contrib import admin
from django.contrib import admin # noqa

# Register your models here.
Empty file added app/core/management/__init__.py
Empty file.
Empty file.
22 changes: 22 additions & 0 deletions app/core/management/commands/wait_for_db.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
import time

from psycopg2 import OperationalError as Psycopg2OpError

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


class Command(BaseCommand):

def handle(self, *args, **options):
self.stdout.write('Waiting for db...')
db_up = False
while db_up is False:
try:
self.check(databases=['default'])
db_up = True
except (Psycopg2OpError, OperationalError):
self.stdout.write('db unavailable, waiting 1 second...')
time.sleep(1)

self.stdout.write(self.style.SUCCESS('db available!'))
2 changes: 1 addition & 1 deletion app/core/models.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
from django.db import models
from django.db import models # noqa

# Create your models here.
31 changes: 31 additions & 0 deletions app/core/tests/test_commands.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
"""
Test custom Django management commands.
"""
from unittest.mock import patch

from psycopg2 import OperationalError as Psycopg2OpError

from django.core.management import call_command
from django.db.utils import OperationalError
from django.test import SimpleTestCase


@patch('core.management.commands.wait_for_db.Command.check')
class CommandTests(SimpleTestCase):

def test_wait_for_db_ready(self, patched_check):
patched_check.return_value = True

call_command('wait_for_db')

patched_check.assert_called_once_with(databases=['default'])

@patch('time.sleep')
def test_wait_for_db_delay(self, patched_sleep, patched_check):
patched_check.side_effect = [Psycopg2OpError] * 2 + \
[OperationalError] * 3 + [True]

call_command('wait_for_db')

self.assertEqual(patched_check.call_count, 6)
patched_check.assert_called_with(databases=['default'])
2 changes: 1 addition & 1 deletion docker-compose.yml
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ services:
- DB_HOST=db
- DB_NAME=devdb
- DB_USER=devuser
- DB_PASS=Password1
- DB_PASSWORD=Password1
depends_on:
- db

Expand Down

0 comments on commit fa3e5f3

Please sign in to comment.