fetch_daily.py 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740
  1. #!/usr/bin/env python
  2. # coding=utf-8
  3. """
  4. 按天增量获取数据 - 通用版本
  5. 支持并发获取,自动跳过已有数据
  6. 用法:
  7. python fetch_daily.py tasks/xxx/query.sql # 获取最近7天
  8. python fetch_daily.py tasks/xxx/query.sql --days 30 # 获取最近30天
  9. python fetch_daily.py tasks/xxx/query.sql --start 20260101 --end 20260107
  10. python fetch_daily.py tasks/xxx/query.sql --date 20260105 # 单天
  11. python fetch_daily.py tasks/xxx/query.sql --date 20260105 --hh 08 # 单天单小时
  12. python fetch_daily.py tasks/xxx/query.sql --force # 强制重新获取
  13. python fetch_daily.py tasks/xxx/query.sql --workers 10 # 设置天级并发数
  14. python fetch_daily.py tasks/xxx/query.sql --parallel 50 # 单天多线程下载(默认50,大数据量推荐)
  15. python fetch_daily.py tasks/xxx/query.sql --parallel 0 # 关闭多线程,使用单线程下载
  16. python fetch_daily.py tasks/xxx/query.sql --feishu # 获取后上传到飞书表格
  17. python fetch_daily.py tasks/xxx/query.sql --feishu TOKEN # 指定飞书表格token
  18. python fetch_daily.py tasks/xxx/query.sql --merge --feishu # 仅合并并上传飞书
  19. python fetch_daily.py tasks/xxx/query.sql --config piaoquan_api # 切换 ODPS 配置
  20. """
  21. import argparse
  22. import sys
  23. from datetime import datetime, timedelta
  24. from pathlib import Path
  25. from concurrent.futures import ThreadPoolExecutor, as_completed
  26. import threading
  27. sys.path.insert(0, str(Path(__file__).parent / "lib"))
  28. from odps_module import ODPSClient
  29. import csv
  30. # 线程安全的计数器
  31. counter_lock = threading.Lock()
  32. success_count = 0
  33. fail_count = 0
  34. def get_existing_dates(daily_dir, hh=None):
  35. """获取已下载的日期列表(可选指定小时)"""
  36. existing = set()
  37. if not daily_dir.exists():
  38. return existing
  39. for f in daily_dir.glob("*.csv"):
  40. try:
  41. stem = f.stem
  42. if hh is not None:
  43. # 带小时格式:20250101_08
  44. if len(stem) == 11 and stem[8] == '_':
  45. dt = stem[:8]
  46. file_hh = stem[9:11]
  47. if dt.isdigit() and file_hh == hh:
  48. existing.add(dt)
  49. else:
  50. # 仅日期格式:20250101
  51. if len(stem) == 8 and stem.isdigit():
  52. existing.add(stem)
  53. except:
  54. pass
  55. return existing
  56. def merge_csv_files(daily_dir, output_file=None):
  57. """合并目录下所有日期 CSV 文件,只保留一个表头"""
  58. csv_files = sorted(daily_dir.glob("*.csv"))
  59. if not csv_files:
  60. print("没有找到 CSV 文件")
  61. return None
  62. if output_file is None:
  63. output_file = daily_dir.parent / f"{daily_dir.name}_merged.csv"
  64. with open(output_file, "w", encoding="utf-8") as out:
  65. header_written = False
  66. total_rows = 0
  67. for csv_file in csv_files:
  68. with open(csv_file, "r", encoding="utf-8") as f:
  69. lines = f.readlines()
  70. if not lines:
  71. continue
  72. if not header_written:
  73. out.write(lines[0])
  74. header_written = True
  75. for line in lines[1:]:
  76. out.write(line)
  77. total_rows += 1
  78. print(f"合并完成: {len(csv_files)} 个文件, {total_rows} 行数据")
  79. print(f"输出文件: {output_file}")
  80. return output_file
  81. def infer_column_types(rows):
  82. """推断每列的类型:int, float, 或 str"""
  83. if not rows:
  84. return []
  85. num_cols = len(rows[0])
  86. col_types = []
  87. for col_idx in range(num_cols):
  88. has_float = False
  89. all_numeric = True
  90. for row in rows:
  91. if col_idx >= len(row):
  92. continue
  93. v = row[col_idx].strip() if row[col_idx] else ""
  94. if not v: # 空值不影响类型判断
  95. continue
  96. try:
  97. if '.' in v or 'e' in v.lower():
  98. float(v)
  99. has_float = True
  100. else:
  101. int(v)
  102. except ValueError:
  103. all_numeric = False
  104. break
  105. if all_numeric:
  106. col_types.append('float' if has_float else 'int')
  107. else:
  108. col_types.append('str')
  109. return col_types
  110. def convert_row_by_types(row, col_types):
  111. """按列类型转换一行数据"""
  112. result = []
  113. for i, cell in enumerate(row):
  114. if i >= len(col_types):
  115. result.append(cell)
  116. continue
  117. v = cell.strip() if cell else ""
  118. if not v:
  119. result.append("")
  120. continue
  121. col_type = col_types[i]
  122. if col_type == 'int':
  123. result.append(int(v))
  124. elif col_type == 'float':
  125. result.append(float(v))
  126. else:
  127. result.append(cell)
  128. return result
  129. def load_feishu_config(sql_file):
  130. """加载飞书配置,优先级: {sql名}.json > sql目录/default.json > 根目录/default.json > 默认值"""
  131. import json
  132. defaults = {
  133. "token": "ONZqsxB9BhGH8tt90EScSJT5nHh",
  134. "sheet_id": None,
  135. "sort": "dt:desc",
  136. "cols": None,
  137. }
  138. root_dir = Path(__file__).parent
  139. sql_dir = sql_file.parent
  140. sql_name = sql_file.stem
  141. def load_json(path, name):
  142. if path.exists():
  143. try:
  144. with open(path, "r", encoding="utf-8") as f:
  145. defaults.update(json.load(f))
  146. except Exception as e:
  147. print(f"警告: 读取 {name} 失败: {e}")
  148. # 按优先级从低到高加载(后加载的覆盖先加载的)
  149. load_json(root_dir / "default.json", "根目录/default.json")
  150. load_json(sql_dir / "default.json", "sql目录/default.json")
  151. load_json(sql_dir / f"{sql_name}.json", f"{sql_name}.json")
  152. return defaults
  153. def parse_sort_spec(sort_spec):
  154. """解析排序规格,如 'dt:desc,name:asc' -> [('dt', True), ('name', False)]"""
  155. if not sort_spec:
  156. return []
  157. result = []
  158. for part in sort_spec.split(","):
  159. part = part.strip()
  160. if not part:
  161. continue
  162. if ":" in part:
  163. field, order = part.rsplit(":", 1)
  164. desc = order.lower() != "asc"
  165. else:
  166. field, desc = part, True # 默认逆序
  167. result.append((field.strip(), desc))
  168. return result
  169. def parse_cols_spec(cols_spec):
  170. """解析列映射规格,如 'dt:日期,name,value:数值' -> [('dt', '日期'), ('name', 'name'), ('value', '数值')]"""
  171. if not cols_spec:
  172. return []
  173. result = []
  174. for part in cols_spec.split(","):
  175. part = part.strip()
  176. if not part:
  177. continue
  178. if ":" in part:
  179. old_name, new_name = part.split(":", 1)
  180. result.append((old_name.strip(), new_name.strip()))
  181. else:
  182. result.append((part, part))
  183. return result
  184. def apply_cols_mapping(header, data_rows, cols_spec):
  185. """应用列映射:筛选、排序、重命名"""
  186. col_mapping = parse_cols_spec(cols_spec)
  187. if not col_mapping:
  188. return header, data_rows
  189. # 构建索引映射
  190. header_index = {name: i for i, name in enumerate(header)}
  191. new_header = []
  192. col_indices = []
  193. for old_name, new_name in col_mapping:
  194. if old_name in header_index:
  195. col_indices.append(header_index[old_name])
  196. new_header.append(new_name)
  197. else:
  198. print(f"警告: 字段 '{old_name}' 不存在,已跳过")
  199. if not col_indices:
  200. print("警告: 没有有效的列映射,保持原样")
  201. return header, data_rows
  202. # 应用映射
  203. new_rows = []
  204. for row in data_rows:
  205. new_row = [row[i] if i < len(row) else "" for i in col_indices]
  206. new_rows.append(new_row)
  207. print(f"列映射: {len(col_indices)} 列")
  208. return new_header, new_rows
  209. def column_index_to_letter(col_idx):
  210. """列索引转字母,如 1->A, 26->Z, 27->AA"""
  211. result = ""
  212. while col_idx > 0:
  213. col_idx, remainder = divmod(col_idx - 1, 26)
  214. result = chr(65 + remainder) + result
  215. return result
  216. def upload_to_feishu(csv_file, sheet_token, sheet_id=None, sort_spec="dt:desc", cols_spec=None):
  217. """上传 CSV 文件到飞书表格(通过模板行继承样式)
  218. 第1行: 表头
  219. 第2行: 样式模板(用于继承,最后删除)
  220. 第3行起: 数据
  221. Args:
  222. csv_file: CSV 文件路径
  223. sheet_token: 飞书表格 token
  224. sheet_id: 工作表 ID(None 时自动获取第一个)
  225. sort_spec: 排序规格,如 "dt:desc,name:asc"
  226. cols_spec: 列映射规格,如 "dt:日期,name,value:数值"
  227. """
  228. from feishu import Client, LARK_HOST, APP_ID, APP_SECRET, request
  229. # 读取 CSV
  230. with open(csv_file, "r", encoding="utf-8") as f:
  231. reader = csv.reader(f)
  232. rows = list(reader)
  233. if len(rows) < 2:
  234. print("CSV 文件为空,跳过上传")
  235. return
  236. header = rows[0]
  237. data_rows = rows[1:]
  238. # 排序(在列映射之前,使用原始列名)
  239. sort_fields = parse_sort_spec(sort_spec)
  240. if sort_fields:
  241. applied = []
  242. for field, desc in reversed(sort_fields):
  243. if field in header:
  244. idx = header.index(field)
  245. data_rows.sort(key=lambda row: row[idx] if idx < len(row) else "", reverse=desc)
  246. applied.append(f"{field}:{'desc' if desc else 'asc'}")
  247. if applied:
  248. print(f"排序: {', '.join(reversed(applied))}")
  249. # 列映射(排序之后)
  250. header, data_rows = apply_cols_mapping(header, data_rows, cols_spec)
  251. # 按列推断类型并转换
  252. col_types = infer_column_types(data_rows)
  253. converted_rows = [convert_row_by_types(row, col_types) for row in data_rows]
  254. # 初始化飞书客户端
  255. client = Client(LARK_HOST)
  256. access_token = client.get_tenant_access_token(APP_ID, APP_SECRET)
  257. # 获取 sheet_id
  258. if sheet_id is None:
  259. sheet_id = client.get_sheetid(access_token, sheet_token)
  260. print(f"Sheet ID: {sheet_id}")
  261. # 获取表格信息
  262. sheet_props = client.get_sheet_properties(access_token, sheet_token, sheet_id)
  263. current_cols = sheet_props['column_count'] if sheet_props else 26
  264. header_end_col = column_index_to_letter(current_cols)
  265. # 扩展列数(CSV 列数超过当前 sheet 列数时)
  266. num_csv_cols = len(header)
  267. if num_csv_cols > current_cols:
  268. add_cols = num_csv_cols - current_cols
  269. expand_headers = {
  270. 'Content-Type': 'application/json; charset=utf-8',
  271. 'Authorization': f'Bearer {access_token}'
  272. }
  273. expand_payload = {
  274. "dimension": {
  275. "sheetId": sheet_id,
  276. "majorDimension": "COLUMNS",
  277. "length": add_cols
  278. }
  279. }
  280. try:
  281. request("POST", f"{LARK_HOST}/open-apis/sheets/v2/spreadsheets/{sheet_token}/dimension_range",
  282. expand_headers, expand_payload)
  283. print(f"扩展列数: {current_cols} -> {num_csv_cols} (+{add_cols}列)")
  284. current_cols = num_csv_cols
  285. header_end_col = column_index_to_letter(current_cols)
  286. except Exception as e:
  287. print(f" 扩展列数失败: {e}")
  288. # 读取飞书表头(获取所有列)
  289. feishu_header = client.read_range_values(access_token, sheet_token, f"{sheet_id}!A1:{header_end_col}1")
  290. feishu_cols = []
  291. if feishu_header and feishu_header[0]:
  292. feishu_cols = [c for c in feishu_header[0] if c] # 过滤 None 和空字符串
  293. if feishu_cols:
  294. print(f"飞书表头: {feishu_cols}")
  295. print(f"CSV表头: {header}")
  296. # 校验字段一致性(警告但继续,以飞书表头为准)
  297. feishu_set = set(feishu_cols)
  298. csv_set = set(header)
  299. missing_in_csv = feishu_set - csv_set
  300. missing_in_feishu = csv_set - feishu_set
  301. if missing_in_csv:
  302. print(f"警告: CSV缺少字段(将填空值): {missing_in_csv}")
  303. if missing_in_feishu:
  304. print(f"警告: 飞书缺少字段(将忽略): {missing_in_feishu}")
  305. # 按飞书表头顺序重排数据
  306. csv_col_index = {name: i for i, name in enumerate(header)}
  307. new_converted_rows = []
  308. for row in converted_rows:
  309. new_row = []
  310. for col_name in feishu_cols:
  311. if col_name in csv_col_index:
  312. new_row.append(row[csv_col_index[col_name]])
  313. else:
  314. new_row.append("") # CSV缺少的字段填空
  315. new_converted_rows.append(new_row)
  316. converted_rows = new_converted_rows
  317. header = feishu_cols
  318. print(f"已按飞书表头顺序重排数据")
  319. else:
  320. # 飞书表头为空,用 CSV 表头写入(飞书单次最多写100列,需分批)
  321. print(f"飞书表头为空,使用 CSV 表头写入")
  322. col_batch = 100
  323. for start in range(0, len(header), col_batch):
  324. end = min(start + col_batch, len(header))
  325. start_col = column_index_to_letter(start + 1)
  326. end_col = column_index_to_letter(end)
  327. batch_range = f"{sheet_id}!{start_col}1:{end_col}1"
  328. client.batch_update_values(access_token, sheet_token, {
  329. "valueRanges": [{"range": batch_range, "values": [header[start:end]]}]
  330. })
  331. total_rows = len(converted_rows)
  332. num_cols = len(header)
  333. end_col = column_index_to_letter(num_cols)
  334. print(f"上传到飞书: {total_rows} 行数据")
  335. batch_size = 500
  336. # 获取当前行数(复用之前获取的 sheet_props)
  337. current_rows = sheet_props['row_count'] if sheet_props else 2
  338. print(f"当前行数: {current_rows}, 需要数据行: {total_rows}")
  339. headers = {
  340. 'Content-Type': 'application/json; charset=utf-8',
  341. 'Authorization': f'Bearer {access_token}'
  342. }
  343. # 第1步:删除旧数据行(保留第1行表头 + 第2行样式模板),分批删除
  344. if current_rows > 2:
  345. print(f"清理旧数据({current_rows - 2}行)...")
  346. rows_to_delete = current_rows - 2
  347. delete_batch = 5000
  348. while rows_to_delete > 0:
  349. # 每次从第3行开始删除,删除后行号会自动调整
  350. batch = min(rows_to_delete, delete_batch)
  351. try:
  352. client.delete_rows(access_token, sheet_token, sheet_id, 3, 2 + batch)
  353. rows_to_delete -= batch
  354. if rows_to_delete > 0:
  355. print(f" 已删除 {current_rows - 2 - rows_to_delete}/{current_rows - 2}")
  356. except Exception as e:
  357. print(f" 清理失败: {e}")
  358. break
  359. # 第2步:扩展表格容量(insert 不会自动扩展)
  360. # 删除后当前只有2行(表头+模板),需要扩展到 2 + total_rows 行
  361. needed_rows = 2 + total_rows
  362. add_url = f"{LARK_HOST}/open-apis/sheets/v2/spreadsheets/{sheet_token}/dimension_range"
  363. add_payload = {
  364. "dimension": {
  365. "sheetId": sheet_id,
  366. "majorDimension": "ROWS",
  367. "length": total_rows # 添加数据行数
  368. }
  369. }
  370. try:
  371. request("POST", add_url, headers, add_payload)
  372. print(f"扩展容量: +{total_rows} 行")
  373. except Exception as e:
  374. print(f" 扩展容量失败: {e}")
  375. # 第3步:分批插入空行(继承第2行样式)并写入数据
  376. print(f"插入并写入 {total_rows} 行...")
  377. insert_url = f"{LARK_HOST}/open-apis/sheets/v2/spreadsheets/{sheet_token}/insert_dimension_range"
  378. # 反向处理批次(从最后一批开始),因为每次都在第3行前插入
  379. batches = [converted_rows[i:i + batch_size] for i in range(0, total_rows, batch_size)]
  380. processed = 0
  381. for batch in reversed(batches):
  382. batch_count = len(batch)
  383. # 在第3行前插入空行(继承第2行样式)
  384. insert_payload = {
  385. "dimension": {
  386. "sheetId": sheet_id,
  387. "majorDimension": "ROWS",
  388. "startIndex": 2, # 0-indexed, 第3行位置
  389. "endIndex": 2 + batch_count
  390. },
  391. "inheritStyle": "BEFORE"
  392. }
  393. try:
  394. request("POST", insert_url, headers, insert_payload)
  395. except Exception as e:
  396. print(f" 插入行失败: {e}")
  397. break
  398. # 写入数据到插入的行(第3行开始,飞书单次最多100列,需按列分批)
  399. col_batch = 100
  400. value_ranges = []
  401. for col_start in range(0, num_cols, col_batch):
  402. col_end = min(col_start + col_batch, num_cols)
  403. sc = column_index_to_letter(col_start + 1)
  404. ec = column_index_to_letter(col_end)
  405. col_range = f"{sheet_id}!{sc}3:{ec}{2 + batch_count}"
  406. col_values = [row[col_start:col_end] for row in batch]
  407. value_ranges.append({"range": col_range, "values": col_values})
  408. client.batch_update_values(access_token, sheet_token, {
  409. "valueRanges": value_ranges
  410. })
  411. processed += batch_count
  412. print(f" 处理: {processed}/{total_rows}")
  413. # 第4步:删除末尾多余的空行(扩展容量时添加的)
  414. # 重新查询实际行数,避免因初始行数不同导致计算偏差
  415. sheet_props_after = client.get_sheet_properties(access_token, sheet_token, sheet_id)
  416. actual_row_count = sheet_props_after['row_count'] if sheet_props_after else 0
  417. # 期望行数:初始保留行数 + 插入的数据行数
  418. rows_after_cleanup = 2 if current_rows > 2 else current_rows
  419. final_row_count = rows_after_cleanup + total_rows
  420. if actual_row_count > final_row_count:
  421. print(f"清理多余空行({actual_row_count - final_row_count}行)...")
  422. try:
  423. client.delete_rows(access_token, sheet_token, sheet_id, final_row_count + 1, actual_row_count)
  424. except Exception as e:
  425. print(f" 清理失败: {e}")
  426. # 第5步:删除模板行(第2行),仅当初始存在模板行时
  427. if current_rows >= 2:
  428. print(f"删除模板行...")
  429. try:
  430. client.delete_rows(access_token, sheet_token, sheet_id, 2, 2)
  431. except Exception as e:
  432. print(f" 删除模板行失败: {e}")
  433. print(f"飞书上传完成: {sheet_token}")
  434. def get_date_range(start_str, end_str):
  435. """生成日期范围列表"""
  436. start = datetime.strptime(start_str, "%Y%m%d")
  437. end = datetime.strptime(end_str, "%Y%m%d")
  438. dates = []
  439. current = start
  440. while current <= end:
  441. dates.append(current.strftime("%Y%m%d"))
  442. current += timedelta(days=1)
  443. return dates
  444. def fetch_single_day(dt, sql_template, daily_dir, parallel_threads=0, config="default", hh=None):
  445. """获取单天数据(可选指定小时)"""
  446. global success_count, fail_count
  447. try:
  448. client = ODPSClient(config=config)
  449. sql = sql_template.replace("${dt}", dt)
  450. if hh is not None:
  451. sql = sql.replace("${hh}", hh)
  452. output_file = daily_dir / f"{dt}_{hh}.csv"
  453. else:
  454. output_file = daily_dir / f"{dt}.csv"
  455. # 下载到文件
  456. if parallel_threads > 0:
  457. # 多线程并行下载(适合大数据量)
  458. client.execute_sql_result_save_file_parallel(sql, str(output_file), workers=parallel_threads)
  459. else:
  460. # 单线程下载
  461. client.execute_sql_result_save_file(sql, str(output_file))
  462. # 检查结果
  463. if output_file.exists():
  464. row_count = sum(1 for _ in open(output_file)) - 1 # 减去表头
  465. with counter_lock:
  466. success_count += 1
  467. if row_count > 0:
  468. return (dt, "success", row_count)
  469. else:
  470. return (dt, "empty", 0)
  471. else:
  472. with counter_lock:
  473. fail_count += 1
  474. return (dt, "fail", 0)
  475. except Exception as e:
  476. with counter_lock:
  477. fail_count += 1
  478. return (dt, "error", str(e))
  479. def main():
  480. global success_count, fail_count
  481. parser = argparse.ArgumentParser(description="按天增量获取数据")
  482. parser.add_argument("sql_file", type=str, help="SQL文件路径")
  483. parser.add_argument("--days", type=int, default=7, help="获取最近N天 (默认7)")
  484. parser.add_argument("--start", type=str, help="开始日期 YYYYMMDD")
  485. parser.add_argument("--end", type=str, help="结束日期 YYYYMMDD")
  486. parser.add_argument("--date", type=str, help="单天日期 YYYYMMDD")
  487. parser.add_argument("--hh", type=str, default=None, help="小时 HH (00-23),需配合 --date 使用")
  488. parser.add_argument("--force", action="store_true", help="强制重新获取")
  489. parser.add_argument("--workers", type=int, default=5, help="天级并发数 (默认5)")
  490. parser.add_argument("--parallel", type=int, default=50, help="单天多线程下载 (默认50, 大数据量推荐)")
  491. parser.add_argument("--merge", action="store_true", help="合并所有日期数据到一个文件")
  492. parser.add_argument("--feishu", nargs="?", const="__USE_CONFIG__",
  493. help="上传到飞书表格")
  494. parser.add_argument("--sheet-id", type=str, default=None, help="飞书工作表ID")
  495. parser.add_argument("--sort", type=str, default=None, help="排序: 字段:asc/desc")
  496. parser.add_argument("--cols", type=str, default=None, help="列映射: 原名:新名,...")
  497. parser.add_argument("--config", type=str, default="default", help="ODPS配置: default 或 piaoquan_api")
  498. args = parser.parse_args()
  499. # 解析 SQL 文件路径
  500. sql_file = Path(args.sql_file).resolve()
  501. if not sql_file.exists():
  502. print(f"错误: 找不到 {sql_file}")
  503. return
  504. # 加载飞书配置(优先级: 命令行 > {sql名}.json > sql目录/default.json > 根目录/default.json > 默认值)
  505. feishu_config = load_feishu_config(sql_file)
  506. if args.feishu == "__USE_CONFIG__":
  507. args.feishu = feishu_config["token"]
  508. elif args.feishu is None:
  509. pass # 未启用飞书上传
  510. # 命令行参数覆盖配置文件
  511. if args.sheet_id is None:
  512. args.sheet_id = feishu_config["sheet_id"]
  513. if args.sort is None:
  514. args.sort = feishu_config["sort"]
  515. if args.cols is None:
  516. args.cols = feishu_config["cols"]
  517. # 打印飞书配置
  518. if args.feishu:
  519. print(f"飞书配置: token={args.feishu}, sheet_id={args.sheet_id}, sort={args.sort}, cols={args.cols}")
  520. # 输出目录:SQL 同目录下的 output/SQL文件名/
  521. output_dir = sql_file.parent / "output"
  522. daily_dir = output_dir / sql_file.stem
  523. daily_dir.mkdir(parents=True, exist_ok=True)
  524. print(f"SQL文件: {sql_file}")
  525. print(f"数据目录: {daily_dir}")
  526. # 仅合并模式:不获取数据,直接合并已有文件
  527. if args.merge:
  528. existing_dates = get_existing_dates(daily_dir)
  529. print(f"已有数据: {len(existing_dates)}天")
  530. if existing_dates:
  531. merged_file = merge_csv_files(daily_dir)
  532. # 如果指定了飞书上传
  533. if args.feishu and merged_file:
  534. upload_to_feishu(merged_file, args.feishu, args.sheet_id, args.sort, args.cols)
  535. else:
  536. print("没有可合并的数据")
  537. return
  538. # 确定日期范围
  539. if args.date:
  540. target_dates = [args.date]
  541. elif args.start and args.end:
  542. target_dates = get_date_range(args.start, args.end)
  543. else:
  544. today = datetime.now()
  545. end_date = (today - timedelta(days=1)).strftime("%Y%m%d")
  546. start_date = (today - timedelta(days=args.days)).strftime("%Y%m%d")
  547. target_dates = get_date_range(start_date, end_date)
  548. print(f"目标日期: {target_dates[0]} ~ {target_dates[-1]} ({len(target_dates)}天)")
  549. # 检查已有数据
  550. existing_dates = get_existing_dates(daily_dir, args.hh)
  551. if args.hh:
  552. print(f"已有数据: {len(existing_dates)}天 (hh={args.hh})")
  553. else:
  554. print(f"已有数据: {len(existing_dates)}天")
  555. # 确定需要获取的日期
  556. if args.force:
  557. missing_dates = target_dates
  558. print(f"强制模式: 重新获取所有 {len(missing_dates)} 天")
  559. else:
  560. missing_dates = [d for d in target_dates if d not in existing_dates]
  561. print(f"需要获取: {len(missing_dates)}天")
  562. if not missing_dates:
  563. print("没有需要获取的数据,退出")
  564. return
  565. # 读取 SQL 模板
  566. sql_template = sql_file.read_text(encoding="utf-8")
  567. # 检测 SQL 中是否包含 ${dt} 变量
  568. has_dt_var = "${dt}" in sql_template
  569. # 重置计数器
  570. success_count = 0
  571. fail_count = 0
  572. # 如果 SQL 中没有 ${dt},只需执行一次
  573. if not has_dt_var:
  574. print("\n检测到 SQL 中不含 ${dt} 变量,只执行一次...")
  575. target_dates = ["20000101"] # 用虚拟日期
  576. missing_dates = target_dates
  577. output_file = output_dir / f"{sql_file.stem}.csv"
  578. output_file.parent.mkdir(parents=True, exist_ok=True)
  579. try:
  580. client = ODPSClient(config=args.config)
  581. if args.parallel > 0:
  582. client.execute_sql_result_save_file_parallel(sql_template, str(output_file), workers=args.parallel)
  583. else:
  584. client.execute_sql_result_save_file(sql_template, str(output_file))
  585. print(f"数据目录: {output_file}")
  586. # 如果指定了飞书上传
  587. if args.feishu and output_file.exists():
  588. upload_to_feishu(output_file, args.feishu, args.sheet_id, args.sort, args.cols)
  589. except Exception as e:
  590. print(f"✗ 执行失败: {e}")
  591. return
  592. # 并发获取
  593. print(f"目标日期: {target_dates[0]} ~ {target_dates[-1]} ({len(target_dates)}天)")
  594. workers = min(args.workers, len(missing_dates))
  595. if args.parallel > 0:
  596. print(f"\n开始获取 (天级并发: {workers}, 单天多线程: {args.parallel})...")
  597. else:
  598. print(f"\n开始获取 (并发数: {workers})...")
  599. with ThreadPoolExecutor(max_workers=workers) as executor:
  600. futures = {
  601. executor.submit(fetch_single_day, dt, sql_template, daily_dir, args.parallel, args.config, args.hh): dt
  602. for dt in missing_dates
  603. }
  604. completed = 0
  605. for future in as_completed(futures):
  606. completed += 1
  607. dt, status, info = future.result()
  608. if status == "success":
  609. print(f" [{completed}/{len(missing_dates)}] ✓ {dt}: {info} 行")
  610. elif status == "empty":
  611. print(f" [{completed}/{len(missing_dates)}] ⚠ {dt}: 无数据")
  612. elif status == "error":
  613. print(f" [{completed}/{len(missing_dates)}] ✗ {dt}: {info}")
  614. else:
  615. print(f" [{completed}/{len(missing_dates)}] ✗ {dt}: 失败")
  616. print(f"\n完成! 成功: {success_count}, 失败: {fail_count}")
  617. print(f"数据目录: {daily_dir}")
  618. # 如果指定了飞书上传,先合并再上传
  619. if args.feishu:
  620. merged_file = merge_csv_files(daily_dir)
  621. if merged_file:
  622. upload_to_feishu(merged_file, args.feishu, args.sheet_id, args.sort, args.cols)
  623. if __name__ == "__main__":
  624. main()