asktable-advisor 1.0.1__py3-none-any.whl

asktable_advisor/database/data_generator.py

@@ -0,0 +1,143 @@
+"""Database data generator using LLM."""
+
+import logging
+from typing import Dict, Any, Optional
+
+from ..agent.llm_client import LLMClient
+
+logger = logging.getLogger(__name__)
+
+
+class DataGenerator:
+    """
+    Generate realistic sample data using LLM.
+
+    Creates INSERT statements with business-logic-aware data.
+    """
+
+    def __init__(self, llm_client: LLMClient):
+        """
+        Initialize data generator.
+
+        Args:
+            llm_client: LLM client for AI-powered generation
+        """
+        self.llm_client = llm_client
+
+    def generate_data(
+        self,
+        schema_sql: str,
+        scenario_context: str,
+        data_volume: Dict[str, int],
+    ) -> str:
+        """
+        Generate sample data for database tables.
+
+        Args:
+            schema_sql: Database schema (CREATE TABLE statements)
+            scenario_context: Business context for generating realistic data
+            data_volume: Dict mapping table names to desired row counts
+
+        Returns:
+            SQL INSERT statements
+        """
+        logger.info(f"Generating sample data for {len(data_volume)} tables")
+
+        # Build prompt for data generation
+        prompt = self._build_data_prompt(
+            schema_sql,
+            scenario_context,
+            data_volume,
+        )
+
+        # Call LLM to generate data
+        try:
+            response = self.llm_client.create_message(
+                messages=[{"role": "user", "content": prompt}],
+                system=self._get_data_system_prompt(),
+            )
+
+            sql = self.llm_client.extract_text(response)
+
+            # Extract SQL from markdown if needed
+            sql = self._extract_sql_from_markdown(sql)
+
+            logger.info("Sample data generated successfully")
+            return sql
+
+        except Exception as e:
+            logger.error(f"Failed to generate data: {e}")
+            raise
+
+    def _build_data_prompt(
+        self,
+        schema_sql: str,
+        scenario_context: str,
+        data_volume: Dict[str, int],
+    ) -> str:
+        """Build prompt for data generation."""
+        prompt = f"""请为以下数据库表结构生成真实感的示例数据：
+
+场景上下文：
+{scenario_context}
+
+表结构：
+```sql
+{schema_sql}
+```
+
+数据量要求：
+"""
+
+        for table, count in data_volume.items():
+            prompt += f"- {table} 表：生成 {count} 条数据\n"
+
+        prompt += """
+生成要求：
+1. 数据要符合业务逻辑和真实场景
+2. 确保外键关系正确（引用的 ID 必须存在）
+3. 姓名、地址、商品名等使用真实感的中文内容
+4. 数值数据要有合理的分布（不要全部相同）
+5. 时间数据要分布在合理的时间范围内
+6. 使用批量 INSERT 语句（每条语句插入多行）
+
+只返回 INSERT INTO 语句，不要包含其他解释。
+注意：先插入父表数据，再插入子表数据（按外键依赖顺序）。
+"""
+
+        return prompt
+
+    def _get_data_system_prompt(self) -> str:
+        """Get system prompt for data generation."""
+        return """你是一位数据生成专家，擅长创建真实感、符合业务逻辑的示例数据。
+
+你的特点：
+- 生成的数据具有真实性和多样性
+- 严格遵守外键约束和业务规则
+- 理解各种业务场景的数据特征
+- 确保数据分布合理（符合真实业务规律）
+
+生成策略：
+1. 姓名：使用常见的中文姓名，多样化
+2. 日期：分布在合理的时间范围，考虑业务周期
+3. 金额：符合业务场景（如订单金额有大有小）
+4. 状态：不同状态合理分布（如订单状态：待付款、已付款、已发货等）
+5. 关联数据：确保外键 ID 正确关联（如订单的 user_id 必须是真实存在的用户）
+
+请根据表结构和业务场景，生成高质量的示例数据。
+"""
+
+    def _extract_sql_from_markdown(self, text: str) -> str:
+        """Extract SQL from markdown code blocks."""
+        if "```sql" in text:
+            start = text.find("```sql") + 6
+            end = text.find("```", start)
+            if end > start:
+                return text[start:end].strip()
+        elif "```" in text:
+            start = text.find("```") + 3
+            end = text.find("```", start)
+            if end > start:
+                return text[start:end].strip()
+
+        return text.strip()

asktable_advisor/database/manager.py

