Improve logic of wait_for_database for alembic

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