Add upload and download backup APIs

Author: mdegat01

aiohasupervisor/backups.py

@@ -1,5 +1,10 @@
 """Backups client for supervisor."""
 
+from collections.abc import AsyncIterator
+
+from aiohttp import MultipartWriter
+from multidict import MultiDict
+
 from .client import _SupervisorComponentClient
 from .const import ResponseType
 from .models.backups import (
@@ -15,6 +20,8 @@
     NewBackup,
     PartialBackupOptions,
     PartialRestoreOptions,
+    UploadBackupOptions,
+    UploadedBackup,
 )
 
 
@@ -102,4 +109,29 @@ async def partial_restore(
         )
         return BackupJob.from_dict(result.data)
 
-    # Omitted for now - Upload and download backup
+    async def upload_backup(
+        self, stream: AsyncIterator[bytes], options: UploadBackupOptions | None = None
+    ) -> str:
+        """Upload backup by stream and return slug."""
+        params = MultiDict()
+        if options and options.location:
+            for location in options.location:
+                params.add("location", location or "")
+
+        with MultipartWriter("form-data") as mp:
+            mp.append(stream)
+            result = await self._client.post(
+                "backups/new/upload",
+                params=params,
+                data=mp,
+                response_type=ResponseType.JSON,
+            )
+
+        return UploadedBackup.from_dict(result.data).slug
+
+    async def download_backup(self, backup: str) -> AsyncIterator[bytes]:
+        """Download backup and return stream."""
+        result = await self._client.get(
+            f"backups/{backup}/download", response_type=ResponseType.STREAM
+        )
+        return result.data

aiohasupervisor/client.py

@@ -12,6 +12,7 @@
     ClientSession,
     ClientTimeout,
 )
+from multidict import MultiDict
 from yarl import URL
 
 from .const import DEFAULT_TIMEOUT, ResponseType
@@ -27,6 +28,7 @@
     SupervisorTimeoutError,
 )
 from .models.base import Response, ResultType
+from .utils.aiohttp import ChunkAsyncStreamIterator
 
 VERSION = metadata.version(__package__)
 
@@ -53,12 +55,33 @@ class _SupervisorClient:
     session: ClientSession | None = None
     _close_session: bool = field(default=False, init=False)
 
