DM-Agent / src /loader.py
wen Qiu
feat: upgrade Meta API to ad-level (素材级别) data pull
381a123
Raw History Blame Contribute Delete
18.5 kB
from __future__ import annotations
"""
数据加载:支持 Google Sheets(优先)和本地 Excel 两种数据源。
处理编码异常列名、过滤汇总行、从 Date 派生 year/week/month。
"""
import logging
import os
from pathlib import Path
from typing import Any
import pandas as pd
import yaml
log = logging.getLogger("ad-agent")
def _project_root() -> Path:
return Path(__file__).resolve().parent.parent
def load_config(config_path: Path | None = None) -> dict[str, Any]:
path = config_path or (_project_root() / "config.yaml")
with open(path, "r", encoding="utf-8") as f:
return yaml.safe_load(f)
# ─────────────────────────────────────────────
# Google Sheets 数据源
# ─────────────────────────────────────────────
def _build_google_creds(config: dict[str, Any]):
"""
构建 Google API 凭证,支持三种认证方式(按优先级):
1. GOOGLE_CREDENTIALS 环境变量(服务账号 JSON 字符串,云端部署推荐)
2. GOOGLE_OAUTH_REFRESH_TOKEN 环境变量(OAuth refresh token,GitHub Actions 等场景)
3. credentials.json 本地文件(本地开发)
"""
from google.oauth2.service_account import Credentials as SACredentials
from google.oauth2.credentials import Credentials as OAuthCredentials
import json
scopes = [
"https://www.googleapis.com/auth/spreadsheets",
"https://www.googleapis.com/auth/drive",
]
creds_json = os.environ.get("GOOGLE_CREDENTIALS", "").strip()
if creds_json:
info = json.loads(creds_json)
creds = SACredentials.from_service_account_info(info, scopes=scopes)
log.info("使用环境变量 GOOGLE_CREDENTIALS(服务账号)认证")
return creds
refresh_token = os.environ.get("GOOGLE_OAUTH_REFRESH_TOKEN", "").strip()
client_id = os.environ.get("GOOGLE_OAUTH_CLIENT_ID", "").strip()
client_secret = os.environ.get("GOOGLE_OAUTH_CLIENT_SECRET", "").strip()
if refresh_token and client_id and client_secret:
creds = OAuthCredentials(
token=None,
refresh_token=refresh_token,
client_id=client_id,
client_secret=client_secret,
token_uri="https://oauth2.googleapis.com/token",
)
log.info("使用 OAuth refresh token 认证")
return creds
gs_cfg = config.get("google_sheets", {})
creds_file = gs_cfg.get("credentials_file", "credentials.json")
creds_path = Path(creds_file)
if not creds_path.is_absolute():
creds_path = _project_root() / creds_path
if creds_path.exists():
creds = SACredentials.from_service_account_file(str(creds_path), scopes=scopes)
log.info(f"使用凭证文件 {creds_path} 认证")
return creds
return None
def _load_from_google_sheets(config: dict[str, Any]) -> pd.DataFrame | None:
"""
从 Google Sheets 分平台表(Meta/Google/Criteo/Bing)读取 API 同步数据。
单一数据源,不读 Export 表,避免多源合并导致的重复。
"""
gs_cfg = config.get("google_sheets", {})
if not gs_cfg.get("enabled"):
return None
url = gs_cfg.get("spreadsheet_url", "").strip()
if not url:
log.warning("Google Sheets 已启用但未配置 spreadsheet_url,回退到本地 Excel")
return None
try:
import gspread
creds = _build_google_creds(config)
if creds is None:
log.warning("无可用 Google 认证凭证,回退到本地 Excel")
return None
gc = gspread.authorize(creds)
sh = gc.open_by_url(url)
platform_sheets = gs_cfg.get("platform_sheets", ["Meta", "Google", "Criteo", "Bing"])
col_cfg = config.get("columns", {})
idx_rename = {
col_cfg.get("country_idx", 0): "country",
col_cfg.get("platform_idx", 1): "platform",
col_cfg.get("campaign_type_idx", 2): "campaign_type",
col_cfg.get("campaign_name_idx", 3): "campaign_name",
col_cfg.get("date_idx", 4): "date",
}
def _read_one_sheet(ws_name: str) -> pd.DataFrame | None:
try:
ws = sh.worksheet(ws_name)
records = ws.get_all_values()
if len(records) < 2:
return None
header = records[0]
seen: dict[str, int] = {}
unique_header = []
for col in header:
if col in seen:
seen[col] += 1
unique_header.append(f"{col}_{seen[col]}")
else:
seen[col] = 0
unique_header.append(col)
sdf = pd.DataFrame(records[1:], columns=unique_header)
rename = {}
for idx, name in idx_rename.items():
if idx < len(sdf.columns):
rename[sdf.columns[idx]] = name
sdf = sdf.rename(columns=rename)
log.info(f"从 Sheet [{ws_name}] 读取 {len(sdf)} 行")
return sdf
except Exception:
return None
all_dfs: list[pd.DataFrame] = []
for ps in platform_sheets:
ps_df = _read_one_sheet(ps)
if ps_df is not None and not ps_df.empty:
all_dfs.append(ps_df)
if not all_dfs:
log.warning("Google Sheets 所有分平台工作表均无数据")
return None
df = pd.concat(all_dfs, ignore_index=True)
log.info(f"从 Google Sheets 分平台表共读取 {len(df)} 行数据({', '.join(platform_sheets)})")
return df
except Exception as e:
log.error(f"Google Sheets 读取失败: {e},回退到本地 Excel")
return None
# ─────────────────────────────────────────────
# 本地 Excel 数据源
# ─────────────────────────────────────────────
def _load_from_excel(config: dict[str, Any]) -> pd.DataFrame:
"""从 data 目录读取最新的 .xlsx 文件。"""
root = _project_root()
data_dir = Path(config.get("data_dir", "data"))
if not data_dir.is_absolute():
data_dir = root / data_dir
files = sorted(
(f for f in data_dir.glob("*.xlsx") if not f.name.startswith("~$")),
key=lambda p: p.stat().st_mtime, reverse=True,
)
if not files:
raise FileNotFoundError(f"在 {data_dir} 中未找到 .xlsx 文件")
sheet = config.get("excel", {}).get("sheet", 0)
df = pd.read_excel(files[0], sheet_name=sheet, engine="openpyxl")
log.info(f"从本地 Excel 读取: {files[0].name}({len(df)} 行)")
return df
# ─────────────────────────────────────────────
# 写回 Google Sheets
# ─────────────────────────────────────────────
def _get_spreadsheet(config: dict[str, Any]):
"""获取 gspread spreadsheet 对象(读写权限)。"""
import gspread
gs_cfg = config.get("google_sheets", {})
url = gs_cfg.get("spreadsheet_url", "").strip()
if not url:
raise RuntimeError("未配置 spreadsheet_url")
creds = _build_google_creds(config)
if creds is None:
raise RuntimeError("无可用 Google 认证凭证(需要 GOOGLE_CREDENTIALS、GOOGLE_OAUTH_REFRESH_TOKEN 或 credentials.json)")
gc = gspread.authorize(creds)
return gc.open_by_url(url)
def _get_or_create_worksheet(sh, sheet_name: str, header: list[str]):
"""获取或创建指定名称的 worksheet,确保有表头。"""
import gspread
try:
ws = sh.worksheet(sheet_name)
except gspread.WorksheetNotFound:
ws = sh.add_worksheet(title=sheet_name, rows=1000, cols=len(header))
ws.append_row(header, value_input_option="RAW")
log.info(f"创建新 Sheet: {sheet_name}")
return ws
def _build_header(config: dict[str, Any]) -> list[str]:
"""根据 config 构建标准表头。"""
col_cfg = config.get("columns", {})
dim_count = 5
if col_cfg.get("ad_id_idx") is not None:
dim_count = max(dim_count, col_cfg["ad_id_idx"] + 1)
dims = [""] * dim_count
dims[col_cfg.get("country_idx", 0)] = "Country"
dims[col_cfg.get("platform_idx", 1)] = "Platform"
dims[col_cfg.get("campaign_type_idx", 2)] = "Campaign Type"
dims[col_cfg.get("campaign_name_idx", 3)] = "Campaign Name"
dims[col_cfg.get("date_idx", 4)] = "Date"
if col_cfg.get("adset_name_idx") is not None:
dims[col_cfg["adset_name_idx"]] = "Ad Set Name"
if col_cfg.get("ad_name_idx") is not None:
dims[col_cfg["ad_name_idx"]] = "Ad Name"
if col_cfg.get("ad_id_idx") is not None:
dims[col_cfg["ad_id_idx"]] = "Ad ID"
return dims + config.get("metrics", [])
def write_api_data_to_sheets(
df: pd.DataFrame,
config_path: Path | None = None,
) -> dict[str, int]:
"""
将 API 拉取的数据按平台写入不同的 Google Sheets worksheet。
每个平台一个 sheet(如 Meta、Google、Criteo),自动创建不存在的 sheet。
按 (country, date, campaign_name, ad_name) 去重,只追加新行。
"""
config = load_config(config_path)
sh = _get_spreadsheet(config)
col_cfg = config.get("columns", {})
country_idx = col_cfg.get("country_idx", 0)
date_idx = col_cfg.get("date_idx", 4)
campaign_name_idx = col_cfg.get("campaign_name_idx", 3)
metrics_list = config.get("metrics", [])
header = _build_header(config)
key_indices = [country_idx, date_idx, campaign_name_idx]
platforms = df["platform"].unique() if "platform" in df.columns else []
total_added = 0
for platform in platforms:
p_df = df[df["platform"] == platform]
ws = _get_or_create_worksheet(sh, platform, header)
all_values = ws.get_all_values()
col_order = all_values[0] if all_values else header
existing_rows = all_values[1:] if len(all_values) > 1 else []
existing_keys = set()
for row in existing_rows:
if len(row) > max(key_indices):
key = tuple(row[i].strip() for i in key_indices)
existing_keys.add(key)
new_rows = []
for _, row_data in p_df.iterrows():
dt = row_data.get("date")
date_str = dt.strftime("%Y-%m-%d") if hasattr(dt, "strftime") else str(dt).split(" ")[0]
campaign_name = str(row_data.get("campaign_name", ""))
ad_name = str(row_data.get("ad_name", ""))
country = str(row_data.get("country", ""))
key = (country, date_str, campaign_name, ad_name)
if key in existing_keys:
continue
out_row = [""] * len(col_order)
dim_map = {
col_cfg.get("country_idx", 0): str(row_data.get("country", "")),
col_cfg.get("platform_idx", 1): platform,
col_cfg.get("campaign_type_idx", 2): str(row_data.get("campaign_type", "")),
campaign_name_idx: campaign_name,
date_idx: date_str,
}
adset_name_idx = col_cfg.get("adset_name_idx")
ad_name_idx = col_cfg.get("ad_name_idx")
ad_id_idx = col_cfg.get("ad_id_idx")
if adset_name_idx is not None:
dim_map[adset_name_idx] = str(row_data.get("adset_name", ""))
if ad_name_idx is not None:
dim_map[ad_name_idx] = ad_name
if ad_id_idx is not None:
dim_map[ad_id_idx] = str(row_data.get("ad_id", ""))
for idx, val in dim_map.items():
if idx < len(out_row):
out_row[idx] = val
for col_name in col_order:
if col_name in metrics_list and col_name in row_data.index:
val = row_data[col_name]
if val is None or (isinstance(val, float) and val != val):
continue
col_idx = col_order.index(col_name)
v = float(val)
out_row[col_idx] = str(int(v)) if v == int(v) else str(v)
new_rows.append(out_row)
existing_keys.add(key)
if new_rows:
ws.append_rows(new_rows, value_input_option="USER_ENTERED")
log.info(f"写入 Sheet [{platform}]: {len(new_rows)} 行新数据")
total_added += len(new_rows)
return {"added": total_added}
# ─────────────────────────────────────────────
# 统一加载入口
# ─────────────────────────────────────────────
_data_source: str = "unknown"
def load_raw(config_path: Path | None = None) -> pd.DataFrame:
"""优先 Google Sheets 分平台表,回退本地 Excel。"""
global _data_source
config = load_config(config_path)
df = _load_from_google_sheets(config)
if df is not None:
_data_source = "sheets"
else:
df = _load_from_excel(config)
_data_source = "excel"
return df
def normalize_columns(df: pd.DataFrame, config: dict[str, Any]) -> pd.DataFrame:
"""用列索引将编码异常的列重命名为可读字段名。"""
col_cfg = config.get("columns", {})
idx_map = {
col_cfg.get("country_idx", 0): "country",
col_cfg.get("platform_idx", 1): "platform",
col_cfg.get("campaign_type_idx", 2): "campaign_type",
col_cfg.get("campaign_name_idx", 3): "campaign_name",
col_cfg.get("date_idx", 4): "date",
}
rename = {}
for idx, name in idx_map.items():
if idx < len(df.columns):
rename[df.columns[idx]] = name
return df.rename(columns=rename)
def _clean_numeric(s: pd.Series) -> pd.Series:
"""
清洗 Google Sheets 格式化字符串。
带 % 的值先标记,去掉符号后除以 100 还原为小数。
"""
raw = s.astype(str)
is_pct = raw.str.contains('%', na=False)
cleaned = raw.str.replace(r'[$,%\s]', '', regex=True)
values = pd.to_numeric(cleaned, errors="coerce")
values[is_pct] = values[is_pct] / 100.0
return values
def load_daily_data(config_path: Path | None = None) -> pd.DataFrame:
"""
加载并清洗数据,只保留日度明细行。
自动从 date 列派生 year、week(ISO week)、month。
兼容本地 Excel(数值)和 Google Sheets(格式化字符串)。
"""
config = load_config(config_path)
df = load_raw(config_path)
df = normalize_columns(df, config)
# Excel 层级结构中,子行的维度列为空字符串,需要前向填充
# API 写入的 Sheets 分平台表每行都有完整字段,不需要 ffill
if _data_source == "excel":
for dim in ["country", "platform", "campaign_type", "campaign_name"]:
if dim in df.columns:
df[dim] = df[dim].replace("", pd.NA).ffill()
df["date"] = pd.to_datetime(df["date"], format="mixed", errors="coerce")
df = df[df["date"].notna()].copy()
df = df[df["country"].notna() & (~df["country"].astype(str).str.contains("Total|Applied filters", na=False))].copy()
df = df[df["platform"].notna() & (df["platform"].astype(str).str.strip() != "Total")].copy()
df = df[df["campaign_type"].notna() & (df["campaign_type"].astype(str).str.strip() != "Total")].copy()
df = df[df["campaign_name"].notna() & (df["campaign_name"].astype(str).str.strip() != "Total")].copy()
df["year"] = df["date"].dt.isocalendar().year.astype(int)
df["week"] = df["date"].dt.isocalendar().week.astype(int)
df["month"] = df["date"].dt.month
metrics = config.get("metrics", [])
for col in metrics:
if col in df.columns:
df[col] = _clean_numeric(df[col])
df = df.reset_index(drop=True)
df = _fix_campaign_type_from_name(df)
dedup_cols = ["date", "platform", "campaign_name", "country"]
if all(c in df.columns for c in dedup_cols):
before = len(df)
df = df.drop_duplicates(subset=dedup_cols, keep="last")
dropped = before - len(df)
if dropped:
log.info(f"去重删除 {dropped} 行重复数据")
df = df.reset_index(drop=True)
return df
_VALID_FUNNELS = {"Awareness", "Consideration", "Conversion"}
def _fix_campaign_type_from_name(df: pd.DataFrame) -> pd.DataFrame:
"""根据 campaign_name 中的第二段修正 campaign_type(如名称含明确漏斗标识)。"""
if "campaign_name" not in df.columns or "campaign_type" not in df.columns:
return df
parts = df["campaign_name"].astype(str).str.split("_", n=2)
name_funnel = parts.str[1]
mask = name_funnel.isin(_VALID_FUNNELS) & (name_funnel != df["campaign_type"])
if mask.any():
log.info(f"根据 campaign_name 修正 {mask.sum()} 行 campaign_type")
df.loc[mask, "campaign_type"] = name_funnel[mask]
return df
# 兼容旧接口
load_weekly_data = load_daily_data
def get_available_weeks(df: pd.DataFrame) -> list[tuple[int, int]]:
"""返回可用的 (year, week) 列表,按时间排序。"""
pairs = df[["year", "week"]].drop_duplicates().sort_values(["year", "week"])
return list(pairs.itertuples(index=False, name=None))
def get_latest_week(df: pd.DataFrame) -> tuple[int, int]:
"""返回数据中最新一周的 (year, week)。"""
weeks = get_available_weeks(df)
if not weeks:
raise ValueError("数据中无可用周数据")
return weeks[-1]
def get_available_months(df: pd.DataFrame) -> list[tuple[int, int]]:
"""返回可用的 (year, month) 列表,按时间排序。"""
cal_year = df["date"].dt.year
pairs = pd.DataFrame({"year": cal_year, "month": df["month"]}).drop_duplicates().sort_values(["year", "month"])
return list(pairs.itertuples(index=False, name=None))