diff --git a/README.rst b/README.rst index 886756d..fe888ca 100644 --- a/README.rst +++ b/README.rst @@ -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'`` diff --git a/django_probes/management/commands/wait_for_database.py b/django_probes/management/commands/wait_for_database.py index 318c197..6d80876 100644 --- a/django_probes/management/commands/wait_for_database.py +++ b/django_probes/management/commands/wait_for_database.py @@ -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:]) diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index cc02c67..d232c2b 100644 --- a/tests/unit/test_cli.py +++ b/tests/unit/test_cli.py @@ -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 @@ -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