improve progress update when many dependent ressources

Author: aletournel

datadog_sync/model/synthetics_mobile_applications_versions.py

@@ -52,8 +52,14 @@ async def get_resources(self, client: CustomClient) -> List[Dict]:
             total_versions,
         )
 
+        # If import progress bar is active, add sub-steps so the bar advances during fetch
+        add_total = getattr(self.config, "_import_progress_add_total", None)
+        if callable(add_total):
+            add_total(total_versions)
+
         resources = []
         progress_interval = max(1, total_versions // 20)  # log at INFO roughly every 5%
+        progress_update = getattr(self.config, "_import_progress_update", None)
         for app_idx, application in enumerate(applications):
             versions = application.get("versions", [])
             for ver_idx, version in enumerate(versions):
@@ -67,6 +73,8 @@ async def get_resources(self, client: CustomClient) -> List[Dict]:
                 )
                 resource = await client.get(self.resource_config.base_path + f"/{_id}")
                 resources.append(resource)
+                if callable(progress_update):
+                    progress_update(1)
                 if current % progress_interval == 0 or current == total_versions:
                     self.config.logger.info(
                         "synthetics_mobile_applications_versions: progress %d/%d",

datadog_sync/utils/workers.py

@@ -95,13 +95,30 @@ async def schedule_workers(self, additional_coros: List = []) -> Future:
     async def schedule_workers_with_pbar(self, total, additional_coros: List = []) -> Future:
         self.pbar = tqdm(total=total)
 
-        self._shutdown_workers = False
-        with logging_redirect_tqdm():
-            additional_coros.append(self._refresh_pbar())
-            await self.schedule_workers(additional_coros)
+        # Allow long-running get_resources (e.g. mobile app versions) to update the bar
+        def _update(n: int = 1) -> None:
+            if self.pbar:
+                self.pbar.update(n)
 
-        self.pbar.close()
-        self.pbar = None
+        def _add_total(n: int) -> None:
+            if self.pbar:
+                self.pbar.total += n
+                self.pbar.refresh()
+
+        setattr(self.config, "_import_progress_update", _update)
+        setattr(self.config, "_import_progress_add_total", _add_total)
+
+        self._shutdown_workers = False
+        try:
+            with logging_redirect_tqdm():
+                additional_coros.append(self._refresh_pbar())
+                await self.schedule_workers(additional_coros)
+        finally:
+            for attr in ("_import_progress_update", "_import_progress_add_total"):
+                if hasattr(self.config, attr):
+                    delattr(self.config, attr)
+            self.pbar.close()
+            self.pbar = None
 
 
 @dataclass