Add a --command argument, allowing to directly call a Django command …

…when databases are ready, without patching the command itself and without loading the django machinery twice.

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'``

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:])

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