Add location to backup download and remove APIs (#5482)

Author: mdegat01

supervisor/api/backups.py

@@ -126,6 +126,14 @@ def _ensure_list(item: Any) -> list:
     }
 )
 
+SCHEMA_REMOVE = vol.Schema(
+    {
+        vol.Optional(ATTR_LOCATION): vol.All(
+            _ensure_list, [vol.Maybe(str)], vol.Unique()
+        ),
+    }
+)
+
 
 class APIBackups(CoreSysAttributes):
     """Handle RESTful API for backups functions."""
@@ -411,17 +419,29 @@ async def thaw(self, request: web.Request):
     async def remove(self, request: web.Request):
         """Remove a backup."""
         backup = self._extract_slug(request)
-        self._validate_cloud_backup_location(request, backup.location)
-        return self.sys_backups.remove(backup)
+        body = await api_validate(SCHEMA_REMOVE, request)
+        locations: list[LOCATION_TYPE] | None = None
+
+        if ATTR_LOCATION in body:
+            self._validate_cloud_backup_location(request, body[ATTR_LOCATION])
+            locations = [self._location_to_mount(name) for name in body[ATTR_LOCATION]]
+        else:
+            self._validate_cloud_backup_location(request, backup.location)
+
+        return self.sys_backups.remove(backup, locations=locations)
 
     @api_process
     async def download(self, request: web.Request):
         """Download a backup file."""
         backup = self._extract_slug(request)
-        self._validate_cloud_backup_location(request, backup.location)
+        # Query will give us '' for /backups, convert value to None
+        location = request.query.get(ATTR_LOCATION, backup.location) or None
+        self._validate_cloud_backup_location(request, location)
+        if location not in backup.all_locations:
+            raise APIError(f"Backup {backup.slug} is not in location {location}")
 
         _LOGGER.info("Downloading backup %s", backup.slug)
-        response = web.FileResponse(backup.tarfile)
+        response = web.FileResponse(backup.all_locations[location])
         response.content_type = CONTENT_TYPE_TAR
         response.headers[CONTENT_DISPOSITION] = (
             f"attachment; filename={RE_SLUGIFY_NAME.sub('_', backup.name)}.tar"

tests/api/test_backups.py

@@ -747,3 +747,65 @@ async def test_backup_not_found(api_client: TestClient, method: str, url: str):
     assert resp.status == 404
     resp = await resp.json()
     assert resp["message"] == "Backup does not exist"
+
+
+@pytest.mark.usefixtures("tmp_supervisor_data")
+async def test_remove_backup_from_location(api_client: TestClient, coresys: CoreSys):
+    """Test removing a backup from one location of multiple."""
+    backup_file = get_fixture_path("backup_example.tar")
+    location_1 = Path(copy(backup_file, coresys.config.path_backup))
+    location_2 = Path(copy(backup_file, coresys.config.path_core_backup))
+
+    await coresys.backups.reload()
+    assert (backup := coresys.backups.get("7fed74c8"))
+    assert backup.all_locations == {None: location_1, ".cloud_backup": location_2}
+
+    resp = await api_client.delete(
+        "/backups/7fed74c8", json={"location": ".cloud_backup"}
+    )
+    assert resp.status == 200
+
+    assert location_1.exists()
+    assert not location_2.exists()
+    assert coresys.backups.get("7fed74c8")
+    assert backup.all_locations == {None: location_1}
+
+
+async def test_download_backup_from_location(
+    api_client: TestClient, coresys: CoreSys, tmp_supervisor_data: Path
+):
+    """Test downloading a backup from a specific location."""
+    backup_file = get_fixture_path("backup_example.tar")
+    location_1 = Path(copy(backup_file, coresys.config.path_backup))
+    location_2 = Path(copy(backup_file, coresys.config.path_core_backup))
+
+    await coresys.backups.reload()
+    assert (backup := coresys.backups.get("7fed74c8"))
+    assert backup.all_locations == {None: location_1, ".cloud_backup": location_2}
+
+    # The use case of this is user might want to pick a particular mount if one is flaky
+    # To simulate this, remove the file from one location and show one works and the other doesn't
+    assert backup.location is None
+    location_1.unlink()
+
+    resp = await api_client.get("/backups/7fed74c8/download?location=")
+    assert resp.status == 404
+
+    resp = await api_client.get("/backups/7fed74c8/download?location=.cloud_backup")
+    assert resp.status == 200
+    out_file = tmp_supervisor_data / "backup_example.tar"
+    with out_file.open("wb") as out:
+        out.write(await resp.read())
+
+    out_backup = Backup(coresys, out_file, "out", None)
+    await out_backup.load()
+    assert backup == out_backup
+
+
+@pytest.mark.usefixtures("mock_full_backup")
+async def test_download_backup_from_invalid_location(api_client: TestClient):
+    """Test error for invalid download location."""
+    resp = await api_client.get("/backups/test/download?location=.cloud_backup")
+    assert resp.status == 400
+    body = await resp.json()
+    assert body["message"] == "Backup test is not in location .cloud_backup"

tests/backups/test_manager.py

@@ -1977,9 +1977,8 @@ async def test_partial_reload_multiple_locations(
     assert backup.all_locations.keys() == {".cloud_backup", None, "backup_test"}
 
 
-async def test_backup_remove_multiple_locations(
-    coresys: CoreSys, tmp_supervisor_data: Path
-):
+@pytest.mark.usefixtures("tmp_supervisor_data")
+async def test_backup_remove_multiple_locations(coresys: CoreSys):
     """Test removing a backup that exists in multiple locations."""
     backup_file = get_fixture_path("backup_example.tar")
     location_1 = Path(copy(backup_file, coresys.config.path_backup))
@@ -1995,9 +1994,8 @@ async def test_backup_remove_multiple_locations(
     assert not coresys.backups.get("7fed74c8")
 
 
-async def test_backup_remove_one_location_of_multiple(
-    coresys: CoreSys, tmp_supervisor_data: Path
-):
+@pytest.mark.usefixtures("tmp_supervisor_data")
+async def test_backup_remove_one_location_of_multiple(coresys: CoreSys):
     """Test removing a backup that exists in multiple locations from one location."""
     backup_file = get_fixture_path("backup_example.tar")
     location_1 = Path(copy(backup_file, coresys.config.path_backup))