Allow backups to report the upload progress (#163608)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

Author: zweckj

homeassistant/components/aws_s3/backup.py

@@ -14,6 +14,7 @@
     BackupAgent,
     BackupAgentError,
     BackupNotFound,
+    OnProgressCallback,
     suggested_filename,
 )
 from homeassistant.core import HomeAssistant, callback
@@ -132,6 +133,7 @@ async def async_upload_backup(
         *,
         open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]],
         backup: AgentBackup,
+        on_progress: OnProgressCallback,
         **kwargs: Any,
     ) -> None:
         """Upload a backup.

homeassistant/components/azure_storage/backup.py

@@ -16,6 +16,7 @@
     BackupAgent,
     BackupAgentError,
     BackupNotFound,
+    OnProgressCallback,
     suggested_filename,
 )
 from homeassistant.core import HomeAssistant, callback
@@ -129,6 +130,7 @@ async def async_upload_backup(
         *,
         open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]],
         backup: AgentBackup,
+        on_progress: OnProgressCallback,
         **kwargs: Any,
     ) -> None:
         """Upload a backup."""

homeassistant/components/backblaze_b2/backup.py

@@ -17,6 +17,7 @@
     BackupAgent,
     BackupAgentError,
     BackupNotFound,
+    OnProgressCallback,
     suggested_filename,
 )
 from homeassistant.core import HomeAssistant, callback
@@ -230,6 +231,7 @@ async def async_upload_backup(
         *,
         open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]],
         backup: AgentBackup,
+        on_progress: OnProgressCallback,
         **kwargs: Any,
     ) -> None:
         """Upload a backup to Backblaze B2.

homeassistant/components/backup/__init__.py

@@ -17,6 +17,7 @@
     BackupAgentError,
     BackupAgentPlatformProtocol,
     LocalBackupAgent,
+    OnProgressCallback,
 )
 from .config import BackupConfig, CreateBackupParametersDict
 from .const import DATA_MANAGER, DOMAIN
@@ -41,6 +42,7 @@
     RestoreBackupEvent,
     RestoreBackupStage,
     RestoreBackupState,
+    UploadBackupEvent,
     WrittenBackup,
 )
 from .models import AddonInfo, AgentBackup, BackupNotFound, Folder
@@ -72,9 +74,11 @@
     "LocalBackupAgent",
     "ManagerBackup",
     "NewBackup",
+    "OnProgressCallback",
     "RestoreBackupEvent",
     "RestoreBackupStage",
     "RestoreBackupState",
+    "UploadBackupEvent",
     "WrittenBackup",
     "async_get_manager",
     "suggested_filename",

homeassistant/components/backup/agent.py

@@ -14,6 +14,13 @@
 from .models import AgentBackup, BackupAgentError
 
 
+class OnProgressCallback(Protocol):
+    """Protocol for on_progress callback."""
+
+    def __call__(self, *, bytes_uploaded: int, **kwargs: Any) -> None:
+        """Report upload progress."""
+
+
 class BackupAgentUnreachableError(BackupAgentError):
     """Raised when the agent can't reach its API."""
 
@@ -53,12 +60,14 @@ async def async_upload_backup(
         *,
         open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]],
         backup: AgentBackup,
+        on_progress: OnProgressCallback,
         **kwargs: Any,
     ) -> None:
         """Upload a backup.
 
         :param open_stream: A function returning an async iterator that yields bytes.
         :param backup: Metadata about the backup that should be uploaded.
+        :param on_progress: A callback to report the number of uploaded bytes.
         """
 
     @abc.abstractmethod

homeassistant/components/backup/backup.py

@@ -11,7 +11,7 @@
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.hassio import is_hassio
 
-from .agent import BackupAgent, LocalBackupAgent
+from .agent import BackupAgent, LocalBackupAgent, OnProgressCallback
 from .const import DOMAIN, LOGGER
 from .models import AgentBackup, BackupNotFound
 from .util import read_backup, suggested_filename
