feat: Multi-user system with PostgreSQL - WIP temporary save
This commit is contained in:
TIANHE
@@ -1,6 +1,12 @@
|
||||
"""
|
||||
Authentication Utilities
|
||||
|
||||
JWT token generation, verification, and middleware decorators.
|
||||
Supports multi-user authentication with role-based access control.
|
||||
"""
|
||||
import jwt
|
||||
import datetime
|
||||
import os
|
||||
from functools import wraps
|
||||
from flask import request, jsonify, g
|
||||
from app.config.settings import Config
|
||||
@@ -8,13 +14,26 @@ from app.utils.logger import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
def generate_token(username):
|
||||
"""Generate JWT token."""
|
||||
|
||||
def generate_token(user_id: int, username: str, role: str = 'user') -> str:
|
||||
"""
|
||||
Generate JWT token with user information.
|
||||
|
||||
Args:
|
||||
user_id: User ID
|
||||
username: Username
|
||||
role: User role (admin/manager/user/viewer)
|
||||
|
||||
Returns:
|
||||
JWT token string
|
||||
"""
|
||||
try:
|
||||
payload = {
|
||||
'exp': datetime.datetime.utcnow() + datetime.timedelta(days=7),
|
||||
'iat': datetime.datetime.utcnow(),
|
||||
'sub': username
|
||||
'sub': username,
|
||||
'user_id': user_id,
|
||||
'role': role,
|
||||
}
|
||||
return jwt.encode(
|
||||
payload,
|
||||
@@ -25,18 +44,44 @@ def generate_token(username):
|
||||
logger.error(f"Token generation failed: {e}")
|
||||
return None
|
||||
|
||||
def verify_token(token):
|
||||
"""Verify JWT token."""
|
||||
|
||||
def verify_token(token: str) -> dict:
|
||||
"""
|
||||
Verify JWT token and return payload.
|
||||
|
||||
Args:
|
||||
token: JWT token string
|
||||
|
||||
Returns:
|
||||
Token payload dict or None if invalid
|
||||
"""
|
||||
try:
|
||||
payload = jwt.decode(token, Config.SECRET_KEY, algorithms=['HS256'])
|
||||
return payload['sub']
|
||||
return payload
|
||||
except jwt.ExpiredSignatureError:
|
||||
logger.debug("Token expired")
|
||||
return None
|
||||
except jwt.InvalidTokenError:
|
||||
except jwt.InvalidTokenError as e:
|
||||
logger.debug(f"Invalid token: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def get_current_user_id() -> int:
|
||||
"""Get current user ID from flask.g context"""
|
||||
return getattr(g, 'user_id', None)
|
||||
|
||||
|
||||
def get_current_user_role() -> str:
|
||||
"""Get current user role from flask.g context"""
|
||||
return getattr(g, 'user_role', 'user')
|
||||
|
||||
|
||||
def login_required(f):
|
||||
"""Decorator that enforces Bearer token auth."""
|
||||
"""
|
||||
Decorator that enforces Bearer token auth.
|
||||
|
||||
Sets g.user, g.user_id, g.user_role on successful auth.
|
||||
"""
|
||||
@wraps(f)
|
||||
def decorated(*args, **kwargs):
|
||||
token = None
|
||||
@@ -51,13 +96,96 @@ def login_required(f):
|
||||
if not token:
|
||||
return jsonify({'code': 401, 'msg': 'Token missing', 'data': None}), 401
|
||||
|
||||
username = verify_token(token)
|
||||
if not username:
|
||||
payload = verify_token(token)
|
||||
if not payload:
|
||||
return jsonify({'code': 401, 'msg': 'Token invalid or expired', 'data': None}), 401
|
||||
|
||||
# Store user in flask.g
|
||||
g.user = username
|
||||
|
||||
# Store user info in flask.g
|
||||
g.user = payload.get('sub')
|
||||
g.user_id = payload.get('user_id')
|
||||
g.user_role = payload.get('role', 'user')
|
||||
|
||||
return f(*args, **kwargs)
|
||||
|
||||
return decorated
|
||||
|
||||
|
||||
def admin_required(f):
|
||||
"""
|
||||
Decorator that requires admin role.
|
||||
Must be used after @login_required.
|
||||
"""
|
||||
@wraps(f)
|
||||
def decorated(*args, **kwargs):
|
||||
role = getattr(g, 'user_role', None)
|
||||
if role != 'admin':
|
||||
return jsonify({'code': 403, 'msg': 'Admin access required', 'data': None}), 403
|
||||
return f(*args, **kwargs)
|
||||
return decorated
|
||||
|
||||
|
||||
def manager_required(f):
|
||||
"""
|
||||
Decorator that requires manager or admin role.
|
||||
Must be used after @login_required.
|
||||
"""
|
||||
@wraps(f)
|
||||
def decorated(*args, **kwargs):
|
||||
role = getattr(g, 'user_role', None)
|
||||
if role not in ('admin', 'manager'):
|
||||
return jsonify({'code': 403, 'msg': 'Manager access required', 'data': None}), 403
|
||||
return f(*args, **kwargs)
|
||||
return decorated
|
||||
|
||||
|
||||
def permission_required(permission: str):
|
||||
"""
|
||||
Decorator factory that checks for a specific permission.
|
||||
Must be used after @login_required.
|
||||
|
||||
Usage:
|
||||
@login_required
|
||||
@permission_required('strategy')
|
||||
def my_endpoint():
|
||||
...
|
||||
"""
|
||||
def decorator(f):
|
||||
@wraps(f)
|
||||
def decorated(*args, **kwargs):
|
||||
role = getattr(g, 'user_role', 'user')
|
||||
|
||||
# Import here to avoid circular import
|
||||
from app.services.user_service import get_user_service
|
||||
permissions = get_user_service().get_user_permissions(role)
|
||||
|
||||
if permission not in permissions:
|
||||
return jsonify({
|
||||
'code': 403,
|
||||
'msg': f'Permission denied: {permission}',
|
||||
'data': None
|
||||
}), 403
|
||||
|
||||
return f(*args, **kwargs)
|
||||
return decorated
|
||||
return decorator
|
||||
|
||||
|
||||
# Legacy compatibility: single-user mode fallback
|
||||
def _is_single_user_mode() -> bool:
|
||||
"""Check if system is in single-user (legacy) mode"""
|
||||
return os.getenv('SINGLE_USER_MODE', 'false').lower() == 'true'
|
||||
|
||||
|
||||
def authenticate_legacy(username: str, password: str) -> dict:
|
||||
"""
|
||||
Legacy single-user authentication (for backward compatibility).
|
||||
Uses ADMIN_USER and ADMIN_PASSWORD from environment.
|
||||
"""
|
||||
if username == Config.ADMIN_USER and password == Config.ADMIN_PASSWORD:
|
||||
return {
|
||||
'user_id': 1,
|
||||
'username': username,
|
||||
'role': 'admin',
|
||||
'nickname': 'Admin',
|
||||
}
|
||||
return None
|
||||
|
||||
@@ -1,662 +1,65 @@
|
||||
|
||||
"""
|
||||
SQLite 数据库连接工具 (本地化适配版)
|
||||
Database Connection Utility - PostgreSQL Only
|
||||
|
||||
Provides unified interface for PostgreSQL database operations.
|
||||
|
||||
Usage:
|
||||
from app.utils.db import get_db_connection
|
||||
|
||||
with get_db_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT * FROM users WHERE id = %s", (user_id,))
|
||||
row = cursor.fetchone()
|
||||
conn.commit()
|
||||
|
||||
Configuration:
|
||||
DATABASE_URL=postgresql://user:password@host:port/dbname
|
||||
"""
|
||||
import sqlite3
|
||||
import os
|
||||
import threading
|
||||
import shutil
|
||||
from typing import Optional, Any, List, Dict
|
||||
from contextlib import contextmanager
|
||||
from app.utils.logger import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# SQLite 主库路径解析
|
||||
#
|
||||
# 目标默认行为：把主库放到 `backend_api_python/data/quantdinger.db`
|
||||
# 兼容旧行为：老版本会在 `backend_api_python/quantdinger.db` 建库
|
||||
_BASE_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
|
||||
_DEFAULT_DB_FILE = os.path.join(_BASE_DIR, 'data', 'quantdinger.db')
|
||||
_LEGACY_DB_FILE = os.path.join(_BASE_DIR, 'quantdinger.db')
|
||||
# Re-export from PostgreSQL module
|
||||
from app.utils.db_postgres import (
|
||||
get_pg_connection as get_db_connection,
|
||||
get_pg_connection_sync as get_db_connection_sync,
|
||||
is_postgres_available,
|
||||
close_pool as close_db,
|
||||
)
|
||||
|
||||
|
||||
def _get_db_file() -> str:
|
||||
def get_db_type() -> str:
|
||||
"""Get database type (always postgresql)"""
|
||||
return 'postgresql'
|
||||
|
||||
|
||||
def is_postgres() -> bool:
|
||||
"""Check if using PostgreSQL (always True)"""
|
||||
return True
|
||||
|
||||
|
||||
def init_database():
|
||||
"""
|
||||
Resolve SQLite DB file path.
|
||||
|
||||
Priority:
|
||||
- SQLITE_DATABASE_FILE env (Docker 推荐：/app/data/quantdinger.db)
|
||||
- default: backend_api_python/data/quantdinger.db (local recommended)
|
||||
|
||||
Also performs a best-effort one-time migration:
|
||||
If the configured/default path doesn't exist but legacy db exists, copy legacy to configured/default path.
|
||||
Initialize database connection.
|
||||
Schema is created via migrations/init.sql on PostgreSQL container start.
|
||||
"""
|
||||
env_path = os.getenv('SQLITE_DATABASE_FILE')
|
||||
db_path = (env_path or '').strip() or _DEFAULT_DB_FILE
|
||||
if is_postgres_available():
|
||||
from app.utils.logger import get_logger
|
||||
logger = get_logger(__name__)
|
||||
logger.info("PostgreSQL connection verified")
|
||||
else:
|
||||
raise RuntimeError("Cannot connect to PostgreSQL. Check DATABASE_URL.")
|
||||
|
||||
# Ensure parent dir exists
|
||||
parent = os.path.dirname(db_path)
|
||||
if parent and not os.path.exists(parent):
|
||||
try:
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Best-effort migration from legacy path
|
||||
try:
|
||||
if os.path.abspath(db_path) != os.path.abspath(_LEGACY_DB_FILE):
|
||||
if (not os.path.exists(db_path)) and os.path.exists(_LEGACY_DB_FILE):
|
||||
shutil.copy2(_LEGACY_DB_FILE, db_path)
|
||||
logger.info(f"Migrated SQLite DB from legacy path to {db_path}")
|
||||
except Exception as e:
|
||||
logger.warning(f"SQLite DB migration skipped/failed: {e}")
|
||||
|
||||
return db_path
|
||||
|
||||
# 线程锁，用于简单的并发控制（SQLite 对写操作有限制）
|
||||
_db_lock = threading.Lock()
|
||||
|
||||
def _init_db_schema(conn):
|
||||
"""初始化数据库表结构"""
|
||||
cursor = conn.cursor()
|
||||
|
||||
def ensure_columns(table: str, columns: Dict[str, str]) -> None:
|
||||
"""
|
||||
Ensure columns exist for an existing SQLite table (simple migration).
|
||||
columns: {column_name: "TYPE DEFAULT ..."}
|
||||
"""
|
||||
try:
|
||||
cursor.execute(f"PRAGMA table_info({table})")
|
||||
existing = {row[1] for row in cursor.fetchall() or []} # row[1] is column name
|
||||
for col, ddl in columns.items():
|
||||
if col in existing:
|
||||
continue
|
||||
cursor.execute(f"ALTER TABLE {table} ADD COLUMN {col} {ddl}")
|
||||
except Exception as e:
|
||||
logger.warning(f"ensure_columns failed for table={table}: {e}")
|
||||
|
||||
# 1. 策略表
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS qd_strategies_trading (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
strategy_name TEXT NOT NULL,
|
||||
strategy_type TEXT DEFAULT 'IndicatorStrategy',
|
||||
market_category TEXT DEFAULT 'Crypto',
|
||||
execution_mode TEXT DEFAULT 'signal',
|
||||
notification_config TEXT DEFAULT '', -- JSON string
|
||||
status TEXT DEFAULT 'stopped',
|
||||
symbol TEXT,
|
||||
timeframe TEXT,
|
||||
initial_capital REAL DEFAULT 1000,
|
||||
leverage INTEGER DEFAULT 1,
|
||||
market_type TEXT DEFAULT 'swap',
|
||||
exchange_config TEXT, -- JSON string
|
||||
indicator_config TEXT, -- JSON string
|
||||
trading_config TEXT, -- JSON string
|
||||
ai_model_config TEXT, -- JSON string
|
||||
decide_interval INTEGER DEFAULT 300,
|
||||
created_at INTEGER,
|
||||
updated_at INTEGER
|
||||
)
|
||||
""")
|
||||
|
||||
ensure_columns("qd_strategies_trading", {
|
||||
"market_category": "TEXT DEFAULT 'Crypto'",
|
||||
"execution_mode": "TEXT DEFAULT 'signal'",
|
||||
"notification_config": "TEXT DEFAULT ''",
|
||||
"strategy_group_id": "TEXT DEFAULT ''", # 策略组ID，批量创建的策略共享同一个组ID
|
||||
"group_base_name": "TEXT DEFAULT ''" # 策略组基础名称（用于显示）
|
||||
})
|
||||
|
||||
# 2. 持仓表
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS qd_strategy_positions (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
strategy_id INTEGER,
|
||||
symbol TEXT,
|
||||
side TEXT, -- long/short
|
||||
size REAL,
|
||||
entry_price REAL,
|
||||
current_price REAL,
|
||||
highest_price REAL DEFAULT 0,
|
||||
lowest_price REAL DEFAULT 0,
|
||||
unrealized_pnl REAL DEFAULT 0,
|
||||
pnl_percent REAL DEFAULT 0,
|
||||
equity REAL DEFAULT 0,
|
||||
updated_at INTEGER,
|
||||
UNIQUE(strategy_id, symbol, side)
|
||||
)
|
||||
""")
|
||||
|
||||
ensure_columns("qd_strategy_positions", {
|
||||
"highest_price": "REAL DEFAULT 0",
|
||||
"lowest_price": "REAL DEFAULT 0",
|
||||
})
|
||||
|
||||
# 3. 交易记录表
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS qd_strategy_trades (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
strategy_id INTEGER,
|
||||
symbol TEXT,
|
||||
type TEXT, -- open_long, close_short, etc.
|
||||
price REAL,
|
||||
amount REAL,
|
||||
value REAL,
|
||||
commission REAL DEFAULT 0,
|
||||
commission_ccy TEXT DEFAULT '',
|
||||
profit REAL DEFAULT 0,
|
||||
created_at INTEGER
|
||||
)
|
||||
""")
|
||||
|
||||
ensure_columns("qd_strategy_trades", {
|
||||
"commission_ccy": "TEXT DEFAULT ''",
|
||||
})
|
||||
|
||||
# NOTE:
|
||||
# We intentionally do not persist runtime logs in DB for local deployments.
|
||||
# Use console logs / stdout prints instead.
|
||||
|
||||
# 3.1 Pending orders queue (signal dispatch / live execution)
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS pending_orders (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
strategy_id INTEGER,
|
||||
symbol TEXT NOT NULL,
|
||||
signal_type TEXT NOT NULL, -- open_long/close_long/open_short/close_short/add_long/add_short
|
||||
signal_ts INTEGER, -- candle timestamp (seconds). used for strict de-dup per candle
|
||||
market_type TEXT DEFAULT 'swap',
|
||||
order_type TEXT DEFAULT 'market',
|
||||
amount REAL DEFAULT 0, -- base amount (or stake amount depending on execution)
|
||||
price REAL DEFAULT 0, -- reference price at enqueue time
|
||||
execution_mode TEXT DEFAULT 'signal', -- signal/live
|
||||
status TEXT DEFAULT 'pending', -- pending/processing/sent/failed/deferred
|
||||
priority INTEGER DEFAULT 0,
|
||||
attempts INTEGER DEFAULT 0,
|
||||
max_attempts INTEGER DEFAULT 10,
|
||||
last_error TEXT DEFAULT '',
|
||||
payload_json TEXT DEFAULT '', -- JSON string for dispatcher
|
||||
-- Live execution result fields (best-effort)
|
||||
dispatch_note TEXT DEFAULT '',
|
||||
exchange_id TEXT DEFAULT '',
|
||||
exchange_order_id TEXT DEFAULT '',
|
||||
exchange_response_json TEXT DEFAULT '',
|
||||
filled REAL DEFAULT 0,
|
||||
avg_price REAL DEFAULT 0,
|
||||
executed_at INTEGER,
|
||||
created_at INTEGER,
|
||||
updated_at INTEGER,
|
||||
processed_at INTEGER,
|
||||
sent_at INTEGER
|
||||
)
|
||||
""")
|
||||
|
||||
ensure_columns("pending_orders", {
|
||||
"signal_ts": "INTEGER",
|
||||
"market_type": "TEXT DEFAULT 'swap'",
|
||||
"order_type": "TEXT DEFAULT 'market'",
|
||||
"price": "REAL DEFAULT 0",
|
||||
"execution_mode": "TEXT DEFAULT 'signal'",
|
||||
"status": "TEXT DEFAULT 'pending'",
|
||||
"priority": "INTEGER DEFAULT 0",
|
||||
"attempts": "INTEGER DEFAULT 0",
|
||||
"max_attempts": "INTEGER DEFAULT 10",
|
||||
"last_error": "TEXT DEFAULT ''",
|
||||
"payload_json": "TEXT DEFAULT ''",
|
||||
"dispatch_note": "TEXT DEFAULT ''",
|
||||
"exchange_id": "TEXT DEFAULT ''",
|
||||
"exchange_order_id": "TEXT DEFAULT ''",
|
||||
"exchange_response_json": "TEXT DEFAULT ''",
|
||||
"filled": "REAL DEFAULT 0",
|
||||
"avg_price": "REAL DEFAULT 0",
|
||||
"executed_at": "INTEGER",
|
||||
"created_at": "INTEGER",
|
||||
"updated_at": "INTEGER",
|
||||
"processed_at": "INTEGER",
|
||||
"sent_at": "INTEGER",
|
||||
})
|
||||
|
||||
# 3.2 Strategy notifications (browser polling / audit trail)
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS qd_strategy_notifications (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
strategy_id INTEGER,
|
||||
symbol TEXT DEFAULT '',
|
||||
signal_type TEXT DEFAULT '',
|
||||
channels TEXT DEFAULT '',
|
||||
title TEXT DEFAULT '',
|
||||
message TEXT DEFAULT '',
|
||||
payload_json TEXT DEFAULT '',
|
||||
created_at INTEGER
|
||||
)
|
||||
""")
|
||||
|
||||
ensure_columns("qd_strategy_notifications", {
|
||||
"strategy_id": "INTEGER",
|
||||
"symbol": "TEXT DEFAULT ''",
|
||||
"signal_type": "TEXT DEFAULT ''",
|
||||
"channels": "TEXT DEFAULT ''",
|
||||
"title": "TEXT DEFAULT ''",
|
||||
"message": "TEXT DEFAULT ''",
|
||||
"payload_json": "TEXT DEFAULT ''",
|
||||
"created_at": "INTEGER",
|
||||
"is_read": "INTEGER DEFAULT 0",
|
||||
})
|
||||
|
||||
# 4. 指标代码表（参考 MySQL: qd_indicator_codes）
|
||||
# 说明：
|
||||
# - 本地化后统一使用 SQLite，但字段保持与 MySQL 结构接近，便于前端/业务复用。
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS qd_indicator_codes (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL DEFAULT 1,
|
||||
is_buy INTEGER NOT NULL DEFAULT 0,
|
||||
end_time INTEGER NOT NULL DEFAULT 1,
|
||||
name TEXT NOT NULL DEFAULT '',
|
||||
code TEXT,
|
||||
description TEXT DEFAULT '',
|
||||
publish_to_community INTEGER NOT NULL DEFAULT 0,
|
||||
pricing_type TEXT NOT NULL DEFAULT 'free',
|
||||
price REAL NOT NULL DEFAULT 0,
|
||||
is_encrypted INTEGER NOT NULL DEFAULT 0,
|
||||
preview_image TEXT DEFAULT '',
|
||||
createtime INTEGER,
|
||||
updatetime INTEGER,
|
||||
-- legacy local columns (kept for backward compatibility)
|
||||
created_at INTEGER,
|
||||
updated_at INTEGER
|
||||
)
|
||||
""")
|
||||
|
||||
# Migrate older local DBs (missing columns) to the new schema shape.
|
||||
ensure_columns("qd_indicator_codes", {
|
||||
"user_id": "INTEGER NOT NULL DEFAULT 1",
|
||||
"is_buy": "INTEGER NOT NULL DEFAULT 0",
|
||||
"end_time": "INTEGER NOT NULL DEFAULT 1",
|
||||
"publish_to_community": "INTEGER NOT NULL DEFAULT 0",
|
||||
"pricing_type": "TEXT NOT NULL DEFAULT 'free'",
|
||||
"price": "REAL NOT NULL DEFAULT 0",
|
||||
"is_encrypted": "INTEGER NOT NULL DEFAULT 0",
|
||||
"preview_image": "TEXT DEFAULT ''",
|
||||
"createtime": "INTEGER",
|
||||
"updatetime": "INTEGER"
|
||||
})
|
||||
|
||||
# 4.1 策略代码表（indicator-analysis 本地策略；与交易执行器的 qd_strategies_trading 区分）
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS qd_strategy_codes (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL DEFAULT 1,
|
||||
name TEXT NOT NULL DEFAULT '',
|
||||
code TEXT,
|
||||
description TEXT DEFAULT '',
|
||||
createtime INTEGER,
|
||||
updatetime INTEGER
|
||||
)
|
||||
""")
|
||||
|
||||
ensure_columns("qd_strategy_codes", {
|
||||
"user_id": "INTEGER NOT NULL DEFAULT 1",
|
||||
"name": "TEXT NOT NULL DEFAULT ''",
|
||||
"code": "TEXT",
|
||||
"description": "TEXT DEFAULT ''",
|
||||
"createtime": "INTEGER",
|
||||
"updatetime": "INTEGER"
|
||||
})
|
||||
|
||||
# 5. AI决策记录表
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS qd_ai_decisions (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
strategy_id INTEGER,
|
||||
decision_data TEXT, -- JSON
|
||||
context_data TEXT, -- JSON
|
||||
created_at INTEGER
|
||||
)
|
||||
""")
|
||||
|
||||
# 6. 插件/系统配置表（原来由 MySQL 提供，这里用 SQLite 本地化）
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS qd_addon_config (
|
||||
config_key TEXT PRIMARY KEY,
|
||||
config_value TEXT,
|
||||
type TEXT DEFAULT 'string'
|
||||
)
|
||||
""")
|
||||
|
||||
# 7. Watchlist (local-only, single-user by default)
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS qd_watchlist (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER DEFAULT 1,
|
||||
market TEXT NOT NULL,
|
||||
symbol TEXT NOT NULL,
|
||||
name TEXT DEFAULT '',
|
||||
created_at INTEGER,
|
||||
updated_at INTEGER,
|
||||
UNIQUE(user_id, market, symbol)
|
||||
)
|
||||
""")
|
||||
|
||||
# 8. Analysis tasks / history (local-only)
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS qd_analysis_tasks (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER DEFAULT 1,
|
||||
market TEXT NOT NULL,
|
||||
symbol TEXT NOT NULL,
|
||||
model TEXT DEFAULT '',
|
||||
language TEXT DEFAULT 'en-US',
|
||||
status TEXT DEFAULT 'completed', -- completed/failed/processing/pending
|
||||
result_json TEXT DEFAULT '',
|
||||
error_message TEXT DEFAULT '',
|
||||
created_at INTEGER,
|
||||
completed_at INTEGER
|
||||
)
|
||||
""")
|
||||
|
||||
# 9. Backtest runs (for AI optimization / history)
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS qd_backtest_runs (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL DEFAULT 1,
|
||||
indicator_id INTEGER,
|
||||
market TEXT NOT NULL,
|
||||
symbol TEXT NOT NULL,
|
||||
timeframe TEXT NOT NULL,
|
||||
start_date TEXT NOT NULL, -- YYYY-MM-DD
|
||||
end_date TEXT NOT NULL, -- YYYY-MM-DD
|
||||
initial_capital REAL DEFAULT 10000,
|
||||
commission REAL DEFAULT 0.001,
|
||||
slippage REAL DEFAULT 0,
|
||||
leverage INTEGER DEFAULT 1,
|
||||
trade_direction TEXT DEFAULT 'long',
|
||||
strategy_config TEXT DEFAULT '', -- JSON string
|
||||
status TEXT DEFAULT 'success', -- success/failed
|
||||
error_message TEXT DEFAULT '',
|
||||
result_json TEXT DEFAULT '', -- JSON string
|
||||
created_at INTEGER
|
||||
)
|
||||
""")
|
||||
|
||||
ensure_columns("qd_backtest_runs", {
|
||||
"user_id": "INTEGER NOT NULL DEFAULT 1",
|
||||
"indicator_id": "INTEGER",
|
||||
"market": "TEXT NOT NULL DEFAULT ''",
|
||||
"symbol": "TEXT NOT NULL DEFAULT ''",
|
||||
"timeframe": "TEXT NOT NULL DEFAULT ''",
|
||||
"start_date": "TEXT NOT NULL DEFAULT ''",
|
||||
"end_date": "TEXT NOT NULL DEFAULT ''",
|
||||
"initial_capital": "REAL DEFAULT 10000",
|
||||
"commission": "REAL DEFAULT 0.001",
|
||||
"slippage": "REAL DEFAULT 0",
|
||||
"leverage": "INTEGER DEFAULT 1",
|
||||
"trade_direction": "TEXT DEFAULT 'long'",
|
||||
"strategy_config": "TEXT DEFAULT ''",
|
||||
"status": "TEXT DEFAULT 'success'",
|
||||
"error_message": "TEXT DEFAULT ''",
|
||||
"result_json": "TEXT DEFAULT ''",
|
||||
"created_at": "INTEGER"
|
||||
})
|
||||
|
||||
# 10. Exchange credentials vault (local-only)
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS qd_exchange_credentials (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL DEFAULT 1,
|
||||
name TEXT DEFAULT '',
|
||||
exchange_id TEXT NOT NULL,
|
||||
api_key_hint TEXT DEFAULT '',
|
||||
encrypted_config TEXT NOT NULL, -- encrypted JSON string
|
||||
created_at INTEGER,
|
||||
updated_at INTEGER
|
||||
)
|
||||
""")
|
||||
|
||||
ensure_columns("qd_exchange_credentials", {
|
||||
"user_id": "INTEGER NOT NULL DEFAULT 1",
|
||||
"name": "TEXT DEFAULT ''",
|
||||
"exchange_id": "TEXT NOT NULL DEFAULT ''",
|
||||
"api_key_hint": "TEXT DEFAULT ''",
|
||||
"encrypted_config": "TEXT NOT NULL DEFAULT ''",
|
||||
"created_at": "INTEGER",
|
||||
"updated_at": "INTEGER"
|
||||
})
|
||||
|
||||
# 11. Manual positions (user's existing holdings outside the system)
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS qd_manual_positions (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL DEFAULT 1,
|
||||
market TEXT NOT NULL, -- Crypto/USStock/AShare/HShare/Forex/Futures
|
||||
symbol TEXT NOT NULL,
|
||||
name TEXT DEFAULT '', -- Display name
|
||||
side TEXT DEFAULT 'long', -- long/short
|
||||
quantity REAL NOT NULL DEFAULT 0,
|
||||
entry_price REAL NOT NULL DEFAULT 0,
|
||||
entry_time INTEGER, -- Position open timestamp
|
||||
notes TEXT DEFAULT '', -- User notes
|
||||
tags TEXT DEFAULT '', -- JSON array of tags
|
||||
group_name TEXT DEFAULT '', -- Group name for categorization
|
||||
created_at INTEGER,
|
||||
updated_at INTEGER,
|
||||
UNIQUE(user_id, market, symbol, side)
|
||||
)
|
||||
""")
|
||||
|
||||
ensure_columns("qd_manual_positions", {
|
||||
"user_id": "INTEGER NOT NULL DEFAULT 1",
|
||||
"market": "TEXT NOT NULL DEFAULT ''",
|
||||
"symbol": "TEXT NOT NULL DEFAULT ''",
|
||||
"name": "TEXT DEFAULT ''",
|
||||
"side": "TEXT DEFAULT 'long'",
|
||||
"quantity": "REAL NOT NULL DEFAULT 0",
|
||||
"entry_price": "REAL NOT NULL DEFAULT 0",
|
||||
"entry_time": "INTEGER",
|
||||
"notes": "TEXT DEFAULT ''",
|
||||
"tags": "TEXT DEFAULT ''",
|
||||
"group_name": "TEXT DEFAULT ''",
|
||||
"created_at": "INTEGER",
|
||||
"updated_at": "INTEGER"
|
||||
})
|
||||
|
||||
# 11.1 Position alerts (price alerts and PnL alerts for manual positions)
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS qd_position_alerts (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL DEFAULT 1,
|
||||
position_id INTEGER, -- FK to qd_manual_positions (NULL = symbol-level alert)
|
||||
market TEXT DEFAULT '',
|
||||
symbol TEXT DEFAULT '',
|
||||
alert_type TEXT NOT NULL, -- price_above / price_below / pnl_above / pnl_below
|
||||
threshold REAL NOT NULL DEFAULT 0,
|
||||
notification_config TEXT DEFAULT '', -- JSON: channels, targets
|
||||
is_active INTEGER DEFAULT 1,
|
||||
is_triggered INTEGER DEFAULT 0,
|
||||
last_triggered_at INTEGER,
|
||||
trigger_count INTEGER DEFAULT 0,
|
||||
repeat_interval INTEGER DEFAULT 0, -- 0=once, >0=repeat every N seconds
|
||||
notes TEXT DEFAULT '',
|
||||
created_at INTEGER,
|
||||
updated_at INTEGER
|
||||
)
|
||||
""")
|
||||
|
||||
ensure_columns("qd_position_alerts", {
|
||||
"user_id": "INTEGER NOT NULL DEFAULT 1",
|
||||
"position_id": "INTEGER",
|
||||
"market": "TEXT DEFAULT ''",
|
||||
"symbol": "TEXT DEFAULT ''",
|
||||
"alert_type": "TEXT NOT NULL DEFAULT 'price_above'",
|
||||
"threshold": "REAL NOT NULL DEFAULT 0",
|
||||
"notification_config": "TEXT DEFAULT ''",
|
||||
"is_active": "INTEGER DEFAULT 1",
|
||||
"is_triggered": "INTEGER DEFAULT 0",
|
||||
"last_triggered_at": "INTEGER",
|
||||
"trigger_count": "INTEGER DEFAULT 0",
|
||||
"repeat_interval": "INTEGER DEFAULT 0",
|
||||
"notes": "TEXT DEFAULT ''",
|
||||
"created_at": "INTEGER",
|
||||
"updated_at": "INTEGER"
|
||||
})
|
||||
|
||||
# 12. Position monitors (AI analysis tasks for manual positions)
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS qd_position_monitors (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL DEFAULT 1,
|
||||
name TEXT DEFAULT '', -- Monitor name
|
||||
position_ids TEXT DEFAULT '', -- JSON array of position IDs (empty = all positions)
|
||||
monitor_type TEXT DEFAULT 'ai', -- ai / price_alert / pnl_alert
|
||||
config TEXT DEFAULT '', -- JSON config (interval_minutes, prompt, thresholds, etc.)
|
||||
notification_config TEXT DEFAULT '', -- JSON notification config (channels, targets)
|
||||
is_active INTEGER DEFAULT 1,
|
||||
last_run_at INTEGER,
|
||||
next_run_at INTEGER,
|
||||
last_result TEXT DEFAULT '', -- Last analysis result (JSON)
|
||||
run_count INTEGER DEFAULT 0,
|
||||
created_at INTEGER,
|
||||
updated_at INTEGER
|
||||
)
|
||||
""")
|
||||
|
||||
ensure_columns("qd_position_monitors", {
|
||||
"user_id": "INTEGER NOT NULL DEFAULT 1",
|
||||
"name": "TEXT DEFAULT ''",
|
||||
"position_ids": "TEXT DEFAULT ''",
|
||||
"monitor_type": "TEXT DEFAULT 'ai'",
|
||||
"config": "TEXT DEFAULT ''",
|
||||
"notification_config": "TEXT DEFAULT ''",
|
||||
"is_active": "INTEGER DEFAULT 1",
|
||||
"last_run_at": "INTEGER",
|
||||
"next_run_at": "INTEGER",
|
||||
"last_result": "TEXT DEFAULT ''",
|
||||
"run_count": "INTEGER DEFAULT 0",
|
||||
"created_at": "INTEGER",
|
||||
"updated_at": "INTEGER"
|
||||
})
|
||||
|
||||
conn.commit()
|
||||
logger.info("Database schema initialized (SQLite)")
|
||||
|
||||
# 初始化一次（按 db_file 维度）
|
||||
_has_initialized = False
|
||||
_initialized_db_file = None
|
||||
|
||||
class SQLiteCursor:
|
||||
"""模拟 pymysql DictCursor"""
|
||||
def __init__(self, cursor):
|
||||
self._cursor = cursor
|
||||
|
||||
def execute(self, query: str, args: Any = None):
|
||||
# 适配 MySQL -> SQLite 语法
|
||||
# 1. 替换占位符: %s -> ?
|
||||
query = query.replace('%s', '?')
|
||||
# 2. 替换 INSERT IGNORE -> INSERT OR IGNORE
|
||||
query = query.replace('INSERT IGNORE', 'INSERT OR IGNORE')
|
||||
# 3. 替换 ON DUPLICATE KEY UPDATE -> 简化为 UPSERT (SQLite 3.24+)
|
||||
# 注意：复杂的 ON DUPLICATE KEY UPDATE 很难自动转换，建议业务代码改写
|
||||
# 这里做一个简单的替换尝试，如果失败则需要人工介入代码
|
||||
if 'ON DUPLICATE KEY UPDATE' in query:
|
||||
# 简单的正则替换很难完美，这里记录日志提醒
|
||||
logger.warning(f"Complex SQL may require manual SQLite adaptation: {query}")
|
||||
# 尝试转换为 SQLite 的 ON CONFLICT (id) DO UPDATE SET ...
|
||||
# 但由于不知道主键冲突列，很难自动转换。
|
||||
# 临时方案：如果遇到这种 SQL，可能报错。我们假设主要业务逻辑已经重构。
|
||||
pass
|
||||
|
||||
if args:
|
||||
return self._cursor.execute(query, args)
|
||||
return self._cursor.execute(query)
|
||||
|
||||
def fetchone(self):
|
||||
row = self._cursor.fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
# Convert sqlite3.Row to dict
|
||||
return dict(row)
|
||||
|
||||
def fetchall(self):
|
||||
rows = self._cursor.fetchall()
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
def close(self):
|
||||
self._cursor.close()
|
||||
|
||||
@property
|
||||
def lastrowid(self):
|
||||
return self._cursor.lastrowid
|
||||
|
||||
class SQLiteConnection:
|
||||
"""数据库连接包装类"""
|
||||
def __init__(self, db_path):
|
||||
self._conn = sqlite3.connect(db_path, check_same_thread=False, timeout=30.0)
|
||||
# 设置 Row factory 以支持字段名访问
|
||||
self._conn.row_factory = sqlite3.Row
|
||||
|
||||
def cursor(self):
|
||||
return SQLiteCursor(self._conn.cursor())
|
||||
|
||||
def commit(self):
|
||||
self._conn.commit()
|
||||
|
||||
def rollback(self):
|
||||
self._conn.rollback()
|
||||
|
||||
def close(self):
|
||||
self._conn.close()
|
||||
|
||||
@contextmanager
|
||||
def get_db_connection():
|
||||
"""
|
||||
获取数据库连接 (Context Manager)
|
||||
"""
|
||||
global _has_initialized, _initialized_db_file
|
||||
|
||||
# 简单的连接创建，不使用连接池（SQLite 文件锁机制决定了连接池意义不大）
|
||||
# 使用线程锁防止写冲突（虽然 SQLite 有 WAL 模式，但稳妥起见）
|
||||
# 注意：这里加锁粒度较大，如果是高并发场景可能会慢，但对于个人量化系统足够。
|
||||
|
||||
# 初始化表结构（确保每个 db_file 都被初始化过）
|
||||
db_file = _get_db_file()
|
||||
if (not _has_initialized) or (_initialized_db_file != db_file):
|
||||
try:
|
||||
conn_init = sqlite3.connect(db_file)
|
||||
_init_db_schema(conn_init)
|
||||
conn_init.close()
|
||||
_has_initialized = True
|
||||
_initialized_db_file = db_file
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize database: {e}")
|
||||
|
||||
conn = SQLiteConnection(db_file)
|
||||
try:
|
||||
# with _db_lock: # SQLite 内部有锁，这里如果不跨线程共享连接其实不用强加锁
|
||||
yield conn
|
||||
except Exception as e:
|
||||
logger.error(f"Database operation error: {e}")
|
||||
conn.rollback()
|
||||
raise e
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def get_db_connection_sync():
|
||||
"""兼容旧接口"""
|
||||
global _has_initialized, _initialized_db_file
|
||||
db_file = _get_db_file()
|
||||
if (not _has_initialized) or (_initialized_db_file != db_file):
|
||||
try:
|
||||
conn_init = sqlite3.connect(db_file)
|
||||
_init_db_schema(conn_init)
|
||||
conn_init.close()
|
||||
_has_initialized = True
|
||||
_initialized_db_file = db_file
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize database: {e}")
|
||||
|
||||
return SQLiteConnection(db_file)
|
||||
|
||||
# Legacy alias
|
||||
def close_db_connection():
|
||||
"""Legacy alias for close_db"""
|
||||
pass
|
||||
|
||||
|
||||
__all__ = [
|
||||
'get_db_connection',
|
||||
'get_db_connection_sync',
|
||||
'close_db_connection',
|
||||
'init_database',
|
||||
'close_db',
|
||||
'get_db_type',
|
||||
'is_postgres',
|
||||
]
|
||||
|
||||
@@ -0,0 +1,297 @@
|
||||
"""
|
||||
PostgreSQL Database Connection Utility
|
||||
|
||||
Supports multi-user mode with connection pooling and SQLite compatibility layer.
|
||||
"""
|
||||
import os
|
||||
import threading
|
||||
from typing import Optional, Any, List, Dict
|
||||
from contextlib import contextmanager
|
||||
from app.utils.logger import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Try to import psycopg2
|
||||
try:
|
||||
import psycopg2
|
||||
from psycopg2 import pool
|
||||
from psycopg2.extras import RealDictCursor
|
||||
HAS_PSYCOPG2 = True
|
||||
except ImportError:
|
||||
HAS_PSYCOPG2 = False
|
||||
logger.warning("psycopg2 not installed. PostgreSQL support disabled.")
|
||||
|
||||
# Connection pool (global singleton)
|
||||
_connection_pool: Optional[Any] = None
|
||||
_pool_lock = threading.Lock()
|
||||
|
||||
|
||||
def _get_database_url() -> str:
|
||||
"""Get database connection URL from environment"""
|
||||
return os.getenv('DATABASE_URL', '').strip()
|
||||
|
||||
|
||||
def _parse_database_url(url: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Parse DATABASE_URL format: postgresql://user:password@host:port/dbname
|
||||
"""
|
||||
if not url:
|
||||
return {}
|
||||
|
||||
# Remove protocol prefix
|
||||
if url.startswith('postgresql://'):
|
||||
url = url[13:]
|
||||
elif url.startswith('postgres://'):
|
||||
url = url[11:]
|
||||
else:
|
||||
return {}
|
||||
|
||||
result = {}
|
||||
|
||||
# Split user:password@host:port/dbname
|
||||
if '@' in url:
|
||||
auth, hostpart = url.rsplit('@', 1)
|
||||
if ':' in auth:
|
||||
result['user'], result['password'] = auth.split(':', 1)
|
||||
else:
|
||||
result['user'] = auth
|
||||
else:
|
||||
hostpart = url
|
||||
|
||||
# Split host:port/dbname
|
||||
if '/' in hostpart:
|
||||
hostport, result['dbname'] = hostpart.split('/', 1)
|
||||
else:
|
||||
hostport = hostpart
|
||||
|
||||
if ':' in hostport:
|
||||
result['host'], port_str = hostport.split(':', 1)
|
||||
result['port'] = int(port_str)
|
||||
else:
|
||||
result['host'] = hostport
|
||||
result['port'] = 5432
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _get_connection_pool():
|
||||
"""Get or create connection pool"""
|
||||
global _connection_pool
|
||||
|
||||
if _connection_pool is not None:
|
||||
return _connection_pool
|
||||
|
||||
with _pool_lock:
|
||||
if _connection_pool is not None:
|
||||
return _connection_pool
|
||||
|
||||
if not HAS_PSYCOPG2:
|
||||
raise RuntimeError("psycopg2 is not installed. Cannot use PostgreSQL.")
|
||||
|
||||
db_url = _get_database_url()
|
||||
if not db_url:
|
||||
raise RuntimeError("DATABASE_URL environment variable is not set.")
|
||||
|
||||
params = _parse_database_url(db_url)
|
||||
if not params:
|
||||
raise RuntimeError(f"Invalid DATABASE_URL format: {db_url}")
|
||||
|
||||
try:
|
||||
_connection_pool = pool.ThreadedConnectionPool(
|
||||
minconn=2,
|
||||
maxconn=20,
|
||||
host=params.get('host', 'localhost'),
|
||||
port=params.get('port', 5432),
|
||||
user=params.get('user', 'quantdinger'),
|
||||
password=params.get('password', ''),
|
||||
dbname=params.get('dbname', 'quantdinger'),
|
||||
connect_timeout=10,
|
||||
)
|
||||
logger.info(f"PostgreSQL connection pool created: {params.get('host')}:{params.get('port')}/{params.get('dbname')}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create PostgreSQL connection pool: {e}")
|
||||
raise
|
||||
|
||||
return _connection_pool
|
||||
|
||||
|
||||
class PostgresCursor:
|
||||
"""PostgreSQL cursor wrapper with SQLite placeholder compatibility"""
|
||||
|
||||
def __init__(self, cursor):
|
||||
self._cursor = cursor
|
||||
self._last_insert_id = None
|
||||
|
||||
def _convert_placeholders(self, query: str) -> str:
|
||||
"""
|
||||
Convert SQLite-style ? placeholders to PostgreSQL %s
|
||||
Also handle some SQL syntax differences
|
||||
"""
|
||||
# Replace ? -> %s
|
||||
query = query.replace('?', '%s')
|
||||
|
||||
# SQLite: INSERT OR IGNORE -> PostgreSQL: INSERT ... ON CONFLICT DO NOTHING
|
||||
query = query.replace('INSERT OR IGNORE', 'INSERT')
|
||||
|
||||
return query
|
||||
|
||||
def execute(self, query: str, args: Any = None):
|
||||
"""Execute SQL statement"""
|
||||
query = self._convert_placeholders(query)
|
||||
|
||||
# Check if this is an INSERT and add RETURNING id if not present
|
||||
is_insert = query.strip().upper().startswith('INSERT')
|
||||
if is_insert and 'RETURNING' not in query.upper():
|
||||
query = query.rstrip(';').rstrip() + ' RETURNING id'
|
||||
|
||||
if args:
|
||||
if not isinstance(args, (tuple, list)):
|
||||
args = (args,)
|
||||
result = self._cursor.execute(query, args)
|
||||
else:
|
||||
result = self._cursor.execute(query)
|
||||
|
||||
# Capture last insert id for INSERT statements
|
||||
if is_insert:
|
||||
try:
|
||||
row = self._cursor.fetchone()
|
||||
if row and 'id' in row:
|
||||
self._last_insert_id = row['id']
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return result
|
||||
|
||||
def fetchone(self) -> Optional[Dict[str, Any]]:
|
||||
"""Fetch single row"""
|
||||
row = self._cursor.fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
return dict(row) if row else None
|
||||
|
||||
def fetchall(self) -> List[Dict[str, Any]]:
|
||||
"""Fetch all rows"""
|
||||
rows = self._cursor.fetchall()
|
||||
return [dict(row) for row in rows] if rows else []
|
||||
|
||||
def close(self):
|
||||
"""Close cursor"""
|
||||
self._cursor.close()
|
||||
|
||||
@property
|
||||
def lastrowid(self) -> Optional[int]:
|
||||
"""Get last inserted row ID"""
|
||||
return self._last_insert_id
|
||||
|
||||
@property
|
||||
def rowcount(self) -> int:
|
||||
"""Get affected row count"""
|
||||
return self._cursor.rowcount
|
||||
|
||||
|
||||
class PostgresConnection:
|
||||
"""PostgreSQL connection wrapper"""
|
||||
|
||||
def __init__(self, conn):
|
||||
self._conn = conn
|
||||
self._pool = _get_connection_pool()
|
||||
|
||||
def cursor(self) -> PostgresCursor:
|
||||
"""Create cursor"""
|
||||
return PostgresCursor(self._conn.cursor(cursor_factory=RealDictCursor))
|
||||
|
||||
def commit(self):
|
||||
"""Commit transaction"""
|
||||
self._conn.commit()
|
||||
|
||||
def rollback(self):
|
||||
"""Rollback transaction"""
|
||||
self._conn.rollback()
|
||||
|
||||
def close(self):
|
||||
"""Return connection to pool"""
|
||||
if self._pool and self._conn:
|
||||
try:
|
||||
self._pool.putconn(self._conn)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to return connection to pool: {e}")
|
||||
|
||||
|
||||
@contextmanager
|
||||
def get_pg_connection():
|
||||
"""
|
||||
Get PostgreSQL database connection (Context Manager)
|
||||
"""
|
||||
pool = _get_connection_pool()
|
||||
conn = None
|
||||
try:
|
||||
conn = pool.getconn()
|
||||
pg_conn = PostgresConnection(conn)
|
||||
yield pg_conn
|
||||
except Exception as e:
|
||||
if conn:
|
||||
try:
|
||||
conn.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
logger.error(f"PostgreSQL operation error: {e}")
|
||||
raise
|
||||
finally:
|
||||
if conn:
|
||||
try:
|
||||
pool.putconn(conn)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def get_pg_connection_sync() -> PostgresConnection:
|
||||
"""
|
||||
Get connection synchronously (caller must close)
|
||||
"""
|
||||
pool = _get_connection_pool()
|
||||
conn = pool.getconn()
|
||||
return PostgresConnection(conn)
|
||||
|
||||
|
||||
def execute_sql(sql: str, params: tuple = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Execute SQL and return results (convenience function)
|
||||
"""
|
||||
with get_pg_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(sql, params)
|
||||
if sql.strip().upper().startswith('SELECT'):
|
||||
return cursor.fetchall()
|
||||
conn.commit()
|
||||
return []
|
||||
|
||||
|
||||
def is_postgres_available() -> bool:
|
||||
"""Check if PostgreSQL is available"""
|
||||
if not HAS_PSYCOPG2:
|
||||
return False
|
||||
|
||||
db_url = _get_database_url()
|
||||
if not db_url:
|
||||
return False
|
||||
|
||||
try:
|
||||
with get_pg_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT 1")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.debug(f"PostgreSQL not available: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def close_pool():
|
||||
"""Close connection pool (call on app shutdown)"""
|
||||
global _connection_pool
|
||||
if _connection_pool:
|
||||
try:
|
||||
_connection_pool.closeall()
|
||||
_connection_pool = None
|
||||
logger.info("PostgreSQL connection pool closed")
|
||||
except Exception as e:
|
||||
logger.warning(f"Error closing connection pool: {e}")
|
||||
Reference in New Issue
Block a user