Add sensors to Google Drive

homeassistant/components/google_drive/__init__.py

@@ -6,7 +6,7 @@
 
 from google_drive_api.exceptions import GoogleDriveApiError
 
-from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import instance_id
@@ -19,13 +19,13 @@
 
 from .api import AsyncConfigEntryAuth, DriveClient
 from .const import DOMAIN
+from .coordinator import GoogleDriveConfigEntry, GoogleDriveDataUpdateCoordinator
 
 DATA_BACKUP_AGENT_LISTENERS: HassKey[list[Callable[[], None]]] = HassKey(
     f"{DOMAIN}.backup_agent_listeners"
 )
 
-
-type GoogleDriveConfigEntry = ConfigEntry[DriveClient]
+_PLATFORMS = (Platform.SENSOR,)
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: GoogleDriveConfigEntry) -> bool:
@@ -41,11 +41,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: GoogleDriveConfigEntry)
     await auth.async_get_access_token()
 
     client = DriveClient(await instance_id.async_get(hass), auth)
-    entry.runtime_data = client
 
     # Test we can access Google Drive and raise if not
     try:
-        await client.async_create_ha_root_folder_if_not_exists()
+        folder_id, _ = await client.async_create_ha_root_folder_if_not_exists()
     except GoogleDriveApiError as err:
         raise ConfigEntryNotReady from err
 
@@ -55,11 +54,20 @@ def async_notify_backup_listeners() -> None:
 
     entry.async_on_unload(entry.async_on_state_change(async_notify_backup_listeners))
 
+    entry.runtime_data = GoogleDriveDataUpdateCoordinator(
+        hass, entry=entry, client=client, backup_folder_id=folder_id
+    )
+    await entry.runtime_data.async_config_entry_first_refresh()
+
+    await hass.config_entries.async_forward_entry_setups(entry, _PLATFORMS)
+
     return True
 
 
 async def async_unload_entry(
     hass: HomeAssistant, entry: GoogleDriveConfigEntry
 ) -> bool:
     """Unload a config entry."""
+    await hass.config_entries.async_unload_platforms(entry, _PLATFORMS)
+
     return True

homeassistant/components/google_drive/api.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from collections.abc import AsyncIterator, Callable, Coroutine
+from dataclasses import dataclass
 import json
 import logging
 from typing import Any
@@ -27,6 +28,16 @@
 _LOGGER = logging.getLogger(__name__)
 
 
+@dataclass
+class StorageQuotaData:
+    """Class to represent storage quota data."""
+
+    limit: int | None
+    usage: int
+    usage_in_drive: int
+    usage_in_trash: int
+
+
 class AsyncConfigEntryAuth(AbstractAuth):
     """Provide Google Drive authentication tied to an OAuth2 based config entry."""
 
@@ -95,6 +106,19 @@ async def async_get_email_address(self) -> str:
         res = await self._api.get_user(params={"fields": "user(emailAddress)"})
         return str(res["user"]["emailAddress"])
 
+    async def async_get_storage_quota(self) -> StorageQuotaData:
+        """Get storage quota of the current user."""
+        res = await self._api.get_user(params={"fields": "storageQuota"})
+
+        storageQuota = res["storageQuota"]
+        limit = storageQuota.get("limit")
+        return StorageQuotaData(
+            limit=int(limit) if limit is not None else None,
+            usage=int(storageQuota.get("usage", 0)),
+            usage_in_drive=int(storageQuota.get("usageInDrive", 0)),
+            usage_in_trash=int(storageQuota.get("usageInTrash", 0)),
+        )
+
     async def async_create_ha_root_folder_if_not_exists(self) -> tuple[str, str]:
         """Create Home Assistant folder if it doesn't exist."""
         fields = "id,name"
@@ -178,6 +202,12 @@ async def async_list_backups(self) -> list[AgentBackup]:
             backups.append(backup)
         return backups
 
