利用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 datastd::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 separatorstd::string header_sep = "+";for (size_t width : col_widths) {header_sep += std::string(width + 2, '-') + "+";}std::cout << header_sep << "\n";// Print headerstd::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 rowsfor (size_t i = 0; i < std::min(print_count / 2, total_rows); ++i) {print_row(rows[i], col_widths);}// Print ellipsis if truncatedif (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 truncatedif (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 truncatedfor (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 namesvoid 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 emptyif (!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;
}