apps/worker/run_backfill.py

import asyncio
import sys
import os
import argparse
from datetime import datetime, timedelta
from typing import Iterator, Tuple

sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))

from packages.core.logger import app_logger


def iter_month_windows(start_dt: datetime, end_dt: datetime) -> Iterator[Tuple[datetime, datetime]]:
    """
    将日期范围拆分为逐月窗口，供回补任务按月执行。
    """
    current_dt = start_dt
    while current_dt <= end_dt:
        next_month = current_dt.replace(day=28) + timedelta(days=4)
        next_month_start = next_month.replace(day=1)
        batch_end_dt = min(next_month_start - timedelta(days=1), end_dt)
        yield current_dt, batch_end_dt
        current_dt = next_month_start

async def run_backfill(country_code: str, start_date: str, end_date: str, local_data_dir: str = None):
    """
    历史数据回溯脚本 (Backfill)
    允许指定国家和时间范围，批量补录过去的数据。
    """
    app_logger.info(f"--- Starting BACKFILL Job for {country_code} from {start_date} to {end_date} ---")
    
    # 如果环境变量中没有 DATABASE_URL，则设置一个默认值（兼容宿主机直接运行）
    if "DATABASE_URL" not in os.environ:
        os.environ["DATABASE_URL"] = "postgresql+asyncpg://postgres:postgres@localhost:5433/customs_data"

    from packages.core.database import AsyncSessionLocal
    from packages.connectors.mock.us_mock import MockUSConnector
    from packages.connectors.brazil import BrazilComexStatConnector
    from packages.connectors.chile import ChileAduanasConnector
    from packages.connectors.mock.extended_mock import ExtendedMockConnector, COUNTRY_CONFIGS
    
    start_dt = datetime.strptime(start_date, "%Y-%m-%d")
    end_dt = datetime.strptime(end_date, "%Y-%m-%d")
    
    async with AsyncSessionLocal() as session:
        # 巴西真实源支持按月范围回补，直接一次性下发给 connector。
        if country_code == "BR":
            connector = BrazilComexStatConnector(
                session=session,
                start_period=start_dt.strftime("%Y-%m"),
                end_period=end_dt.strftime("%Y-%m"),
                local_data_dir=local_data_dir,
            )
            connector.batch_no = f"BACKFILL_{country_code}_{start_dt.strftime('%Y%m')}_{end_dt.strftime('%Y%m')}"
            await connector.run()
            app_logger.info(f"--- Finished BACKFILL Job for {country_code} ---")
            return
            
        if country_code == "CL":
            connector = ChileAduanasConnector(
                session=session,
                start_period=start_dt.strftime("%Y-%m"),
                end_period=end_dt.strftime("%Y-%m"),
                local_data_dir=local_data_dir,
            )
            connector.batch_no = f"BACKFILL_{country_code}_{start_dt.strftime('%Y%m')}_{end_dt.strftime('%Y%m')}"
            await connector.run()
            app_logger.info(f"--- Finished BACKFILL Job for {country_code} ---")
            return

        # 1. 实例化对应的 Connector
        if country_code == "US":
            connector = MockUSConnector(session=session)
        elif country_code in COUNTRY_CONFIGS:
            connector = ExtendedMockConnector(country_code=country_code, session=session)
        else:
            app_logger.error(f"Unsupported country code for backfill: {country_code}")
            return

        # 2. 模拟按月拆分任务进行回溯
        # 真实场景下，海关接口通常限制单次查询跨度，比如只能查一个月
        for window_start, batch_end_dt in iter_month_windows(start_dt, end_dt):
            app_logger.info(f"[{country_code}] Backfilling for period: {window_start.strftime('%Y-%m-%d')} to {batch_end_dt.strftime('%Y-%m-%d')}")
            
            # TODO: 真实场景下，这里要将 current_dt 和 batch_end_dt 传给 fetch/discover
            # 目前 MVP 阶段调用 run() 模拟执行一轮抓取
            # 为了防止冲突，强行注入一个特定的 batch_no 前缀
            connector.batch_no = f"BACKFILL_{country_code}_{window_start.strftime('%Y%m')}"
            await connector.run()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Backfill historical customs data")
    parser.add_argument("--country", type=str, required=True, help="Country Code (e.g., US, BR, ID)")
    parser.add_argument("--start", type=str, required=True, help="Start Date (YYYY-MM-DD)")
    parser.add_argument("--end", type=str, required=True, help="End Date (YYYY-MM-DD)")
    parser.add_argument("--local-data-dir", type=str, required=False, help="Local directory containing historical CSVs (e.g., /data/brazil_csvs)")
    
    args = parser.parse_args()
    
    asyncio.run(run_backfill(
        country_code=args.country, 
        start_date=args.start, 
        end_date=args.end,
        local_data_dir=args.local_data_dir
    ))