利用DeepSeek为chdb命令行客户端添加输出重定向和执行SQL脚本功能

在重定向时，csv格式不输出表格，也不重定向统计信息/计时信息。

#include <iostream>
#include <string>
#include <vector>
#include <chrono>
#include <algorithm>
#include <iomanip>
#include <map>
#include <fstream>  // 添加文件流支持
#include "chdb.h"

class ChDBClient {
private:
    chdb_connection* conn;
    bool timer_enabled;
    std::vector<std::string> history;
    std::string current_format;
    std::map<std::string, std::string> format_aliases;
    std::ofstream output_file;  // 文件输出流
    std::string output_filename;
    bool output_redirected = false;

    void initialize_format_aliases() {
        format_aliases = {
            {"csv", "CSVWithNames"},
            {"json", "JSONEachRow"},
            {"pretty", "Pretty"},
            {"vertical", "Vertical"},
            {"markdown", "Markdown"},
            {"tsv", "TabSeparatedWithNames"},
            {"table", "PrettyCompact"}
        };
    }

    chdb_result* execute_query(const std::string& query) {

        chdb_result* result = chdb_query(*conn, query.c_str(), current_format.c_str());
        return result;
    }
        void print_table(const std::vector<std::vector<std::string>>& rows, const std::vector<std::string>& headers) {
        if (rows.empty()) return;

        // Calculate column widths based on both headers and data
        std::vector<size_t> col_widths(headers.size(), 0);
        for (size_t i = 0; i < headers.size(); ++i) {
            col_widths[i] = headers[i].size();
        }
        for (const auto& row : rows) {
            for (size_t i = 0; i < row.size() && i < col_widths.size(); ++i) {
                col_widths[i] = std::max(col_widths[i], row[i].size());
            }
        }

        // Print header separator
        std::string header_sep = "+";
        for (size_t width : col_widths) {
            header_sep += std::string(width + 2, '-') + "+";
        }
        std::cout << header_sep << "\n";

        // Print header
        std::cout << "|";
        for (size_t i = 0; i < headers.size(); ++i) {
            std::cout << " " << std::left << std::setw(col_widths[i]) << headers[i] << " |";
        }
        std::cout << "\n" << header_sep << "\n";

        // Print rows (with potential truncation)
        size_t total_rows = rows.size();
        bool truncated = false;
        size_t print_count = total_rows;

        if (total_rows > 10) {
            print_count = 10;
            truncated = true;
        }

        // Print first 5 rows
        for (size_t i = 0; i < std::min(print_count / 2, total_rows); ++i) {
            print_row(rows[i], col_widths);
        }

        // Print ellipsis if truncated
        if (truncated) {
            std::cout << "|";
            for (size_t i = 0; i < col_widths.size(); ++i) {
                std::cout << " " << std::left << std::setw(col_widths[i]) << "..." << " |";
            }
            std::cout << "\n";
        }

        // Print last 5 rows if truncated
        if (truncated) {
            for (size_t i = std::max(total_rows - print_count / 2, print_count / 2); i < total_rows; ++i) {
                print_row(rows[i], col_widths);
            }
        } else {
            // Print remaining rows if not truncated
            for (size_t i = print_count / 2; i < total_rows; ++i) {
                print_row(rows[i], col_widths);
            }
        }

        std::cout << header_sep << "\n";

        if (truncated) {
            std::cout << "(" << total_rows << " rows total, showing first and last 5 rows)\n";
        } /*else if (total_rows > 0) {
            std::cout << "(" << total_rows << " rows)\n";
        }*/
    }

    void print_row(const std::vector<std::string>& row, const std::vector<size_t>& col_widths) {
        std::cout << "|";
        for (size_t i = 0; i < row.size() && i < col_widths.size(); ++i) {
            std::cout << " " << std::left << std::setw(col_widths[i]) << row[i] << " |";
        }
        std::cout << "\n";
    }

