"""Small wrappers around Google Drive and Sheets API operations.""" from __future__ import annotations import logging import time from pathlib import Path from typing import Any from google.oauth2.service_account import Credentials from googleapiclient.discovery import build from googleapiclient.errors import HttpError LOGGER = logging.getLogger(__name__) DRIVE_FOLDER_MIME = "application/vnd.google-apps.folder" GOOGLE_SHEET_MIME = "application/vnd.google-apps.spreadsheet" SCOPES = ( "https://www.googleapis.com/auth/drive", "https://www.googleapis.com/auth/spreadsheets", ) class GoogleWorkspaceClient: """Drive and Sheets helper for fixed HACCP spreadsheets.""" def __init__(self, credentials_file: Path) -> None: """Create authenticated Drive and Sheets clients from a service account.""" if not credentials_file.exists(): raise RuntimeError(f"Missing Google credentials file: {credentials_file}") credentials = Credentials.from_service_account_file( str(credentials_file), scopes=SCOPES ) self.drive = build("drive", "v3", credentials=credentials, cache_discovery=False) self.sheets = build("sheets", "v4", credentials=credentials, cache_discovery=False) def find_spreadsheet_recursive(self, root_folder_id: str, name: str) -> str: """Find a precreated spreadsheet by exact name anywhere below a root folder.""" matches: list[dict[str, Any]] = [] folders = [root_folder_id] while folders: folder_id = folders.pop(0) for item in self.list_children(folder_id): if item["mimeType"] == DRIVE_FOLDER_MIME: folders.append(item["id"]) if item["mimeType"] == GOOGLE_SHEET_MIME and item["name"] == name: matches.append(item) if not matches: raise RuntimeError(f"Could not find precreated spreadsheet: {name}") if len(matches) > 1: LOGGER.warning("Found %s spreadsheet duplicates for %s; using the first", len(matches), name) LOGGER.info("Using spreadsheet %s (%s)", matches[0]["name"], matches[0]["id"]) return matches[0]["id"] def list_children(self, parent_id: str) -> list[dict[str, Any]]: """List visible non-trashed direct children of a Drive folder.""" query = f"'{self._query_value(parent_id)}' in parents and trashed = false" result = self._execute( self.drive.files().list( q=query, spaces="drive", pageSize=1000, fields="files(id,name,mimeType,webViewLink,createdTime)", supportsAllDrives=True, includeItemsFromAllDrives=True, ), f"list Drive children for {parent_id}", ) return result.get("files", []) def find_file(self, parent_id: str, name: str, mime_type: str) -> dict[str, Any] | None: """Find the first non-trashed file with an exact name, parent, and MIME type.""" query = " and ".join( [ f"'{self._query_value(parent_id)}' in parents", f"name = '{self._query_value(name)}'", f"mimeType = '{self._query_value(mime_type)}'", "trashed = false", ] ) result = self._execute( self.drive.files().list( q=query, spaces="drive", pageSize=10, fields="files(id,name,mimeType,webViewLink,createdTime)", supportsAllDrives=True, includeItemsFromAllDrives=True, ), f"find Drive file {name}", ) files = result.get("files", []) if len(files) > 1: LOGGER.warning("Found %s duplicates for %s; using the first", len(files), name) return files[0] if files else None def update_values(self, spreadsheet_id: str, data: list[dict[str, Any]]) -> None: """Batch-update spreadsheet values while preserving sheet formatting.""" if not data: return self._execute( self.sheets.spreadsheets().values().batchUpdate( spreadsheetId=spreadsheet_id, body={"valueInputOption": "RAW", "data": data}, ), f"update spreadsheet {spreadsheet_id}", ) def get_sheet_id(self, spreadsheet_id: str, title: str) -> int | None: """Return a tab's numeric sheet ID by title, or None when missing.""" for properties in self.list_sheets(spreadsheet_id): if properties.get("title") == title: return properties["sheetId"] return None def list_sheets(self, spreadsheet_id: str) -> list[dict[str, Any]]: """Return sheet tab properties in their current left-to-right order.""" metadata = self._execute( self.sheets.spreadsheets().get( spreadsheetId=spreadsheet_id, fields="sheets(properties(sheetId,title,index))", ), f"read spreadsheet metadata {spreadsheet_id}", ) sheets = [sheet.get("properties", {}) for sheet in metadata.get("sheets", [])] return sorted(sheets, key=lambda properties: properties.get("index", 0)) def ensure_sheet( self, spreadsheet_id: str, title: str, rows: int, columns: int, index: int | None = None, ) -> int: """Return a tab ID, creating the month tab when it does not exist.""" sheet_id = self.get_sheet_id(spreadsheet_id, title) if sheet_id is not None: return sheet_id response = self.batch_update( spreadsheet_id, [ { "addSheet": { "properties": { "title": title, **({"index": index} if index is not None else {}), "gridProperties": { "rowCount": rows, "columnCount": columns, }, } } } ], f"create tab {title}", ) added = response["replies"][0]["addSheet"]["properties"] LOGGER.info("Created tab %s (%s)", title, added["sheetId"]) return added["sheetId"] def delete_sheet(self, spreadsheet_id: str, sheet_id: int) -> None: """Delete one tab from a spreadsheet.""" self.batch_update( spreadsheet_id, [{"deleteSheet": {"sheetId": sheet_id}}], f"delete tab {sheet_id}", ) def set_sheet_order(self, spreadsheet_id: str, sheet_ids: list[int]) -> None: """Set tab order from left to right using a list of sheet IDs.""" requests = [ { "updateSheetProperties": { "properties": {"sheetId": sheet_id, "index": index}, "fields": "index", } } for index, sheet_id in enumerate(sheet_ids) ] self.batch_update( spreadsheet_id, requests, f"reorder spreadsheet {spreadsheet_id}" ) def move_sheet(self, spreadsheet_id: str, sheet_id: int, index: int) -> None: """Move one tab to a specific left-to-right index.""" if index == 0: other_sheet_ids = [ sheet["sheetId"] for sheet in self.list_sheets(spreadsheet_id) if sheet["sheetId"] != sheet_id ] self.set_sheet_order(spreadsheet_id, [sheet_id, *other_sheet_ids]) return self.batch_update( spreadsheet_id, [ { "updateSheetProperties": { "properties": {"sheetId": sheet_id, "index": index}, "fields": "index", } } ], f"move tab {sheet_id}", ) def reset_sheet(self, spreadsheet_id: str, sheet_id: int, rows: int, columns: int) -> None: """Clear values, formatting, and merges before rerendering a month tab.""" full_range = {"sheetId": sheet_id} self.batch_update( spreadsheet_id, [ {"unmergeCells": {"range": full_range}}, { "updateSheetProperties": { "properties": { "sheetId": sheet_id, "gridProperties": { "rowCount": rows, "columnCount": columns, "frozenRowCount": min(rows, 6), }, }, "fields": "gridProperties(rowCount,columnCount,frozenRowCount)", } }, { "repeatCell": { "range": full_range, "cell": {}, "fields": "userEnteredValue,userEnteredFormat,note", } }, ], f"reset tab {sheet_id}", ) def batch_update( self, spreadsheet_id: str, requests: list[dict[str, Any]], description: str ) -> dict[str, Any]: """Run a Sheets batchUpdate request when there is at least one request.""" if not requests: return {} return self._execute( self.sheets.spreadsheets().batchUpdate( spreadsheetId=spreadsheet_id, body={"requests": requests}, ), description, ) @staticmethod def _query_value(value: str) -> str: """Escape a value for Drive query string literals.""" return value.replace("\\", "\\\\").replace("'", "\\'") def _execute(self, request: Any, description: str) -> Any: """Execute a Google API request with simple retry handling.""" for attempt in range(1, 6): try: return request.execute() except HttpError as exc: status = getattr(exc.resp, "status", None) retryable = status in {429, 500, 502, 503, 504} if not retryable or attempt == 5: LOGGER.error("Google API request failed during %s: %s", description, exc) raise sleep_for = min(2**attempt, 30) LOGGER.warning( "Retrying %s after HTTP %s in %ss", description, status, sleep_for ) time.sleep(sleep_for)