"""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秒）")