Validate KB and backup bucket exist before a restore

nucliadb/src/nucliadb/backups/const.py

@@ -32,6 +32,7 @@ class StorageKeys:
     Defines the key templates used to store backup files in the backups bucket of the storage.
     """
 
+    BACKUP_FOLDER = "backups/{backup_id}"
     BACKUP_PREFIX = "backups/{backup_id}/"
     RESOURCES_PREFIX = "backups/{backup_id}/resources/"
     RESOURCE = "backups/{backup_id}/resources/{resource_id}.tar"

nucliadb/src/nucliadb/backups/restore.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-
 import asyncio
 import functools
 import json
@@ -32,6 +31,8 @@
 from nucliadb.backups.const import MaindbKeys, StorageKeys
 from nucliadb.backups.models import RestoreBackupRequest
 from nucliadb.backups.settings import settings
+from nucliadb.backups.utils import exists_backup
+from nucliadb.common import datamanagers
 from nucliadb.common.context import ApplicationContext
 from nucliadb.export_import.utils import (
     import_binary,
@@ -74,6 +75,20 @@ async def restore_kb(context: ApplicationContext, kbid: str, backup_id: str):
     """
     Downloads the backup files from the cloud storage and imports them into the KB.
     """
+    kb_exists = await datamanagers.atomic.kb.exists_kb(kbid=kbid)
+    if not kb_exists:
+        logger.warning(
+            "Trying to restore on a KB that doesn't exist",
+            extra={"kbid": kbid, "backup_id": backup_id},
+        )
+        return
+    if not await exists_backup(context.blob_storage, backup_id):
+        logger.warning(
+            "Trying to restore from a bucket that doesn't exist",
+            extra={"kbid": kbid, "backup_id": backup_id},
+        )
+        return
+
     await restore_resources(context, kbid, backup_id)
     await restore_labels(context, kbid, backup_id)
     await restore_synonyms(context, kbid, backup_id)

nucliadb/tests/nucliadb/integration/test_backups.py

@@ -23,6 +23,7 @@
 
 import pytest
 from httpx import AsyncClient
+from pytest import LogCaptureFixture
 
 from nucliadb.backups.const import StorageKeys
 from nucliadb.backups.create import backup_kb_task, get_metadata, set_metadata
@@ -379,3 +380,29 @@ def test_all_broker_message_fields_are_backed_up():
             or field_name in BM_FIELDS["common"]
             or field_name in BM_FIELDS["deprecated"]
         ), f"Field {field_name} is not being taken into account in the backup!"
+
+
+@pytest.mark.deploy_modes("standalone")
+async def test_restore_with_nonexistent_kb(
+    nucliadb_reader: AsyncClient,
+    settings: BackupSettings,
+    context: ApplicationContext,
+    caplog: LogCaptureFixture,
+):
+    fake_kbid = str(uuid.uuid4())
+    fake_backup_id = str(uuid.uuid4())
+    await restore_kb_task(context, RestoreBackupRequest(kb_id=fake_kbid, backup_id=fake_backup_id))
+    assert caplog.records[0].message == "Trying to restore on a KB that doesn't exist"
+
+
+@pytest.mark.deploy_modes("standalone")
+async def test_restore_with_nonexistent_backup_bucket(
+    nucliadb_reader: AsyncClient,
+    dst_kb: str,
+    settings: BackupSettings,
+    context: ApplicationContext,
+    caplog: LogCaptureFixture,
+):
+    fake_backup_id = str(uuid.uuid4())
+    await restore_kb_task(context, RestoreBackupRequest(kb_id=dst_kb, backup_id=fake_backup_id))
+    assert caplog.records[0].message == "Trying to restore from a bucket that doesn't exist"