Initial commit
This commit is contained in:
freedakgmail
@@ -0,0 +1,587 @@
|
||||
"""
|
||||
数据库连接和用户认证
|
||||
"""
|
||||
import psycopg2
|
||||
from psycopg2.extras import RealDictCursor
|
||||
from werkzeug.security import generate_password_hash, check_password_hash
|
||||
from flask import session
|
||||
import functools
|
||||
from config import Config
|
||||
|
||||
|
||||
def get_db():
|
||||
"""获取数据库连接"""
|
||||
try:
|
||||
conn = psycopg2.connect(
|
||||
host=Config.DB_HOST,
|
||||
port=Config.DB_PORT,
|
||||
database=Config.DB_NAME,
|
||||
user=Config.DB_USER,
|
||||
password=Config.DB_PASSWORD
|
||||
)
|
||||
return conn
|
||||
except Exception as e:
|
||||
print(f"数据库连接失败: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def login_required(f):
|
||||
"""登录验证装饰器"""
|
||||
@functools.wraps(f)
|
||||
def decorated_function(*args, **kwargs):
|
||||
if 'user_id' not in session:
|
||||
from flask import jsonify
|
||||
return jsonify({'success': False, 'error': '请先登录'}), 401
|
||||
return f(*args, **kwargs)
|
||||
return decorated_function
|
||||
|
||||
|
||||
def get_current_user_id():
|
||||
"""获取当前登录用户ID"""
|
||||
return session.get('user_id')
|
||||
|
||||
|
||||
def get_current_username():
|
||||
"""获取当前登录用户名"""
|
||||
return session.get('username')
|
||||
|
||||
|
||||
# ========== 可用资金操作 ==========
|
||||
|
||||
def db_get_available_cash(user_id):
|
||||
"""获取用户可用资金"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return 0
|
||||
|
||||
try:
|
||||
cur = conn.cursor(cursor_factory=RealDictCursor)
|
||||
cur.execute("SELECT available_cash FROM users WHERE id = %s", (user_id,))
|
||||
result = cur.fetchone()
|
||||
return float(result['available_cash'] or 0) if result else 0
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def db_update_available_cash(user_id, amount):
|
||||
"""更新用户可用资金"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return False, '数据库连接失败'
|
||||
|
||||
try:
|
||||
cur = conn.cursor()
|
||||
cur.execute("UPDATE users SET available_cash = %s WHERE id = %s", (amount, user_id))
|
||||
conn.commit()
|
||||
return True, None
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
return False, str(e)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
# ========== 用户操作 ==========
|
||||
|
||||
def create_user(email, password):
|
||||
"""创建用户（使用邮箱）"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return None, '数据库连接失败'
|
||||
|
||||
try:
|
||||
cur = conn.cursor(cursor_factory=RealDictCursor)
|
||||
|
||||
# 检查邮箱是否已存在
|
||||
cur.execute("SELECT id FROM users WHERE email = %s OR username = %s", (email, email))
|
||||
if cur.fetchone():
|
||||
return None, '该邮箱已注册'
|
||||
|
||||
# 创建用户（username和email都存邮箱）
|
||||
password_hash = generate_password_hash(password)
|
||||
cur.execute(
|
||||
"INSERT INTO users (username, email, password_hash) VALUES (%s, %s, %s) RETURNING id, username, email",
|
||||
(email, email, password_hash)
|
||||
)
|
||||
user = cur.fetchone()
|
||||
conn.commit()
|
||||
|
||||
return user, None
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
return None, str(e)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def verify_user(email, password):
|
||||
"""验证用户登录（使用邮箱）"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return None, '数据库连接失败'
|
||||
|
||||
try:
|
||||
cur = conn.cursor(cursor_factory=RealDictCursor)
|
||||
# 同时检查email和username字段（兼容旧数据）
|
||||
cur.execute("SELECT * FROM users WHERE email = %s OR username = %s", (email, email))
|
||||
user = cur.fetchone()
|
||||
|
||||
if not user or not check_password_hash(user['password_hash'], password):
|
||||
return None, '邮箱或密码错误'
|
||||
|
||||
return {'id': user['id'], 'username': user.get('email') or user['username']}, None
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def change_user_password(user_id, old_password, new_password):
|
||||
"""修改用户密码"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return False, '数据库连接失败'
|
||||
|
||||
try:
|
||||
cur = conn.cursor(cursor_factory=RealDictCursor)
|
||||
cur.execute("SELECT * FROM users WHERE id = %s", (user_id,))
|
||||
user = cur.fetchone()
|
||||
|
||||
if not user:
|
||||
return False, '用户不存在'
|
||||
|
||||
if not check_password_hash(user['password_hash'], old_password):
|
||||
return False, '当前密码错误'
|
||||
|
||||
new_hash = generate_password_hash(new_password)
|
||||
cur.execute("UPDATE users SET password_hash = %s WHERE id = %s", (new_hash, user_id))
|
||||
conn.commit()
|
||||
return True, None
|
||||
except Exception as e:
|
||||
return False, str(e)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
# ========== 交易记录操作（数据库版） ==========
|
||||
|
||||
def db_get_trades(user_id):
|
||||
"""从数据库获取用户交易记录"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return []
|
||||
|
||||
try:
|
||||
cur = conn.cursor(cursor_factory=RealDictCursor)
|
||||
cur.execute("""
|
||||
SELECT id, stock_code, stock_name, trade_type, price, quantity,
|
||||
trade_date::text, reason, result, profit_amount, stop_loss_price, notes,
|
||||
created_at::text
|
||||
FROM trades
|
||||
WHERE user_id = %s
|
||||
ORDER BY trade_date DESC, created_at DESC
|
||||
""", (user_id,))
|
||||
return cur.fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def db_get_trade(user_id, trade_id):
|
||||
"""获取单条交易记录"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return None
|
||||
try:
|
||||
cur = conn.cursor(cursor_factory=RealDictCursor)
|
||||
cur.execute("""
|
||||
SELECT id, stock_code, stock_name, trade_type, price, quantity,
|
||||
trade_date::text, reason, result, profit_amount, stop_loss_price, notes,
|
||||
created_at::text
|
||||
FROM trades WHERE id = %s AND user_id = %s
|
||||
""", (trade_id, user_id))
|
||||
return cur.fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def db_add_trade(user_id, data):
|
||||
"""添加交易记录到数据库"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return None, '数据库连接失败'
|
||||
|
||||
try:
|
||||
cur = conn.cursor(cursor_factory=RealDictCursor)
|
||||
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)
|
||||
RETURNING id, stock_code, stock_name, trade_type, price, quantity,
|
||||
trade_date::text, reason, result, profit_amount, stop_loss_price,
|
||||
notes, created_at::text
|
||||
""", (
|
||||
user_id,
|
||||
data.get('stock_code'),
|
||||
data.get('stock_name'),
|
||||
data.get('trade_type'),
|
||||
data.get('price'),
|
||||
data.get('quantity'),
|
||||
data.get('trade_date'),
|
||||
data.get('reason'),
|
||||
data.get('result'),
|
||||
data.get('profit_amount'),
|
||||
data.get('stop_loss_price'),
|
||||
data.get('notes')
|
||||
))
|
||||
trade = cur.fetchone()
|
||||
conn.commit()
|
||||
return trade, None
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
return None, str(e)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def db_update_trade(user_id, trade_id, data):
|
||||
"""更新交易记录"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return None, '数据库连接失败'
|
||||
|
||||
try:
|
||||
cur = conn.cursor(cursor_factory=RealDictCursor)
|
||||
cur.execute("""
|
||||
UPDATE trades SET
|
||||
stock_code = COALESCE(%s, stock_code),
|
||||
stock_name = COALESCE(%s, stock_name),
|
||||
trade_type = COALESCE(%s, trade_type),
|
||||
price = COALESCE(%s, price),
|
||||
quantity = COALESCE(%s, quantity),
|
||||
trade_date = COALESCE(%s, trade_date),
|
||||
reason = COALESCE(%s, reason),
|
||||
result = COALESCE(%s, result),
|
||||
profit_amount = COALESCE(%s, profit_amount),
|
||||
stop_loss_price = COALESCE(%s, stop_loss_price),
|
||||
notes = COALESCE(%s, notes)
|
||||
WHERE id = %s AND user_id = %s
|
||||
RETURNING id, stock_code, stock_name, trade_type, price, quantity,
|
||||
trade_date::text, reason, result, profit_amount, stop_loss_price,
|
||||
notes, created_at::text
|
||||
""", (
|
||||
data.get('stock_code'),
|
||||
data.get('stock_name'),
|
||||
data.get('trade_type'),
|
||||
data.get('price'),
|
||||
data.get('quantity'),
|
||||
data.get('trade_date'),
|
||||
data.get('reason'),
|
||||
data.get('result'),
|
||||
data.get('profit_amount'),
|
||||
data.get('stop_loss_price'),
|
||||
data.get('notes'),
|
||||
trade_id,
|
||||
user_id
|
||||
))
|
||||
trade = cur.fetchone()
|
||||
conn.commit()
|
||||
return trade, None
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
return None, str(e)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def db_delete_trade(user_id, trade_id):
|
||||
"""删除交易记录"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return False
|
||||
|
||||
try:
|
||||
cur = conn.cursor()
|
||||
cur.execute("DELETE FROM trades WHERE id = %s AND user_id = %s", (trade_id, user_id))
|
||||
conn.commit()
|
||||
return cur.rowcount > 0
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
# ========== 关注列表操作（数据库版） ==========
|
||||
|
||||
def db_get_watchlist(user_id):
|
||||
"""从数据库获取用户关注列表"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return []
|
||||
|
||||
try:
|
||||
cur = conn.cursor(cursor_factory=RealDictCursor)
|
||||
cur.execute("""
|
||||
SELECT stock_code as code, stock_name as name, added_time::text
|
||||
FROM watchlist
|
||||
WHERE user_id = %s
|
||||
ORDER BY added_time DESC
|
||||
""", (user_id,))
|
||||
return cur.fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def db_add_to_watchlist(user_id, code, name):
|
||||
"""添加到关注列表"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return None, '数据库连接失败'
|
||||
|
||||
try:
|
||||
cur = conn.cursor(cursor_factory=RealDictCursor)
|
||||
cur.execute("""
|
||||
INSERT INTO watchlist (user_id, stock_code, stock_name)
|
||||
VALUES (%s, %s, %s)
|
||||
ON CONFLICT (user_id, stock_code) DO NOTHING
|
||||
RETURNING stock_code as code, stock_name as name
|
||||
""", (user_id, code, name))
|
||||
conn.commit()
|
||||
return db_get_watchlist(user_id), None
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
return None, str(e)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def db_remove_from_watchlist(user_id, code):
|
||||
"""从关注列表移除"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return None
|
||||
|
||||
try:
|
||||
cur = conn.cursor()
|
||||
cur.execute("DELETE FROM watchlist WHERE user_id = %s AND stock_code = %s", (user_id, code))
|
||||
conn.commit()
|
||||
return db_get_watchlist(user_id)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
# ========== 分析缓存操作（数据库版） ==========
|
||||
|
||||
def db_get_alerts_cache(user_id):
|
||||
"""从数据库获取分析缓存"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return None
|
||||
|
||||
try:
|
||||
cur = conn.cursor(cursor_factory=RealDictCursor)
|
||||
cur.execute("""
|
||||
SELECT data, updated_at::text as lastUpdate
|
||||
FROM alerts_cache
|
||||
WHERE user_id = %s
|
||||
""", (user_id,))
|
||||
result = cur.fetchone()
|
||||
if result:
|
||||
return {
|
||||
'alerts': result['data'] or [],
|
||||
'lastUpdate': result['lastupdate']
|
||||
}
|
||||
return {'alerts': [], 'lastUpdate': None}
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def db_save_alerts_cache(user_id, alerts):
|
||||
"""保存分析缓存到数据库"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return False
|
||||
|
||||
try:
|
||||
import json
|
||||
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()
|
||||
return True
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
print(f"保存分析缓存失败: {e}")
|
||||
return False
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
# ========== 基本面数据操作（数据库版） ==========
|
||||
|
||||
def db_get_fundamental(code):
|
||||
"""从数据库获取基本面数据（当日缓存）"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return None
|
||||
|
||||
try:
|
||||
from datetime import date
|
||||
cur = conn.cursor(cursor_factory=RealDictCursor)
|
||||
cur.execute("""
|
||||
SELECT code, name, pe, pb, total_market_cap, industry,
|
||||
latest_price, change_pct, update_date::text, updated_at::text,
|
||||
roe, eps, bps, revenue_yoy, profit_yoy, gross_margin, net_margin
|
||||
FROM stock_fundamental
|
||||
WHERE code = %s AND update_date = %s
|
||||
""", (code, date.today()))
|
||||
return cur.fetchone()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def db_save_fundamental(code, data):
|
||||
"""保存基本面数据到数据库"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return False
|
||||
|
||||
try:
|
||||
from datetime import date
|
||||
cur = conn.cursor()
|
||||
cur.execute("""
|
||||
INSERT INTO stock_fundamental
|
||||
(code, name, pe, pb, total_market_cap, industry, latest_price, change_pct, update_date,
|
||||
roe, eps, bps, revenue_yoy, profit_yoy, gross_margin, net_margin)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||||
ON CONFLICT (code) DO UPDATE SET
|
||||
name = EXCLUDED.name,
|
||||
pe = EXCLUDED.pe,
|
||||
pb = EXCLUDED.pb,
|
||||
total_market_cap = EXCLUDED.total_market_cap,
|
||||
industry = EXCLUDED.industry,
|
||||
latest_price = EXCLUDED.latest_price,
|
||||
change_pct = EXCLUDED.change_pct,
|
||||
update_date = EXCLUDED.update_date,
|
||||
roe = EXCLUDED.roe,
|
||||
eps = EXCLUDED.eps,
|
||||
bps = EXCLUDED.bps,
|
||||
revenue_yoy = EXCLUDED.revenue_yoy,
|
||||
profit_yoy = EXCLUDED.profit_yoy,
|
||||
gross_margin = EXCLUDED.gross_margin,
|
||||
net_margin = EXCLUDED.net_margin,
|
||||
updated_at = NOW()
|
||||
""", (
|
||||
code,
|
||||
data.get('name') or data.get('stock_name'),
|
||||
data.get('pe') or data.get('pe_ttm'),
|
||||
data.get('pb'),
|
||||
data.get('total_market_cap'),
|
||||
data.get('industry'),
|
||||
data.get('latest_price'),
|
||||
data.get('change_pct'),
|
||||
date.today(),
|
||||
data.get('roe'),
|
||||
data.get('eps'),
|
||||
data.get('bps'),
|
||||
data.get('revenue_yoy'),
|
||||
data.get('profit_yoy'),
|
||||
data.get('gross_margin'),
|
||||
data.get('net_margin'),
|
||||
))
|
||||
conn.commit()
|
||||
return True
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
print(f"保存基本面数据失败: {e}")
|
||||
return False
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
# ========== 资金流向历史数据操作（数据库版） ==========
|
||||
|
||||
def db_get_fund_flow_history(code):
|
||||
"""获取股票的资金流向历史数据"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return None, None
|
||||
|
||||
try:
|
||||
cur = conn.cursor(cursor_factory=RealDictCursor)
|
||||
cur.execute("""
|
||||
SELECT code, trade_date::text, close_price, change_pct,
|
||||
main_net_inflow, main_net_inflow_pct,
|
||||
super_net_inflow, super_net_inflow_pct,
|
||||
big_net_inflow, big_net_inflow_pct
|
||||
FROM stock_fund_flow_history
|
||||
WHERE code = %s
|
||||
ORDER BY trade_date DESC
|
||||
""", (code,))
|
||||
rows = cur.fetchall()
|
||||
|
||||
# 获取最新日期
|
||||
latest_date = rows[0]['trade_date'] if rows else None
|
||||
|
||||
return [dict(row) for row in rows], latest_date
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def db_save_fund_flow_history(code, records):
|
||||
"""保存资金流向历史数据到数据库"""
|
||||
conn = get_db()
|
||||
if not conn:
|
||||
return False
|
||||
|
||||
try:
|
||||
cur = conn.cursor()
|
||||
for r in records:
|
||||
cur.execute("""
|
||||
INSERT INTO stock_fund_flow_history
|
||||
(code, trade_date, close_price, change_pct,
|
||||
main_net_inflow, main_net_inflow_pct,
|
||||
super_net_inflow, super_net_inflow_pct,
|
||||
big_net_inflow, big_net_inflow_pct,
|
||||
mid_net_inflow, mid_net_inflow_pct,
|
||||
small_net_inflow, small_net_inflow_pct)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||||
ON CONFLICT (code, trade_date) DO UPDATE SET
|
||||
close_price = EXCLUDED.close_price,
|
||||
change_pct = EXCLUDED.change_pct,
|
||||
main_net_inflow = EXCLUDED.main_net_inflow,
|
||||
main_net_inflow_pct = EXCLUDED.main_net_inflow_pct,
|
||||
super_net_inflow = EXCLUDED.super_net_inflow,
|
||||
super_net_inflow_pct = EXCLUDED.super_net_inflow_pct,
|
||||
big_net_inflow = EXCLUDED.big_net_inflow,
|
||||
big_net_inflow_pct = EXCLUDED.big_net_inflow_pct,
|
||||
mid_net_inflow = EXCLUDED.mid_net_inflow,
|
||||
mid_net_inflow_pct = EXCLUDED.mid_net_inflow_pct,
|
||||
small_net_inflow = EXCLUDED.small_net_inflow,
|
||||
small_net_inflow_pct = EXCLUDED.small_net_inflow_pct,
|
||||
updated_at = NOW()
|
||||
""", (
|
||||
code,
|
||||
r.get('日期') or r.get('trade_date'),
|
||||
r.get('收盘价') or r.get('close_price'),
|
||||
r.get('涨跌幅') or r.get('change_pct'),
|
||||
r.get('主力净流入-净额') or r.get('main_net_inflow'),
|
||||
r.get('主力净流入-净占比') or r.get('main_net_inflow_pct'),
|
||||
r.get('超大单净流入-净额') or r.get('super_net_inflow'),
|
||||
r.get('超大单净流入-净占比') or r.get('super_net_inflow_pct'),
|
||||
r.get('大单净流入-净额') or r.get('big_net_inflow'),
|
||||
r.get('大单净流入-净占比') or r.get('big_net_inflow_pct'),
|
||||
r.get('中单净流入-净额') or r.get('mid_net_inflow'),
|
||||
r.get('中单净流入-净占比') or r.get('mid_net_inflow_pct'),
|
||||
r.get('小单净流入-净额') or r.get('small_net_inflow'),
|
||||
r.get('小单净流入-净占比') or r.get('small_net_inflow_pct')
|
||||
))
|
||||
conn.commit()
|
||||
return True
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
print(f"保存资金流向历史失败: {e}")
|
||||
return False
|
||||
finally:
|
||||
conn.close()
|
||||
Reference in New Issue
Block a user