From 84f3537bc39128fa9bce5c4b21f8b4f39db971b9 Mon Sep 17 00:00:00 2001
From: Stanislav Khlud <stanislav.khlud@saritasa.com>
Date: Wed, 8 Nov 2023 16:42:29 +0700
Subject: [PATCH] Improve logic of wait_for_database for alembic

---
 saritasa_invocations/alembic.py | 30 ++++++++++++++++++++++++------
 1 file changed, 24 insertions(+), 6 deletions(-)

diff --git a/saritasa_invocations/alembic.py b/saritasa_invocations/alembic.py
index a80ce05..78b6fe6 100644
--- a/saritasa_invocations/alembic.py
+++ b/saritasa_invocations/alembic.py
@@ -1,4 +1,5 @@
 import pathlib
+import time
 
 import invoke
 
@@ -23,7 +24,7 @@ def wait_for_database(context: invoke.Context) -> None:
             "hide": "out",
         },
     ) as context:
-        for _ in range(config.alembic.connect_attempts):
+        for _ in range(config.alembic.connect_attempts - 1):
             try:
                 # Doing it manually to avoid loop
                 python.run(
@@ -33,13 +34,30 @@ def wait_for_database(context: invoke.Context) -> None:
                 wait_for_database._called = True  # type: ignore
                 return
             except invoke.UnexpectedExit:
+                time.sleep(1)
                 continue
 
-    printing.print_error(
-        "Failed to connect to db, "
-        f"after {config.alembic.connect_attempts} attempts",
-    )
-    raise invoke.Exit(code=1)
+    with _config.context_override(
+        context,
+        run={
+            "echo": True,
+            "hide": None,
+        },
+    ) as context:
+        try:
+            # Do it one more time but without hiding the terminal output
+            python.run(
+                context,
+                command=f"{config.alembic.command} current",
+            )
+            wait_for_database._called = True  # type: ignore
+            return
+        except invoke.UnexpectedExit:
+            printing.print_error(
+                "Failed to connect to db, "
+                f"after {config.alembic.connect_attempts} attempts",
+            )
+            raise invoke.Exit(code=1)
 
 
 @invoke.task