Skip to content

Commit

Permalink
Add a --command argument, allowing to directly call a Django command …
Browse files Browse the repository at this point in the history
…when databases are ready, without patching the command itself and without loading the django machinery twice.
  • Loading branch information
d9pouces committed Oct 16, 2023
1 parent 3995f1e commit 1c5e83d
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 5 deletions.
3 changes: 3 additions & 0 deletions README.rst
Original file line number Diff line number Diff line change
Expand Up @@ -109,3 +109,6 @@ if needed:
delay between checks when database is up (seconds), default: ``1``
:--database:
which database of ``settings.DATABASES`` to wait for, default: ``default``
:--command, -c:
execute this Django command when the database is ready.
This option can be used multiple times: ``wait_for_database -c 'migrate' -c 'runserver --skip-checks'``
10 changes: 5 additions & 5 deletions django_probes/management/commands/wait_for_database.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,20 +90,20 @@ def add_arguments(self, parser):
help='which database of `settings.DATABASES` '
'to wait for. Defaults to the "default" '
'database.')
parser.add_argument("--command", "-c", default=None,
action="store", dest='command',
help='execute this command when database is up')
parser.add_argument("--command", "-c", default=[],
action="append", dest='command',
help='execute this command when database is up (can be repeated multiple times).')

def handle(self, *args, **options):
"""
Wait for a database connection to come up. Exit with error
status when a timeout threshold is surpassed.
"""
command = options.pop("command", None)
commands = options.pop("command", [])
try:
wait_for_database(**options)
except TimeoutError as err:
raise CommandError(err) from err
if command:
for command in commands:
command_list = shlex.split(command)
call_command(command_list[0], *command_list[1:])
20 changes: 20 additions & 0 deletions tests/unit/test_cli.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
"""
Verify that ``python manage.py wait_for_database`` works fine.
"""
import os
import tempfile
from unittest.mock import patch

import pytest
Expand Down Expand Up @@ -57,3 +59,21 @@ def test_command_error_raised_when_connection_absent(mock_db_cursor):
call_command('wait_for_database', stable=0, timeout=0)

assert mock_db_cursor.called


@patch('django.db.connection.cursor')
def test_can_call_through_management_with_commands(mock_db_cursor):
"""
Executing the management command works (w/o operational errors).
"""
with tempfile.TemporaryDirectory() as dirname:
call_command('wait_for_database',
"--stable=0",
"--timeout=0",
"-c",
f"shell -c 'open(\"{dirname}/0\", \"w\").close()'",
"-c",
f"shell -c 'open(\"{dirname}/1\", \"w\").close()'",
)
assert set(os.listdir(dirname)) == {"0", "1"}
assert mock_db_cursor.called

0 comments on commit 1c5e83d

Please sign in to comment.