601 lines
19 KiB
Python
601 lines
19 KiB
Python
#!/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 datetime import datetime
|
|
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.last_shard_col = None
|
|
self.total_rows_col = None
|
|
self.skip_descriptions_col = None
|
|
self.last_updated_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"]
|
|
)
|
|
self.last_shard_col = self._find_column_index(
|
|
["last completed shard", "last_completed_shard",
|
|
"last shard", "last_shard", "shard"]
|
|
)
|
|
self.total_rows_col = self._find_column_index(
|
|
["total rows", "rows", "total_rows"]
|
|
)
|
|
self.skip_descriptions_col = self._find_column_index(
|
|
["skip descriptions", "skip_descriptions", "descriptions"]
|
|
)
|
|
self.last_updated_col = self._find_column_index(
|
|
["last updated", "last_updated", "updated", "timestamp"]
|
|
)
|
|
|
|
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:
|
|
# Calculate maximum column index needed for status data
|
|
max_col = -1
|
|
for col in [
|
|
self.dataset_name_col,
|
|
self.status_col,
|
|
self.error_col,
|
|
]:
|
|
if col is not None and col > max_col:
|
|
max_col = col
|
|
|
|
# Use dynamic range: A to max_column (or Z if no columns found)
|
|
if max_col >= 0:
|
|
# +1 because _col_letter is 1-indexed
|
|
end_col = self._col_letter(max_col + 1)
|
|
range_str = f"{self.worksheet_name}!A:{end_col}"
|
|
else:
|
|
# Fallback to Z if no status columns found
|
|
range_str = f"{self.worksheet_name}!A:Z"
|
|
|
|
result = (
|
|
self.sheet.values()
|
|
.get(
|
|
spreadsheetId=self.spreadsheet_id,
|
|
range=range_str,
|
|
)
|
|
.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
|
|
|
|
# No columns to update (both status_col and error_col are None)
|
|
return False
|
|
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 get_checkpoint(self, dataset_id: str) -> dict[str, Any] | None:
|
|
"""Get checkpoint data for a dataset.
|
|
|
|
Args:
|
|
dataset_id: Dataset identifier
|
|
|
|
Returns:
|
|
Dictionary with checkpoint data or None if not found
|
|
"""
|
|
if not self.authenticated or self.sheet is None:
|
|
return None
|
|
|
|
try:
|
|
dataset_status = self.get_dataset_status(dataset_id)
|
|
row_num = dataset_status.get("row")
|
|
|
|
if row_num is None:
|
|
return None
|
|
|
|
# Calculate maximum column index needed for checkpoint data
|
|
max_col = -1
|
|
for col in [
|
|
self.last_shard_col,
|
|
self.total_rows_col,
|
|
self.skip_descriptions_col,
|
|
self.last_updated_col,
|
|
]:
|
|
if col is not None and col > max_col:
|
|
max_col = col
|
|
|
|
# Use dynamic range: A to max_column (or Z if no columns found)
|
|
if max_col >= 0:
|
|
# +1 because _col_letter is 1-indexed
|
|
end_col = self._col_letter(max_col + 1)
|
|
range_str = f"{self.worksheet_name}!A{row_num}:{end_col}{row_num}"
|
|
else:
|
|
# Fallback to Z if no checkpoint columns found
|
|
range_str = f"{self.worksheet_name}!A{row_num}:Z{row_num}"
|
|
|
|
result = (
|
|
self.sheet.values()
|
|
.get(
|
|
spreadsheetId=self.spreadsheet_id,
|
|
range=range_str,
|
|
)
|
|
.execute()
|
|
)
|
|
values = result.get("values", [])
|
|
|
|
if not values or not values[0]:
|
|
return None
|
|
|
|
row = values[0]
|
|
checkpoint = {}
|
|
|
|
if self.last_shard_col is not None and len(row) > self.last_shard_col:
|
|
val = row[self.last_shard_col].strip()
|
|
checkpoint["last_completed_shard"] = int(val) if val else 0
|
|
|
|
if self.total_rows_col is not None and len(row) > self.total_rows_col:
|
|
val = row[self.total_rows_col].strip()
|
|
checkpoint["total_rows"] = int(val) if val else 0
|
|
|
|
if (
|
|
self.skip_descriptions_col is not None
|
|
and len(row) > self.skip_descriptions_col
|
|
):
|
|
val = row[self.skip_descriptions_col].strip()
|
|
checkpoint["skip_descriptions"] = int(val) if val else 0
|
|
|
|
if self.last_updated_col is not None and len(row) > self.last_updated_col:
|
|
checkpoint["timestamp"] = row[self.last_updated_col].strip()
|
|
|
|
return checkpoint if checkpoint else None
|
|
|
|
except Exception:
|
|
return None
|
|
|
|
def update_checkpoint(
|
|
self,
|
|
dataset_id: str,
|
|
last_shard: int,
|
|
total_rows: int,
|
|
skip_descriptions: int,
|
|
) -> bool:
|
|
"""Update checkpoint data for a dataset.
|
|
|
|
Args:
|
|
dataset_id: Dataset identifier
|
|
last_shard: Last completed shard index
|
|
total_rows: Total rows processed
|
|
skip_descriptions: Total descriptions processed
|
|
|
|
Returns:
|
|
True if update successful
|
|
"""
|
|
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.last_shard_col is not None:
|
|
shard_range = (
|
|
f"{self.worksheet_name}!"
|
|
f"{self._col_letter(self.last_shard_col + 1)}{row_num}"
|
|
)
|
|
updates.append({"range": shard_range, "values": [[last_shard]]})
|
|
|
|
if self.total_rows_col is not None:
|
|
rows_range = (
|
|
f"{self.worksheet_name}!"
|
|
f"{self._col_letter(self.total_rows_col + 1)}{row_num}"
|
|
)
|
|
updates.append({"range": rows_range, "values": [[total_rows]]})
|
|
|
|
if self.skip_descriptions_col is not None:
|
|
desc_range = (
|
|
f"{self.worksheet_name}!"
|
|
f"{self._col_letter(self.skip_descriptions_col + 1)}{row_num}"
|
|
)
|
|
updates.append(
|
|
{"range": desc_range, "values": [[skip_descriptions]]}
|
|
)
|
|
|
|
if self.last_updated_col is not None:
|
|
timestamp = datetime.now().isoformat()
|
|
updated_range = (
|
|
f"{self.worksheet_name}!"
|
|
f"{self._col_letter(self.last_updated_col + 1)}{row_num}"
|
|
)
|
|
updates.append({"range": updated_range, "values": [[timestamp]]})
|
|
|
|
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 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
|
|
|