|
|
# -*- coding: utf-8 -*-
|
|
|
import os
|
|
|
import random
|
|
|
import time
|
|
|
import sqlite3
|
|
|
import httpx
|
|
|
from httpx import BasicAuth
|
|
|
|
|
|
|
|
|
class DataSetDownloader:
|
|
|
def __init__(self):
|
|
|
self.base_api_url = 'https://api.worldquantbrain.com'
|
|
|
self.db_path = os.path.join(os.getcwd(), 'data_sets.db')
|
|
|
self.client = None
|
|
|
self.login_time = None
|
|
|
self.total_tasks = 0
|
|
|
self.completed_tasks = 0
|
|
|
self.current_task_index = 0
|
|
|
self.login()
|
|
|
|
|
|
def login(self):
|
|
|
"""登录并返回客户端实例"""
|
|
|
username, password = "jack0210_@hotmail.com", "!QAZ2wsx+0913"
|
|
|
client = httpx.Client(auth=BasicAuth(username, password))
|
|
|
|
|
|
try:
|
|
|
response = client.post(f'{self.base_api_url}/authentication')
|
|
|
print(f"登录状态: {response.status_code}")
|
|
|
|
|
|
if response.status_code in [200, 201]:
|
|
|
print("登录成功!")
|
|
|
self.client = client
|
|
|
self.login_time = time.time()
|
|
|
return client
|
|
|
else:
|
|
|
print(f"登录失败: {response.json()}")
|
|
|
exit(1)
|
|
|
except Exception as e:
|
|
|
print(f"登录过程中出现错误: {e}")
|
|
|
exit(1)
|
|
|
|
|
|
def check_and_renew_login(self):
|
|
|
"""检查登录状态，超过3.5小时则重新登录"""
|
|
|
if not self.login_time:
|
|
|
return
|
|
|
|
|
|
current_time = time.time()
|
|
|
elapsed_hours = (current_time - self.login_time) / 3600
|
|
|
|
|
|
if elapsed_hours > 3.5:
|
|
|
print(f"\n⚠️ 登录已超过{elapsed_hours:.1f}小时，正在重新登录...")
|
|
|
self.login()
|
|
|
print("✅ 重新登录成功，继续执行任务")
|
|
|
|
|
|
def get_download_stats(self):
|
|
|
"""获取下载统计信息"""
|
|
|
conn = sqlite3.connect(self.db_path)
|
|
|
cursor = conn.cursor()
|
|
|
|
|
|
# 获取总任务数
|
|
|
cursor.execute('SELECT COUNT(*) FROM category WHERE downloaded = 0')
|
|
|
remaining_tasks = cursor.fetchone()[0]
|
|
|
|
|
|
# 获取已完成任务数
|
|
|
cursor.execute('SELECT COUNT(*) FROM category WHERE downloaded = 1')
|
|
|
completed_tasks = cursor.fetchone()[0]
|
|
|
|
|
|
conn.close()
|
|
|
|
|
|
total_tasks = remaining_tasks + completed_tasks
|
|
|
progress_percentage = (completed_tasks / total_tasks * 100) if total_tasks > 0 else 0
|
|
|
|
|
|
return {
|
|
|
'total_tasks': total_tasks,
|
|
|
'remaining_tasks': remaining_tasks,
|
|
|
'completed_tasks': completed_tasks,
|
|
|
'progress_percentage': progress_percentage
|
|
|
}
|
|
|
|
|
|
def print_progress_bar(self, current, total, prefix='', suffix='', length=50, fill='█'):
|
|
|
"""打印进度条"""
|
|
|
percent = f"{100 * (current / float(total)):.1f}"
|
|
|
filled_length = int(length * current // total)
|
|
|
bar = fill * filled_length + '-' * (length - filled_length)
|
|
|
|
|
|
progress_text = f'\r{prefix} |{bar}| {percent}% {suffix}'
|
|
|
if current == total:
|
|
|
print(progress_text)
|
|
|
else:
|
|
|
print(progress_text, end='', flush=True)
|
|
|
|
|
|
def test_download(self):
|
|
|
"""测试每个category分类是否能成功下载数据"""
|
|
|
if not self.client:
|
|
|
print("❌ 客户端未初始化")
|
|
|
exit(1)
|
|
|
|
|
|
print("\n" + "=" * 60)
|
|
|
print("开始测试所有category分类...")
|
|
|
print("=" * 60)
|
|
|
|
|
|
# 从数据库获取每个category分类的一条测试记录
|
|
|
conn = sqlite3.connect(self.db_path)
|
|
|
cursor = conn.cursor()
|
|
|
|
|
|
# 使用DISTINCT获取每个category.name的第一条记录
|
|
|
cursor.execute('''
|
|
|
SELECT DISTINCT name, category_id, region, universe, instrumentType, delay
|
|
|
FROM category
|
|
|
WHERE downloaded = 0
|
|
|
ORDER BY name
|
|
|
''')
|
|
|
|
|
|
test_categories = cursor.fetchall()
|
|
|
conn.close()
|
|
|
|
|
|
print(f"找到 {len(test_categories)} 个不同的category分类进行测试")
|
|
|
|
|
|
success_count = 0
|
|
|
fail_count = 0
|
|
|
|
|
|
for idx, category in enumerate(test_categories):
|
|
|
# 检查登录状态
|
|
|
self.check_and_renew_login()
|
|
|
|
|
|
category_name, category_id, region, universe, instrumentType, delay = category
|
|
|
|
|
|
print(f"\n🔍 测试 [{idx + 1}/{len(test_categories)}]: {category_name}")
|
|
|
print(f" category_id: {category_id}")
|
|
|
print(f" region: {region}, universe: {universe}")
|
|
|
|
|
|
try:
|
|
|
# 测试下载（只下载第一页的1条数据）
|
|
|
result = self._test_single_category(
|
|
|
category_name, category_id, region, universe, instrumentType, delay
|
|
|
)
|
|
|
|
|
|
if result:
|
|
|
print(f" ✅ 测试成功")
|
|
|
success_count += 1
|
|
|
else:
|
|
|
print(f" ❌ 测试失败")
|
|
|
fail_count += 1
|
|
|
|
|
|
except Exception as e:
|
|
|
print(f" ❌ 测试异常: {e}")
|
|
|
fail_count += 1
|
|
|
|
|
|
# 打印测试进度
|
|
|
self.print_progress_bar(idx + 1, len(test_categories),
|
|
|
prefix='测试进度:',
|
|
|
suffix=f'完成 {success_count} 成功, {fail_count} 失败')
|
|
|
|
|
|
# 测试间隔
|
|
|
time.sleep(random.uniform(3, 5))
|
|
|
|
|
|
# 输出测试结果
|
|
|
print("\n" + "=" * 60)
|
|
|
print("测试结果汇总:")
|
|
|
print("=" * 60)
|
|
|
print(f"总测试分类数: {len(test_categories)}")
|
|
|
print(f"成功: {success_count}")
|
|
|
print(f"失败: {fail_count}")
|
|
|
|
|
|
if fail_count == 0:
|
|
|
print("\n🎉 所有category分类测试通过！")
|
|
|
else:
|
|
|
print(f"\n⚠️ 有 {fail_count} 个分类测试失败")
|
|
|
|
|
|
return fail_count == 0
|
|
|
|
|
|
def _test_single_category(self, category_name, category_id, region, universe, instrumentType, delay):
|
|
|
"""测试单个category是否能成功下载"""
|
|
|
endpoint = 'data-fields'
|
|
|
url = f"{self.base_api_url}/{endpoint}"
|
|
|
|
|
|
# 构建参数，只请求1条数据
|
|
|
params = {
|
|
|
'dataset.id': category_id,
|
|
|
'delay': delay,
|
|
|
'instrumentType': instrumentType,
|
|
|
'limit': 1,
|
|
|
'offset': 0,
|
|
|
'region': region,
|
|
|
'universe': universe
|
|
|
}
|
|
|
|
|
|
try:
|
|
|
response = self.client.get(url, params=params, timeout=30)
|
|
|
|
|
|
if response.status_code != 200:
|
|
|
print(f" HTTP状态码: {response.status_code}")
|
|
|
return False
|
|
|
|
|
|
data = response.json()
|
|
|
total_count = data.get('count', 0)
|
|
|
results = data.get('results', [])
|
|
|
|
|
|
print(f" 数据总数: {total_count}, 本页获取: {len(results)}")
|
|
|
|
|
|
if total_count > 0 and len(results) > 0:
|
|
|
# 打印第一条数据的部分信息用于验证
|
|
|
item = results[0]
|
|
|
print(f" 示例数据字段 - id: {item.get('id', 'N/A')}, name: {item.get('name', 'N/A')}")
|
|
|
return True
|
|
|
else:
|
|
|
print(f" ⚠️ 无数据返回")
|
|
|
return total_count == 0 # 如果count=0也算测试成功（只是没数据）
|
|
|
|
|
|
except httpx.TimeoutException:
|
|
|
print(" ⏰ 请求超时")
|
|
|
return False
|
|
|
except Exception as e:
|
|
|
print(f" ❌ 请求异常: {e}")
|
|
|
return False
|
|
|
|
|
|
def get_categories_to_download(self):
|
|
|
"""从数据库获取需要下载的category记录"""
|
|
|
conn = sqlite3.connect(self.db_path)
|
|
|
cursor = conn.cursor()
|
|
|
|
|
|
cursor.execute('''
|
|
|
SELECT name, base_url, category_id, region, universe, instrumentType, delay
|
|
|
FROM category
|
|
|
WHERE downloaded = 0
|
|
|
ORDER BY rowid
|
|
|
''')
|
|
|
|
|
|
categories = cursor.fetchall()
|
|
|
conn.close()
|
|
|
|
|
|
# 更新统计信息
|
|
|
self.total_tasks = len(categories)
|
|
|
self.completed_tasks = 0
|
|
|
self.current_task_index = 0
|
|
|
|
|
|
print(f"找到 {self.total_tasks} 个需要下载的category")
|
|
|
|
|
|
# 打印总体进度
|
|
|
if self.total_tasks > 0:
|
|
|
stats = self.get_download_stats()
|
|
|
print(f"总体进度: {stats['completed_tasks']}/{stats['total_tasks']} ({stats['progress_percentage']:.1f}%)")
|
|
|
|
|
|
return categories
|
|
|
|
|
|
def _build_params(self, category_id, region, universe, instrumentType, delay, offset=0, limit=50):
|
|
|
"""构建data-fields请求参数"""
|
|
|
return {
|
|
|
'dataset.id': category_id,
|
|
|
'delay': delay,
|
|
|
'instrumentType': instrumentType,
|
|
|
'limit': limit,
|
|
|
'offset': offset,
|
|
|
'region': region,
|
|
|
'universe': universe
|
|
|
}
|
|
|
|
|
|
def _process_item(self, item):
|
|
|
"""处理单个data-field项 - name字段对应返回的id字段"""
|
|
|
return {
|
|
|
'name': item.get('id', ''), # name字段对应返回的id字段
|
|
|
'description': item.get('description', ''),
|
|
|
'dataset_id': item.get('dataset', {}).get('id', ''),
|
|
|
'dataset_name': item.get('dataset', {}).get('name', ''),
|
|
|
'category_id': item.get('category', {}).get('id', ''),
|
|
|
'category_name': item.get('category', {}).get('name', ''),
|
|
|
'region': item.get('region', ''),
|
|
|
'delay': str(item.get('delay', '')),
|
|
|
'universe': item.get('universe', ''),
|
|
|
'type': item.get('type', '')
|
|
|
}
|
|
|
|
|
|
def save_to_data_sets_table(self, data_items):
|
|
|
"""批量保存到data_sets表"""
|
|
|
conn = sqlite3.connect(self.db_path)
|
|
|
cursor = conn.cursor()
|
|
|
|
|
|
for item in data_items:
|
|
|
cursor.execute('''
|
|
|
INSERT INTO data_sets (name, description, dataset_id, dataset_name,
|
|
|
category_id, category_name, region, delay,
|
|
|
universe, type)
|
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
|
''', (
|
|
|
item['name'], item['description'], item['dataset_id'], item['dataset_name'],
|
|
|
item['category_id'], item['category_name'], item['region'], item['delay'],
|
|
|
item['universe'], item['type']
|
|
|
))
|
|
|
|
|
|
conn.commit()
|
|
|
conn.close()
|
|
|
print(f"✅ 保存了 {len(data_items)} 条记录到 data_sets 表")
|
|
|
|
|
|
def update_category_status(self, category_id, region, universe):
|
|
|
"""更新category表的下载状态"""
|
|
|
conn = sqlite3.connect(self.db_path)
|
|
|
cursor = conn.cursor()
|
|
|
|
|
|
cursor.execute('''
|
|
|
UPDATE category
|
|
|
SET downloaded = 1
|
|
|
WHERE category_id = ?
|
|
|
AND region = ?
|
|
|
AND universe = ?
|
|
|
''', (category_id, region, universe))
|
|
|
|
|
|
conn.commit()
|
|
|
conn.close()
|
|
|
|
|
|
# 更新进度统计
|
|
|
self.completed_tasks += 1
|
|
|
stats = self.get_download_stats()
|
|
|
print(f"✅ 已更新category状态: {category_id} - {region} - {universe}")
|
|
|
print(f"📊 总体进度: {stats['completed_tasks']}/{stats['total_tasks']} ({stats['progress_percentage']:.1f}%)")
|
|
|
|
|
|
def download_category_data_fields(self, category_name, category_id, region, universe, instrumentType, delay):
|
|
|
"""下载指定category的所有data-fields"""
|
|
|
endpoint = 'data-fields'
|
|
|
url = f"{self.base_api_url}/{endpoint}"
|
|
|
|
|
|
print(f"\n开始下载: {category_name} - {category_id} - {region} - {universe}")
|
|
|
|
|
|
# 打印当前任务进度
|
|
|
self.current_task_index += 1
|
|
|
print(f"📋 任务进度: {self.current_task_index}/{self.total_tasks}")
|
|
|
|
|
|
# 获取数据总数
|
|
|
params = self._build_params(category_id, region, universe, instrumentType, delay, limit=1)
|
|
|
response = self.client.get(url, params=params)
|
|
|
|
|
|
if response.status_code != 200:
|
|
|
print(f"❌ 获取数据总数失败: {response.status_code}")
|
|
|
exit(1)
|
|
|
|
|
|
data = response.json()
|
|
|
total_count = data.get('count', 0)
|
|
|
|
|
|
print(f"📊 数据总数: {total_count}")
|
|
|
|
|
|
if total_count == 0:
|
|
|
print("⚠️ 没有找到数据，跳过")
|
|
|
return
|
|
|
|
|
|
# 下载所有数据
|
|
|
limit = 50
|
|
|
all_data_items = []
|
|
|
|
|
|
for offset in range(0, total_count, limit):
|
|
|
# 每次循环开始前检查登录状态
|
|
|
self.check_and_renew_login()
|
|
|
|
|
|
time.sleep(random.uniform(2, 2.5))
|
|
|
|
|
|
params = self._build_params(category_id, region, universe, instrumentType, delay, offset, limit)
|
|
|
|
|
|
# 打印数据下载进度
|
|
|
data_progress = f" {offset}/{total_count} ({offset / total_count * 100:.1f}%)"
|
|
|
print(f"📥 数据下载进度:{data_progress}")
|
|
|
|
|
|
response = self.client.get(url, params=params)
|
|
|
|
|
|
if response.status_code != 200:
|
|
|
print(f"❌ 下载失败: {response.status_code}")
|
|
|
exit(1)
|
|
|
|
|
|
data = response.json()
|
|
|
results = data.get('results', [])
|
|
|
|
|
|
# 处理数据
|
|
|
processed_items = [self._process_item(item) for item in results]
|
|
|
all_data_items.extend(processed_items)
|
|
|
|
|
|
print(f"✅ 本页获取到 {len(results)} 条记录")
|
|
|
|
|
|
if len(results) < limit:
|
|
|
print("🎯 到达数据末尾")
|
|
|
break
|
|
|
|
|
|
# 保存到数据库
|
|
|
if all_data_items:
|
|
|
self.save_to_data_sets_table(all_data_items)
|
|
|
|
|
|
def run(self):
|
|
|
"""主运行函数"""
|
|
|
if not self.client:
|
|
|
print("❌ 客户端未初始化")
|
|
|
exit(1)
|
|
|
|
|
|
# 打印开始时的总体进度
|
|
|
print("\n" + "=" * 60)
|
|
|
print("开始下载任务...")
|
|
|
stats = self.get_download_stats()
|
|
|
print(f"初始状态: 已完成 {stats['completed_tasks']}/{stats['total_tasks']} ({stats['progress_percentage']:.1f}%)")
|
|
|
print("=" * 60)
|
|
|
|
|
|
categories = self.get_categories_to_download()
|
|
|
|
|
|
if not categories:
|
|
|
print("✅ 所有category都已下载完成")
|
|
|
return
|
|
|
|
|
|
for category in categories:
|
|
|
# 每次循环开始前检查登录状态
|
|
|
self.check_and_renew_login()
|
|
|
|
|
|
category_name, base_url, category_id, region, universe, instrumentType, delay = category
|
|
|
|
|
|
print(f"\n{'=' * 60}")
|
|
|
print(f"处理category: {category_name}")
|
|
|
print(f"category_id: {category_id}")
|
|
|
print(f"region: {region}, universe: {universe}")
|
|
|
print(f"{'=' * 60}")
|
|
|
|
|
|
try:
|
|
|
# 下载data-fields
|
|
|
self.download_category_data_fields(
|
|
|
category_name, category_id, region, universe, instrumentType, delay
|
|
|
)
|
|
|
|
|
|
# 更新下载状态
|
|
|
self.update_category_status(category_id, region, universe)
|
|
|
|
|
|
print(f"✅ 完成: {category_name} - {region} - {universe}")
|
|
|
|
|
|
except Exception as e:
|
|
|
print(f"❌ 处理失败: {e}")
|
|
|
return
|
|
|
|
|
|
sleep_time = random.uniform(20, 30)
|
|
|
print(f"\n等待 {sleep_time} 秒后继续下一个...")
|
|
|
time.sleep(sleep_time)
|
|
|
|
|
|
# 打印最终进度
|
|
|
print("\n" + "=" * 60)
|
|
|
print("🎉 所有category处理完成！")
|
|
|
final_stats = self.get_download_stats()
|
|
|
print(f"最终状态: 已完成 {final_stats['completed_tasks']}/{final_stats['total_tasks']} ({final_stats['progress_percentage']:.1f}%)")
|
|
|
print("=" * 60)
|
|
|
exit(0)
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
# category_list = [
|
|
|
# 'analyst',
|
|
|
# 'broker',
|
|
|
# 'earnings',
|
|
|
# 'fundamental',
|
|
|
# 'imbalance',
|
|
|
# 'insiders',
|
|
|
# 'institutions',
|
|
|
# 'macro',
|
|
|
# 'model',
|
|
|
# 'news',
|
|
|
# 'option',
|
|
|
# 'other',
|
|
|
# 'pv',
|
|
|
# 'risk',
|
|
|
# 'sentiment',
|
|
|
# 'shortinterest',
|
|
|
# 'socialmedia'
|
|
|
# ]
|
|
|
downloader = DataSetDownloader()
|
|
|
|
|
|
while True:
|
|
|
downloader.run()
|
|
|
time.sleep(30) |
