Cloud backup support (#33)

* Add upload and download backup APIs

* Add model changes

* Fix fixtures

* Apply suggestions from code review

---------

Co-authored-by: Stefan Agner <[email protected]>

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] | 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] | 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] | 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] | 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] | 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

@@ -47,7 +47,9 @@ class BackupBaseFields(ABC):
     date: datetime
     type: BackupType
     size: float
+    size_bytes: int
     location: str | None
+    locations: set[str | None]
     protected: bool
     compressed: bool
 
@@ -73,12 +75,13 @@ 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]
     homeassistant_exclude_database: bool | None
+    extra: dict | None
 
 
 @dataclass(frozen=True, slots=True)
@@ -132,9 +135,10 @@ class FullBackupOptions(Request):
     name: str | None = None
     password: str | None = None
     compressed: bool | None = None
-    location: str | None = None
+    location: set[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 +171,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/models/mounts.py

@@ -79,6 +79,7 @@ class MountResponse(ABC):
     name: str
     read_only: bool
     state: MountState | None
+    user_path: PurePath | None
 
 
 @dataclass(frozen=True)

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_info.json

@@ -6,11 +6,13 @@
     "name": "addon_core_mosquitto_6.4.0",
     "date": "2024-05-31T00:00:00.000000+00:00",
     "size": 0.01,
+    "size_bytes": 10123,
     "compressed": true,
     "protected": false,
     "supervisor_version": "2024.05.0",
     "homeassistant": null,
     "location": null,
+    "locations": [null],
     "addons": [
       {
         "slug": "core_mosquitto",
@@ -27,6 +29,7 @@
       "https://github.com/hassio-addons/repository"
     ],
     "folders": [],
-    "homeassistant_exclude_database": null
+    "homeassistant_exclude_database": null,
+    "extra": null
   }
 }