@@ -0,0 +1,228 @@
+"""Database manager for AskTable Advisor."""
+
+import logging
+from typing import Optional, List
+
+from sqlalchemy import create_engine, text
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import sessionmaker, Session
+
+from ..config import AdvisorSettings
+
+logger = logging.getLogger(__name__)
+
+
+class DatabaseManager:
+    """
+    Database connection and operation manager.
+
+    Handles MySQL database connections, SQL execution, and database creation.
+    """
+
+    def __init__(self, settings: AdvisorSettings):
+        """
+        Initialize database manager.
+
+        Args:
+            settings: Application settings with database credentials
+        """
+        self.settings = settings
+        self._engine: Optional[Engine] = None
+        self._session_maker: Optional[sessionmaker] = None
+
+        logger.info(
+            f"Database manager initialized: "
+            f"{settings.mysql_host}:{settings.mysql_port}/{settings.mysql_database}"
+        )
+
+    @property
+    def engine(self) -> Engine:
+        """Get or create database engine."""
+        if self._engine is None:
+            self._engine = create_engine(
+                self.settings.mysql_url,
+                pool_pre_ping=True,
+                pool_size=self.settings.mysql_pool_size,
+                max_overflow=self.settings.mysql_max_overflow,
+                echo=False,
+            )
+        return self._engine
+
+    @property
+    def session_maker(self) -> sessionmaker:
+        """Get or create session maker."""
+        if self._session_maker is None:
+            self._session_maker = sessionmaker(
+                autocommit=False,
+                autoflush=False,
+                bind=self.engine,
+            )
+        return self._session_maker
+
+    def get_session(self) -> Session:
+        """Create a new database session."""
+        return self.session_maker()
+
+    def test_connection(self) -> bool:
+        """
+        Test database connection.
+
+        Returns:
+            True if connection successful, False otherwise
+        """
+        try:
+            with self.engine.connect() as conn:
+                result = conn.execute(text("SELECT 1"))
+                result.fetchone()
+            logger.info("Database connection successful")
+            return True
+        except Exception as e:
+            logger.error(f"Database connection failed: {e}")
+            return False
+
+    def create_database(self, database_name: Optional[str] = None) -> None:
+        """
+        Create database if it doesn't exist.
+
+        Args:
+            database_name: Database name (default: use settings)
+        """
+        db_name = database_name or self.settings.mysql_database
+
+        # Create engine without database selection
+        engine_without_db = create_engine(
+            f"mysql+pymysql://{self.settings.mysql_user}:{self.settings.mysql_password}"
+            f"@{self.settings.mysql_host}:{self.settings.mysql_port}",
+            pool_pre_ping=True,
+        )
+
+        try:
+            with engine_without_db.connect() as conn:
+                # Use CREATE DATABASE with proper MySQL syntax
+                conn.execute(
+                    text(f"CREATE DATABASE IF NOT EXISTS `{db_name}` "
+                         f"CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci")
+                )
+                conn.commit()
+            logger.info(f"Database '{db_name}' ready")
+        except Exception as e:
+            logger.error(f"Failed to create database: {e}")
+            raise
+        finally:
+            engine_without_db.dispose()
+
+    def execute_sql(self, sql: str, description: Optional[str] = None) -> int:
+        """
+        Execute SQL statements.
+
+        Args:
+            sql: SQL content (may contain multiple statements separated by ;)
+            description: Optional description of the SQL operation
+
+        Returns:
+            Number of statements executed
+        """
+        if description:
+            logger.info(f"Executing SQL: {description}")
+
+        with self.get_session() as session:
+            try:
+                # Split by semicolon and filter empty statements
+                statements = [
+                    s.strip()
+                    for s in sql.split(";")
+                    if s.strip() and not s.strip().startswith("--")
+                ]
+
+                for statement in statements:
+                    session.execute(text(statement))
+
+                session.commit()
+                logger.info(f"Successfully executed {len(statements)} SQL statements")
+                return len(statements)
+
+            except Exception as e:
+                session.rollback()
+                logger.error(f"SQL execution failed: {e}")
+                logger.error(f"Failed SQL:\n{sql[:500]}...")  # Log first 500 chars
+                raise
+
+    def execute_sql_statements(self, statements: List[str]) -> int:
+        """
+        Execute a list of SQL statements.
+
+        Args:
+            statements: List of SQL statements
+
+        Returns:
+            Number of statements executed
+        """
+        with self.get_session() as session:
+            try:
+                for statement in statements:
+                    if statement.strip():
+                        session.execute(text(statement))
+
+                session.commit()
+                logger.info(f"Successfully executed {len(statements)} SQL statements")
+                return len(statements)
+
+            except Exception as e:
+                session.rollback()
+                logger.error(f"SQL execution failed: {e}")
+                raise
+
+    def database_exists(self, database_name: Optional[str] = None) -> bool:
+        """
+        Check if database exists.
+
+        Args:
+            database_name: Database name (default: use settings)
+
+        Returns:
+            True if database exists
+        """
+        db_name = database_name or self.settings.mysql_database
+
+        engine_without_db = create_engine(
+            f"mysql+pymysql://{self.settings.mysql_user}:{self.settings.mysql_password}"
+            f"@{self.settings.mysql_host}:{self.settings.mysql_port}",
+            pool_pre_ping=True,
+        )
+
+        try:
+            with engine_without_db.connect() as conn:
+                result = conn.execute(
+                    text("SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA "
+                         "WHERE SCHEMA_NAME = :db_name"),
+                    {"db_name": db_name}
+                )
+                exists = result.fetchone() is not None
+                return exists
+        finally:
+            engine_without_db.dispose()
+
+    def list_tables(self, database_name: Optional[str] = None) -> List[str]:
+        """
+        List all tables in the database.
+
+        Args:
+            database_name: Database name (default: use settings)
+
+        Returns:
+            List of table names
+        """
+        db_name = database_name or self.settings.mysql_database
+
+        with self.engine.connect() as conn:
+            result = conn.execute(
+                text(f"SHOW TABLES FROM `{db_name}`")
+            )
+            tables = [row[0] for row in result.fetchall()]
+            return tables
+
+    def close(self) -> None:
+        """Close database connections."""
+        if self._engine:
+            self._engine.dispose()
+            logger.info("Database connections closed")

