Initial commit
This commit is contained in:
freedakgmail
@@ -0,0 +1,242 @@
|
||||
"""
|
||||
数据迁移脚本 - 将JSON数据迁移到PostgreSQL数据库
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import psycopg2
|
||||
from psycopg2.extras import RealDictCursor
|
||||
from werkzeug.security import generate_password_hash
|
||||
from config import Config
|
||||
|
||||
def get_db():
|
||||
"""获取数据库连接"""
|
||||
return psycopg2.connect(
|
||||
host=Config.DB_HOST,
|
||||
port=Config.DB_PORT,
|
||||
database=Config.DB_NAME,
|
||||
user=Config.DB_USER,
|
||||
password=Config.DB_PASSWORD
|
||||
)
|
||||
|
||||
|
||||
def create_default_user(conn, email=None, password=None):
|
||||
"""创建用户"""
|
||||
cur = conn.cursor(cursor_factory=RealDictCursor)
|
||||
|
||||
# 使用指定的或默认的邮箱密码
|
||||
user_email = email or 'admin@admin.com'
|
||||
user_password = password or 'admin123'
|
||||
|
||||
# 检查是否已存在用户
|
||||
cur.execute("SELECT id FROM users WHERE email = %s OR username = %s", (user_email, user_email))
|
||||
user = cur.fetchone()
|
||||
|
||||
if user:
|
||||
print(f"用户 {user_email} 已存在，使用现有用户")
|
||||
return user['id']
|
||||
|
||||
# 创建用户
|
||||
password_hash = generate_password_hash(user_password)
|
||||
cur.execute(
|
||||
"INSERT INTO users (username, email, password_hash) VALUES (%s, %s, %s) RETURNING id",
|
||||
(user_email, user_email, password_hash)
|
||||
)
|
||||
user_id = cur.fetchone()['id']
|
||||
conn.commit()
|
||||
print(f"创建用户: {user_email}")
|
||||
return user_id
|
||||
|
||||
|
||||
def migrate_trades(conn, user_id):
|
||||
"""迁移交易记录"""
|
||||
if not os.path.exists(Config.TRADES_FILE):
|
||||
print("trades.json 不存在，跳过")
|
||||
return 0
|
||||
|
||||
with open(Config.TRADES_FILE, 'r', encoding='utf-8') as f:
|
||||
trades = json.load(f)
|
||||
|
||||
if not trades:
|
||||
print("trades.json 为空，跳过")
|
||||
return 0
|
||||
|
||||
cur = conn.cursor()
|
||||
count = 0
|
||||
|
||||
for trade in trades:
|
||||
try:
|
||||
# 处理日期
|
||||
trade_date = trade.get('trade_date')
|
||||
if trade_date and len(trade_date) > 10:
|
||||
trade_date = trade_date[:10]
|
||||
|
||||
# 处理数值
|
||||
price = trade.get('price')
|
||||
if price and price != '':
|
||||
price = float(price)
|
||||
else:
|
||||
price = None
|
||||
|
||||
quantity = trade.get('quantity')
|
||||
if quantity and quantity != '':
|
||||
quantity = int(quantity)
|
||||
else:
|
||||
quantity = None
|
||||
|
||||
profit_amount = trade.get('profit_amount')
|
||||
if profit_amount and profit_amount != '':
|
||||
profit_amount = float(profit_amount)
|
||||
else:
|
||||
profit_amount = None
|
||||
|
||||
stop_loss_price = trade.get('stop_loss_price')
|
||||
if stop_loss_price and stop_loss_price != '':
|
||||
stop_loss_price = float(stop_loss_price)
|
||||
else:
|
||||
stop_loss_price = None
|
||||
|
||||
cur.execute("""
|
||||
INSERT INTO trades (user_id, stock_code, stock_name, trade_type, price,
|
||||
quantity, trade_date, reason, result, profit_amount,
|
||||
stop_loss_price, notes)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||||
""", (
|
||||
user_id,
|
||||
trade.get('stock_code'),
|
||||
trade.get('stock_name'),
|
||||
trade.get('trade_type'),
|
||||
price,
|
||||
quantity,
|
||||
trade_date,
|
||||
trade.get('reason'),
|
||||
trade.get('result'),
|
||||
profit_amount,
|
||||
stop_loss_price,
|
||||
trade.get('notes')
|
||||
))
|
||||
count += 1
|
||||
except Exception as e:
|
||||
print(f"迁移交易记录失败: {e}, 数据: {trade}")
|
||||
|
||||
conn.commit()
|
||||
print(f"迁移交易记录: {count} 条")
|
||||
return count
|
||||
|
||||
|
||||
def migrate_watchlist(conn, user_id):
|
||||
"""迁移关注列表"""
|
||||
if not os.path.exists(Config.WATCHLIST_FILE):
|
||||
print("watchlist.json 不存在，跳过")
|
||||
return 0
|
||||
|
||||
with open(Config.WATCHLIST_FILE, 'r', encoding='utf-8') as f:
|
||||
watchlist = json.load(f)
|
||||
|
||||
if not watchlist:
|
||||
print("watchlist.json 为空，跳过")
|
||||
return 0
|
||||
|
||||
cur = conn.cursor()
|
||||
count = 0
|
||||
|
||||
for item in watchlist:
|
||||
try:
|
||||
cur.execute("""
|
||||
INSERT INTO watchlist (user_id, stock_code, stock_name)
|
||||
VALUES (%s, %s, %s)
|
||||
ON CONFLICT (user_id, stock_code) DO NOTHING
|
||||
""", (
|
||||
user_id,
|
||||
item.get('code'),
|
||||
item.get('name')
|
||||
))
|
||||
count += 1
|
||||
except Exception as e:
|
||||
print(f"迁移关注列表失败: {e}, 数据: {item}")
|
||||
|
||||
conn.commit()
|
||||
print(f"迁移关注列表: {count} 条")
|
||||
return count
|
||||
|
||||
|
||||
def migrate_alerts_cache(conn, user_id):
|
||||
"""迁移分析缓存"""
|
||||
if not os.path.exists(Config.ALERTS_CACHE_FILE):
|
||||
print("alerts_cache.json 不存在，跳过")
|
||||
return 0
|
||||
|
||||
with open(Config.ALERTS_CACHE_FILE, 'r', encoding='utf-8') as f:
|
||||
cache = json.load(f)
|
||||
|
||||
alerts = cache.get('alerts', [])
|
||||
if not alerts:
|
||||
print("alerts_cache.json 为空，跳过")
|
||||
return 0
|
||||
|
||||
cur = conn.cursor()
|
||||
cur.execute("""
|
||||
INSERT INTO alerts_cache (user_id, data, updated_at)
|
||||
VALUES (%s, %s, NOW())
|
||||
ON CONFLICT (user_id) DO UPDATE SET
|
||||
data = EXCLUDED.data,
|
||||
updated_at = NOW()
|
||||
""", (user_id, json.dumps(alerts)))
|
||||
|
||||
conn.commit()
|
||||
print(f"迁移分析缓存: {len(alerts)} 条")
|
||||
return len(alerts)
|
||||
|
||||
|
||||
def main():
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser(description='数据迁移工具')
|
||||
parser.add_argument('--email', default=None, help='用户邮箱')
|
||||
parser.add_argument('--password', default=None, help='用户密码')
|
||||
args = parser.parse_args()
|
||||
|
||||
print("=" * 60)
|
||||
print("数据迁移 - JSON -> PostgreSQL")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
conn = get_db()
|
||||
print("数据库连接成功")
|
||||
except Exception as e:
|
||||
print(f"数据库连接失败: {e}")
|
||||
print("\n请先执行: psql -U postgres -f init_db.sql")
|
||||
sys.exit(1)
|
||||
|
||||
try:
|
||||
# 创建用户
|
||||
user_id = create_default_user(conn, args.email, args.password)
|
||||
|
||||
# 先清空现有数据
|
||||
cur = conn.cursor()
|
||||
cur.execute("DELETE FROM trades WHERE user_id = %s", (user_id,))
|
||||
cur.execute("DELETE FROM watchlist WHERE user_id = %s", (user_id,))
|
||||
cur.execute("DELETE FROM alerts_cache WHERE user_id = %s", (user_id,))
|
||||
conn.commit()
|
||||
print("清空现有数据")
|
||||
|
||||
# 迁移数据
|
||||
migrate_trades(conn, user_id)
|
||||
migrate_watchlist(conn, user_id)
|
||||
migrate_alerts_cache(conn, user_id)
|
||||
|
||||
print("=" * 60)
|
||||
print("迁移完成！")
|
||||
print(f"登录邮箱: {args.email or 'admin@admin.com'}")
|
||||
print(f"登录密码: {args.password or 'admin123'}")
|
||||
print("=" * 60)
|
||||
|
||||
except Exception as e:
|
||||
print(f"迁移失败: {e}")
|
||||
conn.rollback()
|
||||
raise
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user