sync_job.py

import os, threading, time, subprocess, sys, logging, importlib, re
from contextlib import asynccontextmanager
from fastapi import FastAPI
import uvicorn
from datasets import load_dataset
from supabase import create_client
from bs4 import BeautifulSoup

# --- 1. 环境自修复 ---
try:
    import bleach
except ImportError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "bleach"])
    import bleach

# 屏蔽无用日志
logging.getLogger("uvicorn.access").setLevel(logging.ERROR)

# --- 2. 动态加载配置 ---
CONFIG_MODULE_NAME = os.getenv("CONFIG_FILE", "fields_config_plants")
TASK_KEY = os.getenv("SYNC_TASK", "all")  # 默认跑全量，收割 title

try:
    cfg_module = importlib.import_module(CONFIG_MODULE_NAME)
    CFG = cfg_module.CONFIG
except Exception as e:
    print(f"❌ 配置加载失败: {e}")
    sys.exit(1)

URL = os.getenv("SUPABASE_URL")
KEY = os.getenv("SUPABASE_SERVICE_ROLE_KEY")

# 状态追踪
status = {"total_todo": 0, "scanned": 0, "hits": 0, "speed": 0, "task": TASK_KEY}

# --- 3. 增强版清洗函数 ---
def extract_and_clean_v2(html, rule):
    if not html: 
        return "[Empty:NoHtml]"
    try:
        soup = BeautifulSoup(html, 'lxml')
        col_name = rule["col"]
        
        # A. 正常寻找配置中的 ID
        el = soup.find('span', id=rule["html_id"])
        
        # B. 针对 Title 的专项保底逻辑 (处理 id="Head1" 问题)
        if not el and col_name == "title":
            title_tag = soup.find('title')
            if title_tag:
                raw_t = title_tag.get_text()
                # 剔除 PFAF 常见后缀
                return raw_t.replace("PFAF Plant Database", "").strip()
        
        # C. 针对 Common Name 的拼写纠错保底
        if not el and col_name == "common_name":
            # 尝试 PFAF 常见的拼写错误 ID
            el = soup.find('span', id="ContentPlaceHolder1_lblCommanName")
            
        if not el: 
            return "[Empty:NoTag]"
        
        # D. 深度内容提取逻辑
        # 药用、食用、俗名等字段强制穿透所有内嵌标签
        force_text_fields = ['medicinal_uses', 'edible_uses', 'common_name', 'title']
        
        if col_name in force_text_fields or not rule.get("use_bleach"):
            content = el.get_text(separator=" ", strip=True)
            
            # 剔除药用免责声明
            disclaimer = "Plants For A Future can not take any responsibility for any adverse effects from the use of plants. Always seek advice from a professional before using a plant medicinally."
            if disclaimer in content:
                content = content.replace(disclaimer, "").strip()
            
            # 清理多余空格
            content = re.sub(r'\s+', ' ', content)
        else:
            # 需要保留基础格式的字段 (如栽培详情)
            content = bleach.clean(str(el), tags=['p', 'br', 'b', 'i'], strip=True)
            
        return content if content.strip() else "[Empty:Blank]"
    except Exception as e:
        return f"[Error:{str(e)[:20]}]"

# --- 4. 同步主逻辑 ---
def run_sync_logic():
    global status
    try:
        client = create_client(URL, KEY)
    except Exception as e:
        print(f"❌ Supabase 客户端初始化失败: {e}")
        return

    # 确定任务
    if TASK_KEY.lower() == "all":
        active_tasks = CFG["FIELD_RULES"]
    else:
        active_tasks = {TASK_KEY: CFG["FIELD_RULES"][TASK_KEY]}

    print(f"🔍 模式: {TASK_KEY} | 正在检索缺失项...")
    
    try:
        res = client.table(CFG["TARGET_TABLE"]).select("*").execute()
        db_rows = res.data if res.data else []
    except Exception as e:
        print(f"❌ 读取数据库失败: {e}")
        return
    
    todo_ids = set()
    db_lookup = {} 
    
    for r in db_rows:
        p_id = r[CFG["TARGET_ID_COL"]]
        needs_update = False
        for t_key, rule in active_tasks.items():
            val = r.get(rule["col"])
            # 识别 Null, 空值, 或之前的 [Empty:...] 标记
            if val is None or val == "" or str(val).startswith("[Empty:"): 
                needs_update = True
                break
        if needs_update:
            todo_ids.add(p_id)
            db_lookup[p_id] = r

    status["total_todo"] = len(todo_ids)
    print(f"📊 待修复/待补齐: {status['total_todo']:,} 条")

    if not todo_ids:
        print("✅ 状态完美，无需更新。")
        return

    # 连接湖泊
    ds = load_dataset(CFG["REPO_ID"], data_dir=CFG["DATA_DIR"], split="train", streaming=True)
    start_time = time.time()

    for row in ds:
        status["scanned"] += 1
        p_id = row.get(CFG["TARGET_ID_COL"])
        
        if p_id in todo_ids:
            update_payload = {}
            html = row.get('raw_html')
            current_db_row = db_lookup.get(p_id, {})

            for t_key, rule in active_tasks.items():
                col_name = rule["col"]
                db_val = current_db_row.get(col_name)
                if db_val is None or db_val == "" or str(db_val).startswith("[Empty:"):
                    update_payload[col_name] = extract_and_clean_v2(html, rule)
            
            if update_payload:
                try:
                    update_payload["updated_at"] = "now()"
                    client.table(CFG["TARGET_TABLE"]).update(update_payload).eq(CFG["TARGET_ID_COL"], p_id).execute()
                    status["hits"] += 1
                except Exception:
                    continue

            if status["hits"] > 0 and status["hits"] % 50 == 0:
                elapsed = time.time() - start_time
                status["speed"] = status["hits"] / elapsed
                print(f"🎯 修复进度: {status['hits']}/{status['total_todo']} | 速度: {status['speed']:.1f}条/秒")

            if status["hits"] >= status["total_todo"]:
                print(f"🎉 修复完成！共处理 {status['hits']} 条。")
                os._exit(0)

# --- 5. FastAPI 入口 ---
@asynccontextmanager
async def lifespan(app: FastAPI):
    threading.Timer(1.0, run_sync_logic).start()
    yield

app = FastAPI(lifespan=lifespan)

@app.get("/")
def health():
    return {"status": "running", "progress": f"{status['hits']}/{status['total_todo']}", "task": status["task"]}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860, access_log=False)