asktable_advisor/database/schema_generator.py

@@ -0,0 +1,148 @@
+"""Database schema generator using LLM."""
+
+import logging
+from typing import Dict, List, Any, Optional
+
+from ..agent.llm_client import LLMClient
+
+logger = logging.getLogger(__name__)
+
+
+class SchemaGenerator:
+    """
+    Generate database schema using LLM.
+
+    Analyzes business scenarios and creates appropriate table structures.
+    """
+
+    def __init__(self, llm_client: LLMClient):
+        """
+        Initialize schema generator.
+
+        Args:
+            llm_client: LLM client for AI-powered generation
+        """
+        self.llm_client = llm_client
+
+    def generate_schema(
+        self,
+        scenario_name: str,
+        scenario_description: str,
+        requirements: Optional[List[str]] = None,
+        scale_info: Optional[Dict[str, Any]] = None,
+    ) -> str:
+        """
+        Generate database schema SQL for a scenario.
+
+        Args:
+            scenario_name: Name of the scenario
+            scenario_description: Detailed scenario description
+            requirements: List of specific requirements
+            scale_info: Information about data scale
+
+        Returns:
+            SQL DDL statements (CREATE TABLE)
+        """
+        logger.info(f"Generating schema for scenario: {scenario_name}")
+
+        # Build prompt for schema generation
+        prompt = self._build_schema_prompt(
+            scenario_name,
+            scenario_description,
+            requirements,
+            scale_info,
+        )
+
+        # Call LLM to generate schema
+        try:
+            response = self.llm_client.create_message(
+                messages=[{"role": "user", "content": prompt}],
+                system=self._get_schema_system_prompt(),
+            )
+
+            sql = self.llm_client.extract_text(response)
+
+            # Extract SQL from markdown code blocks if present
+            sql = self._extract_sql_from_markdown(sql)
+
+            logger.info(f"Schema generated successfully")
+            return sql
+
+        except Exception as e:
+            logger.error(f"Failed to generate schema: {e}")
+            raise
+
+    def _build_schema_prompt(
+        self,
+        scenario_name: str,
+        scenario_description: str,
+        requirements: Optional[List[str]],
+        scale_info: Optional[Dict[str, Any]],
+    ) -> str:
+        """Build prompt for schema generation."""
+        prompt = f"""请为以下业务场景设计 MySQL 数据库表结构：
+
+场景名称：{scenario_name}
+场景描述：{scenario_description}
+"""
+
+        if requirements:
+            prompt += f"\n具体需求：\n"
+            for req in requirements:
+                prompt += f"- {req}\n"
+
+        if scale_info:
+            prompt += f"\n数据规模：\n"
+            for key, value in scale_info.items():
+                prompt += f"- {key}: {value}\n"
+
+        prompt += """
+请生成完整的 CREATE TABLE 语句。要求：
+1. 每个表都要有主键（id INT AUTO_INCREMENT PRIMARY KEY）
+2. 正确设置外键关系（使用 FOREIGN KEY）
+3. 包含 created_at 和 updated_at 时间戳字段
+4. 选择合适的数据类型（INT、VARCHAR、DECIMAL、DATETIME 等）
+5. 添加必要的索引
+6. 使用 InnoDB 引擎
+7. 字符集使用 utf8mb4
+
+只返回 SQL 语句，不要包含其他解释文字。
+"""
+
+        return prompt
+
+    def _get_schema_system_prompt(self) -> str:
+        """Get system prompt for schema generation."""
+        return """你是一位经验丰富的数据库设计专家。
+
+你精通：
+- 各种业务场景的数据建模（电商、CRM、教育、医疗等）
+- MySQL 数据库设计最佳实践
+- 数据库规范化原则（1NF、2NF、3NF）
+- 性能优化（索引设计、查询优化）
+
+设计原则：
+1. 表结构清晰、字段命名规范
+2. 正确建立表关系（一对一、一对多、多对多）
+3. 适当冗余以提升查询性能
+4. 考虑数据一致性和完整性
+5. 为常用查询添加索引
+
+请根据业务场景，设计合理、规范的数据库表结构。
+"""
+
+    def _extract_sql_from_markdown(self, text: str) -> str:
+        """Extract SQL from markdown code blocks."""
+        # Check if text contains markdown SQL code block
+        if "```sql" in text:
+            start = text.find("```sql") + 6
+            end = text.find("```", start)
+            if end > start:
+                return text[start:end].strip()
+        elif "```" in text:
+            start = text.find("```") + 3
+            end = text.find("```", start)
+            if end > start:
+                return text[start:end].strip()
+
+        return text.strip()