@@ -73,6 +73,7 @@ async def async_upload_backup(
         *,
         open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]],
         backup: AgentBackup,
+        on_progress: OnProgressCallback,
         **kwargs: Any,
     ) -> None:
         """Upload a backup."""

homeassistant/components/backup/manager.py

@@ -252,6 +252,15 @@ class BlockedEvent(ManagerStateEvent):
     manager_state: BackupManagerState = BackupManagerState.BLOCKED
 
 
+@dataclass(frozen=True, kw_only=True, slots=True)
+class UploadBackupEvent(ManagerStateEvent):
+    """Backup agent upload progress event."""
+
+    agent_id: str
+    uploaded_bytes: int
+    total_bytes: int
+
+
 class BackupPlatformProtocol(Protocol):
     """Define the format that backup platforms can have."""
 
@@ -579,9 +588,24 @@ async def upload_backup_to_agent(agent_id: str) -> None:
                 _backup = replace(
                     backup, protected=should_encrypt, size=streamer.size()
                 )
-            await self.backup_agents[agent_id].async_upload_backup(
+            agent = self.backup_agents[agent_id]
+
+            @callback
+            def on_upload_progress(*, bytes_uploaded: int, **kwargs: Any) -> None:
+                """Handle upload progress."""
+                self.async_on_backup_event(
+                    UploadBackupEvent(
+                        manager_state=self.state,
+                        agent_id=agent_id,
+                        uploaded_bytes=bytes_uploaded,
+                        total_bytes=_backup.size,
+                    )
+                )
+
+            await agent.async_upload_backup(
                 open_stream=open_stream_func,
                 backup=_backup,
+                on_progress=on_upload_progress,
             )
             if streamer:
                 await streamer.wait()
@@ -1374,9 +1398,10 @@ def async_on_backup_event(
         """Forward event to subscribers."""
         if (current_state := self.state) != (new_state := event.manager_state):
             LOGGER.debug("Backup state: %s -> %s", current_state, new_state)
-        self.last_event = event
-        if not isinstance(event, (BlockedEvent, IdleEvent)):
-            self.last_action_event = event
+        if not isinstance(event, UploadBackupEvent):
+            self.last_event = event
+            if not isinstance(event, (BlockedEvent, IdleEvent)):
+                self.last_action_event = event
         for subscription in self._backup_event_subscriptions:
             subscription(event)
 

homeassistant/components/cloud/backup.py

@@ -18,6 +18,7 @@
     BackupAgent,
     BackupAgentError,
     BackupNotFound,
+    OnProgressCallback,
 )
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.aiohttp_client import ChunkAsyncStreamIterator
@@ -106,6 +107,7 @@ async def async_upload_backup(
         *,
         open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]],
         backup: AgentBackup,
+        on_progress: OnProgressCallback,
         **kwargs: Any,
     ) -> None:
         """Upload a backup.

homeassistant/components/cloudflare_r2/backup.py

@@ -14,6 +14,7 @@
     BackupAgent,
     BackupAgentError,
     BackupNotFound,
+    OnProgressCallback,
     suggested_filename,
 )
 from homeassistant.core import HomeAssistant, callback
@@ -129,6 +130,7 @@ async def async_upload_backup(
         *,
         open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]],
         backup: AgentBackup,
+        on_progress: OnProgressCallback,
         **kwargs: Any,
     ) -> None:
         """Upload a backup.

homeassistant/components/google_drive/backup.py

@@ -13,6 +13,7 @@
     BackupAgent,
     BackupAgentError,
     BackupNotFound,
+    OnProgressCallback,
 )
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
@@ -75,6 +76,7 @@ async def async_upload_backup(
         *,
         open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]],
         backup: AgentBackup,
+        on_progress: OnProgressCallback,
         **kwargs: Any,
     ) -> None:
         """Upload a backup.