    // Parse CSV results and extract headers from column names
    void parse_csv(const std::string& csv_data, std::vector<std::string>& headers, std::vector<std::vector<std::string>>& rows) {
        headers.clear();
        rows.clear();
        
        std::vector<std::string> current_row;
        std::string current_field;
        bool in_quotes = false;
        bool is_first_row = true;

        for (char c : csv_data) {
            if (c == '"') {
                in_quotes = !in_quotes;
            } else if (c == ',' && !in_quotes) {
                current_row.push_back(current_field);
                current_field.clear();
            } else if (c == '\n' && !in_quotes) {
                current_row.push_back(current_field);
                
                if (is_first_row) {
                    headers = current_row;
                    is_first_row = false;
                } else {
                    rows.push_back(current_row);
                }
                
                current_row.clear();
                current_field.clear();
            } else {
                current_field += c;
            }
        }

        // Add the last field if not empty
        if (!current_field.empty()) {
            current_row.push_back(current_field);
        }
        if (!current_row.empty()) {
            if (is_first_row) {
                headers = current_row;
            } else {
                rows.push_back(current_row);
            }
        }
    }
    void print_result(chdb_result* result) {
        const char* error = chdb_result_error(result);
        if (error) {
            std::string error_msg = "Query error: " + std::string(error) + "\n";
            /*if (output_redirected) {
                output_file << error_msg;
            } else */{
                std::cerr << error_msg;
            }
            return;
        }

        char* data = chdb_result_buffer(result);
        size_t data_len = chdb_result_length(result);
        size_t rows_read = chdb_result_rows_read(result);
        double elapsed = chdb_result_elapsed(result);

        if (data_len > 0) {
            std::string result_data(data, data_len);
            if (output_redirected) {
                output_file << result_data;
            } else {
            // 恢复表格显示逻辑
            if (current_format == "CSV" || current_format == "CSVWithNames") {
                std::vector<std::string> headers;
                std::vector<std::vector<std::string>> rows;
                parse_csv(result_data, headers, rows);
                print_table(rows, headers);
            } else {
                std::cout << result_data;
            }
            }

            std::string stats;
            if (timer_enabled) {
                stats = "\n(" + std::to_string(rows_read) + " rows, " + 
                       std::to_string(elapsed) + " sec)\n";
            } else if (rows_read > 0) {
                stats = "\n(" + std::to_string(rows_read) + " rows)\n";
            }

            if (!stats.empty()) {
                /*if (output_redirected) {
                    output_file << stats;
                } else */{
                    std::cout << stats;
                }
            }
        } else {
            std::string msg = "Query executed successfully. No results returned.\n";
            /*if (output_redirected) {
                output_file << msg;
            } else */{
                std::cout << msg;
            }
        }
    }

    void handle_output_command(const std::string& filename) {
        if (filename.empty()) {
            if (output_redirected) {
                output_file.close();
                std::cout << "Output redirected to console\n";
                output_redirected = false;
                output_filename.clear();
            }
        } else {
            if (output_redirected) {
                output_file.close();
            }
            output_filename = filename;
            output_file.open(output_filename);
            if (!output_file.is_open()) {
                std::cerr << "Failed to open file: " << output_filename << "\n";
                output_redirected = false;
                output_filename.clear();
            } else {
                output_redirected = true;
                std::cout << "Output redirected to file: " << output_filename << "\n";
            }
        }
    }
    void read_and_execute_sql_file(const std::string& filename) {
    std::ifstream file(filename);
    if (!file.is_open()) {
        std::cerr << "Failed to open file: " << filename << "\n";
        return;
    }

    std::string buffer;
    std::string line;
    
    while (std::getline(file, line)) {
        // 去除行尾回车和空白
        line.erase(line.find_last_not_of(" \t\n\r") + 1);
        
        if (line.empty() || line.find("--") == 0) {
            continue;
        }
        
        buffer += line;
        
        // 检查是否以分号结尾（完整命令/语句）
        if (!buffer.empty() && buffer.back() == ';') {
            buffer.pop_back(); // 移除分号
            buffer.erase(0, buffer.find_first_not_of(" \t\n\r"));
            
            if (!buffer.empty()) {
                if (!process_command(buffer)) {
                    break;
                }
            }
            buffer.clear();
        }
        else {
            buffer += "\n"; // 不是完整语句则保留换行
        }
    }
    
    // 处理文件末尾未以分号结束的语句
    if (!buffer.empty()) {
        buffer.erase(0, buffer.find_first_not_of(" \t\n\r"));
        if (!buffer.empty()) {
            process_command(buffer);
        }
    }
    
    file.close();
}
    // 提取命令处理逻辑到单独函数
    bool process_command(const std::string& command) {
        if (command.empty()) return true;
    std::string cmd = command;
    cmd.erase(0, cmd.find_first_not_of(" \t\n\r"));
    cmd.erase(cmd.find_last_not_of(" \t\n\r") + 1);
    
    if (cmd.empty()) return true;

    if (cmd[0] == '.') {
        // 统一处理带/不带分号的命令
        if (cmd.back() == ';') cmd.pop_back();
        if (command[0] == '.') {
        	
            if (command == ".exit") {
                return false;
            } else if (command == ".timer on") {
                timer_enabled = true;
                std::cout << "Query timer enabled\n";
            } else if (command == ".timer off") {
                timer_enabled = false;
                std::cout << "Query timer disabled\n";
            } else if (command.substr(0, 8) == ".format ") {
                set_format(command.substr(8));
            } else if (command == ".output" || command == ".output;") {
                handle_output_command("");
            } else if (command.substr(0, 8) == ".output ") {
                handle_output_command(command.substr(8));
            } else if (command.substr(0, 6) == ".read ") {
                std::string filename = command.substr(6);
                filename.erase(std::remove(filename.begin(), filename.end(), '\"'), filename.end());
                filename.erase(std::remove(filename.begin(), filename.end(), '\''), filename.end());
                read_and_execute_sql_file(filename);
            } else if (command == ".help") {
                print_help();
            } else if (command == ".history") {
                print_history();
            } else {
                std::cerr << "Unknown command: " << command << "\n";
            }
            return true;
        }
}
        // 处理SQL查询
        chdb_result* result = execute_query(command);
        if (result) {
            print_result(result);
            chdb_destroy_query_result(result);
        }
        return true;
    }
public:
    ChDBClient() : conn(nullptr), timer_enabled(false), current_format("CSVWithNames") {
        initialize_format_aliases();
    }

