From 4663afcd975c64a14db4288ce504b850cf180dd1 Mon Sep 17 00:00:00 2001
From: Vladyslav Tarasenko <vladislavstargate@gmail.com>
Date: Wed, 19 Jun 2024 01:55:50 +0300
Subject: [PATCH] Format.

---
 src/backups.py                    | 33 ++++++++++++++-----------------
 tests/integration/test_backups.py |  3 ++-
 2 files changed, 17 insertions(+), 19 deletions(-)

diff --git a/src/backups.py b/src/backups.py
index 5e8ea06eb4..217c17b8aa 100644
--- a/src/backups.py
+++ b/src/backups.py
@@ -798,24 +798,23 @@ def _on_restore_action(self, event):
             f" has been requested on the unit"
         )
 
-        if backup_id:
-            # Validate the provided backup id.
-            logger.info("Validating provided backup-id")
-            try:
-                backups = self._list_backups(show_failed=False)
-                if backup_id not in backups.keys():
-                    error_message = f"Invalid backup-id: {backup_id}"
-                    logger.error(f"Restore failed: {error_message}")
-                    event.fail(error_message)
-                    return
-            except ListBackupsError as e:
-                logger.exception(e)
-                error_message = "Failed to retrieve backup id"
+        # Validate the provided backup id and restore to time.
+        logger.info("Validating provided backup-id and restore-to-time")
+        try:
+            backups = self._list_backups(show_failed=False)
+            if backup_id and backup_id not in backups.keys():
+                error_message = f"Invalid backup-id: {backup_id}"
                 logger.error(f"Restore failed: {error_message}")
                 event.fail(error_message)
                 return
-        elif not self._list_backups(show_failed=False):
-            error_message = "Cannot restore PITR without any backups created"
+            if not backup_id and restore_to_time and not backups:
+                error_message = "Cannot restore PITR without any backups created"
+                logger.error(f"Restore failed: {error_message}")
+                event.fail(error_message)
+                return
+        except ListBackupsError as e:
+            logger.exception(e)
+            error_message = "Failed to retrieve backups list"
             logger.error(f"Restore failed: {error_message}")
             event.fail(error_message)
             return
@@ -863,9 +862,7 @@ def _on_restore_action(self, event):
         # Mark the cluster as in a restoring backup state and update the Patroni configuration.
         logger.info("Configuring Patroni to restore the backup")
         self.charm.app_peer_data.update({
-            "restoring-backup": self._fetch_backup_from_id(backup_id)
-            if backup_id
-            else "",
+            "restoring-backup": self._fetch_backup_from_id(backup_id) if backup_id else "",
             "restore-stanza": backups[backup_id]
             if backup_id
             else self.charm.app_peer_data.get("stanza", self.stanza_name),
diff --git a/tests/integration/test_backups.py b/tests/integration/test_backups.py
index cf2048d310..8e25f23a8d 100644
--- a/tests/integration/test_backups.py
+++ b/tests/integration/test_backups.py
@@ -277,7 +277,8 @@ async def test_backup(ops_test: OpsTest, cloud_configs: Tuple[Dict, Dict], charm
         # Wait for the restore to complete.
         async with ops_test.fast_forward():
             await ops_test.model.block_until(
-                lambda: remaining_unit.workload_status_message == MOVE_RESTORED_CLUSTER_TO_ANOTHER_BUCKET,
+                lambda: remaining_unit.workload_status_message
+                == MOVE_RESTORED_CLUSTER_TO_ANOTHER_BUCKET,
                 timeout=1000,
             )