+    async def async_get_size_of_all_backups(self) -> int:
+        """Get size of all backups."""
+        backups = await self.async_list_backups()
+
+        return sum(backup.size for backup in backups)
+
     async def async_get_backup_file_id(self, backup_id: str) -> str | None:
         """Get file_id of backup if it exists."""
         query = " and ".join(

homeassistant/components/google_drive/backup.py

@@ -68,7 +68,7 @@ def __init__(self, config_entry: GoogleDriveConfigEntry) -> None:
         assert config_entry.unique_id
         self.name = config_entry.title
         self.unique_id = slugify(config_entry.unique_id)
-        self._client = config_entry.runtime_data
+        self._client = config_entry.runtime_data.client
 
     async def async_upload_backup(
         self,

homeassistant/components/google_drive/config_flow.py

@@ -14,10 +14,9 @@
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
 from .api import AsyncConfigFlowAuth, DriveClient
-from .const import DOMAIN
+from .const import DOMAIN, DRIVE_FOLDER_URL_PREFIX
 
 DEFAULT_NAME = "Google Drive"
-DRIVE_FOLDER_URL_PREFIX = "https://drive.google.com/drive/folders/"
 OAUTH2_SCOPES = [
     "https://www.googleapis.com/auth/drive.file",
 ]

homeassistant/components/google_drive/const.py

@@ -2,4 +2,9 @@
 
 from __future__ import annotations
 
+from datetime import timedelta
+
 DOMAIN = "google_drive"
+
+SCAN_INTERVAL = timedelta(hours=6)
+DRIVE_FOLDER_URL_PREFIX = "https://drive.google.com/drive/folders/"

homeassistant/components/google_drive/coordinator.py

@@ -0,0 +1,76 @@
+"""DataUpdateCoordinator for Google Drive."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import logging
+
+from google_drive_api.exceptions import GoogleDriveApiError
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .api import DriveClient, StorageQuotaData
+from .const import DOMAIN, SCAN_INTERVAL
+
+type GoogleDriveConfigEntry = ConfigEntry[GoogleDriveDataUpdateCoordinator]
+
+_LOGGER = logging.getLogger(__name__)
+
+
+@dataclass
+class SensorData:
+    """Class to represent sensor data."""
+
+    storage_quota: StorageQuotaData
+    all_backups_size: int
+
+
+class GoogleDriveDataUpdateCoordinator(DataUpdateCoordinator[SensorData]):
+    """Class to manage fetching Google Drive data from single endpoint."""
+
+    client: DriveClient
+    config_entry: GoogleDriveConfigEntry
+    email_address: str
+    backup_folder_id: str
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        *,
+        client: DriveClient,
+        backup_folder_id: str,
+        entry: GoogleDriveConfigEntry,
+    ) -> None:
+        """Initialize Google Drive data updater."""
+        self.client = client
+        self.backup_folder_id = backup_folder_id
+
+        super().__init__(
+            hass,
+            _LOGGER,
+            config_entry=entry,
+            name=DOMAIN,
+            update_interval=SCAN_INTERVAL,
+        )
+
+    async def _async_setup(self) -> None:
+        """Do initialization logic."""
+        self.email_address = await self.client.async_get_email_address()
+
+    async def _async_update_data(self) -> SensorData:
+        """Fetch data from Google Drive."""
+        try:
+            storage_quota = await self.client.async_get_storage_quota()
+            all_backups_size = await self.client.async_get_size_of_all_backups()
+            return SensorData(
+                storage_quota=storage_quota,
+                all_backups_size=all_backups_size,
+            )
+        except GoogleDriveApiError as error:
+            raise UpdateFailed(
+                translation_domain=DOMAIN,
+                translation_key="invalid_response_google_drive_error",
+                translation_placeholders={"error": str(error)},
+            ) from error

homeassistant/components/google_drive/diagnostics.py

@@ -0,0 +1,48 @@
+"""Diagnostics support for Google Drive."""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Any
+
+from homeassistant.components.backup import (
+    DATA_MANAGER as BACKUP_DATA_MANAGER,
+    BackupManager,
+)
+from homeassistant.components.diagnostics import async_redact_data
+from homeassistant.const import CONF_ACCESS_TOKEN
+from homeassistant.core import HomeAssistant
+
+from .const import DOMAIN
+from .coordinator import GoogleDriveConfigEntry
+
+TO_REDACT = (CONF_ACCESS_TOKEN, "refresh_token")
+
+
+async def async_get_config_entry_diagnostics(
+    hass: HomeAssistant,
+    entry: GoogleDriveConfigEntry,
+) -> dict[str, Any]:
+    """Return diagnostics for a config entry."""
+
+    coordinator = entry.runtime_data
+    backup_manager: BackupManager = hass.data[BACKUP_DATA_MANAGER]
+
+    backups = await coordinator.client.async_list_backups()
+
+    data = {
+        "coordinator_data": dataclasses.asdict(coordinator.data),
+        "config": {
+            **entry.data,
+            **entry.options,
+        },
+        "backup_folder_id": coordinator.backup_folder_id,
+        "backup_agents": [
+            {"name": agent.name}
+            for agent in backup_manager.backup_agents.values()
+            if agent.domain == DOMAIN
+        ],
+        "backup": [backup.as_dict() for backup in backups],
+    }
+
+    return async_redact_data(data, TO_REDACT)

homeassistant/components/google_drive/entity.py

@@ -0,0 +1,25 @@
+"""Define the Google Drive entity."""
+
+from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
+
+from .const import DOMAIN, DRIVE_FOLDER_URL_PREFIX
+from .coordinator import GoogleDriveDataUpdateCoordinator
+
+
+class GoogleDriveEntity(CoordinatorEntity[GoogleDriveDataUpdateCoordinator]):
+    """Defines a base Google Drive entity."""
+
+    _attr_has_entity_name = True
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return device information about this Google Drive device."""
+        return DeviceInfo(
+            identifiers={(DOMAIN, str(self.coordinator.config_entry.unique_id))},
+            name=self.coordinator.email_address,
+            manufacturer="Google",
+            model="Google Drive",
+            configuration_url=f"{DRIVE_FOLDER_URL_PREFIX}{self.coordinator.backup_folder_id}",
+            entry_type=DeviceEntryType.SERVICE,
+        )

homeassistant/components/google_drive/icons.json

@@ -0,0 +1,21 @@
+{
+  "entity": {
+    "sensor": {
+      "backups_size": {
+        "default": "mdi:database"
+      },
+      "storage_total": {
+        "default": "mdi:database"
+      },
+      "storage_used": {
+        "default": "mdi:database"
+      },
+      "storage_used_in_drive": {
+        "default": "mdi:database"
+      },
+      "storage_used_in_drive_trash": {
+        "default": "mdi:database"
+      }
+    }
+  }
+}