    ~ChDBClient() {
        if (output_redirected) {
            output_file.close();
        }
        if (conn) {
            chdb_close_conn(conn);
        }
    }

    bool connect() {
        char *argv[] = {
            "chdb_client",
            "--path=:memory:"
        };
        int argc = 2;

        conn = chdb_connect(argc, argv);
        if (!conn) {
            std::cerr << "Failed to create database connection\n";
            return false;
        }
        std::cout << "Connected to in-memory database successfully\n";
        return true;
    }
void run() {
    std::cout << "chDB client (type '.exit' to quit, '.help' for help)\n";

    while (true) {
        std::string input;
        std::cout << "chdb> ";
        
        std::string multiline_sql;
        while (true) {
            std::getline(std::cin, input);
            
            if (input.empty() && multiline_sql.empty()) {
                continue;
            }
            
            size_t semicolon_pos = input.find(';');
            if (semicolon_pos != std::string::npos) {
                multiline_sql += input.substr(0, semicolon_pos);
                break;
            } else {
                multiline_sql += input + "\n";
                //std::cout << "   ...> ";
            }
        }

        multiline_sql.erase(0, multiline_sql.find_first_not_of(" \t\n\r\f\v"));
        multiline_sql.erase(multiline_sql.find_last_not_of(" \t\n\r\f\v") + 1);

        if (multiline_sql.empty()) continue;

        history.push_back(multiline_sql + ";");
        
        if (!process_command(multiline_sql)) {
            break; // 处理.exit命令
        }
    }
}
 

    void set_format(const std::string& format) {
        std::string fmt_lower = format;
        std::transform(fmt_lower.begin(), fmt_lower.end(), fmt_lower.begin(), ::tolower);
        
        if (format_aliases.find(fmt_lower) != format_aliases.end()) {
            current_format = format_aliases[fmt_lower];
            std::cout << "Output format set to: " << current_format << "\n";
        } else {
            current_format = format;
            std::cout << "Output format set to: " << current_format << "\n";
        }
    }

    void print_help() {
        std::cout << "Available commands:\n"
                  << "  .exit          - Exit the client\n"
                  << "  .timer on/off  - Enable/disable query timing\n"
                  << "  .format <fmt>  - Set output format (csv, json, pretty, vertical, markdown, tsv)\n"
                  << "  .output [file] - Redirect output to file (no filename to cancel)\n"
                  << "  .history       - Show command history\n"
                  << "  .read <file>   - Execute SQL commands from file\n"
                  << "  .help          - Show this help message\n";
    }

    void print_history() {
        std::cout << "Command history:\n";
        for (size_t i = 0; i < history.size(); ++i) {
            std::cout << "  " << i + 1 << ": " << history[i] << "\n";
        }
    }
};

int main() {
    ChDBClient client;
    if (!client.connect()) {
        return EXIT_FAILURE;
    }

    client.run();
    std::cout << "Goodbye!\n";
    return EXIT_SUCCESS;
}