feat(cli): improve quickstart stability (#7839)

datahub-project · Apr 18, 2023 · 399e333 · 399e333
1 parent 80650f0
commit 399e333
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 13 deletions.
diff --git a/metadata-ingestion/src/datahub/cli/docker_check.py b/metadata-ingestion/src/datahub/cli/docker_check.py
@@ -137,6 +137,17 @@ def errors(self) -> List[str]:
     def is_ok(self) -> bool:
         return not self.errors()
 
+    def needs_up(self) -> bool:
+        return any(
+            container.status
+            in {
+                ContainerStatus.EXITED_WITH_FAILURE,
+                ContainerStatus.DIED,
+                ContainerStatus.MISSING,
+            }
+            for container in self.containers
+        )
+
     def to_exception(
         self, header: str, footer: Optional[str] = None
     ) -> QuickstartError:

diff --git a/metadata-ingestion/src/datahub/cli/docker_cli.py b/metadata-ingestion/src/datahub/cli/docker_cli.py
@@ -25,6 +25,7 @@
     DATAHUB_COMPOSE_LEGACY_VOLUME_FILTERS,
     DATAHUB_COMPOSE_PROJECT_FILTER,
     DockerComposeVersionError,
+    QuickstartStatus,
     check_docker_quickstart,
     get_docker_client,
     run_quickstart_preflight_checks,
@@ -60,6 +61,10 @@
 )
 
 
+_QUICKSTART_MAX_WAIT_TIME = datetime.timedelta(minutes=10)
+_QUICKSTART_STATUS_CHECK_INTERVAL = datetime.timedelta(seconds=2)
+
+
 class Architectures(Enum):
     x86 = "x86"
     arm64 = "arm64"
@@ -714,27 +719,21 @@ def quickstart(
     if build_locally:
         logger.info("Building docker images locally...")
         subprocess.run(
-            [
-                *base_command,
-                "build",
-                "--pull",
-                "-q",
-            ],
+            base_command + ["build", "--pull", "-q"],
             check=True,
             env=_docker_subprocess_env(),
         )
         logger.info("Finished building docker images!")
 
     # Start it up! (with retries)
     click.echo("\nStarting up DataHub...")
-    max_wait_time = datetime.timedelta(minutes=8)
     start_time = datetime.datetime.now()
-    sleep_interval = datetime.timedelta(seconds=2)
-    up_interval = datetime.timedelta(seconds=30)
+    status: Optional[QuickstartStatus] = None
     up_attempts = 0
-    while (datetime.datetime.now() - start_time) < max_wait_time:
-        # Attempt to run docker compose up every `up_interval`.
-        if (datetime.datetime.now() - start_time) > up_attempts * up_interval:
+    while (datetime.datetime.now() - start_time) < _QUICKSTART_MAX_WAIT_TIME:
+        # We must run docker-compose up at least once.
+        # Beyond that, we should run it again if something goes wrong.
+        if up_attempts == 0 or (status and status.needs_up()):
             if up_attempts > 0:
                 click.echo()
             subprocess.run(
@@ -750,8 +749,10 @@ def quickstart(
 
         # Wait until next iteration.
         click.echo(".", nl=False)
-        time.sleep(sleep_interval.total_seconds())
+        time.sleep(_QUICKSTART_STATUS_CHECK_INTERVAL.total_seconds())
     else:
+        assert status
+
         # Falls through if the while loop doesn't exit via break.
         click.echo()
         with tempfile.NamedTemporaryFile(suffix=".log", delete=False) as log_file: