"""Support for Google Sheets.""" from typing import TYPE_CHECKING from google.auth.exceptions import RefreshError from google.oauth2.credentials import Credentials from gspread import Client, Spreadsheet, Worksheet from gspread.exceptions import APIError from gspread.utils import ValueInputOption import voluptuous as vol from homeassistant.const import CONF_ACCESS_TOKEN, CONF_TOKEN from homeassistant.core import ( HomeAssistant, ServiceCall, ServiceResponse, SupportsResponse, callback, ) from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import config_validation as cv, service from homeassistant.helpers.selector import ConfigEntrySelector from homeassistant.util import dt as dt_util from homeassistant.util.json import JsonObjectType from .const import DOMAIN if TYPE_CHECKING: from . import GoogleSheetsConfigEntry ADD_CREATED_COLUMN = "add_created_column" DATA = "data" DATA_CONFIG_ENTRY = "config_entry" ROWS = "rows" WORKSHEET = "worksheet" SERVICE_APPEND_SHEET = "append_sheet" SERVICE_GET_SHEET = "get_sheet" SHEET_SERVICE_SCHEMA = vol.All( { vol.Required(DATA_CONFIG_ENTRY): ConfigEntrySelector({"integration": DOMAIN}), vol.Optional(WORKSHEET): cv.string, vol.Optional(ADD_CREATED_COLUMN, default=True): cv.boolean, vol.Required(DATA): vol.Any(cv.ensure_list, [dict]), }, ) get_SHEET_SERVICE_SCHEMA = vol.All( { vol.Required(DATA_CONFIG_ENTRY): ConfigEntrySelector({"integration": DOMAIN}), vol.Optional(WORKSHEET): cv.string, vol.Required(ROWS): cv.positive_int, }, ) def _get_worksheet(sheet: Spreadsheet, name: str | None) -> Worksheet: """Return the requested worksheet, or the first one when none was given. Looking the name up eagerly would fetch the document metadata even when a worksheet was requested and the first one is discarded. """ if name is None: return sheet.sheet1 return sheet.worksheet(name) def _append_to_sheet(call: ServiceCall, entry: GoogleSheetsConfigEntry) -> None: """Run append in the executor.""" client = Client(Credentials(entry.data[CONF_TOKEN][CONF_ACCESS_TOKEN])) # type: ignore[no-untyped-call] sheet = client.open_by_key(entry.unique_id) worksheet = _get_worksheet(sheet, call.data.get(WORKSHEET)) columns: list[str] = next(iter(worksheet.get_values("A1:ZZ1")), []) add_created_column = call.data[ADD_CREATED_COLUMN] now = str(dt_util.now()) rows = [] for d in call.data[DATA]: row_data = ({"created": now} | d) if add_created_column else d row = [row_data.get(column, "") for column in columns] for key, value in row_data.items(): if key not in columns: columns.append(key) worksheet.update_cell(1, len(columns), key) row.append(value) rows.append(row) worksheet.append_rows(rows, value_input_option=ValueInputOption.user_entered) def _get_from_sheet( call: ServiceCall, entry: GoogleSheetsConfigEntry ) -> JsonObjectType: """Run get in the executor.""" client = Client(Credentials(entry.data[CONF_TOKEN][CONF_ACCESS_TOKEN])) # type: ignore[no-untyped-call] sheet = client.open_by_key(entry.unique_id) worksheet = _get_worksheet(sheet, call.data.get(WORKSHEET)) all_values = worksheet.get_values() return {"range": all_values[-call.data[ROWS] :]} async def _async_append_to_sheet(call: ServiceCall) -> None: """Append new line of data to a Google Sheets document.""" entry: GoogleSheetsConfigEntry = service.async_get_config_entry( call.hass, DOMAIN, call.data[DATA_CONFIG_ENTRY] ) await entry.runtime_data.async_ensure_token_valid() try: await call.hass.async_add_executor_job(_append_to_sheet, call, entry) except RefreshError: entry.async_start_reauth(call.hass) raise except APIError as ex: raise HomeAssistantError( translation_domain=DOMAIN, translation_key="append_failed" ) from ex async def _async_get_from_sheet(call: ServiceCall) -> ServiceResponse: """Get lines of data from a Google Sheets document.""" entry: GoogleSheetsConfigEntry = service.async_get_config_entry( call.hass, DOMAIN, call.data[DATA_CONFIG_ENTRY] ) await entry.runtime_data.async_ensure_token_valid() try: return await call.hass.async_add_executor_job(_get_from_sheet, call, entry) except RefreshError: entry.async_start_reauth(call.hass) raise except APIError as ex: raise HomeAssistantError( translation_domain=DOMAIN, translation_key="get_failed" ) from ex @callback def async_setup_services(hass: HomeAssistant) -> None: """Add the services for Google Sheets.""" hass.services.async_register( DOMAIN, SERVICE_APPEND_SHEET, _async_append_to_sheet, schema=SHEET_SERVICE_SCHEMA, ) hass.services.async_register( DOMAIN, SERVICE_GET_SHEET, _async_get_from_sheet, schema=get_SHEET_SERVICE_SCHEMA, supports_response=SupportsResponse.ONLY, )