+    async def _raise_on_status(self, response: ClientResponse) -> None:
+        """Raise appropriate exception on status."""
+        if response.status >= HTTPStatus.BAD_REQUEST.value:
+            exc_type: type[SupervisorError] = SupervisorError
+            match response.status:
+                case HTTPStatus.BAD_REQUEST:
+                    exc_type = SupervisorBadRequestError
+                case HTTPStatus.UNAUTHORIZED:
+                    exc_type = SupervisorAuthenticationError
+                case HTTPStatus.FORBIDDEN:
+                    exc_type = SupervisorForbiddenError
+                case HTTPStatus.NOT_FOUND:
+                    exc_type = SupervisorNotFoundError
+                case HTTPStatus.SERVICE_UNAVAILABLE:
+                    exc_type = SupervisorServiceUnavailableError
+
+            if is_json(response):
+                result = Response.from_json(await response.text())
+                raise exc_type(result.message, result.job_id)
+            raise exc_type()
+
     async def _request(
         self,
         method: HTTPMethod,
         uri: str,
         *,
-        params: dict[str, str] | None,
+        params: dict[str, str] | MultiDict[str, str] | None,
         response_type: ResponseType,
         json: dict[str, Any] | None = None,
         data: Any = None,
@@ -94,42 +117,28 @@ async def _request(
             self._close_session = True
 
         try:
-            async with self.session.request(
+            response = await self.session.request(
                 method.value,
                 url,
                 timeout=timeout,
                 headers=headers,
                 params=params,
                 json=json,
                 data=data,
-            ) as response:
-                if response.status >= HTTPStatus.BAD_REQUEST.value:
-                    exc_type: type[SupervisorError] = SupervisorError
-                    match response.status:
-                        case HTTPStatus.BAD_REQUEST:
-                            exc_type = SupervisorBadRequestError
-                        case HTTPStatus.UNAUTHORIZED:
-                            exc_type = SupervisorAuthenticationError
-                        case HTTPStatus.FORBIDDEN:
-                            exc_type = SupervisorForbiddenError
-                        case HTTPStatus.NOT_FOUND:
-                            exc_type = SupervisorNotFoundError
-                        case HTTPStatus.SERVICE_UNAVAILABLE:
-                            exc_type = SupervisorServiceUnavailableError
-
-                    if is_json(response):
-                        result = Response.from_json(await response.text())
-                        raise exc_type(result.message, result.job_id)
-                    raise exc_type()
-
-                match response_type:
-                    case ResponseType.JSON:
-                        is_json(response, raise_on_fail=True)
-                        return Response.from_json(await response.text())
-                    case ResponseType.TEXT:
-                        return Response(ResultType.OK, await response.text())
-                    case _:
-                        return Response(ResultType.OK)
+            )
+            await self._raise_on_status(response)
+            match response_type:
+                case ResponseType.JSON:
+                    is_json(response, raise_on_fail=True)
+                    return Response.from_json(await response.text())
+                case ResponseType.TEXT:
+                    return Response(ResultType.OK, await response.text())
+                case ResponseType.STREAM:
+                    return Response(
+                        ResultType.OK, ChunkAsyncStreamIterator(response.content)
+                    )
+                case _:
+                    return Response(ResultType.OK)
 
         except (UnicodeDecodeError, ClientResponseError) as err:
             raise SupervisorResponseError(
@@ -146,7 +155,7 @@ async def get(
         self,
         uri: str,
         *,
-        params: dict[str, str] | None = None,
+        params: dict[str, str] | MultiDict[str, str] | None = None,
         response_type: ResponseType = ResponseType.JSON,
         timeout: ClientTimeout | None = DEFAULT_TIMEOUT,
     ) -> Response:
@@ -163,7 +172,7 @@ async def post(
         self,
         uri: str,
         *,
-        params: dict[str, str] | None = None,
+        params: dict[str, str] | MultiDict[str, str] | None = None,
         response_type: ResponseType = ResponseType.NONE,
         json: dict[str, Any] | None = None,
         data: Any = None,
@@ -184,7 +193,7 @@ async def put(
         self,
         uri: str,
         *,
-        params: dict[str, str] | None = None,
+        params: dict[str, str] | MultiDict[str, str] | None = None,
         json: dict[str, Any] | None = None,
         timeout: ClientTimeout | None = DEFAULT_TIMEOUT,
     ) -> Response:
@@ -202,7 +211,7 @@ async def delete(
         self,
         uri: str,
         *,
-        params: dict[str, str] | None = None,
+        params: dict[str, str] | MultiDict[str, str] | None = None,
         timeout: ClientTimeout | None = DEFAULT_TIMEOUT,
     ) -> Response:
         """Handle a DELETE request to Supervisor."""

aiohasupervisor/const.py

@@ -13,4 +13,5 @@ class ResponseType(StrEnum):
 
     NONE = "none"
     JSON = "json"
+    STREAM = "stream"
     TEXT = "text"

aiohasupervisor/models/__init__.py

@@ -40,6 +40,7 @@
     NewBackup,
     PartialBackupOptions,
     PartialRestoreOptions,
+    UploadBackupOptions,
 )
 from aiohasupervisor.models.discovery import (
     Discovery,
@@ -215,6 +216,7 @@
     "NewBackup",
     "PartialBackupOptions",
     "PartialRestoreOptions",
+    "UploadBackupOptions",
     "Discovery",
     "DiscoveryConfig",
     "AccessPoint",

aiohasupervisor/models/backups.py

@@ -73,8 +73,8 @@ class BackupAddon(ResponseData):
 class BackupComplete(BackupBaseFields, ResponseData):
     """BackupComplete model."""
 
-    supervisor_version: str | None
-    homeassistant: str
+    supervisor_version: str
+    homeassistant: str | None
     addons: list[BackupAddon]
     repositories: list[str]
     folders: list[Folder]
@@ -132,9 +132,10 @@ class FullBackupOptions(Request):
     name: str | None = None
     password: str | None = None
     compressed: bool | None = None
-    location: str | None = None
+    location: list[str | None] | str | None = None
     homeassistant_exclude_database: bool | None = None
     background: bool | None = None
+    extra: dict | None = None
 
 
 @dataclass(frozen=True, slots=True)
@@ -167,3 +168,17 @@ class FullRestoreOptions(Request):
 @dataclass(frozen=True, slots=True)
 class PartialRestoreOptions(FullRestoreOptions, PartialBackupRestoreOptions):
     """PartialRestoreOptions model."""
+
+
+@dataclass(frozen=True, slots=True)
+class UploadBackupOptions(Request):
+    """UploadBackupOptions model."""
+
+    location: set[str | None] = None
+
+
+@dataclass(frozen=True, slots=True)
+class UploadedBackup(ResponseData):
+    """UploadedBackup model."""
+
+    slug: str

aiohasupervisor/utils/__init__.py

@@ -0,0 +1 @@
+"""Utilities used internally in library."""

aiohasupervisor/utils/aiohttp.py

@@ -0,0 +1,31 @@
+"""Utilities for interacting with aiohttp."""
+
+from typing import Self
+
+from aiohttp import StreamReader
+
+
+class ChunkAsyncStreamIterator:
+    """Async iterator for chunked streams.
+
+    Based on aiohttp.streams.ChunkTupleAsyncStreamIterator, but yields
+    bytes instead of tuple[bytes, bool].
+    Borrowed from home-assistant/core.
+    """
+
+    __slots__ = ("_stream",)
+
+    def __init__(self, stream: StreamReader) -> None:
+        """Initialize."""
+        self._stream = stream
+
+    def __aiter__(self) -> Self:
+        """Iterate."""
+        return self
+
+    async def __anext__(self) -> bytes:
+        """Yield next chunk."""
+        rv = await self._stream.readchunk()
+        if rv == (b"", False):
+            raise StopAsyncIteration
+        return rv[0]

tests/__init__.py

@@ -3,7 +3,12 @@
 from pathlib import Path
 
 
+def get_fixture_path(filename: str) -> Path:
+    """Get fixture path."""
+    return Path(__package__) / "fixtures" / filename
+
+
 def load_fixture(filename: str) -> str:
     """Load a fixture."""
-    fixture = Path(__package__) / "fixtures" / filename
+    fixture = get_fixture_path(filename)
     return fixture.read_text(encoding="utf-8")

tests/fixtures/backup_uploaded.json

@@ -0,0 +1,4 @@
+{
+  "result": "ok",
+  "data": { "slug": "7fed74c8" }
+}

tests/test_backups.py

@@ -1,5 +1,7 @@
 """Test backups supervisor client."""
 
+import asyncio
+from collections.abc import AsyncIterator
 from datetime import UTC, datetime
 from typing import Any
 
@@ -15,6 +17,7 @@
     FullBackupOptions,
     PartialBackupOptions,
     PartialRestoreOptions,
+    UploadBackupOptions,
 )
 
 from . import load_fixture
@@ -265,3 +268,52 @@ async def test_partial_restore(
         "abc123", PartialRestoreOptions(addons={"core_ssh"})
     )
     assert result.job_id == "dc9dbc16f6ad4de592ffa72c807ca2bf"
+
+
+async def test_upload_backup(
+    responses: aioresponses, supervisor_client: SupervisorClient
+) -> None:
+    """Test upload backup API."""
+    responses.post(
+        f"{SUPERVISOR_URL}/backups/new/upload",
+        status=200,
+        body=load_fixture("backup_uploaded.json"),
+    )
+    data = asyncio.StreamReader(loop=asyncio.get_running_loop())
+    data.feed_data(b"backup test")
+    data.feed_eof()
+
+    result = await supervisor_client.backups.upload_backup(data)
+    assert result == "7fed74c8"
+
+
+async def test_upload_backup_to_locations(
+    responses: aioresponses, supervisor_client: SupervisorClient
+) -> None:
+    """Test upload backup API with multiple locations."""
+    responses.post(
+        f"{SUPERVISOR_URL}/backups/new/upload?location=&location=test",
+        status=200,
+        body=load_fixture("backup_uploaded.json"),
+    )
+    data = asyncio.StreamReader(loop=asyncio.get_running_loop())
+    data.feed_data(b"backup test")
+    data.feed_eof()
+
+    result = await supervisor_client.backups.upload_backup(
+        data, UploadBackupOptions(location={None, "test"})
+    )
+    assert result == "7fed74c8"
+
+
+async def test_download_backup(
+    responses: aioresponses, supervisor_client: SupervisorClient
+) -> None:
+    """Test download backup API."""
+    responses.get(
+        f"{SUPERVISOR_URL}/backups/7fed74c8/download", status=200, body=b"backup test"
+    )
+    result = await supervisor_client.backups.download_backup("7fed74c8")
+    assert isinstance(result, AsyncIterator)
+    async for chunk in result:
+        assert chunk == b"backup test"