simple_sql_test.py

#!/usr/bin/env python3
"""
简单的SQL测试脚本 - 直接使用DuckDB测试SQL查询
不依赖AI模型，只测试SQL执行
"""

import asyncio
import sys
import os
from pathlib import Path
import pandas as pd
import duckdb
import json
from typing import Dict, List, Any


class DirectSQLTester:
    """直接SQL测试器 - 不依赖AI模型"""
    
    def __init__(self, data_dir: str = "src/sample-data"):
        self.data_dir = Path(data_dir)
        self.connection = None
        self.tables = {}
        
    def setup(self):
        """初始化DuckDB连接和加载数据"""
        print("🔧 初始化DuckDB连接...")
        
        try:
            # 创建DuckDB连接
            self.connection = duckdb.connect()
            
            # 配置DuckDB
            self.connection.execute("SET memory_limit='1GB'")
            self.connection.execute("SET threads=4")
            
            # 加载Parquet文件
            self.load_parquet_files()
            
            print("✅ 初始化完成！")
            
        except Exception as e:
            print(f"❌ 初始化失败: {e}")
            raise
    
    def load_parquet_files(self):
        """加载所有Parquet文件为表"""
        if not self.data_dir.exists():
            print(f"❌ 数据目录不存在: {self.data_dir}")
            return
        
        print(f"📥 从 {self.data_dir} 加载Parquet文件...")
        
        parquet_files = list(self.data_dir.glob("*.parquet"))
        
        if not parquet_files:
            print(f"⚠️  在 {self.data_dir} 中没有找到Parquet文件")
            return
        
        for file_path in parquet_files:
            table_name = file_path.stem  # 文件名不含扩展名
            
            try:
                # 创建视图来注册表
                sql = f'CREATE OR REPLACE VIEW "{table_name}" AS SELECT * FROM read_parquet(\'{file_path}\')'
                self.connection.execute(sql)
                
                # 获取行数
                count_sql = f'SELECT COUNT(*) FROM "{table_name}"'
                row_count = self.connection.execute(count_sql).fetchone()[0]
                
                self.tables[table_name] = {
                    "file_path": str(file_path),
                    "row_count": row_count
                }
                
                print(f"  ✅ {table_name:<20} ({row_count:,} 行)")
                
            except Exception as e:
                print(f"  ❌ {table_name:<20} 加载失败: {e}")
        
        print(f"📋 成功加载 {len(self.tables)} 个表")
    
    def show_tables(self):
        """显示所有可用的表"""
        if not self.tables:
            print("⚠️  没有可用的表")
            return
        
        print(f"\n📋 可用的表 ({len(self.tables)}个):")
        print("-" * 50)
        
        for table_name, info in sorted(self.tables.items()):
            row_count = info["row_count"]
            print(f"  • {table_name:<25} {row_count:,} 行")
        
        print("-" * 50)
    
    def describe_table(self, table_name: str):
        """显示表结构"""
        if table_name not in self.tables:
            print(f"❌ 表 '{table_name}' 不存在")
            available = ", ".join(self.tables.keys())
            print(f"可用的表: {available}")
            return
        
        try:
            # 获取表结构
            schema_sql = f'DESCRIBE "{table_name}"'
            schema_result = self.connection.execute(schema_sql).fetchdf()
            
            print(f"\n📊 表 '{table_name}' 的结构:")
            print("-" * 70)
            print(f"行数: {self.tables[table_name]['row_count']:,}")
            print(f"文件: {self.tables[table_name]['file_path']}")
            print("\n列信息:")
            
            for _, row in schema_result.iterrows():
                col_name = row['column_name']
                col_type = row['column_type']
                nullable = "NULL" if row['null'] == 'YES' else "NOT NULL"
                print(f"  • {col_name:<25} {col_type:<20} {nullable}")
            
            print("-" * 70)
            
        except Exception as e:
            print(f"❌ 获取表结构失败: {e}")
    
    def test_sql(self, sql: str, limit_rows: int = 10):
        """测试SQL查询"""
        if not self.connection:
            print("❌ 数据库连接未初始化")
            return None
        
        print(f"\n🔍 执行SQL查询:")
        print("=" * 80)
        print(sql)
        print("=" * 80)
        
        import time
        start_time = time.time()
        
        try:
            # 执行查询
            result_df = self.connection.execute(sql).fetchdf()
            execution_time = (time.time() - start_time) * 1000
            
            print(f"✅ 查询成功!")
            print(f"   • 返回行数: {len(result_df):,}")
            print(f"   • 列数: {len(result_df.columns)}")
            print(f"   • 执行时间: {execution_time:.2f}ms")
            
            # 显示结果
            if len(result_df) > 0:
                print(f"\n📄 查询结果 (显示前{min(limit_rows, len(result_df))}行):")
                print("-" * 80)
                
                # 限制显示的行数
                display_df = result_df.head(limit_rows)
                print(display_df.to_string(index=False, max_cols=10))
                
                if len(result_df) > limit_rows:
                    print(f"\n... (还有 {len(result_df) - limit_rows} 行未显示)")
            else:
                print("\n📄 查询结果为空")
            
            print("-" * 80)
            
            return {
                "success": True,
                "data": result_df,
                "row_count": len(result_df),
                "execution_time_ms": execution_time
            }
            
        except Exception as e:
            execution_time = (time.time() - start_time) * 1000
            print(f"❌ 查询失败: {e}")
            print(f"   • 执行时间: {execution_time:.2f}ms")
            
            return {
                "success": False,
                "error": str(e),
                "execution_time_ms": execution_time
            }
    
    def sample_queries(self):
        """显示一些示例查询"""
        if not self.tables:
            print("⚠️  没有可用的表来生成示例查询")
            return
        
        print("\n💡 示例查询:")
        print("-" * 50)
        
        table_names = list(self.tables.keys())
        
        # 基本查询示例
        examples = [
            f'SELECT * FROM "{table_names[0]}" LIMIT 5',
            f'SELECT COUNT(*) FROM "{table_names[0]}"',
        ]
        
        # 如果有多个表，添加JOIN示例
        if len(table_names) >= 2:
            examples.append(f'SELECT COUNT(*) FROM "{table_names[0]}" a, "{table_names[1]}" b')
        
        for i, query in enumerate(examples, 1):
            print(f"{i}. {query}")
        
        print("-" * 50)
    
    def interactive_mode(self):
        user_input = """WITH CustomerRevenue AS (
    SELECT
        a.Id AS AccountId,
        a.Name AS CustomerName,
        a.BillingCountry AS Country,
        SUM(CASE
            WHEN o.CurrencyIsoCode = 'USD' THEN o.Amount
            WHEN dcr.ConversionRate IS NOT NULL THEN o.Amount / dcr.ConversionRate
            ELSE o.Amount -- Assume original amount if no conversion rate found and not USD
        END) AS TotalCustomerRevenueUSD
    FROM
        Account AS a
    JOIN
        Opportunity AS o
        ON a.Id = o.AccountId
    LEFT JOIN
        DatedConversionRate AS dcr
        ON o.CurrencyIsoCode = dcr.IsoCode
        AND o.CloseDate >= dcr.StartDate
        AND o.CloseDate < dcr.NextStartDate
    WHERE
        o.IsWon = TRUE
        AND o.IsClosed = TRUE
        AND a.BillingCountry IS NOT NULL
    GROUP BY
        a.Id,
        a.Name,
        a.BillingCountry
),
RankedCustomerRevenue AS (
    SELECT
        CustomerName,
        Country,
        TotalCustomerRevenueUSD,
        ROW_NUMBER() OVER (PARTITION BY Country ORDER BY TotalCustomerRevenueUSD DESC) AS rn
    FROM
        CustomerRevenue
)
SELECT
    CustomerName,
    Country,
    TotalCustomerRevenueUSD
FROM
    RankedCustomerRevenue
WHERE
    rn <= 10
ORDER BY
    Country,
    TotalCustomerRevenueUSD DESC;"""

        # 执行SQL
        self.test_sql(user_input)
    
    def cleanup(self):
        """清理资源"""
        if self.connection:
            self.connection.close()
            self.connection = None
        print("🧹 资源已清理")


def main():
    """主函数"""
    print("🚀 简单SQL测试工具 (DuckDB)")
    print("=" * 50)
    
    # 创建测试器
    tester = DirectSQLTester()
    
    try:
        # 初始化
        tester.setup()
        
        # 显示可用表
        tester.show_tables()
        
        # 检查命令行参数
        if len(sys.argv) > 1:
            # 如果提供了SQL参数，直接执行
            sql = " ".join(sys.argv[1:])
            tester.test_sql(sql)
        else:
            # 显示示例查询
            tester.sample_queries()
            
            # 进入交互模式
            tester.interactive_mode()
    
    except KeyboardInterrupt:
        print("\n👋 程序被中断")
    except Exception as e:
        print(f"💥 程序异常: {e}")
        import traceback
        traceback.print_exc()
    finally:
        tester.cleanup()


if __name__ == "__main__":
    main()