From 0c0e98fb0cb40c191fa63ecc61b345a745c35194 Mon Sep 17 00:00:00 2001
From: Rommel Layco <rommel.layco@aiven.io>
Date: Mon, 21 Oct 2024 11:14:34 +0100
Subject: [PATCH] restore: Make max stale seconds configurable

---
 pghoard/restore.py | 24 ++++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

diff --git a/pghoard/restore.py b/pghoard/restore.py
index 4cfe839f..4eb7b873 100644
--- a/pghoard/restore.py
+++ b/pghoard/restore.py
@@ -275,6 +275,12 @@ def target_args():
                 dest="cancel_preserve_on_success",
                 action="store_false",
             )
+            cmd.add_argument(
+                "--max-stale-seconds",
+                help="Maximum time in seconds a download can stall before aborting",
+                type=int,
+                default=120
+            )
 
         cmd = add_cmd(self.list_basebackups_http)
         host_port_user_args()
@@ -341,6 +347,7 @@ def get_basebackup(self, arg):
                 tablespace_base_dir=arg.tablespace_base_dir,
                 preserve_until=arg.preserve_until,
                 cancel_preserve_on_success=arg.cancel_preserve_on_success,
+                max_stale_seconds=arg.max_stale_seconds,
             )
         except RestoreError:  # pylint: disable=try-except-raise
             # Pass RestoreErrors thru
@@ -462,6 +469,7 @@ def _get_basebackup(
         tablespace_base_dir=None,
         preserve_until: Optional[str] = None,
         cancel_preserve_on_success: bool = True,
+        max_stale_seconds: int = 120,
     ):
         targets = [recovery_target_name, recovery_target_time, recovery_target_xid]
         if sum(0 if flag is None else 1 for flag in targets) > 1:
@@ -613,6 +621,7 @@ def _get_basebackup(
             pgdata=pgdata,
             site=site,
             tablespaces=tablespaces,
+            max_stale_seconds=max_stale_seconds
         )
         fetcher.fetch_all()
 
@@ -657,7 +666,18 @@ def run(self, args=None):
 
 
 class BasebackupFetcher:
-    def __init__(self, *, app_config, debug, site, pgdata, tablespaces, data_files: List[FileInfo], status_output_file=None):
+    def __init__(
+        self,
+        *,
+        app_config,
+        debug,
+        site,
+        pgdata,
+        tablespaces,
+        data_files: List[FileInfo],
+        status_output_file=None,
+        max_stale_seconds=120
+    ):
         self.log = logging.getLogger(self.__class__.__name__)
         self.completed_jobs: Set[str] = set()
         self.config = app_config
@@ -669,7 +689,7 @@ def __init__(self, *, app_config, debug, site, pgdata, tablespaces, data_files:
         self.last_total_downloaded = 0
         self.lock = RLock()
         self.manager_class = multiprocessing.Manager if self._process_count() > 1 else ThreadingManager
-        self.max_stale_seconds = 120
+        self.max_stale_seconds = max_stale_seconds
         self.pending_jobs: Set[str] = set()
         self.jobs_to_retry: Set[str] = set()
         self.pgdata = pgdata