201 lines
7.0 KiB
Python
Executable File
201 lines
7.0 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""
|
||
通过数据采集模块API导入测试数据到正式数据库
|
||
自动化脚本 - 一键导入所有测试数据
|
||
"""
|
||
|
||
import requests
|
||
import time
|
||
import json
|
||
from typing import Dict, Any
|
||
|
||
API_BASE = "http://localhost:8000/api/v1/data"
|
||
|
||
class DataImporter:
|
||
def __init__(self):
|
||
self.api_base = API_BASE
|
||
self.uploaded_files = {}
|
||
|
||
def upload_file(self, file_path: str, import_type: str) -> Dict[str, Any]:
|
||
"""上传文件并返回upload信息"""
|
||
print(f"\n📤 正在上传: {file_path}")
|
||
print(f" 类型: {import_type}")
|
||
|
||
with open(file_path, 'rb') as f:
|
||
files = {'file': (file_path.split('/')[-1], f)}
|
||
data = {'import_type': import_type}
|
||
|
||
response = requests.post(f"{self.api_base}/upload", files=files, data=data)
|
||
|
||
if response.status_code == 200:
|
||
result = response.json()
|
||
upload_id = result['upload_id']
|
||
total_rows = result['total_rows']
|
||
field_mapping = result.get('field_mapping', {})
|
||
print(f" ✅ 上传成功 (upload_id: {upload_id}, {total_rows}条记录)")
|
||
return {
|
||
'upload_id': upload_id,
|
||
'field_mapping': field_mapping,
|
||
'total_rows': total_rows
|
||
}
|
||
else:
|
||
print(f" ❌ 上传失败: {response.text}")
|
||
raise Exception(f"上传失败: {response.status_code}")
|
||
|
||
def confirm_import(self, upload_id: str, import_type: str, field_mapping: Dict[str, str] = None) -> Dict[str, Any]:
|
||
"""确认导入并返回结果"""
|
||
print(f"\n🔄 正在导入数据...")
|
||
|
||
# 如果没有提供 field_mapping，创建一个空的映射
|
||
if field_mapping is None:
|
||
field_mapping = {}
|
||
|
||
response = requests.post(
|
||
f"{self.api_base}/confirm",
|
||
json={
|
||
'upload_id': upload_id,
|
||
'import_type': import_type,
|
||
'field_mapping': field_mapping
|
||
}
|
||
)
|
||
|
||
if response.status_code == 200:
|
||
result = response.json()
|
||
success_rows = result['success_rows']
|
||
failed_rows = result['failed_rows']
|
||
total_rows = result['total_rows']
|
||
task_id = result['task_id']
|
||
|
||
print(f" ✅ 导入完成!")
|
||
print(f" 📊 总计: {total_rows}, 成功: {success_rows}, 失败: {failed_rows}")
|
||
print(f" 🆔 任务ID: {task_id}")
|
||
|
||
if failed_rows > 0:
|
||
print(f" ⚠️ 错误详情:")
|
||
for error in result['error_log'][:3]: # 只显示前3个错误
|
||
print(f" 行{error['row']}: {error['errorMsg']}")
|
||
|
||
return result
|
||
else:
|
||
print(f" ❌ 导入失败: {response.text}")
|
||
raise Exception(f"导入失败: {response.status_code}")
|
||
|
||
def import_data_type(self, file_path: str, import_type: str, name: str):
|
||
"""导入单个数据类型"""
|
||
print("\n" + "="*80)
|
||
print(f"📦 正在导入: {name}")
|
||
print("="*80)
|
||
|
||
try:
|
||
# 上传文件
|
||
upload_info = self.upload_file(file_path, import_type)
|
||
|
||
# 确认导入
|
||
result = self.confirm_import(
|
||
upload_info['upload_id'],
|
||
import_type,
|
||
upload_info['field_mapping']
|
||
)
|
||
|
||
print(f"\n✅ {name} 导入成功!")
|
||
return True
|
||
|
||
except Exception as e:
|
||
print(f"\n❌ {name} 导入失败: {str(e)}")
|
||
return False
|
||
|
||
def check_import_history(self):
|
||
"""查看导入历史"""
|
||
print("\n" + "="*80)
|
||
print("📋 导入历史记录")
|
||
print("="*80)
|
||
|
||
response = requests.get(f"{self.api_base}/history?page=1&size=10")
|
||
|
||
if response.status_code == 200:
|
||
result = response.json()
|
||
records = result.get('records', [])
|
||
|
||
for record in records[-5:]: # 显示最近5条
|
||
status_icon = "✅" if record['status'] == 'success' else "⚠️"
|
||
print(f"{status_icon} {record['import_type_name']}")
|
||
print(f" 文件: {record['file_name']}")
|
||
print(f" 状态: {record['status']} | 成功: {record['success_rows']} | 失败: {record['failed_rows']}")
|
||
print(f" 时间: {record['created_at']}")
|
||
print()
|
||
else:
|
||
print(f"❌ 查询失败: {response.text}")
|
||
|
||
def run(self):
|
||
"""执行完整导入流程"""
|
||
print("="*80)
|
||
print("🚀 数据采集模块 - 测试数据导入工具")
|
||
print("="*80)
|
||
|
||
# 定义要导入的文件
|
||
files_to_import = [
|
||
{
|
||
'file': '/Users/liulujian/Documents/code/deeprisk-claude-1/backend/test_data/revenue_test/主播信息导入模板.xlsx',
|
||
'type': 'streamer',
|
||
'name': '主播信息'
|
||
},
|
||
{
|
||
'file': '/Users/liulujian/Documents/code/deeprisk-claude-1/backend/test_data/revenue_test/分成协议导入模板.xlsx',
|
||
'type': 'contract',
|
||
'name': '分成协议'
|
||
},
|
||
{
|
||
'file': '/Users/liulujian/Documents/code/deeprisk-claude-1/backend/test_data/revenue_test/充值记录导入模板.xlsx',
|
||
'type': 'recharge',
|
||
'name': '充值记录'
|
||
},
|
||
{
|
||
'file': '/Users/liulujian/Documents/code/deeprisk-claude-1/backend/test_data/revenue_test/税务申报导入模板.xlsx',
|
||
'type': 'tax',
|
||
'name': '税务申报'
|
||
},
|
||
]
|
||
|
||
success_count = 0
|
||
failed_count = 0
|
||
|
||
# 逐个导入
|
||
for item in files_to_import:
|
||
if self.import_data_type(item['file'], item['type'], item['name']):
|
||
success_count += 1
|
||
else:
|
||
failed_count += 1
|
||
|
||
time.sleep(1) # 等待1秒避免请求过快
|
||
|
||
# 查看历史
|
||
self.check_import_history()
|
||
|
||
# 总结
|
||
print("\n" + "="*80)
|
||
print("📊 导入完成统计")
|
||
print("="*80)
|
||
print(f"✅ 成功导入: {success_count} 个文件")
|
||
print(f"❌ 导入失败: {failed_count} 个文件")
|
||
print(f"📈 总计: {len(files_to_import)} 个文件")
|
||
print("\n🎉 数据导入完成!")
|
||
print("="*80)
|
||
|
||
if failed_count == 0:
|
||
print("\n💡 提示:")
|
||
print(" 现在您可以运行算法测试:")
|
||
print(" python /Users/liulujian/Documents/code/deeprisk-claude-1/backend/scripts/test_revenue_algorithm.py")
|
||
else:
|
||
print("\n⚠️ 部分导入失败，请检查错误信息并重试")
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
importer = DataImporter()
|
||
importer.run()
|
||
except KeyboardInterrupt:
|
||
print("\n\n⏹️ 用户取消操作")
|
||
except Exception as e:
|
||
print(f"\n\n❌ 程序执行失败: {str(e)}")
|
||
import traceback
|
||
traceback.print_exc()
|