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