test_executor.py

#!/usr/bin/env python3
"""
测试 QueryExecutor 和 SQL 安全验证功能

这是一个简单的测试脚本，用于验证数据库连接和查询执行器的基本功能。
"""

import sys
from nl2sql_agent import DatabaseConfig, QueryExecutor, validate_sql, setup_logging


def test_database_config():
    """测试数据库配置加载"""
    print("\n=== 测试 DatabaseConfig ===")
    
    try:
        # 从文件加载配置
        config = DatabaseConfig.from_file("rds_demo/db_config.txt")
        print(f"✓ 成功从文件加载配置")
        print(f"  数据库: {config.database}")
        print(f"  主机: {config.host}")
        print(f"  端口: {config.port}")
        
        # 验证配置
        config.validate()
        print(f"✓ 配置验证通过")
        
        return config
    except Exception as e:
        print(f"✗ 配置加载失败: {e}")
        return None


def test_sql_validation():
    """测试SQL安全验证"""
    print("\n=== 测试 SQL 安全验证 ===")
    
    test_cases = [
        ("SELECT * FROM customers", True, "简单SELECT查询"),
        ("SELECT name, email FROM customers WHERE city = 'Beijing'", True, "带WHERE条件的查询"),
        ("SELECT COUNT(*) FROM orders", True, "聚合查询"),
        ("DROP TABLE customers", False, "DROP语句（危险）"),
        ("DELETE FROM customers", False, "DELETE语句（危险）"),
        ("UPDATE customers SET name = 'test'", False, "UPDATE语句（危险）"),
        ("SELECT * FROM customers; DROP TABLE orders;", False, "多条语句（SQL注入）"),
        ("SELECT * FROM customers -- comment", False, "包含注释"),
        ("SELECT * FROM customers /* comment */", False, "包含多行注释"),
    ]
    
    passed = 0
    failed = 0
    
    for sql, expected_safe, description in test_cases:
        is_safe, message = validate_sql(sql)
        
        if is_safe == expected_safe:
            print(f"✓ {description}: {message}")
            passed += 1
        else:
            print(f"✗ {description}: 预期{'安全' if expected_safe else '不安全'}，实际{'安全' if is_safe else '不安全'}")
            failed += 1
    
    print(f"\n验证测试: {passed} 通过, {failed} 失败")
    return failed == 0


def test_query_executor(config):
    """测试查询执行器"""
    print("\n=== 测试 QueryExecutor ===")
    
    try:
        # 创建执行器
        executor = QueryExecutor(config)
        print(f"✓ 成功创建 QueryExecutor")
        
        # 测试连接
        success, message = executor.test_connection()
        if success:
            print(f"✓ 数据库连接测试成功")
        else:
            print(f"✗ 数据库连接测试失败: {message}")
            return False
        
        # 获取Schema信息
        print("\n获取数据库Schema信息...")
        schema_summary = executor.get_schema_summary()
        print(schema_summary[:500] + "..." if len(schema_summary) > 500 else schema_summary)
        print(f"✓ 成功获取Schema信息")
        
        # 执行简单查询
        print("\n执行测试查询: SELECT COUNT(*) as count FROM customers")
        success, result = executor.execute_query("SELECT COUNT(*) as count FROM customers")
        
        if success:
            print(f"✓ 查询执行成功")
            print(f"  结果: {result}")
        else:
            print(f"✗ 查询执行失败: {result}")
            return False
        
        # 测试不安全的SQL
        print("\n测试SQL安全验证（尝试执行DROP语句）...")
        success, result = executor.execute_query("DROP TABLE customers")
        
        if not success and "安全验证失败" in result:
            print(f"✓ SQL安全验证正常工作: {result}")
        else:
            print(f"✗ SQL安全验证未能阻止危险操作")
            return False
        
        # 清理资源
        executor.close()
        print(f"\n✓ 资源清理完成")
        
        return True
        
    except Exception as e:
        print(f"✗ QueryExecutor测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def main():
    """主测试函数"""
    print("=" * 60)
    print("NL2SQL Agent - 数据库连接和查询执行器测试")
    print("=" * 60)
    
    # 配置日志
    setup_logging(level="INFO")
    
    # 测试数据库配置
    config = test_database_config()
    if not config:
        print("\n❌ 数据库配置测试失败，无法继续")
        sys.exit(1)
    
    # 测试SQL验证
    if not test_sql_validation():
        print("\n❌ SQL安全验证测试失败")
        sys.exit(1)
    
    # 测试查询执行器
    if not test_query_executor(config):
        print("\n❌ QueryExecutor测试失败")
        sys.exit(1)
    
    print("\n" + "=" * 60)
    print("✅ 所有测试通过！")
    print("=" * 60)


if __name__ == "__main__":
    main()