#!/usr/bin/env python3 """ Google Sheets tracker for dataset processing status. This module provides functionality to track dataset processing status in Google Sheets. It's designed to be optional and non-intrusive. """ from __future__ import annotations import logging import os from pathlib import Path from typing import Any from dotenv import load_dotenv from google.oauth2 import service_account from googleapiclient.discovery import build env_path = Path(__file__).resolve().parent.parent / ".env" if env_path.exists(): load_dotenv(dotenv_path=env_path, override=False) logger = logging.getLogger(__name__) class GoogleSheetsTracker: """Track dataset processing status in Google Sheets.""" def __init__( self, spreadsheet_id: str, worksheet_name: str = "Sheet1", credentials_path: str | None = None, ): """Initialize Google Sheets tracker. Args: spreadsheet_id: Google Sheets spreadsheet ID worksheet_name: Name of the worksheet to use credentials_path: Path to service account JSON credentials """ self.spreadsheet_id = spreadsheet_id self.worksheet_name = worksheet_name credentials_path_raw = credentials_path or os.getenv( "GOOGLE_SHEETS_CREDENTIALS_PATH" ) if credentials_path_raw: credentials_path_obj = Path(credentials_path_raw) if not credentials_path_obj.is_absolute(): base_dir = Path(__file__).resolve().parent.parent credentials_path_obj = base_dir / credentials_path_raw self.credentials_path = str(credentials_path_obj.resolve()) else: self.credentials_path = None self.service = None self.sheet = None self.authenticated = False self.dataset_name_col = None self.status_col = None self.error_col = None self.headers = [] def authenticate(self) -> bool: """Authenticate with Google Sheets API. Returns: True if authentication successful, False otherwise """ if not self.credentials_path: return False if not os.path.exists(self.credentials_path): return False try: credentials = service_account.Credentials.from_service_account_file( self.credentials_path, scopes=["https://www.googleapis.com/auth/spreadsheets"], ) self.service = build("sheets", "v4", credentials=credentials) self.sheet = self.service.spreadsheets() self.authenticated = True self._load_headers() return True except Exception: return False def _load_headers(self) -> None: """Load column headers from the sheet.""" if self.sheet is None: return try: result = ( self.sheet.values() .get( spreadsheetId=self.spreadsheet_id, range=f"{self.worksheet_name}!1:1", ) .execute() ) values = result.get("values", []) if values: self.headers = [str(h).strip().casefold() for h in values[0]] self.dataset_name_col = self._find_column_index( ["dataset name", "dataset", "dataset_id", "id", "name"] ) self.status_col = self._find_column_index(["status", "state"]) self.error_col = self._find_column_index( ["error", "error message"] ) if self.dataset_name_col is None: logger.warning("Could not find dataset name column in sheet") if self.status_col is None: logger.warning("Could not find status column in sheet") if self.error_col is None: logger.warning("Could not find error column in sheet") except Exception: self.headers = [] self.dataset_name_col = 0 self.status_col = 1 self.error_col = 2 def _find_column_index(self, possible_names: list[str]) -> int | None: """Find column index by name. Args: possible_names: List of possible column names (will be casefolded) Returns: Column index or None if not found """ for name in possible_names: try: return self.headers.index(name.casefold()) except ValueError: continue return None def get_dataset_status(self, dataset_id: str) -> dict[str, Any]: """Get current status of a dataset from the sheet. Args: dataset_id: Dataset identifier Returns: Dictionary with status information """ if not self.authenticated or self.sheet is None: return {"status": None, "error": None, "row": None} try: result = ( self.sheet.values() .get( spreadsheetId=self.spreadsheet_id, range=f"{self.worksheet_name}!A:Z", ) .execute() ) values = result.get("values", []) if len(values) < 2: return {"status": None, "error": None, "row": None} for row_idx, row in enumerate(values[1:], start=2): if ( self.dataset_name_col is not None and len(row) > self.dataset_name_col and str(row[self.dataset_name_col]).strip().casefold() == dataset_id.casefold() ): status = ( row[self.status_col].strip() if self.status_col is not None and len(row) > self.status_col else None ) error = ( row[self.error_col].strip() if self.error_col is not None and len(row) > self.error_col else None ) return {"status": status, "error": error, "row": row_idx} return {"status": None, "error": None, "row": None} except Exception: return {"status": None, "error": None, "row": None} def update_status( self, dataset_id: str, status: str, error_message: str | None = None, ) -> bool: """Update dataset status in the sheet. Args: dataset_id: Dataset identifier status: Status value (e.g., "IN PROGRESS", "DONE", "NOT STARTED") error_message: Error message to update in ERROR column Returns: True if update successful, False otherwise """ if not self.authenticated or self.sheet is None: return False try: dataset_status = self.get_dataset_status(dataset_id) row_num = dataset_status.get("row") if row_num is None: return False updates = [] if self.status_col is not None: status_range = ( f"{self.worksheet_name}!" f"{self._col_letter(self.status_col + 1)}{row_num}" ) updates.append({"range": status_range, "values": [[status]]}) if self.error_col is not None: if error_message: error_text = str(error_message)[:500] error_range = ( f"{self.worksheet_name}!" f"{self._col_letter(self.error_col + 1)}{row_num}" ) updates.append({"range": error_range, "values": [[error_text]]}) else: error_range = ( f"{self.worksheet_name}!" f"{self._col_letter(self.error_col + 1)}{row_num}" ) updates.append({"range": error_range, "values": [[""]]}) if updates: body = {"valueInputOption": "RAW", "data": updates} ( self.sheet.values() .batchUpdate( spreadsheetId=self.spreadsheet_id, body=body, ) .execute() ) return True except Exception: return False def _col_letter(self, col_num: int) -> str: """Convert column number to letter (1 -> A, 2 -> B, etc.). Args: col_num: Column number (1-indexed) Returns: Column letter """ result = "" while col_num > 0: col_num -= 1 result = chr(65 + (col_num % 26)) + result col_num //= 26 return result def check_should_process( self, dataset_id: str, force: bool = False ) -> tuple[bool, str]: """Check if dataset should be processed based on current status. Args: dataset_id: Dataset identifier force: Force processing even if status is "Done" or "In Progress" Returns: Tuple of (should_process, reason) """ if force: return (True, "Force processing enabled") if not self.authenticated: return (True, "Sheet tracking not available") dataset_status = self.get_dataset_status(dataset_id) status = dataset_status.get("status", "") if not status: return (True, "No status found in sheet") status_upper = status.upper().strip() if status_upper == "DONE": return (False, "Dataset already marked as DONE") if status_upper == "IN PROGRESS": return (False, "Dataset already IN PROGRESS") if status_upper == "NOT STARTED": return (True, "Status is NOT STARTED, processing") return ( False, f"Unknown status '{status_upper}': " f"Not processing until status is 'NOT STARTED'", ) def mark_in_progress(self, dataset_id: str) -> bool: """Mark dataset as in progress. Args: dataset_id: Dataset identifier Returns: True if update successful """ return self.update_status(dataset_id, "IN PROGRESS") def mark_completed(self, dataset_id: str) -> bool: """Mark dataset as completed. Args: dataset_id: Dataset identifier Returns: True if update successful """ return self.update_status(dataset_id, "DONE") def mark_error(self, dataset_id: str, error_message: str) -> bool: """Mark dataset as error with error message. Args: dataset_id: Dataset identifier error_message: Error message to record in ERROR column Returns: True if update successful """ return self.update_status( dataset_id, "NOT STARTED", error_message=error_message ) def reset_to_not_started( self, dataset_id: str, error_message: str | None = None ) -> bool: """Reset dataset status to NOT STARTED. Used when script is interrupted or fails. Args: dataset_id: Dataset identifier error_message: Optional error message to record in ERROR column Returns: True if update successful """ return self.update_status( dataset_id, "NOT STARTED", error_message=error_message ) def create_tracker( spreadsheet_id: str | None = None, worksheet_name: str | None = None, credentials_path: str | None = None, ) -> GoogleSheetsTracker | None: """Create and initialize Google Sheets tracker. Args: spreadsheet_id: Google Sheets spreadsheet ID worksheet_name: Name of the worksheet credentials_path: Path to service account credentials Returns: GoogleSheetsTracker instance or None if initialization fails """ if not spreadsheet_id: spreadsheet_id = os.getenv("GOOGLE_SHEETS_ID") if not spreadsheet_id: return None if not worksheet_name: worksheet_name = os.getenv("GOOGLE_SHEETS_WORKSHEET_NAME", "Sheet1") try: tracker = GoogleSheetsTracker( spreadsheet_id=spreadsheet_id, worksheet_name=worksheet_name, credentials_path=credentials_path, ) if tracker.authenticate(): return tracker except Exception: pass return None