Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Cloud backup support #33

Merged
merged 4 commits into from
Dec 9, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 33 additions & 1 deletion aiohasupervisor/backups.py
Original file line number Diff line number Diff line change
@@ -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 (
Expand All @@ -15,6 +20,8 @@
NewBackup,
PartialBackupOptions,
PartialRestoreOptions,
UploadBackupOptions,
UploadedBackup,
)


Expand Down Expand Up @@ -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
77 changes: 43 additions & 34 deletions aiohasupervisor/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
ClientSession,
ClientTimeout,
)
from multidict import MultiDict
from yarl import URL

from .const import DEFAULT_TIMEOUT, ResponseType
Expand All @@ -27,6 +28,7 @@
SupervisorTimeoutError,
)
from .models.base import Response, ResultType
from .utils.aiohttp import ChunkAsyncStreamIterator

VERSION = metadata.version(__package__)

Expand All @@ -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,
Expand Down Expand Up @@ -94,42 +117,28 @@ async def _request(
self._close_session = True

try:
async with self.session.request(
response = await self.session.request(
Copy link
Contributor

@emontnemery emontnemery Dec 7, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure it matters, but if want to keep the context manager we could separate the cases where we can use the context manager from the cases we don't:

if response_type == ResponseType.RAW:
    response = await self.session.request(
        ...
    )
    ...
else:
    async with self.session.request(
        ...
    ) as response:
        ...

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(
Expand All @@ -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:
Expand All @@ -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,
Expand All @@ -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:
Expand All @@ -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."""
Expand Down
1 change: 1 addition & 0 deletions aiohasupervisor/const.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,4 +13,5 @@ class ResponseType(StrEnum):

NONE = "none"
JSON = "json"
STREAM = "stream"
TEXT = "text"
2 changes: 2 additions & 0 deletions aiohasupervisor/models/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@
NewBackup,
PartialBackupOptions,
PartialRestoreOptions,
UploadBackupOptions,
)
from aiohasupervisor.models.discovery import (
Discovery,
Expand Down Expand Up @@ -215,6 +216,7 @@
"NewBackup",
"PartialBackupOptions",
"PartialRestoreOptions",
"UploadBackupOptions",
"Discovery",
"DiscoveryConfig",
"AccessPoint",
Expand Down
24 changes: 21 additions & 3 deletions aiohasupervisor/models/backups.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
1 change: 1 addition & 0 deletions aiohasupervisor/models/mounts.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,7 @@ class MountResponse(ABC):
name: str
read_only: bool
state: MountState | None
user_path: PurePath | None


@dataclass(frozen=True)
Expand Down
1 change: 1 addition & 0 deletions aiohasupervisor/utils/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
"""Utilities used internally in library."""
31 changes: 31 additions & 0 deletions aiohasupervisor/utils/aiohttp.py
Original file line number Diff line number Diff line change
@@ -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]
7 changes: 6 additions & 1 deletion tests/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
5 changes: 4 additions & 1 deletion tests/fixtures/backup_info.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -27,6 +29,7 @@
"https://github.com/hassio-addons/repository"
],
"folders": [],
"homeassistant_exclude_database": null
"homeassistant_exclude_database": null,
"extra": null
}
}
5 changes: 4 additions & 1 deletion tests/fixtures/backup_info_no_homeassistant.json
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,13 @@
"name": "Studio Code Server",
"date": "2023-08-10T19:37:01.084215+00:00",
"size": 0.12,
"size_bytes": 120123,
"compressed": true,
"protected": false,
"supervisor_version": "2023.08.2.dev1002",
"homeassistant": null,
"location": "Test",
"locations": ["Test"],
"addons": [
{
"slug": "a0d7b954_vscode",
Expand All @@ -27,6 +29,7 @@
"https://github.com/esphome/home-assistant-addon"
],
"folders": [],
"homeassistant_exclude_database": null
"homeassistant_exclude_database": null,
"extra": null
}
}
Loading
Loading