From 1c5e83d20659e053dd94bfb2579a6f7dcb6ec003 Mon Sep 17 00:00:00 2001
From: d9pouces <github@19pouces.net>
Date: Mon, 16 Oct 2023 07:03:38 +0200
Subject: [PATCH] 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                                    |  3 +++
 .../management/commands/wait_for_database.py  | 10 +++++-----
 tests/unit/test_cli.py                        | 20 +++++++++++++++++++
 3 files changed, 28 insertions(+), 5 deletions(-)

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