core/homeassistant/components/google_sheets/services.py

"""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,
    )