94 lines
5.1 KiB
Python
94 lines
5.1 KiB
Python
"""Read-only SQLAlchemy connectors with server-side credential references."""
|
|
import os
|
|
from pathlib import Path
|
|
from datetime import date, datetime
|
|
from decimal import Decimal
|
|
import sqlite3
|
|
|
|
|
|
def normalize(value, exact_decimals=False):
|
|
if isinstance(value, (date, datetime)): return value.isoformat()
|
|
if isinstance(value, Decimal): return str(value) if exact_decimals else float(value)
|
|
if isinstance(value, bytes): raise ValueError("查询包含二进制数据，请返回附件索引而非文件内容")
|
|
return value
|
|
|
|
|
|
def validate_query(sql, allowed_tables, dialect):
|
|
import sqlglot
|
|
from sqlglot import exp
|
|
try:
|
|
statements = sqlglot.parse(sql, read=dialect)
|
|
except Exception:
|
|
raise ValueError("SQL 无法解析，请使用单条只读 SELECT")
|
|
if len(statements) != 1 or not isinstance(statements[0], (exp.Select, exp.Union)):
|
|
raise ValueError("只允许单条 SELECT 查询")
|
|
tree = statements[0]
|
|
if any(tree.find_all(exp.Insert, exp.Update, exp.Delete, exp.Create, exp.Drop, exp.Command, exp.Into, exp.Lock)):
|
|
raise ValueError("不允许写入、锁定或导出语句")
|
|
# User-defined functions may have side effects even when embedded in SELECT.
|
|
if any(tree.find_all(exp.Anonymous)):
|
|
raise ValueError("不允许自定义函数，请使用标准聚合和表达式")
|
|
aliases = {c.alias for c in tree.find_all(exp.CTE)}
|
|
tables = {".".join(part.name for part in t.parts).lower() for t in tree.find_all(exp.Table) if t.name not in aliases}
|
|
if not allowed_tables or not tables.issubset({t.lower() for t in allowed_tables}):
|
|
raise ValueError("查询使用了未授权的库表或视图")
|
|
|
|
|
|
def execute(connection, sql, project_id, exact_decimals=False):
|
|
from sqlalchemy import create_engine, text
|
|
from sqlalchemy.engine import URL
|
|
kind = connection.get("type")
|
|
dialects = {"PostgreSQL": ("postgresql+psycopg", "postgres"), "MySQL": ("mysql+pymysql", "mysql"), "SQLite": ("sqlite", "sqlite")}
|
|
if kind not in dialects: raise ValueError("该数据库连接器暂不支持")
|
|
driver, dialect = dialects[kind]
|
|
validate_query(sql, connection.get("allowed_tables", []), dialect)
|
|
seconds = max(1, min(int(connection.get("timeout", 15)), 60))
|
|
max_rows = max(1, min(int(connection.get("max_rows", 1000)), 10000))
|
|
if kind == "SQLite":
|
|
root = os.getenv("XINGBANG_SQLITE_ROOT")
|
|
if not root: raise ValueError("服务器未配置 XINGBANG_SQLITE_ROOT")
|
|
file = Path(connection.get("database", "")).resolve()
|
|
if not file.is_relative_to(Path(root).resolve()) or not file.is_file():
|
|
raise ValueError("SQLite 文件不存在或不在服务器允许目录内")
|
|
def creator():
|
|
db = sqlite3.connect(file.as_uri()+"?mode=ro", uri=True)
|
|
def authorize(action, arg1, arg2, *_):
|
|
if action in {sqlite3.SQLITE_SELECT, sqlite3.SQLITE_READ, sqlite3.SQLITE_FUNCTION}:
|
|
return sqlite3.SQLITE_OK
|
|
# SQLAlchemy reads the isolation mode during connection initialization.
|
|
if action == sqlite3.SQLITE_PRAGMA and arg1 == "read_uncommitted" and arg2 is None:
|
|
return sqlite3.SQLITE_OK
|
|
return sqlite3.SQLITE_DENY
|
|
db.set_authorizer(authorize)
|
|
import time
|
|
deadline=time.monotonic()+seconds
|
|
db.set_progress_handler(lambda: int(time.monotonic()>deadline),1000)
|
|
return db
|
|
engine = create_engine("sqlite://", creator=creator)
|
|
else:
|
|
ref = connection.get("password_env", "")
|
|
if not ref.startswith("XINGBANG_DB_") or not os.getenv(ref):
|
|
raise ValueError("请配置 XINGBANG_DB_ 前缀的服务器密码引用")
|
|
url = URL.create(driver, username=connection.get("username"),password=os.environ[ref],host=connection.get("host"),port=int(connection.get("port") or (5432 if kind=="PostgreSQL" else 3306)),database=connection.get("database"))
|
|
args = {"connect_timeout": seconds}
|
|
if kind == "PostgreSQL": args["sslmode"]="require"
|
|
else: args.update(read_timeout=seconds,write_timeout=seconds)
|
|
engine = create_engine(url,connect_args=args,pool_pre_ping=True)
|
|
try:
|
|
with engine.connect() as db:
|
|
if kind == "PostgreSQL":
|
|
db.exec_driver_sql("SET TRANSACTION READ ONLY")
|
|
db.exec_driver_sql(f"SET LOCAL statement_timeout = {seconds*1000}")
|
|
elif kind == "MySQL":
|
|
db.exec_driver_sql("SET SESSION TRANSACTION READ ONLY")
|
|
db.exec_driver_sql(f"SET SESSION MAX_EXECUTION_TIME = {seconds*1000}")
|
|
rows = db.execute(text(sql), {"project_id": project_id}).mappings().fetchmany(max_rows+1)
|
|
if len(rows)>max_rows: raise ValueError(f"查询结果超过 {max_rows} 行，请缩小范围")
|
|
return [{k:normalize(v, exact_decimals) for k,v in row.items()} for row in rows]
|
|
except ValueError:
|
|
raise
|
|
except Exception:
|
|
raise ValueError("数据库连接或查询失败，请检查驱动、只读账号、网络与字段配置") from None
|
|
finally:
|
|
engine.dispose()
|