sql-assistant 1.0.0__py3-none-any.whl

sql_assistant/database/connectors/postgresql.py

@@ -0,0 +1,133 @@
+"""PostgreSQL 连接器"""
+
+import asyncio
+from typing import Any
+
+import psycopg2
+import psycopg2.extras
+
+from .base import BaseConnector, QueryResult
+
+
+class PostgreSQLConnector(BaseConnector):
+    """PostgreSQL 数据库连接器"""
+
+    db_type = "postgresql"
+
+    def __init__(self, host: str, port: int, user: str, password: str, database: str):
+        super().__init__(host, port, user, password, database)
+        self._conn: psycopg2.extensions.connection | None = None
+
+    async def connect(self) -> None:
+        loop = asyncio.get_event_loop()
+        self._conn = await loop.run_in_executor(
+            None,
+            lambda: psycopg2.connect(
+                host=self.host,
+                port=self.port,
+                user=self.user,
+                password=self.password,
+                dbname=self.database,
+                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("PostgreSQL 未连接")
+
+        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("PostgreSQL 未连接")
+
+        loop = asyncio.get_event_loop()
+
+        def _run():
+            schema_sql = """
+                SELECT
+                    c.table_name,
+                    c.column_name,
+                    c.data_type,
+                    c.character_maximum_length,
+                    c.is_nullable,
+                    c.column_default,
+                    tc.constraint_type
+                FROM information_schema.columns c
+                LEFT JOIN information_schema.key_column_usage kcu
+                    ON c.table_schema = kcu.table_schema
+                    AND c.table_name = kcu.table_name
+                    AND c.column_name = kcu.column_name
+                LEFT JOIN information_schema.table_constraints tc
+                    ON kcu.constraint_name = tc.constraint_name
+                    AND tc.constraint_type = 'PRIMARY KEY'
+                WHERE c.table_schema NOT IN ('information_schema', 'pg_catalog')
+                ORDER BY c.table_name, c.ordinal_position
+            """
+            tables_dict: dict[str, list] = {}
+            with self._conn.cursor() as cursor:
+                cursor.execute(schema_sql)
+                for row in cursor.fetchall():
+                    tname = row[0]
+                    col_type = row[2]
+                    if row[3]:
+                        col_type = f"{col_type}({row[3]})"
+                    if tname not in tables_dict:
+                        tables_dict[tname] = []
+                    tables_dict[tname].append({
+                        "name": row[1],
+                        "type": col_type,
+                        "nullable": row[4] == "YES",
+                        "key": "PRI" if row[6] == "PRIMARY KEY" else "",
+                        "default": str(row[5]) if row[5] is not None else None,
+                        "comment": "",
+                    })
+
+            tables = [{"name": k, "columns": v} for k, v in tables_dict.items()]
+            return {"db_type": "postgresql", "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:
+                cur = self._conn.cursor()
+                cur.execute("SELECT 1")
+                cur.close()
+            return format_connector_result(True, data={"message": "PostgreSQL 连接成功"}, db_type="postgresql")
+        except Exception as e:
+            return format_connector_result(False, error=str(e), db_type="postgresql", code="CONNECTION_FAILED")
+        finally:
+            await self.disconnect()

sql_assistant/database/connectors/redis.py

@@ -0,0 +1,132 @@
+"""Redis 连接器"""
+
+import asyncio
+from typing import Any
+
+import redis.asyncio as aioredis
+
+from .base import BaseConnector, QueryResult
+
+
+class RedisConnector(BaseConnector):
+    """Redis 数据库连接器"""
+
+    db_type = "redis"
+
+    def __init__(self, host: str, port: int, user: str, password: str, database: str):
+        # Redis database 是数字索引
+        super().__init__(host, port, user, password, database)
+        self._conn: aioredis.Redis | None = None
+
+    async def connect(self) -> None:
+        db_num = int(self.database) if self.database and self.database.isdigit() else 0
+        self._conn = aioredis.Redis(
+            host=self.host,
+            port=self.port,
+            username=self.user or None,
+            password=self.password or None,
+            db=db_num,
+            decode_responses=True,
+            socket_connect_timeout=10,
+        )
+
+    async def disconnect(self) -> None:
+        if self._conn:
+            await self._conn.aclose()
+            self._conn = None
+
+    async def get_schema(self) -> dict:
+        """Redis 无传统 schema，使用 SCAN 返回 key 列表（避免 KEYS 命令阻塞）"""
+        if not self._conn:
+            raise RuntimeError("Redis 未连接")
+
+        try:
+            all_keys = []
+            async for k in self._conn.scan_iter(match="*", count=50):
+                all_keys.append(k)
+                if len(all_keys) >= 50:
+                    break
+
+            key_count = len(all_keys)
+            keys = sorted(all_keys)
+            types = {}
+            for k in keys:
+                try:
+                    t = await self._conn.type(k)
+                    types[k] = t
+                except Exception:
+                    types[k] = "unknown"
+
+            return {
+                "db_type": "redis",
+                "keys": [
+                    {"name": k, "type": types.get(k, "unknown")}
+                    for k in keys
+                ],
+                "key_count": key_count,
+            }
+        except Exception as e:
+            return {"db_type": "redis", "keys": [], "error": str(e)}
+
+    async def execute(self, sql: str) -> QueryResult:
+        """执行 Redis 命令"""
+        if not self._conn:
+            raise RuntimeError("Redis 未连接")
+
+        result = QueryResult(sql_type="OTHER")
+
+        # Redis 使用原生命令格式：命令 参数1 参数2 ...
+        parts = sql.strip().split()
+        if not parts:
+            return result
+
+        command = parts[0].upper()
+        args = parts[1:] if len(parts) > 1 else []
+
+        try:
+            # 调用 Redis 命令
+            cmd_func = getattr(self._conn, command.lower(), None)
+            if cmd_func is None:
+                raise ValueError(f"不支持的 Redis 命令: {command}")
+
+            value = await cmd_func(*args)
+
+            result.sql_type = "SELECT" if command in (
+                "GET", "HGET", "HGETALL", "LRANGE", "SMEMBERS", "ZRANGE",
+                "KEYS", "MGET", "TYPE", "TTL", "EXISTS", "STRLEN",
+            ) else "OTHER"
+
+            # 格式化结果
+            if isinstance(value, list):
+                result.columns = ["result"]
+                result.rows = [[v] for v in value]
+                result.row_count = len(value)
+            elif isinstance(value, dict):
+                result.columns = ["key", "value"]
+                result.rows = [[k, v] for k, v in value.items()]
+                result.row_count = len(value)
+            elif isinstance(value, (int, float)):
+                result.affected_rows = int(value)
+            else:
+                result.columns = ["result"]
+                result.rows = [[str(value)]]
+                result.row_count = 1
+
+        except Exception as e:
+            result.columns = ["error"]
+            result.rows = [[str(e)]]
+            result.row_count = 1
+
+        return result
+
+    async def test_connection(self) -> dict:
+        from .exceptions import format_connector_result
+        try:
+            await self.connect()
+            if self._conn:
+                await self._conn.ping()
+            return format_connector_result(True, data={"message": "Redis 连接成功"}, db_type="redis")
+        except Exception as e:
+            return format_connector_result(False, error=str(e), db_type="redis", code="CONNECTION_FAILED")
+        finally:
+            await self.disconnect()

sql_assistant/database/connectors/sqlserver.py

@@ -0,0 +1,140 @@
+"""SQL Server 连接器"""
+
+import asyncio
+from typing import Any
+
+from .base import BaseConnector, QueryResult
+
+
+class SQLServerConnector(BaseConnector):
+    """SQL Server 数据库连接器 (使用 pymssql)"""
+
+    db_type = "sqlserver"
+
+    def __init__(self, host: str, port: int, user: str, password: str, database: str):
+        super().__init__(host, port, user, password, database)
+        self._conn = None
+
+    async def connect(self) -> None:
+        try:
+            import pymssql
+        except ImportError:
+            raise ImportError(
+                "请安装 pymssql: uv pip install pymssql  或  pip install pymssql"
+            )
+
+        loop = asyncio.get_event_loop()
+        self._conn = await loop.run_in_executor(
+            None,
+            lambda: pymssql.connect(
+                server=self.host,
+                port=self.port,
+                user=self.user,
+                password=self.password,
+                database=self.database,
+                timeout=10,
+                login_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("SQL Server 未连接")
+
+        loop = asyncio.get_event_loop()
+        result = QueryResult()
+
+        def _run():
+            cursor = self._conn.cursor()
+            try:
+                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
+            finally:
+                cursor.close()
+
+        return await loop.run_in_executor(None, _run)
+
+    async def get_schema(self) -> dict:
+        if not self._conn:
+            raise RuntimeError("SQL Server 未连接")
+
+        loop = asyncio.get_event_loop()
+
+        def _run():
+            schema_sql = """
+                SELECT
+                    c.TABLE_NAME,
+                    c.COLUMN_NAME,
+                    c.DATA_TYPE,
+                    c.CHARACTER_MAXIMUM_LENGTH,
+                    CASE WHEN c.IS_NULLABLE = 'YES' THEN 1 ELSE 0 END,
+                    c.COLUMN_DEFAULT,
+                    CASE WHEN pk.COLUMN_NAME IS NOT NULL THEN 'PRI' ELSE '' END
+                FROM INFORMATION_SCHEMA.COLUMNS c
+                LEFT JOIN (
+                    SELECT ku.TABLE_NAME, ku.COLUMN_NAME
+                    FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS tc
+                    JOIN INFORMATION_SCHEMA.KEY_COLUMN_USAGE ku
+                        ON tc.CONSTRAINT_NAME = ku.CONSTRAINT_NAME
+                    WHERE tc.CONSTRAINT_TYPE = 'PRIMARY KEY'
+                ) pk ON c.TABLE_NAME = pk.TABLE_NAME AND c.COLUMN_NAME = pk.COLUMN_NAME
+                WHERE c.TABLE_SCHEMA = 'dbo'
+                ORDER BY c.TABLE_NAME, c.ORDINAL_POSITION
+            """
+            tables_dict: dict[str, list] = {}
+            cursor = self._conn.cursor()
+            try:
+                cursor.execute(schema_sql)
+                for row in cursor.fetchall():
+                    tname = row[0]
+                    col_type = row[2]
+                    if row[3]:
+                        col_type = f"{col_type}({row[3]})"
+                    if tname not in tables_dict:
+                        tables_dict[tname] = []
+                    tables_dict[tname].append({
+                        "name": row[1],
+                        "type": col_type,
+                        "nullable": bool(row[4]),
+                        "key": row[6],
+                        "default": str(row[5]) if row[5] is not None else None,
+                        "comment": "",
+                    })
+            finally:
+                cursor.close()
+
+            tables = [{"name": k, "columns": v} for k, v in tables_dict.items()]
+            return {"db_type": "sqlserver", "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()
+            return format_connector_result(True, data={"message": "SQL Server 连接成功"}, db_type="sqlserver")
+        except Exception as e:
+            return format_connector_result(False, error=str(e), db_type="sqlserver", code="CONNECTION_FAILED")
+        finally:
+            await self.disconnect()

sql_assistant/database/history.py

@@ -0,0 +1,290 @@
+"""查询历史管理 - 使用 SQLite 存储"""
+
+import json
+import asyncio
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+import aiosqlite
+
+from ..config import DEFAULT_CONFIG_DIR
+
+HISTORY_DB_PATH = DEFAULT_CONFIG_DIR / "history.db"
+
+
+class HistoryManager:
+    """查询历史管理器 - 支持对话管理"""
+
+    def __init__(self, db_path: Optional[Path] = None):
+        self.db_path = db_path or HISTORY_DB_PATH
+        self._db: Optional[aiosqlite.Connection] = None
+
+    async def _get_db(self) -> aiosqlite.Connection:
+        if self._db is None:
+            self.db_path.parent.mkdir(parents=True, exist_ok=True)
+            self._db = await aiosqlite.connect(str(self.db_path))
+            self._db.row_factory = aiosqlite.Row
+            await self._init_tables()
+        return self._db
+
+    async def _init_tables(self):
+        db = self._db
+        # 对话表 - 存储对话会话
+        await db.execute("""
+            CREATE TABLE IF NOT EXISTS conversations (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                title TEXT NOT NULL,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+            )
+        """)
+        await db.execute("""
+            CREATE INDEX IF NOT EXISTS idx_conversations_updated_at
+            ON conversations(updated_at DESC)
+        """)
+        
+        # 查询历史表 - 添加 conversation_id 字段关联对话
+        await db.execute("""
+            CREATE TABLE IF NOT EXISTS query_history (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                conversation_id INTEGER,
+                question TEXT NOT NULL,
+                sql TEXT NOT NULL,
+                result_json TEXT,
+                db_type TEXT DEFAULT '',
+                llm_provider TEXT DEFAULT '',
+                success INTEGER DEFAULT 1,
+                error_message TEXT DEFAULT '',
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
+            )
+        """)
+        await db.execute("""
+            CREATE INDEX IF NOT EXISTS idx_history_conversation_id
+            ON query_history(conversation_id)
+        """)
+        await db.execute("""
+            CREATE INDEX IF NOT EXISTS idx_history_created_at
+            ON query_history(created_at DESC)
+        """)
+        await db.commit()
+
+    # ============ 对话管理方法 ============
+
+    async def create_conversation(self, title: str = "") -> int:
+        """创建新对话，返回对话 ID"""
+        db = await self._get_db()
+        if not title:
+            title = "未命名对话"
+        
+        cursor = await db.execute(
+            """INSERT INTO conversations (title, created_at, updated_at)
+               VALUES (?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)""",
+            (title,),
+        )
+        await db.commit()
+        return cursor.lastrowid
+
+    async def get_conversations(self, limit: int = 50, offset: int = 0) -> list[dict]:
+        """获取对话列表"""
+        db = await self._get_db()
+        cursor = await db.execute(
+            """SELECT c.id, c.title, c.created_at, c.updated_at,
+                      COUNT(q.id) as message_count
+               FROM conversations c
+               LEFT JOIN query_history q ON c.id = q.conversation_id
+               GROUP BY c.id
+               ORDER BY c.updated_at DESC
+               LIMIT ? OFFSET ?""",
+            (limit, offset),
+        )
+        rows = await cursor.fetchall()
+        return [dict(row) for row in rows]
+
+    async def get_conversation(self, conversation_id: int) -> Optional[dict]:
+        """获取单个对话详情"""
+        db = await self._get_db()
+        cursor = await db.execute(
+            """SELECT c.id, c.title, c.created_at, c.updated_at,
+                      COUNT(q.id) as message_count
+               FROM conversations c
+               LEFT JOIN query_history q ON c.id = q.conversation_id
+               WHERE c.id = ?
+               GROUP BY c.id""",
+            (conversation_id,),
+        )
+        row = await cursor.fetchone()
+        return dict(row) if row else None
+
+    async def update_conversation(self, conversation_id: int, title: str) -> bool:
+        """更新对话标题"""
+        db = await self._get_db()
+        cursor = await db.execute(
+            """UPDATE conversations SET title = ?, updated_at = CURRENT_TIMESTAMP
+               WHERE id = ?""",
+            (title, conversation_id),
+        )
+        await db.commit()
+        return cursor.rowcount > 0
+
+    async def delete_conversation(self, conversation_id: int) -> bool:
+        """删除对话（会级联删除相关的查询记录）"""
+        db = await self._get_db()
+        cursor = await db.execute(
+            "DELETE FROM conversations WHERE id = ?",
+            (conversation_id,),
+        )
+        await db.commit()
+        return cursor.rowcount > 0
+
+    async def get_conversation_messages(self, conversation_id: int) -> list[dict]:
+        """获取对话中的所有消息"""
+        db = await self._get_db()
+        cursor = await db.execute(
+            """SELECT q.id, q.question, q.sql, q.result_json, q.db_type,
+                      q.llm_provider, q.success, q.error_message, q.created_at
+               FROM query_history q
+               WHERE q.conversation_id = ?
+               ORDER BY q.created_at ASC""",
+            (conversation_id,),
+        )
+        rows = await cursor.fetchall()
+        return [dict(row) for row in rows]
+
+    # ============ 查询历史方法 ============
+
+    async def add_record(
+        self,
+        question: str,
+        sql: str,
+        result: Optional[dict] = None,
+        db_type: str = "",
+        llm_provider: str = "",
+        success: bool = True,
+        error_message: str = "",
+        conversation_id: Optional[int] = None,
+    ) -> int:
+        """添加一条查询记录，返回记录 ID"""
+        db = await self._get_db()
+        result_json = json.dumps(result, ensure_ascii=False, default=str) if result else None
+
+        cursor = await db.execute(
+            """INSERT INTO query_history
+               (conversation_id, question, sql, result_json, db_type, llm_provider, success, error_message)
+               VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
+            (conversation_id, question, sql, result_json, db_type, llm_provider, int(success), error_message),
+        )
+        await db.commit()
+        
+        # 如果关联了对话，更新对话的 updated_at
+        if conversation_id:
+            await db.execute(
+                "UPDATE conversations SET updated_at = CURRENT_TIMESTAMP WHERE id = ?",
+                (conversation_id,),
+            )
+            await db.commit()
+        
+        return cursor.lastrowid
+
+    async def get_records(
+        self,
+        limit: int = 50,
+        offset: int = 0,
+        conversation_id: Optional[int] = None,
+    ) -> list[dict]:
+        """获取查询历史记录"""
+        db = await self._get_db()
+        if conversation_id:
+            cursor = await db.execute(
+                """SELECT id, conversation_id, question, sql, result_json, db_type, llm_provider,
+                          success, error_message, created_at
+                   FROM query_history
+                   WHERE conversation_id = ?
+                   ORDER BY created_at DESC
+                   LIMIT ? OFFSET ?""",
+                (conversation_id, limit, offset),
+            )
+        else:
+            cursor = await db.execute(
+                """SELECT id, conversation_id, question, sql, result_json, db_type, llm_provider,
+                          success, error_message, created_at
+                   FROM query_history
+                   ORDER BY created_at DESC
+                   LIMIT ? OFFSET ?""",
+                (limit, offset),
+            )
+        rows = await cursor.fetchall()
+        return [dict(row) for row in rows]
+
+    async def get_record(self, record_id: int) -> Optional[dict]:
+        """获取单条记录"""
+        db = await self._get_db()
+        cursor = await db.execute(
+            "SELECT * FROM query_history WHERE id = ?",
+            (record_id,),
+        )
+        row = await cursor.fetchone()
+        return dict(row) if row else None
+
+    async def delete_record(self, record_id: int) -> bool:
+        """删除单条记录"""
+        db = await self._get_db()
+        cursor = await db.execute(
+            "DELETE FROM query_history WHERE id = ?",
+            (record_id,),
+        )
+        await db.commit()
+        return cursor.rowcount > 0
+
+    async def clear_history(self) -> int:
+        """清空所有历史记录（包括所有对话）"""
+        db = await self._get_db()
+        cursor = await db.execute("DELETE FROM query_history")
+        await db.commit()
+        await db.execute("DELETE FROM conversations")
+        await db.commit()
+        return cursor.rowcount
+
+    async def get_count(self, conversation_id: Optional[int] = None) -> int:
+        """获取记录总数"""
+        db = await self._get_db()
+        if conversation_id:
+            cursor = await db.execute(
+                "SELECT COUNT(*) FROM query_history WHERE conversation_id = ?",
+                (conversation_id,),
+            )
+        else:
+            cursor = await db.execute("SELECT COUNT(*) FROM query_history")
+        row = await cursor.fetchone()
+        return row[0] if row else 0
+
+    async def get_conversation_count(self) -> int:
+        """获取对话总数"""
+        db = await self._get_db()
+        cursor = await db.execute("SELECT COUNT(*) FROM conversations")
+        row = await cursor.fetchone()
+        return row[0] if row else 0
+
+    async def close(self):
+        if self._db:
+            await self._db.close()
+            self._db = None
+
+
+# 全局单例
+_history_manager: Optional[HistoryManager] = None
+
+
+def get_history_manager() -> HistoryManager:
+    global _history_manager
+    if _history_manager is None:
+        _history_manager = HistoryManager()
+    return _history_manager
+
+
+async def close_history_manager():
+    global _history_manager
+    if _history_manager:
+        await _history_manager.close()
+        _history_manager = None