QuantEngineByItz/tools/automate_routine.py

import json
import os
import requests
import time
import subprocess
import argparse
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
CLASPRC_PATH = ROOT / ".clasprc.json"
CLASP_PATH = ROOT / ".clasp.json"
SPREADSHEET_ID = "1e1TNlLfnT69nvw-I1wU_oBHmEtI2pfbld3e0fFmtrZM"
OUTPUT_XLSX = ROOT / "GatherTradingData.xlsx"
LOCAL_OUTPUT_XLSX = ROOT / "outputs" / "sector_insights_enhanced" / "GatherTradingData_sector_insights.xlsx"

def get_tokens():
    if not CLASPRC_PATH.exists():
        raise FileNotFoundError(".clasprc.json not found")
    with open(CLASPRC_PATH, "r", encoding="utf-8") as f:
        return json.load(f)["tokens"]["default"]

def get_script_id():
    if not CLASP_PATH.exists():
        raise FileNotFoundError(".clasp.json not found")
    with open(CLASP_PATH, "r", encoding="utf-8") as f:
        return json.load(f)["scriptId"]

def refresh_access_token(tokens):
    url = "https://oauth2.googleapis.com/token"
    payload = {
        "grant_type": "refresh_token",
        "refresh_token": tokens["refresh_token"],
        "client_id": tokens["client_id"],
        "client_secret": tokens["client_secret"],
    }
    resp = requests.post(url, data=payload)
    resp.raise_for_status()
    return resp.json()["access_token"]

def run_gas_function(script_id, access_token, function_name):
    url = f"https://script.googleapis.com/v1/scripts/{script_id}:run"
    headers = {"Authorization": f"Bearer {access_token}"}
    payload = {"function": function_name, "devMode": True}
    print(f"Executing GAS function: {function_name} ...")
    resp = requests.post(url, headers=headers, json=payload)
    
    # Handle response
    if resp.status_code != 200:
        print(f"Error executing function: HTTP {resp.status_code}")
        print(resp.text)
        return False
    
    result = resp.json()
    if "error" in result:
        print(f"Function execution failed: {json.dumps(result['error'], indent=2)}")
        # Check if error is because it's not deployed as API Executable (even if user said it is, common issues persist)
        return False
    
    print("Function execution triggered successfully.")
    return True

def download_spreadsheet(spreadsheet_id, access_token, output_path):
    print(f"Downloading spreadsheet {spreadsheet_id} as XLSX...")
    # Using Drive API v3 to export Google Sheet as XLSX
    url = f"https://www.googleapis.com/drive/v3/files/{spreadsheet_id}/export?mimeType=application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
    headers = {"Authorization": f"Bearer {access_token}"}
    resp = requests.get(url, headers=headers)
    
    if resp.status_code == 403:
        print("Error: 403 Forbidden. This usually means the OAuth token lacks the 'https://www.googleapis.com/auth/drive.readonly' or 'drive' scope.")
        print("Please ensure your clasp login was done with proper scopes or manual token has Drive access.")
        return False
        
    resp.raise_for_status()
    with open(output_path, "wb") as f:
        f.write(resp.content)
    print(f"Successfully downloaded to {output_path}")
    return True

def validate_monthly_sector_refresh(xlsx_path: Path) -> bool:
    cmd = [
        "python",
        "tools/validate_sector_universe_monthly_refresh_v1.py",
        "--xlsx",
        str(xlsx_path),
    ]
    print(f"Validating monthly sector refresh: {xlsx_path} ...")
    res = subprocess.run(cmd, cwd=str(ROOT))
    if res.returncode == 0:
        print("Monthly sector refresh validation passed.")
        return True
    print("Monthly sector refresh validation failed.")
    return False

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--function", default="runDataFeed", help="Primary GAS function to execute before download")
    parser.add_argument("--fallback-function", default="run_all", help="Fallback GAS function to execute if primary fails")
    args = parser.parse_args()

    try:
        tokens = get_tokens()
        script_id = get_script_id()
        access_token = refresh_access_token(tokens)
        
        # Step 1: Execute GAS runDataFeed first, then fallback to run_all if needed.
        primary_ok = run_gas_function(script_id, access_token, args.function)
        if not primary_ok and args.fallback_function and args.fallback_function != args.function:
            print(f"Primary function {args.function} failed; trying fallback {args.fallback_function} ...")
            primary_ok = run_gas_function(script_id, access_token, args.fallback_function)

        if primary_ok:
            print("Waiting a bit for GAS processes to finalize (optional)...")
            time.sleep(5) 
            
            # Step 2: Download spreadsheet
            if download_spreadsheet(SPREADSHEET_ID, access_token, OUTPUT_XLSX):
                print("\nRoutine Part 1 & 2 complete.")
                validate_monthly_sector_refresh(OUTPUT_XLSX)
                print("Final step: npm run prepare-upload-zip")
            else:
                print("\nDownload failed. Please check Drive API scopes.")
        else:
            print("\nGAS execution failed. Process aborted.")
            print("Falling back to local workbook sector-insight build...")
            fallback = subprocess.run(["python", "tools/update_workbook_sector_insights.py"], cwd=str(ROOT))
            if fallback.returncode == 0:
                print("Local sector-insight workbook updated.")
                validate_monthly_sector_refresh(LOCAL_OUTPUT_XLSX)
            else:
                print("Local sector-insight workbook build failed.")
            
    except Exception as e:
        print(f"Error: {e}")

if __name__ == "__main__":
    main()