423 lines
13 KiB
Python
423 lines
13 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 os
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
try:
|
|
from dotenv import load_dotenv
|
|
env_path = Path(__file__).resolve().parent.parent / ".env"
|
|
load_dotenv(dotenv_path=env_path, override=False)
|
|
except ImportError:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
from google.oauth2 import service_account
|
|
from googleapiclient.discovery import build
|
|
|
|
GOOGLE_SHEETS_AVAILABLE = True
|
|
except ImportError:
|
|
GOOGLE_SHEETS_AVAILABLE = False
|
|
|
|
|
|
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
|
|
"""
|
|
if not GOOGLE_SHEETS_AVAILABLE:
|
|
raise ImportError(
|
|
"Google Sheets API not available. "
|
|
"Install with: pip install google-api-python-client google-auth"
|
|
)
|
|
|
|
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."""
|
|
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() for h in values[0]]
|
|
self.dataset_name_col = self._find_column_index(
|
|
[
|
|
"DATASET NAME",
|
|
"dataset name",
|
|
"dataset",
|
|
"dataset_id",
|
|
"id",
|
|
"name",
|
|
]
|
|
)
|
|
self.status_col = self._find_column_index(
|
|
["STATUS", "status", "state"]
|
|
)
|
|
self.error_col = self._find_column_index(
|
|
["ERROR", "error", "error message"]
|
|
)
|
|
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
|
|
|
|
Returns:
|
|
Column index or None if not found
|
|
"""
|
|
for name in possible_names:
|
|
try:
|
|
return self.headers.index(name)
|
|
except ValueError:
|
|
try:
|
|
return self.headers.index(name.upper())
|
|
except ValueError:
|
|
try:
|
|
return self.headers.index(name.lower())
|
|
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:
|
|
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 (
|
|
len(row) > self.dataset_name_col
|
|
and str(row[self.dataset_name_col]).strip().lower()
|
|
== dataset_id.lower()
|
|
):
|
|
status = (
|
|
row[self.status_col].strip()
|
|
if self.status_col and len(row) > self.status_col
|
|
else None
|
|
)
|
|
error = (
|
|
row[self.error_col].strip()
|
|
if self.error_col 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,
|
|
notes: 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
|
|
notes: Optional notes (not used)
|
|
|
|
Returns:
|
|
True if update successful, False otherwise
|
|
"""
|
|
if not self.authenticated:
|
|
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 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":
|
|
if force:
|
|
return (True, "Force processing enabled")
|
|
return (False, "Dataset already marked as DONE")
|
|
|
|
if status_upper == "IN PROGRESS":
|
|
if force:
|
|
return (True, "Force processing enabled")
|
|
return (False, "Dataset already IN PROGRESS")
|
|
|
|
if status_upper == "NOT STARTED":
|
|
return (True, "Status is NOT STARTED, processing")
|
|
|
|
return (True, "Status allows processing")
|
|
|
|
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 = "Sheet1",
|
|
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 GOOGLE_SHEETS_AVAILABLE:
|
|
return None
|
|
|
|
if not spreadsheet_id:
|
|
spreadsheet_id = os.getenv("GOOGLE_SHEETS_ID")
|
|
|
|
if not spreadsheet_id:
|
|
return None
|
|
|
|
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
|
|
|