Files
opsxsql/backend/app/services/db_connector.py
T
logan 2cfd45c9d1 feat: run.sh restart/stop/status + real DB connector + GitLab CI
1. run.sh: start/stop/restart/status/logs/build/test
   - PID file management in .pids/
   - Colored output, health check wait, log files
   - Auto alembic migrate on start

2. DB Connector (replaces mock aiomysql):
   - MySQL: aiomysql (connect/query/dml/test)
   - PostgreSQL: asyncpg (connect/query/dml/test)
   - SQLite: sqlite3 via executor (connect/query/dml/test)
   - Unified interface: DBConnector.test_connection/execute_query/execute_dml
   - Connection test returns version + latency

3. GitLab CI/CD pipeline:
   - test: backend pytest + frontend build
   - build: Docker images (manual trigger)
   - deploy: SSH deploy (manual trigger)
   - Cache: pip + node_modules

Tests: 128 passing, frontend build: OK
2026-05-22 15:16:59 +08:00

289 lines
11 KiB
Python

"""Database connector — unified interface for MySQL/PostgreSQL/SQLite.
Provides real database connections for query execution and connection testing.
Replaces mock aiomysql calls with actual driver connections.
"""
import asyncio
import time
import sqlite3
from typing import Optional, List, Dict, Any, Tuple
from app.models.instance import Instance
from app.services.sql_analyzer import SQLAnalyzer
MAX_ROWS = 10000
QUERY_TIMEOUT = 30
DML_TIMEOUT = 60
class ConnectionResult:
"""Result of a connection test."""
def __init__(self, success: bool, message: str, version: str = None, latency_ms: int = 0):
self.success = success
self.message = message
self.version = version
self.latency_ms = latency_ms
def to_dict(self):
return {
"success": self.success,
"message": self.message,
"version": self.version,
"latency_ms": self.latency_ms,
}
class QueryResult:
"""Result of a SQL query execution."""
def __init__(self, columns: List[str], rows: List[Dict[str, Any]],
total: int, elapsed_ms: int, truncated: bool = False):
self.columns = columns
self.rows = rows
self.total = total
self.elapsed_ms = elapsed_ms
self.truncated = truncated
class DBConnector:
"""Unified database connector supporting MySQL, PostgreSQL, SQLite."""
_analyzer = SQLAnalyzer()
@staticmethod
async def test_connection(instance: Instance) -> ConnectionResult:
"""Test database connectivity and return version info."""
start = time.time()
try:
if instance.db_type == "mysql":
return await DBConnector._test_mysql(instance, start)
elif instance.db_type in ("postgres", "postgresql"):
return await DBConnector._test_postgres(instance, start)
elif instance.db_type == "sqlite":
return await DBConnector._test_sqlite(instance, start)
else:
return ConnectionResult(False, f"暂不支持 {instance.db_type} 类型数据源")
except Exception as e:
elapsed = int((time.time() - start) * 1000)
return ConnectionResult(False, f"连接失败: {str(e)}", latency_ms=elapsed)
@staticmethod
async def execute_query(
instance: Instance, database: Optional[str], sql: str, user_id: int
) -> 'QueryResult':
"""Execute a read-only SQL query on the target database."""
# Validate SQL is read-only
is_safe, error_msg = DBConnector._analyzer.validate_readonly(sql)
if not is_safe:
raise ValueError(error_msg)
# Auto-add LIMIT
upper_sql = sql.strip().upper()
import re
if not re.search(r"LIMIT\s+\d+", upper_sql, re.IGNORECASE):
sql = sql.rstrip(";") + f" LIMIT {MAX_ROWS};"
start = time.time()
if instance.db_type == "mysql":
rows, columns = await DBConnector._query_mysql(instance, database, sql)
elif instance.db_type in ("postgres", "postgresql"):
rows, columns = await DBConnector._query_postgres(instance, database, sql)
elif instance.db_type == "sqlite":
rows, columns = await DBConnector._query_sqlite(instance, sql)
else:
raise ValueError(f"暂不支持 {instance.db_type} 查询")
elapsed = int((time.time() - start) * 1000)
return QueryResult(
columns=columns,
rows=rows,
total=len(rows),
elapsed_ms=elapsed,
truncated=len(rows) >= MAX_ROWS,
)
@staticmethod
async def execute_dml(instance: Instance, database: str, sql: str) -> str:
"""Execute a DML statement (INSERT/UPDATE/DELETE)."""
if instance.db_type == "mysql":
return await DBConnector._dml_mysql(instance, database, sql)
elif instance.db_type in ("postgres", "postgresql"):
return await DBConnector._dml_postgres(instance, database, sql)
elif instance.db_type == "sqlite":
return await DBConnector._dml_sqlite(instance, sql)
else:
raise ValueError(f"暂不支持 {instance.db_type} 执行")
# ===== MySQL =====
@staticmethod
async def _test_mysql(instance: Instance, start: float) -> ConnectionResult:
import aiomysql
conn = await asyncio.wait_for(
aiomysql.connect(
host=instance.host, port=instance.port,
user=instance.username, password=instance.password,
db=instance.database,
),
timeout=10,
)
try:
async with conn.cursor() as cur:
await cur.execute("SELECT VERSION()")
version = (await cur.fetchone())[0]
elapsed = int((time.time() - start) * 1000)
return ConnectionResult(True, "连接成功", version=version, latency_ms=elapsed)
finally:
conn.close()
@staticmethod
async def _query_mysql(instance: Instance, database: Optional[str], sql: str) -> Tuple[List[Dict], List[str]]:
import aiomysql
conn = await aiomysql.connect(
host=instance.host, port=instance.port,
user=instance.username, password=instance.password,
db=database or instance.database,
)
try:
async with conn.cursor(aiomysql.DictCursor) as cur:
await asyncio.wait_for(cur.execute(sql), timeout=QUERY_TIMEOUT)
rows = await cur.fetchall()
columns = [desc[0] for desc in cur.description] if cur.description else []
return rows, columns
except asyncio.TimeoutError:
raise ValueError("查询超时(30秒)")
finally:
conn.close()
@staticmethod
async def _dml_mysql(instance: Instance, database: str, sql: str) -> str:
import aiomysql
conn = await aiomysql.connect(
host=instance.host, port=instance.port,
user=instance.username, password=instance.password,
db=database, autocommit=True,
)
try:
async with conn.cursor() as cur:
await asyncio.wait_for(cur.execute(sql), timeout=DML_TIMEOUT)
affected = cur.rowcount
return f"执行成功,影响 {affected} 行"
except asyncio.TimeoutError:
raise ValueError("执行超时(60秒)")
finally:
conn.close()
# ===== PostgreSQL =====
@staticmethod
async def _test_postgres(instance: Instance, start: float) -> ConnectionResult:
import asyncpg
dsn = f"postgresql://{instance.username}:{instance.password}@{instance.host}:{instance.port}/{instance.database or 'postgres'}"
conn = await asyncio.wait_for(asyncpg.connect(dsn), timeout=10)
try:
version = await conn.fetchval("SELECT version()")
elapsed = int((time.time() - start) * 1000)
# version string is long, take first part
ver_short = version.split(",")[0] if version else "unknown"
return ConnectionResult(True, "连接成功", version=ver_short, latency_ms=elapsed)
finally:
await conn.close()
@staticmethod
async def _query_postgres(instance: Instance, database: Optional[str], sql: str) -> Tuple[List[Dict], List[str]]:
import asyncpg
dsn = f"postgresql://{instance.username}:{instance.password}@{instance.host}:{instance.port}/{database or instance.database or 'postgres'}"
conn = await asyncpg.connect(dsn)
try:
# asyncpg returns Records, convert to dicts
stmt = await asyncio.wait_for(conn.fetch(sql), timeout=QUERY_TIMEOUT)
rows = [dict(r) for r in stmt]
columns = list(stmt[0].keys()) if stmt else []
return rows, columns
except asyncio.TimeoutError:
raise ValueError("查询超时(30秒)")
finally:
await conn.close()
@staticmethod
async def _dml_postgres(instance: Instance, database: str, sql: str) -> str:
import asyncpg
dsn = f"postgresql://{instance.username}:{instance.password}@{instance.host}:{instance.port}/{database}"
conn = await asyncpg.connect(dsn)
try:
status = await asyncio.wait_for(conn.execute(sql), timeout=DML_TIMEOUT)
# status like "UPDATE 3"
parts = status.split()
affected = int(parts[1]) if len(parts) >= 2 else 0
return f"执行成功,影响 {affected} 行"
except asyncio.TimeoutError:
raise ValueError("执行超时(60秒)")
finally:
await conn.close()
# ===== SQLite (remote instance) =====
@staticmethod
async def _test_sqlite(instance: Instance, start: float) -> ConnectionResult:
# SQLite uses host as file path
db_path = instance.host
loop = asyncio.get_event_loop()
try:
def _connect():
conn = sqlite3.connect(db_path)
cur = conn.execute("SELECT sqlite_version()")
ver = cur.fetchone()[0]
conn.close()
return ver
version = await loop.run_in_executor(None, _connect)
elapsed = int((time.time() - start) * 1000)
return ConnectionResult(True, "连接成功", version=f"SQLite {version}", latency_ms=elapsed)
except Exception as e:
elapsed = int((time.time() - start) * 1000)
return ConnectionResult(False, f"连接失败: {e}", latency_ms=elapsed)
@staticmethod
async def _query_sqlite(instance: Instance, sql: str) -> Tuple[List[Dict], List[str]]:
db_path = instance.host
loop = asyncio.get_event_loop()
def _query():
conn = sqlite3.connect(db_path)
conn.row_factory = sqlite3.Row
cur = conn.execute(sql)
rows = [dict(r) for r in cur.fetchall()]
columns = [desc[0] for desc in cur.description] if cur.description else []
conn.close()
return rows, columns
try:
return await asyncio.wait_for(
loop.run_in_executor(None, _query),
timeout=QUERY_TIMEOUT,
)
except asyncio.TimeoutError:
raise ValueError("查询超时(30秒)")
@staticmethod
async def _dml_sqlite(instance: Instance, sql: str) -> str:
db_path = instance.host
loop = asyncio.get_event_loop()
def _execute():
conn = sqlite3.connect(db_path)
cur = conn.execute(sql)
affected = cur.rowcount
conn.commit()
conn.close()
return affected
try:
affected = await asyncio.wait_for(
loop.run_in_executor(None, _execute),
timeout=DML_TIMEOUT,
)
return f"执行成功,影响 {affected} 行"
except asyncio.TimeoutError:
raise ValueError("执行超时(60秒)")