#!/usr/bin/env python # -*- coding: utf-8 -*- """ CSV数据导入数据库脚本 功能：将三个CSV文件的数据以author_name为唯一键导入到数据库 - ai_statistics.csv -> ai_statistics表 - ai_statistics_day.csv -> ai_statistics_day表 - ai_statistics_days.csv -> ai_statistics_days表 """ import os import sys import csv from decimal import Decimal from typing import Dict, List, Optional from datetime import datetime from database_config import DatabaseManager, DB_CONFIG from log_config import setup_logger class CSVImporter: """CSV数据导入器""" def __init__(self, db_config: Optional[Dict] = None): """初始化导入器 Args: db_config: 数据库配置，默认使用database_config.DB_CONFIG """ self.script_dir = os.path.dirname(os.path.abspath(__file__)) self.db_manager = DatabaseManager(db_config) # 初始化日志 self.logger = setup_logger('import_csv', os.path.join(self.script_dir, 'logs', 'import_csv.log')) # CSV文件路径 self.csv_files = { 'ai_statistics': os.path.join(self.script_dir, 'ai_statistics.csv'), 'ai_statistics_day': os.path.join(self.script_dir, 'ai_statistics_day.csv'), 'ai_statistics_days': os.path.join(self.script_dir, 'ai_statistics_days.csv'), } def read_csv(self, file_path: str) -> List[Dict]: """读取CSV文件 Args: file_path: CSV文件路径 Returns: 数据行列表 """ if not os.path.exists(file_path): print(f"[X] CSV文件不存在: {file_path}") self.logger.error(f"CSV文件不存在: {file_path}") return [] rows = [] try: with open(file_path, 'r', encoding='utf-8-sig') as f: reader = csv.DictReader(f) for row in reader: rows.append(row) print(f"[OK] 读取CSV文件成功: {file_path}, 共 {len(rows)} 行") self.logger.info(f"读取CSV文件成功: {file_path}, 共 {len(rows)} 行") return rows except Exception as e: print(f"[X] 读取CSV文件失败: {file_path}, 错误: {e}") self.logger.error(f"读取CSV文件失败: {file_path}, 错误: {e}") return [] def convert_value(self, value: str, field_type: str): """转换数据类型 Args: value: 字符串值 field_type: 字段类型 (int, float, decimal, str) Returns: 转换后的值 """ if value == '' or value is None: return None try: if field_type == 'int': return int(float(value)) elif field_type == 'float': return float(value) elif field_type == 'decimal': return Decimal(str(value)) else: return str(value) except Exception as e: print(f"[!] 数据类型转换失败: {value} -> {field_type}, 错误: {e}") self.logger.warning(f"数据类型转换失败: {value} -> {field_type}, 错误: {e}") return None def get_author_id(self, author_name: str, channel: int = 1) -> int: """从 ai_authors 表查询 author_id Args: author_name: 作者名称 channel: 渠道，默认1=百度 Returns: author_id，如果找不到返回0 """ try: sql = "SELECT id FROM ai_authors WHERE author_name = %s AND channel = %s" result = self.db_manager.execute_query(sql, (author_name, channel), fetch_one=True) if result: author_id = result.get('id', 0) if isinstance(result, dict) else result[0] self.logger.debug(f"查询author_id: {author_name} -> {author_id}") return author_id else: self.logger.warning(f"未找到author_id: {author_name}, channel={channel}") return 0 except Exception as e: self.logger.error(f"查询author_id失败: {author_name}, 错误: {e}") return 0 def import_ai_statistics(self, batch_size: int = 50) -> bool: """导入 ai_statistics 表数据（使用批量提交） Args: batch_size: 批量提交大小，默认50条 """ print("\n" + "="*70) print("开始导入 ai_statistics 表数据") print("="*70) csv_file = self.csv_files['ai_statistics'] rows = self.read_csv(csv_file) if not rows: print("[X] 没有数据可导入") self.logger.warning("ai_statistics表没有数据可导入") return False self.logger.info(f"开始导入ai_statistics表数据，共 {len(rows)} 条记录，批量大小: {batch_size}") print(f"\n总计 {len(rows)} 条记录，分批导入（每批 {batch_size} 条）\n") success_count = 0 failed_count = 0 batch_params = [] # 批量参数列表 # 构建SQL（所有记录使用相同SQL） first_record_keys = None sql_template = None for idx, row in enumerate(rows, 1): author_name = row.get('author_name', '').strip() if not author_name: continue try: # 处理slide_ratio值（CSV中已是小数格式） slide_ratio_value = float(self.convert_value(row.get('slide_ratio', '0'), 'float') or 0.0) slide_ratio_value = min(slide_ratio_value, 9.9999) # 获取channel channel = self.convert_value(row.get('channel', '1'), 'int') or 1 # 从数据库查询正确的 author_id（不使用CSV中的0） author_id = self.get_author_id(author_name, int(channel)) if author_id == 0: self.logger.warning(f"跳过无效记录: {author_name}, author_id未找到") failed_count += 1 continue # 准备数据 record = { 'author_id': author_id, # 使用查询到的author_id 'author_name': author_name, 'channel': channel, 'date': row.get('date', '').strip(), 'submission_count': self.convert_value(row.get('submission_count', '0'), 'int') or 0, 'read_count': self.convert_value(row.get('read_count', '0'), 'int') or 0, 'comment_count': self.convert_value(row.get('comment_count', '0'), 'int') or 0, 'comment_rate': min(float(self.convert_value(row.get('comment_rate', '0'), 'float') or 0.0), 9.9999), 'like_count': self.convert_value(row.get('like_count', '0'), 'int') or 0, 'like_rate': min(float(self.convert_value(row.get('like_rate', '0'), 'float') or 0.0), 9.9999), 'favorite_count': self.convert_value(row.get('favorite_count', '0'), 'int') or 0, 'favorite_rate': min(float(self.convert_value(row.get('favorite_rate', '0'), 'float') or 0.0), 9.9999), 'share_count': self.convert_value(row.get('share_count', '0'), 'int') or 0, 'share_rate': min(float(self.convert_value(row.get('share_rate', '0'), 'float') or 0.0), 9.9999), 'slide_ratio': slide_ratio_value, 'baidu_search_volume': self.convert_value(row.get('baidu_search_volume', '0'), 'int') or 0, 'updated_at': datetime.now().strftime('%Y-%m-%d %H:%M:%S'), # 添加更新时间戳，强制更新 } # 第一条记录时构建SQL if sql_template is None: first_record_keys = list(record.keys()) columns = ', '.join(first_record_keys) placeholders = ', '.join(['%s'] * len(first_record_keys)) update_parts = [f"{key} = VALUES({key})" for key in first_record_keys if key not in ['author_name', 'channel', 'date']] sql_template = f""" INSERT INTO ai_statistics ({columns}) VALUES ({placeholders}) ON DUPLICATE KEY UPDATE {', '.join(update_parts)} """ # 添加到批量参数（确保first_record_keys已初始化） if first_record_keys is not None: batch_params.append(tuple(record[key] for key in first_record_keys)) # 当达到批量大小或最后一条记录时，执行批量插入 if len(batch_params) >= batch_size or idx == len(rows): try: # 使用 execute_many 批量提交 result_count = self.db_manager.execute_many(sql_template, batch_params, autocommit=True) success_count += result_count print(f"[批次提交] 已导入 {success_count} 条记录（本批: {result_count}/{len(batch_params)}）") self.logger.info(f"ai_statistics表批量提交: {result_count}/{len(batch_params)} 条记录") batch_params = [] # 清空批次 except Exception as batch_error: failed_count += len(batch_params) print(f" [X] 批次提交失败: {batch_error}") self.logger.error(f"ai_statistics表批量提交失败: {batch_error}") batch_params = [] except Exception as e: failed_count += 1 print(f" [X] 处理失败 ({author_name}): {e}") self.logger.error(f"ai_statistics表处理失败: {author_name}, 错误: {e}") continue print("\n" + "="*70) print(f"[OK] ai_statistics 表数据导入完成") print(f" 成功: {success_count} 条记录") if failed_count > 0: print(f" 失败: {failed_count} 条记录") print("="*70) self.logger.info(f"ai_statistics表数据导入完成: 成功 {success_count} 条，失败 {failed_count} 条") return success_count > 0 def import_ai_statistics_day(self, batch_size: int = 50) -> bool: """导入 ai_statistics_day 表数据（使用批量提交） Args: batch_size: 批量提交大小，默认50条 """ print("\n" + "="*70) print("开始导入 ai_statistics_day 表数据") print("="*70) csv_file = self.csv_files['ai_statistics_day'] rows = self.read_csv(csv_file) if not rows: print("[X] 没有数据可导入") self.logger.warning("ai_statistics_day表没有数据可导入") return False self.logger.info(f"开始导入ai_statistics_day表数据，共 {len(rows)} 条记录，批量大小: {batch_size}") print(f"\n总计 {len(rows)} 条记录，分批导入（每批 {batch_size} 条）\n") success_count = 0 failed_count = 0 batch_params = [] first_record_keys = None sql_template = None for idx, row in enumerate(rows, 1): author_name = row.get('author_name', '').strip() if not author_name: continue try: # 处理avg_slide_ratio值（CSV中已是小数格式） avg_slide_ratio_value = float(self.convert_value(row.get('avg_slide_ratio', '0'), 'float') or 0.0) avg_slide_ratio_value = min(avg_slide_ratio_value, 9.9999) # 获取channel并查询author_id channel = self.convert_value(row.get('channel', '1'), 'int') or 1 author_id = self.get_author_id(author_name, int(channel)) if author_id == 0: self.logger.warning(f"跳过无效记录: {author_name}, author_id未找到") failed_count += 1 continue record = { 'author_id': author_id, 'author_name': author_name, 'channel': channel, 'stat_date': row.get('stat_date', '').strip(), 'total_submission_count': self.convert_value(row.get('total_submission_count', '0'), 'int') or 0, 'total_read_count': self.convert_value(row.get('total_read_count', '0'), 'int') or 0, 'total_comment_count': self.convert_value(row.get('total_comment_count', '0'), 'int') or 0, 'total_like_count': self.convert_value(row.get('total_like_count', '0'), 'int') or 0, 'total_favorite_count': self.convert_value(row.get('total_favorite_count', '0'), 'int') or 0, 'total_share_count': self.convert_value(row.get('total_share_count', '0'), 'int') or 0, 'avg_comment_rate': min(float(self.convert_value(row.get('avg_comment_rate', '0'), 'float') or 0.0), 9.9999), 'avg_like_rate': min(float(self.convert_value(row.get('avg_like_rate', '0'), 'float') or 0.0), 9.9999), 'avg_favorite_rate': min(float(self.convert_value(row.get('avg_favorite_rate', '0'), 'float') or 0.0), 9.9999), 'avg_share_rate': min(float(self.convert_value(row.get('avg_share_rate', '0'), 'float') or 0.0), 9.9999), 'avg_slide_ratio': avg_slide_ratio_value, 'total_baidu_search_volume': self.convert_value(row.get('total_baidu_search_volume', '0'), 'int') or 0, 'updated_at': datetime.now().strftime('%Y-%m-%d %H:%M:%S'), # 添加更新时间戳，强制更新 } if sql_template is None: first_record_keys = list(record.keys()) columns = ', '.join(first_record_keys) placeholders = ', '.join(['%s'] * len(first_record_keys)) update_parts = [f"{key} = VALUES({key})" for key in first_record_keys if key not in ['author_name', 'channel', 'stat_date']] sql_template = f""" INSERT INTO ai_statistics_day ({columns}) VALUES ({placeholders}) ON DUPLICATE KEY UPDATE {', '.join(update_parts)} """ if first_record_keys is not None: batch_params.append(tuple(record[key] for key in first_record_keys)) if len(batch_params) >= batch_size or idx == len(rows): try: result_count = self.db_manager.execute_many(sql_template, batch_params, autocommit=True) success_count += result_count print(f"[批次提交] 已导入 {success_count} 条记录（本批: {result_count}/{len(batch_params)}）") self.logger.info(f"ai_statistics_day表批量提交: {result_count}/{len(batch_params)} 条记录") batch_params = [] except Exception as batch_error: failed_count += len(batch_params) print(f" [X] 批次提交失败: {batch_error}") self.logger.error(f"ai_statistics_day表批量提交失败: {batch_error}") batch_params = [] except Exception as e: failed_count += 1 print(f" [X] 处理失败 ({author_name}): {e}") self.logger.error(f"ai_statistics_day表处理失败: {author_name}, 错误: {e}") continue print("\n" + "="*70) print(f"[OK] ai_statistics_day 表数据导入完成") print(f" 成功: {success_count} 条记录") if failed_count > 0: print(f" 失败: {failed_count} 条记录") print("="*70) self.logger.info(f"ai_statistics_day表数据导入完成: 成功 {success_count} 条，失败 {failed_count} 条") return success_count > 0 def import_ai_statistics_days(self, batch_size: int = 50) -> bool: """导入 ai_statistics_days 表数据（仅当日数据：day_revenue） 同时自动拆分数据到 ai_statistics_weekly 和 ai_statistics_monthly 表 Args: batch_size: 批量提交大小，默认50条 """ print("\n" + "="*70) print("开始导入 ai_statistics_days 表数据（拆分到3个表）") print("="*70) csv_file = self.csv_files['ai_statistics_days'] rows = self.read_csv(csv_file) if not rows: print("[X] 没有数据可导入") self.logger.warning("ai_statistics_days表没有数据可导入") return False self.logger.info(f"开始导入数据，共 {len(rows)} 条记录，批量大小: {batch_size}") print(f"\n总计 {len(rows)} 条记录，将拆分到3个表\n") # 三个表的统计 days_success = 0 weekly_success = 0 monthly_success = 0 failed_count = 0 # 批量参数 days_batch = [] weekly_batch = [] monthly_batch = [] # SQL模板 days_sql = None weekly_sql = None monthly_sql = None days_keys = None weekly_keys = None monthly_keys = None for idx, row in enumerate(rows, 1): author_name = row.get('author_name', '').strip() if not author_name: continue try: # 获取channel并查询author_id channel = self.convert_value(row.get('channel', '1'), 'int') or 1 author_id = self.get_author_id(author_name, int(channel)) if author_id == 0: self.logger.warning(f"跳过无效记录: {author_name}, author_id未找到") failed_count += 1 continue stat_date = row.get('stat_date', '').strip() # 1. ai_statistics_days 表数据（仅当日数据） day_revenue = self.convert_value(row.get('day_revenue', '0'), 'decimal') or Decimal('0') daily_published_count = self.convert_value(row.get('daily_published_count', '0'), 'int') or 0 cumulative_published_count = self.convert_value(row.get('cumulative_published_count', '0'), 'int') or 0 days_record = { 'author_id': author_id, 'author_name': author_name, 'channel': channel, 'stat_date': stat_date, 'daily_published_count': daily_published_count, 'day_revenue': day_revenue, 'updated_at': datetime.now().strftime('%Y-%m-%d %H:%M:%S'), } # 2. ai_statistics_weekly 表数据 weekly_revenue = self.convert_value(row.get('weekly_revenue', '0'), 'decimal') or Decimal('0') revenue_wow_growth_rate = self.convert_value(row.get('revenue_wow_growth_rate', '0'), 'decimal') or Decimal('0') # 计算该日期所在周次（格式：WW，如51） from datetime import datetime as dt, timedelta date_obj = dt.strptime(stat_date, '%Y-%m-%d') # 使用isocalendar()获取ISO周数（周一为一周开始） year, week_num, _ = date_obj.isocalendar() stat_weekly = week_num # 直接使用数字 weekly_record = { 'author_id': author_id, 'author_name': author_name, 'channel': channel, 'stat_weekly': stat_weekly, 'weekly_revenue': weekly_revenue, 'revenue_wow_growth_rate': revenue_wow_growth_rate, 'updated_at': datetime.now().strftime('%Y-%m-%d %H:%M:%S'), } # 3. ai_statistics_monthly 表数据 monthly_revenue = self.convert_value(row.get('monthly_revenue', '0'), 'decimal') or Decimal('0') revenue_mom_growth_rate = self.convert_value(row.get('revenue_mom_growth_rate', '0'), 'decimal') or Decimal('0') # 计算该日期所在月份（格式：YYYY-MM，如2025-12） stat_monthly = date_obj.strftime('%Y-%m') monthly_record = { 'author_id': author_id, 'author_name': author_name, 'channel': channel, 'stat_monthly': stat_monthly, 'monthly_revenue': monthly_revenue, 'revenue_mom_growth_rate': revenue_mom_growth_rate, 'updated_at': datetime.now().strftime('%Y-%m-%d %H:%M:%S'), } # 构建SQL模板（首次） if days_sql is None: days_keys = list(days_record.keys()) columns = ', '.join(days_keys) placeholders = ', '.join(['%s'] * len(days_keys)) update_parts = [f"{key} = VALUES({key})" for key in days_keys if key not in ['author_name', 'channel', 'stat_date']] days_sql = f""" INSERT INTO ai_statistics_days ({columns}) VALUES ({placeholders}) ON DUPLICATE KEY UPDATE {', '.join(update_parts)} """ if weekly_sql is None: weekly_keys = list(weekly_record.keys()) columns = ', '.join(weekly_keys) placeholders = ', '.join(['%s'] * len(weekly_keys)) update_parts = [f"{key} = VALUES({key})" for key in weekly_keys if key not in ['author_name', 'channel', 'stat_weekly']] weekly_sql = f""" INSERT INTO ai_statistics_weekly ({columns}) VALUES ({placeholders}) ON DUPLICATE KEY UPDATE {', '.join(update_parts)} """ if monthly_sql is None: monthly_keys = list(monthly_record.keys()) columns = ', '.join(monthly_keys) placeholders = ', '.join(['%s'] * len(monthly_keys)) update_parts = [f"{key} = VALUES({key})" for key in monthly_keys if key not in ['author_name', 'channel', 'stat_monthly']] monthly_sql = f""" INSERT INTO ai_statistics_monthly ({columns}) VALUES ({placeholders}) ON DUPLICATE KEY UPDATE {', '.join(update_parts)} """ # 添加到批量参数 days_batch.append(tuple(days_record[key] for key in days_keys)) weekly_batch.append(tuple(weekly_record[key] for key in weekly_keys)) monthly_batch.append(tuple(monthly_record[key] for key in monthly_keys)) # 批量提交 if len(days_batch) >= batch_size or idx == len(rows): try: # 提交 ai_statistics_days result = self.db_manager.execute_many(days_sql, days_batch, autocommit=True) days_success += result print(f"[days] 已导入 {days_success} 条") days_batch = [] except Exception as e: print(f" [X] days表提交失败: {e}") self.logger.error(f"ai_statistics_days批量提交失败: {e}") failed_count += len(days_batch) days_batch = [] try: # 提交 ai_statistics_weekly result = self.db_manager.execute_many(weekly_sql, weekly_batch, autocommit=True) weekly_success += result print(f"[weekly] 已导入 {weekly_success} 条") weekly_batch = [] except Exception as e: print(f" [X] weekly表提交失败: {e}") self.logger.error(f"ai_statistics_weekly批量提交失败: {e}") weekly_batch = [] try: # 提交 ai_statistics_monthly result = self.db_manager.execute_many(monthly_sql, monthly_batch, autocommit=True) monthly_success += result print(f"[monthly] 已导入 {monthly_success} 条") monthly_batch = [] except Exception as e: print(f" [X] monthly表提交失败: {e}") self.logger.error(f"ai_statistics_monthly批量提交失败: {e}") monthly_batch = [] except Exception as e: failed_count += 1 print(f" [X] 处理失败 ({author_name}): {e}") self.logger.error(f"数据处理失败: {author_name}, 错误: {e}") continue print("\n" + "="*70) print(f"[OK] 数据导入完成（拆分到3个表）") print(f" ai_statistics_days: {days_success} 条") print(f" ai_statistics_weekly: {weekly_success} 条") print(f" ai_statistics_monthly: {monthly_success} 条") if failed_count > 0: print(f" 失败: {failed_count} 条") print("="*70) self.logger.info(f"数据导入完成: days={days_success}, weekly={weekly_success}, monthly={monthly_success}, failed={failed_count}") return days_success > 0 def import_all(self) -> bool: """导入所有CSV文件""" print("\n" + "="*70) print("CSV数据导入数据库") print("="*70) print("\n将导入以下三个CSV文件到数据库：") print(" 1. ai_statistics.csv -> ai_statistics表") print(" 2. ai_statistics_day.csv -> ai_statistics_day表") print(" 3. ai_statistics_days.csv -> ai_statistics_days表") print("\n唯一键：(author_name, channel, date/stat_date)") print("="*70) self.logger.info("开始执行CSV数据导入任务") # 导入三个表 result1 = self.import_ai_statistics() result2 = self.import_ai_statistics_day() result3 = self.import_ai_statistics_days() if result1 and result2 and result3: print("\n" + "="*70) print("✓ 所有CSV文件导入完成！") print("="*70) self.logger.info("所有CSV文件导入成功") return True else: print("\n" + "="*70) print("✗ 部分CSV文件导入失败") print("="*70) self.logger.warning("部分CSV文件导入失败") return False def main(): """主函数""" logger = None try: # 创建导入器 importer = CSVImporter() logger = importer.logger # 导入所有CSV文件 success = importer.import_all() return 0 if success else 1 except Exception as e: print(f"\n[X] 程序执行出错: {e}") if logger: logger.error(f"程序执行出错: {e}", exc_info=True) return 1 if __name__ == '__main__': sys.exit(main())