sql-assistant 1.0.0__py3-none-any.whl

sql_assistant/database/connectors/base.py

@@ -0,0 +1,185 @@
+"""数据库连接器抽象基类"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any
+
+from .exceptions import format_connector_result
+
+
+@dataclass
+class QueryResult:
+    """统一查询结果"""
+    columns: list[str] = field(default_factory=list)
+    rows: list[list[Any]] = field(default_factory=list)
+    row_count: int = 0
+    affected_rows: int = 0
+    sql_type: str = ""  # SELECT, INSERT, UPDATE, DELETE, OTHER
+
+
+class BaseConnector(ABC):
+    """数据库连接器抽象基类"""
+
+    db_type: str = "unknown"
+
+    def __init__(self, host: str, port: int, user: str, password: str, database: str):
+        self.host = host
+        self.port = port
+        self.user = user
+        self.password = password
+        self.database = database
+
+    @abstractmethod
+    async def connect(self) -> None:
+        """建立数据库连接"""
+        ...
+
+    @abstractmethod
+    async def disconnect(self) -> None:
+        """断开数据库连接"""
+        ...
+
+    @abstractmethod
+    async def execute(self, sql: str) -> QueryResult:
+        """执行 SQL 语句，返回统一结果"""
+        ...
+
+    @abstractmethod
+    async def test_connection(self) -> bool:
+        """测试连接是否可用"""
+        ...
+
+    @abstractmethod
+    async def get_schema(self) -> dict:
+        """获取数据库 schema（表名 + 列信息）。
+        
+        返回格式:
+        {
+            "db_type": "mysql",
+            "tables": [
+                {"name": "users", "columns": [
+                    {"name": "id", "type": "INT", "nullable": False, "key": "PRI"},
+                    {"name": "name", "type": "VARCHAR(100)", "nullable": True, "key": ""},
+                ]},
+            ]
+        }
+        """
+        ...
+
+    @staticmethod
+    def classify_sql(sql: str) -> str:
+        """判断 SQL 语句类型（忽略前导注释和空行）
+        
+        支持多语句 SQL，检查所有子语句的类型。如果任何子语句是增删改或 DDL，
+        则返回该类型，确保危险操作能触发确认对话框。
+        """
+        s = BaseConnector._strip_sql_comments(sql).strip()
+        
+        # 按分号分割 SQL 语句（处理字符串内的分号）
+        statements = BaseConnector._split_sql_statements(s)
+        
+        for stmt in statements:
+            stmt_upper = stmt.strip().upper()
+            if not stmt_upper:
+                continue
+            
+            # 按优先级判断：DDL > DELETE > UPDATE > INSERT > SELECT
+            if stmt_upper.startswith("CREATE") or stmt_upper.startswith("ALTER") or stmt_upper.startswith("DROP"):
+                return "DDL"
+            elif stmt_upper.startswith("DELETE"):
+                return "DELETE"
+            elif stmt_upper.startswith("UPDATE"):
+                return "UPDATE"
+            elif stmt_upper.startswith("INSERT"):
+                return "INSERT"
+        
+        # 如果没有找到增删改或 DDL，检查是否是查询语句
+        for stmt in statements:
+            stmt_upper = stmt.strip().upper()
+            if stmt_upper.startswith("SELECT") or stmt_upper.startswith("SHOW") or stmt_upper.startswith("DESCRIBE") or stmt_upper.startswith("EXPLAIN"):
+                return "SELECT"
+        
+        return "OTHER"
+    
+    @staticmethod
+    def _split_sql_statements(sql: str) -> list:
+        """安全地按分号分割 SQL 语句，处理字符串内的分号"""
+        statements = []
+        current_stmt = []
+        in_single_quote = False
+        in_double_quote = False
+        in_comment = False
+        
+        for i, char in enumerate(sql):
+            # 处理行注释
+            if i > 0 and sql[i-1] == '-' and char == '-':
+                # 跳过到行尾
+                while i < len(sql) and sql[i] != '\n':
+                    i += 1
+                continue
+            
+            # 处理单引号字符串
+            if char == "'" and not in_double_quote and not in_comment:
+                in_single_quote = not in_single_quote
+            
+            # 处理双引号字符串
+            if char == '"' and not in_single_quote and not in_comment:
+                in_double_quote = not in_double_quote
+            
+            # 处理分号（仅在字符串外）
+            if char == ';' and not in_single_quote and not in_double_quote:
+                statements.append(''.join(current_stmt))
+                current_stmt = []
+                continue
+            
+            current_stmt.append(char)
+        
+        # 添加最后一个语句（如果有）
+        if current_stmt:
+            statements.append(''.join(current_stmt))
+        
+        return statements
+
+    @staticmethod
+    def _strip_sql_comments(sql: str) -> str:
+        """去除 SQL 前导注释行和空行，保留实际语句部分。
+        
+        处理两种注释风格：
+        - 行注释：-- ... 或 # ...
+        - 块注释：/* ... */
+        """
+        lines = sql.split("\n")
+        result_lines = []
+        in_block_comment = False
+
+        for line in lines:
+            stripped = line.strip()
+
+            # 处理块注释（可能跨行）
+            if in_block_comment:
+                end_idx = stripped.find("*/")
+                if end_idx != -1:
+                    in_block_comment = False
+                    remaining = stripped[end_idx + 2:].strip()
+                    if remaining:
+                        result_lines.append(remaining)
+                continue
+
+            # 检查行首是否是块注释开始
+            if stripped.startswith("/*"):
+                end_idx = stripped.find("*/", 2)
+                if end_idx != -1:
+                    remaining = stripped[end_idx + 2:].strip()
+                    if remaining:
+                        result_lines.append(remaining)
+                else:
+                    in_block_comment = True
+                continue
+
+            # 跳过行注释和空行
+            if not stripped or stripped.startswith("--") or stripped.startswith("#"):
+                continue
+
+            result_lines.append(line)
+
+        return "\n".join(result_lines)

sql_assistant/database/connectors/exceptions.py

@@ -0,0 +1,88 @@
+"""数据库连接器统一异常"""
+
+from typing import Optional
+
+
+class ConnectorError(Exception):
+    """数据库连接器基础异常"""
+    def __init__(self, message: str, db_type: str = "", code: Optional[str] = None):
+        super().__init__(message)
+        self.db_type = db_type
+        self.code = code
+
+    def to_dict(self) -> dict:
+        return {
+            "success": False,
+            "error": str(self),
+            "db_type": self.db_type,
+            "code": self.code,
+        }
+
+
+class ConnectionError(ConnectorError):
+    """连接失败异常"""
+    def __init__(self, message: str, db_type: str = ""):
+        super().__init__(message, db_type, code="CONNECTION_FAILED")
+
+
+class QueryError(ConnectorError):
+    """查询执行异常"""
+    def __init__(self, message: str, db_type: str = "", sql: str = ""):
+        super().__init__(message, db_type, code="QUERY_FAILED")
+        self.sql = sql
+
+    def to_dict(self) -> dict:
+        result = super().to_dict()
+        result["sql"] = self.sql
+        return result
+
+
+class SchemaError(ConnectorError):
+    """Schema 获取异常"""
+    def __init__(self, message: str, db_type: str = ""):
+        super().__init__(message, db_type, code="SCHEMA_ERROR")
+
+
+class AuthenticationError(ConnectorError):
+    """认证失败异常"""
+    def __init__(self, message: str, db_type: str = ""):
+        super().__init__(message, db_type, code="AUTH_FAILED")
+
+
+class TimeoutError(ConnectorError):
+    """连接超时异常"""
+    def __init__(self, message: str, db_type: str = ""):
+        super().__init__(message, db_type, code="TIMEOUT")
+
+
+def format_connector_result(
+    success: bool,
+    data: any = None,
+    error: Optional[str] = None,
+    db_type: str = "",
+    code: Optional[str] = None,
+) -> dict:
+    """统一格式化连接器返回结果
+
+    Args:
+        success: 是否成功
+        data: 成功时返回的数据
+        error: 失败时的错误信息
+        db_type: 数据库类型
+        code: 错误码
+
+    Returns:
+        统一格式的字典
+    """
+    result = {"success": success}
+
+    if success:
+        result["data"] = data
+    else:
+        result["error"] = error or "未知错误"
+        if code:
+            result["code"] = code
+        if db_type:
+            result["db_type"] = db_type
+
+    return result

sql_assistant/database/connectors/mongodb.py

@@ -0,0 +1,194 @@
+"""MongoDB 连接器"""
+
+import json
+import asyncio
+from typing import Any
+
+from pymongo import MongoClient
+
+from .base import BaseConnector, QueryResult
+
+
+class MongoDBConnector(BaseConnector):
+    """MongoDB 数据库连接器"""
+
+    db_type = "mongodb"
+
+    def __init__(self, host: str, port: int, user: str, password: str, database: str):
+        super().__init__(host, port, user, password, database)
+        self._client: MongoClient | None = None
+        self._db = None
+
+    async def connect(self) -> None:
+        loop = asyncio.get_event_loop()
+        connection_string = f"mongodb://{self.user}:{self.password}@{self.host}:{self.port}/" if self.user else f"mongodb://{self.host}:{self.port}/"
+        self._client = await loop.run_in_executor(
+            None,
+            lambda: MongoClient(connection_string, serverSelectionTimeoutMS=10000),
+        )
+        self._db = self._client[self.database]
+
+    async def disconnect(self) -> None:
+        if self._client:
+            self._client.close()
+            self._client = None
+            self._db = None
+
+    async def get_schema(self) -> dict:
+        """MongoDB 无固定 schema，采样分析每个 collection 的字段"""
+        if not self._client or self._db is None:
+            raise RuntimeError("MongoDB 未连接")
+
+        loop = asyncio.get_event_loop()
+
+        def _run():
+            collections = self._db.list_collection_names()
+            tables = []
+            for coll_name in collections:
+                # 采样最多 5 条文档推断字段
+                docs = list(self._db[coll_name].find().limit(5))
+                columns_set: dict[str, set] = {}
+                for doc in docs:
+                    for key in doc.keys():
+                        if key not in columns_set:
+                            columns_set[key] = set()
+                        val = doc[key]
+                        columns_set[key].add(type(val).__name__ if val is not None else "null")
+
+                columns = []
+                for col_name, types in columns_set.items():
+                    columns.append({
+                        "name": col_name,
+                        "type": ", ".join(sorted(types)),
+                        "nullable": True,
+                        "key": "_id" if col_name == "_id" else "",
+                        "default": None,
+                        "comment": "",
+                    })
+                tables.append({"name": coll_name, "columns": columns})
+
+            return {"db_type": "mongodb", "tables": tables}
+
+        return await loop.run_in_executor(None, _run)
+
+    async def execute(self, sql: str) -> QueryResult:
+        """执行 MongoDB 查询（JSON 格式）"""
+        if not self._client or self._db is None:
+            raise RuntimeError("MongoDB 未连接")
+
+        loop = asyncio.get_event_loop()
+        result = QueryResult()
+
+        def _run():
+            try:
+                # 解析 JSON 查询
+                query = json.loads(sql)
+                collection_name = query.get("collection", "")
+                operation = query.get("operation", "find").lower()
+                collection = self._db[collection_name]
+
+                if operation == "find":
+                    filter_obj = query.get("filter", {})
+                    projection = query.get("projection", None)
+                    sort_list = query.get("sort", None)
+                    limit_val = query.get("limit", 100)
+
+                    cursor = collection.find(filter_obj, projection)
+                    if sort_list:
+                        cursor = cursor.sort(sort_list)
+                    if limit_val:
+                        cursor = cursor.limit(limit_val)
+
+                    docs = list(cursor)
+                    # 提取所有字段名
+                    columns_set = set()
+                    for doc in docs:
+                        columns_set.update(doc.keys())
+                    result.columns = sorted(columns_set)
+                    result.rows = []
+                    for doc in docs:
+                        result.rows.append([doc.get(col, None) for col in result.columns])
+                    result.row_count = len(docs)
+                    result.sql_type = "SELECT"
+
+                elif operation == "insert":
+                    documents = query.get("documents", [])
+                    if isinstance(documents, dict):
+                        documents = [documents]
+                    result_obj = collection.insert_many(documents)
+                    result.affected_rows = len(result_obj.inserted_ids)
+                    result.sql_type = "INSERT"
+
+                elif operation == "update":
+                    filter_obj = query.get("filter", {})
+                    update_obj = query.get("update", {})
+                    many = query.get("many", False)
+                    if many:
+                        update_result = collection.update_many(filter_obj, update_obj)
+                    else:
+                        update_result = collection.update_one(filter_obj, update_obj)
+                    result.affected_rows = update_result.modified_count
+                    result.sql_type = "UPDATE"
+
+                elif operation == "delete":
+                    filter_obj = query.get("filter", {})
+                    many = query.get("many", False)
+                    if many:
+                        delete_result = collection.delete_many(filter_obj)
+                    else:
+                        delete_result = collection.delete_one(filter_obj)
+                    result.affected_rows = delete_result.deleted_count
+                    result.sql_type = "DELETE"
+
+                elif operation == "aggregate":
+                    pipeline = query.get("pipeline", [])
+                    docs = list(collection.aggregate(pipeline))
+                    columns_set = set()
+                    for doc in docs:
+                        columns_set.update(doc.keys())
+                    result.columns = sorted(columns_set)
+                    result.rows = [[doc.get(col, None) for col in result.columns] for doc in docs]
+                    result.row_count = len(docs)
+                    result.sql_type = "SELECT"
+
+                elif operation == "count":
+                    filter_obj = query.get("filter", {})
+                    count = collection.count_documents(filter_obj)
+                    result.columns = ["count"]
+                    result.rows = [[count]]
+                    result.row_count = 1
+                    result.sql_type = "SELECT"
+
+                else:
+                    result.columns = ["error"]
+                    result.rows = [[f"不支持的操作: {operation}"]]
+                    result.row_count = 1
+
+            except json.JSONDecodeError as e:
+                result.columns = ["error"]
+                result.rows = [[f"JSON 解析错误: {e}"]]
+                result.row_count = 1
+            except Exception as e:
+                result.columns = ["error"]
+                result.rows = [[str(e)]]
+                result.row_count = 1
+
+            return result
+
+        return await loop.run_in_executor(None, _run)
+
+    async def test_connection(self) -> dict:
+        from .exceptions import format_connector_result
+        try:
+            await self.connect()
+            if self._client:
+                loop = asyncio.get_event_loop()
+                await loop.run_in_executor(
+                    None,
+                    lambda: self._client.admin.command("ping"),
+                )
+            return format_connector_result(True, data={"message": "MongoDB 连接成功"}, db_type="mongodb")
+        except Exception as e:
+            return format_connector_result(False, error=str(e), db_type="mongodb", code="CONNECTION_FAILED")
+        finally:
+            await self.disconnect()

sql_assistant/database/connectors/mysql.py

@@ -0,0 +1,110 @@
+"""MySQL 连接器"""
+
+import asyncio
+from typing import Any
+
+import pymysql
+
+from .base import BaseConnector, QueryResult
+
+
+class MySQLConnector(BaseConnector):
+    """MySQL 数据库连接器"""
+
+    db_type = "mysql"
+
+    def __init__(self, host: str, port: int, user: str, password: str, database: str):
+        super().__init__(host, port, user, password, database)
+        self._conn: pymysql.Connection | None = None
+
+    async def connect(self) -> None:
+        loop = asyncio.get_event_loop()
+        self._conn = await loop.run_in_executor(
+            None,
+            lambda: pymysql.connect(
+                host=self.host,
+                port=self.port,
+                user=self.user,
+                password=self.password,
+                database=self.database,
+                charset="utf8mb4",
+                cursorclass=pymysql.cursors.Cursor,
+                connect_timeout=10,
+            ),
+        )
+
+    async def disconnect(self) -> None:
+        if self._conn:
+            self._conn.close()
+            self._conn = None
+
+    async def execute(self, sql: str) -> QueryResult:
+        if not self._conn:
+            raise RuntimeError("MySQL 未连接")
+
+        loop = asyncio.get_event_loop()
+        result = QueryResult()
+
+        def _run():
+            with self._conn.cursor() as cursor:
+                # 分割多条语句
+                statements = [s.strip() for s in sql.split(";") if s.strip()]
+                if not statements:
+                    return result
+
+                for stmt in statements:
+                    cursor.execute(stmt)
+                    sql_type = self.classify_sql(stmt)
+                    result.sql_type = sql_type
+
+                    if sql_type == "SELECT":
+                        result.columns = [col[0] for col in cursor.description] if cursor.description else []
+                        result.rows = cursor.fetchall() if cursor.description else []
+                        result.row_count = len(result.rows)
+                    else:
+                        result.affected_rows = cursor.rowcount
+
+                self._conn.commit()
+                return result
+
+        return await loop.run_in_executor(None, _run)
+
+    async def get_schema(self) -> dict:
+        if not self._conn:
+            raise RuntimeError("MySQL 未连接")
+
+        loop = asyncio.get_event_loop()
+
+        def _run():
+            tables = []
+            with self._conn.cursor() as cursor:
+                cursor.execute("SHOW TABLES")
+                table_rows = cursor.fetchall()
+                for (table_name,) in table_rows:
+                    cursor.execute(f"SHOW FULL COLUMNS FROM `{table_name}`")
+                    columns = []
+                    for col in cursor.fetchall():
+                        columns.append({
+                            "name": col[0],
+                            "type": col[1],
+                            "nullable": col[3] == "YES",
+                            "key": col[4] or "",
+                            "default": str(col[5]) if col[5] is not None else None,
+                            "comment": col[8] or "",
+                        })
+                    tables.append({"name": table_name, "columns": columns})
+            return {"db_type": "mysql", "tables": tables}
+
+        return await loop.run_in_executor(None, _run)
+
+    async def test_connection(self) -> dict:
+        from .exceptions import format_connector_result
+        try:
+            await self.connect()
+            if self._conn:
+                self._conn.ping(reconnect=False)
+            return format_connector_result(True, data={"message": "MySQL 连接成功"}, db_type="mysql")
+        except Exception as e:
+            return format_connector_result(False, error=str(e), db_type="mysql", code="CONNECTION_FAILED")
+        finally:
+            await self.disconnect()