#!/usr/bin/env python3 """One-off helper to backfill existing AWS CSV exports into MySQL.""" from __future__ import annotations import argparse import csv import sys from datetime import datetime from pathlib import Path from typing import Dict, Iterable, List, Optional, Sequence REPO_ROOT = Path(__file__).resolve().parents[2] AWS_SRC = REPO_ROOT / "aws" / "src" sys.path.insert(0, str(AWS_SRC)) sys.path.insert(0, str(REPO_ROOT)) try: # prefer local-style imports used in exporter modules from src.common.config import ConfigManager # type: ignore from src.export.environments import EnvironmentManager # type: ignore from src.export.mysql_writer import AWSMySQLWriter # type: ignore except ModuleNotFoundError: # fallback when src package is not resolved from aws.src.common.config import ConfigManager from aws.src.export.environments import EnvironmentManager from aws.src.export.mysql_writer import AWSMySQLWriter DATE_FORMAT = "%Y-%m-%d" def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Backfill AWS billing CSV exports into MySQL", formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) parser.add_argument( "--config", help="Path to aws/config.json (falls back to default resolution)", ) parser.add_argument( "--exports-root", help="Override the root directory that contains environment export folders", ) parser.add_argument( "--account-details", default=str(REPO_ROOT / "aws" / "aws_account_details.csv"), help="CSV file that maps account_id to account_name", ) parser.add_argument( "--env", dest="environments", nargs="+", choices=["nonprod", "prod", "shared"], help="Limit processing to specific environments", ) parser.add_argument( "--account-id", dest="account_ids", action="append", help="Restrict to specific AWS account IDs (repeatable)", ) parser.add_argument( "--limit", type=int, help="Maximum number of files to process", ) parser.add_argument( "--dry-run", action="store_true", help="Print actions without writing to MySQL", ) return parser.parse_args() def resolve_output_dir(base_dir: Optional[str], configured_dir: Path) -> Path: if base_dir: return Path(base_dir).expanduser().resolve() if configured_dir.is_absolute(): return configured_dir return (REPO_ROOT / "aws" / configured_dir).resolve() def discover_files(output_dir: Path) -> Sequence[Path]: if not output_dir.exists(): return [] return sorted(path for path in output_dir.glob("*.csv") if path.is_file()) def load_rows(csv_path: Path) -> List[Dict[str, str]]: with csv_path.open("r", encoding="utf-8") as handle: reader = csv.DictReader(handle) rows = [dict(row) for row in reader] return rows def parse_dates(rows: Iterable[Dict[str, str]]) -> List[datetime.date]: dates: List[datetime.date] = [] for row in rows: raw = (row.get("date") or "").strip() if not raw: continue try: dates.append(datetime.strptime(raw, DATE_FORMAT).date()) except ValueError: continue return dates def infer_account_id(csv_path: Path, rows: List[Dict[str, str]]) -> str: if rows: row_value = rows[0].get("account_id") if row_value: return row_value.strip() parts = csv_path.stem.split("_") if len(parts) < 2: return csv_path.stem return parts[-2] def load_account_names(path: str) -> Dict[str, str]: mapping: Dict[str, str] = {} csv_path = Path(path) if not csv_path.exists(): return mapping with csv_path.open("r", encoding="utf-8") as handle: reader = csv.DictReader(handle) for row in reader: account_id = (row.get("account_id") or "").strip() account_name = (row.get("account_name") or account_id).strip() if account_id: mapping[account_id] = account_name or account_id return mapping def main() -> None: args = parse_args() config = ConfigManager(args.config) env_manager = EnvironmentManager(config) mysql_writer = AWSMySQLWriter(config) if not mysql_writer.enabled: print("MySQL 同步未启用，请先在 aws/config.json 中设置 mysqlConfig.enabled = true 并提供连接信息。") sys.exit(1) selected_envs = args.environments or env_manager.list_environments() account_filter = set(args.account_ids or []) account_names = load_account_names(args.account_details) total_files = 0 synced_files = 0 skipped_files = 0 for env_name in selected_envs: if args.limit and synced_files >= args.limit: break env_config = env_manager.get_environment(env_name) if not env_config: print(f"跳过未知环境: {env_name}") continue output_dir = resolve_output_dir(args.exports_root, env_config["output_dir"]) csv_files = discover_files(output_dir) file_prefix = env_manager.get_file_prefix(env_name) if file_prefix: csv_files = [ path for path in csv_files if path.name.startswith(f"{file_prefix}_") ] if not csv_files: print(f"环境 {env_name} 未找到 CSV 文件 (目录: {output_dir})") continue print(f"处理环境 {env_name} ({len(csv_files)} 个文件) -> {output_dir}") for csv_path in csv_files: if args.limit and synced_files >= args.limit: break account_id = infer_account_id(csv_path, []) if account_filter and account_id not in account_filter: continue rows = load_rows(csv_path) if not rows: skipped_files += 1 print(f" - 跳过 {csv_path.name} (空文件)") continue account_id = infer_account_id(csv_path, rows) if account_filter and account_id not in account_filter: continue dates = parse_dates(rows) if not dates: skipped_files += 1 print(f" - 跳过 {csv_path.name} (缺少有效日期)") continue start_date = min(dates) end_date = max(dates) account_name = account_names.get(account_id, account_id) total_files += 1 summary = ( f" - {csv_path.name}: {len(rows)} 行, {start_date} 至 {end_date}, 账户 {account_name}" ) if args.dry_run: print(f"{summary} [DRY RUN]") continue try: mysql_writer.sync( account_identifier=account_id, account_name=account_name, environment=rows[0].get("environment", env_name), records=rows, start_date=start_date, end_date=end_date, source_file=csv_path, metadata={"notes": "aws/scripts/backfill_mysql.py"}, ) synced_files += 1 print(f"{summary} -> 已同步") except Exception as exc: # pylint: disable=broad-except skipped_files += 1 print(f" ✗ {csv_path.name} 同步失败: {exc}") print("\n完成掃描") print(f"总文件: {total_files}") print(f"已同步: {synced_files}") print(f"已跳过: {skipped_files}") if __name__ == "__main__": main()