From cc000775c9c7d068b8245227d3b95354d7c47676 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jes=C3=BAs=20Garc=C3=ADa=20Crespo?= <jesus@sevein.com>
Date: Fri, 16 Aug 2024 06:40:58 +0000
Subject: [PATCH] Make install wait for the db

---
 hack/Makefile                                 |  1 +
 .../main/management/commands/install.py       | 42 +++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/hack/Makefile b/hack/Makefile
index 0867d388..14bc90ca 100644
--- a/hack/Makefile
+++ b/hack/Makefile
@@ -100,6 +100,7 @@ bootstrap:  ## Bootstrap the database.
 		--entrypoint /src/worker/manage.py \
 			archivematica-worker \
 				install \
+					--wait-for-db \
 					--migrate \
 					--username="test" \
 					--password="test" \
diff --git a/worker/worker/main/management/commands/install.py b/worker/worker/main/management/commands/install.py
index 9307f64a..5210a292 100644
--- a/worker/worker/main/management/commands/install.py
+++ b/worker/worker/main/management/commands/install.py
@@ -1,9 +1,14 @@
 import uuid
 from argparse import BooleanOptionalAction
+from time import sleep
+from time import time
 
 from django.contrib.auth import get_user_model
 from django.core.management import call_command
 from django.core.management.base import BaseCommand
+from django.db import DEFAULT_DB_ALIAS
+from django.db import connections
+from django.db.utils import OperationalError
 from tastypie.models import ApiKey
 
 from worker.main.models import Agent
@@ -12,6 +17,12 @@
 
 class Command(BaseCommand):
     def add_arguments(self, parser):
+        parser.add_argument(
+            "--wait-for-db",
+            default=False,
+            action=BooleanOptionalAction,
+            help="Wait for the database.",
+        )
         parser.add_argument(
             "--migrate",
             default=False,
@@ -33,6 +44,10 @@ def add_arguments(self, parser):
         parser.add_argument("--site-url", required=False)
 
     def handle(self, *args, **options):
+        self.print_banner("Wait for the database")
+        self.wait_for_db()
+        self.print_done()
+
         if options["migrate"] is True:
             self.print_banner("Apply migrations")
             call_command("migrate", interactive=False)
@@ -64,6 +79,33 @@ def print_banner(self, text):
         middle = f"+ {text} +"
         self.stdout.write(f"\n{border}\n{middle}\n{border}\n")
 
+    def wait_for_db(self, timeout_seconds=120, stable_for_seconds=3):
+        connection = connections[DEFAULT_DB_ALIAS]
+        sleep_seconds = 1
+        seen_alive_at = None
+        started_at = time()
+        while True:
+            while True:
+                try:
+                    connection.cursor().execute("SELECT 1")
+                    if not seen_alive_at:
+                        seen_alive_at = time()
+                    break
+                except OperationalError as err:
+                    seen_alive_at = None
+                    elapsed_time = int(time() - started_at)
+                    if elapsed_time >= timeout_seconds:
+                        raise TimeoutError from err
+                    err_message = str(err).strip()
+                    self.stdout.write(
+                        f"Waiting for database: {err_message}... ({elapsed_time}s)\n"
+                    )
+                    sleep(sleep_seconds)
+            uptime = int(time() - seen_alive_at)
+            if uptime >= stable_for_seconds:
+                break
+            sleep(sleep_seconds)
+
 
 def get_setting(setting, default=""):
     try: