#include "config.h"

#include <Client/ClientBase.h>

#include <Client/ClientBaseHelpers.h>

#include <Client/InternalTextLogs.h>

#include <Client/LineReader.h>

#include <Client/TerminalKeystrokeInterceptor.h>

#include <Client/TestHint.h>

#include <Client/TestTags.h>

#include <Core/SortDescription.h>

#include <Interpreters/sortBlock.h>

#if USE_CLIENT_AI

#include <Client/AI/AISQLGenerator.h>

#include <Client/AI/AIClientFactory.h>

#include <Client/AI/AIConfiguration.h>

#endif

#include <Core/Block.h>

#include <Core/Protocol.h>

#include <Core/Settings.h>

#include <Common/DateLUT.h>

#include <Common/MemoryTracker.h>

#include <Common/formatReadable.h>

#include <Common/scope_guard_safe.h>

#include <Common/Exception.h>

#include <Common/ErrorCodes.h>

#include <Common/getNumberOfCPUCoresToUse.h>

#include <Common/logger_useful.h>

#include <Common/typeid_cast.h>

#include <Common/TerminalSize.h>

#include <Common/StringUtils.h>

#include <Common/filesystemHelpers.h>

#include <Common/NetException.h>

#include <Common/SignalHandlers.h>

#include <Common/tryGetFileNameByFileDescriptor.h>

#include <Columns/ColumnString.h>

#include <Columns/ColumnsNumber.h>

#include <Formats/FormatFactory.h>

#include <Parsers/parseQuery.h>

#include <Parsers/ParserQuery.h>

#include <Parsers/ASTInsertQuery.h>

#include <Parsers/ASTCreateQuery.h>

#include <Parsers/ASTCreateFunctionQuery.h>

#include <Parsers/Access/ASTCreateUserQuery.h>

#include <Parsers/ASTDropQuery.h>

#include <Parsers/ASTExplainQuery.h>

#include <Parsers/ASTSetQuery.h>

#include <Parsers/ASTUseQuery.h>

#include <Parsers/ASTSelectWithUnionQuery.h>

#include <Parsers/ASTQueryWithOutput.h>

#include <Parsers/ASTLiteral.h>

#include <Parsers/ASTIdentifier.h>

#include <Parsers/ASTColumnDeclaration.h>

#include <Parsers/ASTFunction.h>

#include <Parsers/PRQL/ParserPRQLQuery.h>

#include <Parsers/Kusto/ParserKQLStatement.h>

#include <Parsers/Kusto/parseKQLQuery.h>

#include <Parsers/Prometheus/ParserPrometheusQuery.h>

#include <IO/Ask.h>

#include <IO/CompressionMethod.h>

#include <IO/ForkWriteBuffer.h>

#include <IO/ReadHelpers.h>

#include <IO/SharedThreadPools.h>

#include <IO/WriteBufferFromFileDescriptor.h>

#include <IO/WriteBufferFromOStream.h>

#include <Interpreters/InterpreterSetQuery.h>

#include <Interpreters/ProfileEventsExt.h>

#include <Interpreters/ReplaceQueryParameterVisitor.h>

#include <Interpreters/processColumnTransformers.h>

#include <Processors/Executors/PullingAsyncPipelineExecutor.h>

#include <Processors/Formats/IInputFormat.h>

#include <Processors/Formats/Impl/NullFormat.h>

#include <Processors/Formats/Impl/ValuesBlockInputFormat.h>

#include <Processors/QueryPlan/BuildQueryPipelineSettings.h>

#include <Processors/QueryPlan/Optimizations/QueryPlanOptimizationSettings.h>

#include <Processors/QueryPlan/QueryPlan.h>

#include <Processors/Transforms/AddingDefaultsTransform.h>

#include <QueryPipeline/QueryPipeline.h>

#include <QueryPipeline/QueryPipelineBuilder.h>

#include <Storages/MergeTree/MergeTreeSettings.h>

#include <Access/AccessControl.h>

#include <Storages/ColumnsDescription.h>

#include <Storages/SelectQueryInfo.h>

#include <TableFunctions/ITableFunction.h>

#include <filesystem>

#include <iostream>

#include <limits>

#include <map>

#include <memory>

#include <mutex>

#include <string_view>

#include <unordered_map>

#include <csignal>

#include <boost/algorithm/string/case_conv.hpp>

#include <boost/algorithm/string/replace.hpp>

#include <boost/algorithm/string/split.hpp>

#include <Common/config_version.h>

#include <Common/XDGBaseDirectories.h>

#include <base/find_symbols.h>

namespace fs = std::filesystem;

using namespace std::literals;

namespace DB

{

namespace Setting

{

extern const SettingsBool allow_settings_after_format_in_insert;

extern const SettingsBool async_insert;

extern const SettingsDialect dialect;

extern const SettingsNonZeroUInt64 max_block_size;

extern const SettingsNonZeroUInt64 max_insert_block_size;

extern const SettingsUInt64 max_insert_block_size_bytes;

extern const SettingsUInt64 min_insert_block_size_rows;

extern const SettingsUInt64 min_insert_block_size_bytes;

extern const SettingsUInt64 max_parser_backtracks;

extern const SettingsUInt64 max_parser_depth;

extern const SettingsUInt64 max_query_size;

extern const SettingsUInt64 output_format_pretty_max_rows;

extern const SettingsUInt64 output_format_pretty_max_value_width;

extern const SettingsBool partial_result_on_first_cancel;

extern const SettingsBool throw_if_no_data_to_insert;

extern const SettingsBool implicit_select;

extern const SettingsBool apply_settings_from_server;

extern const SettingsString promql_database;

extern const SettingsString promql_table;

extern const SettingsFloatAuto promql_evaluation_time;

extern const SettingsBool into_outfile_create_parent_directories;

}

namespace ErrorCodes

{

extern const int BAD_ARGUMENTS;

extern const int SYNTAX_ERROR;

extern const int DEADLOCK_AVOIDED;

extern const int CLIENT_OUTPUT_FORMAT_SPECIFIED;

extern const int UNKNOWN_PACKET_FROM_SERVER;

extern const int NO_DATA_TO_INSERT;

extern const int UNEXPECTED_PACKET_FROM_SERVER;

extern const int INVALID_USAGE_OF_INPUT;

extern const int CANNOT_SET_SIGNAL_HANDLER;

extern const int LOGICAL_ERROR;

extern const int CANNOT_OPEN_FILE;

extern const int FILE_ALREADY_EXISTS;

extern const int USER_SESSION_LIMIT_EXCEEDED;

extern const int NOT_IMPLEMENTED;

extern const int CANNOT_READ_FROM_FILE_DESCRIPTOR;

extern const int USER_EXPIRED;

extern const int SUPPORT_IS_DISABLED;

extern const int CANNOT_WRITE_TO_FILE;

extern const int CANNOT_CREATE_DIRECTORY;

extern const int TIMEOUT_EXCEEDED;

}

}

namespace ProfileEvents

{

extern const Event UserTimeMicroseconds;

extern const Event SystemTimeMicroseconds;

}

namespace

{

constexpr UInt64 THREAD_GROUP_ID = 0;

void cleanupTempFile(const DB::ASTPtr & parsed_query, const String & tmp_file)

{

if (const auto * query_with_output = dynamic_cast<const DB::ASTQueryWithOutput *>(parsed_query.get()))

{

if (query_with_output->is_outfile_truncate && query_with_output->out_file)

{

if (fs::exists(tmp_file))

fs::remove(tmp_file);

}

}

}

void performAtomicRename(const DB::ASTPtr & parsed_query, const String & out_file)

{

if (const auto * query_with_output = dynamic_cast<const DB::ASTQueryWithOutput *>(parsed_query.get()))

{

if (query_with_output->is_outfile_truncate && query_with_output->out_file)

{

const auto & tmp_file_node = query_with_output->out_file->as<DB::ASTLiteral &>();

String tmp_file = tmp_file_node.value.safeGet<std::string>();

try

{

fs::rename(tmp_file, out_file);

}

catch (const fs::filesystem_error & e)

{

/// Clean up temporary file

if (fs::exists(tmp_file))

fs::remove(tmp_file);

throw DB::Exception(DB::ErrorCodes::CANNOT_WRITE_TO_FILE,

"Cannot rename temporary file {} to {}: {}",

tmp_file, out_file, e.what());

}

}

}

}

}

namespace DB

{

class ClientEmbedded;

ProgressOption toProgressOption(std::string progress)

{

boost::to_upper(progress);

if (progress == "OFF" || progress == "FALSE" || progress == "0" || progress == "NO")

return ProgressOption::OFF;

if (progress == "TTY" || progress == "ON" || progress == "TRUE" || progress == "1" || progress == "YES")

return ProgressOption::TTY;

if (progress == "ERR")

return ProgressOption::ERR;

if (progress == "DEFAULT")

return ProgressOption::DEFAULT;

throw boost::program_options::validation_error(boost::program_options::validation_error::invalid_option_value);

}

std::istream& operator>> (std::istream & in, ProgressOption & progress)

{

std::string token;

in >> token;

progress = toProgressOption(token);

return in;

}

static void incrementProfileEventsBlock(Block & dst, const Block & src)

{

if (dst.empty())

{

dst = src.cloneEmpty();

}

assertBlocksHaveEqualStructure(src, dst, "ProfileEvents");

std::unordered_map<String, size_t> name_pos;

for (size_t i = 0; i < dst.columns(); ++i)

name_pos[dst.getByPosition(i).name] = i;

size_t dst_rows = dst.rows();

MutableColumns mutable_columns = dst.mutateColumns();

auto & dst_column_host_name = typeid_cast<ColumnString &>(*mutable_columns[name_pos["host_name"]]);

auto & dst_array_current_time = typeid_cast<ColumnUInt32 &>(*mutable_columns[name_pos["current_time"]]).getData();

auto & dst_array_type = typeid_cast<ColumnInt8 &>(*mutable_columns[name_pos["type"]]).getData();

auto & dst_column_name = typeid_cast<ColumnString &>(*mutable_columns[name_pos["name"]]);

auto & dst_array_value = typeid_cast<ColumnInt64 &>(*mutable_columns[name_pos["value"]]).getData();

const auto & src_column_host_name = typeid_cast<const ColumnString &>(*src.getByName("host_name").column);

const auto & src_array_current_time = typeid_cast<const ColumnUInt32 &>(*src.getByName("current_time").column).getData();

const auto & src_array_thread_id = typeid_cast<const ColumnUInt64 &>(*src.getByName("thread_id").column).getData();

const auto & src_column_name = typeid_cast<const ColumnString &>(*src.getByName("name").column);

const auto & src_array_value = typeid_cast<const ColumnInt64 &>(*src.getByName("value").column).getData();

struct Id

{

std::string_view name;

std::string_view host_name;

bool operator<(const Id & rhs) const

{

return std::tie(name, host_name)

< std::tie(rhs.name, rhs.host_name);

}

};

std::map<Id, UInt64> rows_by_name;

for (size_t src_row = 0; src_row < src.rows(); ++src_row)

{

/// Filter out threads stats, use stats from thread group

/// Exactly stats from thread group is stored to the table system.query_log

/// The stats from threads are less useful.

/// They take more records, they need to be combined,

/// there even could be several records from one thread.

/// Server doesn't send it any more to the clients, so this code left for compatible

auto thread_id = src_array_thread_id[src_row];

if (thread_id != THREAD_GROUP_ID)

continue;

Id id{

src_column_name.getDataAt(src_row),

src_column_host_name.getDataAt(src_row),

};

rows_by_name[id] = src_row;

}

/// Merge src into dst.

for (size_t dst_row = 0; dst_row < dst_rows; ++dst_row)

{

Id id{

dst_column_name.getDataAt(dst_row),

dst_column_host_name.getDataAt(dst_row),

};

if (auto it = rows_by_name.find(id); it != rows_by_name.end())

{

size_t src_row = it->second;

dst_array_current_time[dst_row] = src_array_current_time[src_row];

switch (static_cast<ProfileEvents::Type>(dst_array_type[dst_row]))

{

case ProfileEvents::Type::INCREMENT:

dst_array_value[dst_row] += src_array_value[src_row];

break;

case ProfileEvents::Type::GAUGE:

dst_array_value[dst_row] = src_array_value[src_row];

break;

}

rows_by_name.erase(it);

}

}

/// Copy rows from src that dst does not contains.

for (const auto & [id, pos] : rows_by_name)

{

for (size_t col = 0; col < src.columns(); ++col)

{

mutable_columns[col]->insert((*src.getByPosition(col).column)[pos]);

}

}

dst.setColumns(std::move(mutable_columns));

}

/// To cancel the query on local format error.

class LocalFormatError : public Exception

{

public:

using Exception::Exception;

LocalFormatError * clone() const override { return new LocalFormatError(*this); }

void rethrow() const override { throw *this; } /// NOLINT(cert-err60-cpp)

};

ClientBase::~ClientBase() = default;

ClientBase::ClientBase(

int in_fd_,

int out_fd_,

int err_fd_,

std::istream & input_stream_,

std::ostream & output_stream_,

std::ostream & error_stream_

)

: stdin_fd(in_fd_)

, stdout_fd(out_fd_)

, stderr_fd(err_fd_)

, cmd_settings(std::make_unique<Settings>())

, cmd_merge_tree_settings(std::make_unique<MergeTreeSettings>())

, std_in(std::make_unique<ReadBufferFromFileDescriptor>(in_fd_))

, std_out(std::make_unique<AutoCanceledWriteBuffer<WriteBufferFromFileDescriptor>>(out_fd_))

, progress_indication(output_stream_, in_fd_, err_fd_)

, progress_table(in_fd_, err_fd_)

, input_stream(input_stream_)

, output_stream(output_stream_)

, error_stream(error_stream_)

{

stdin_is_a_tty = isatty(in_fd_);

stdout_is_a_tty = isatty(out_fd_);

stderr_is_a_tty = isatty(err_fd_);

terminal_width = getTerminalWidth(in_fd_, err_fd_);

}

ASTPtr ClientBase::parseQuery(const char *& pos, const char * end, const Settings & settings, bool allow_multi_statements)

{

std::unique_ptr<IParserBase> parser;

ASTPtr res;

size_t max_length = 0;

if (!allow_multi_statements)

max_length = settings[Setting::max_query_size];

const Dialect dialect = settings[Setting::dialect];

if (dialect == Dialect::kusto)

parser = std::make_unique<ParserKQLStatement>(end, settings[Setting::allow_settings_after_format_in_insert]);

else if (dialect == Dialect::prql)

parser = std::make_unique<ParserPRQLQuery>(max_length, settings[Setting::max_parser_depth], settings[Setting::max_parser_backtracks]);

else if (dialect == Dialect::promql)

parser = std::make_unique<ParserPrometheusQuery>(settings[Setting::promql_database], settings[Setting::promql_table], Field{settings[Setting::promql_evaluation_time]});

else

parser = std::make_unique<ParserQuery>(end, settings[Setting::allow_settings_after_format_in_insert], settings[Setting::implicit_select]);

if (is_interactive || ignore_error)

{

String message;

try

{

if (dialect == Dialect::kusto)

res = tryParseKQLQuery(*parser, pos, end, message, true, "", allow_multi_statements, max_length, settings[Setting::max_parser_depth], settings[Setting::max_parser_backtracks], true);

else

res = tryParseQuery(*parser, pos, end, message, true, "", allow_multi_statements, max_length, settings[Setting::max_parser_depth], settings[Setting::max_parser_backtracks], true);

}

catch (const Exception & e)

{

error_stream << "Exception on client:" << std::endl << getExceptionMessage(e, print_stack_trace, true) << std::endl << std::endl;

client_exception.reset(e.clone());

return nullptr;

}

if (!res)

{

error_stream << std::endl << message << std::endl << std::endl;

return nullptr;

}

}

else

{

if (dialect == Dialect::kusto)

res = parseKQLQueryAndMovePosition(*parser, pos, end, "", allow_multi_statements, max_length, settings[Setting::max_parser_depth], settings[Setting::max_parser_backtracks]);

else

res = parseQueryAndMovePosition(*parser, pos, end, "", allow_multi_statements, max_length, settings[Setting::max_parser_depth], settings[Setting::max_parser_backtracks]);

}

if (is_interactive)

{

WriteBufferFromOwnString res_buf;

IAST::FormatSettings format_settings(/* one_line */ false);

format_settings.show_secrets = true;

format_settings.print_pretty_type_names = true;

res->format(res_buf, format_settings);

res_buf.finalize();

output_stream << std::endl;

#if USE_REPLXX

output_stream << highlighted(res_buf.str(), *client_context);

#else

output_stream << res_buf.str();

#endif

output_stream << std::endl << std::endl;

}

return res;

}

/// Consumes trailing semicolons and tries to consume the same-line trailing comment.

void ClientBase::adjustQueryEnd(

const char *& this_query_end, const char * all_queries_end, uint32_t max_parser_depth, uint32_t max_parser_backtracks)

{

// We have to skip the trailing semicolon that might be left

// after VALUES parsing or just after a normal semicolon-terminated query.

Tokens after_query_tokens(this_query_end, all_queries_end);

IParser::Pos after_query_iterator(after_query_tokens, max_parser_depth, max_parser_backtracks);

while (after_query_iterator.isValid() && after_query_iterator->type == TokenType::Semicolon)

{

this_query_end = after_query_iterator->end;

++after_query_iterator;

}

// Now we have to do some extra work to add the trailing

// same-line comment to the query, but preserve the leading

// comments of the next query. The trailing comment is important

// because the test hints are usually written this way, e.g.:

// select nonexistent_column; -- { serverError 12345 }.

// The token iterator skips comments and whitespace, so we have

// to find the newline in the string manually. If it's earlier

// than the next significant token, it means that the text before

// newline is some trailing whitespace or comment, and we should

// add it to our query. There are also several special cases

// that are described below.

const auto * newline = find_first_symbols<'\n'>(this_query_end, all_queries_end);

const char * next_query_begin = after_query_iterator->begin;

// We include the entire line if the next query starts after

// it. This is a generic case of trailing in-line comment.

// The "equals" condition is for case of end of input (they both equal

// all_queries_end);

if (newline <= next_query_begin)

{

assert(newline >= this_query_end);

this_query_end = newline;

}

else

{

// Many queries on one line, can't do anything. By the way, this

// syntax is probably going to work as expected:

// select nonexistent /* { serverError 12345 } */; select 1

}

}

/// Convert external tables to ExternalTableData and send them using the connection.

void ClientBase::sendExternalTables(ASTPtr parsed_query)

{

const auto * select = parsed_query->as<ASTSelectWithUnionQuery>();

if (!select && !external_tables.empty())

throw Exception(ErrorCodes::BAD_ARGUMENTS, "External tables could be sent only with select query");

if (isEmbeeddedClient() && !external_tables.empty())

throw Exception(ErrorCodes::SUPPORT_IS_DISABLED, "External tables are not allowed in embedded more");

std::vector<ExternalTableDataPtr> data;

for (auto & table : external_tables)

data.emplace_back(table.getData(client_context));

connection->sendExternalTablesData(data);

}

void ClientBase::onData(Block & block, ASTPtr parsed_query)

{

if (block.empty())

return;

processed_rows_from_blocks += block.rows();

/// Even if all blocks are empty, we still need to initialize the output stream to write empty resultset.

initOutputFormat(block, parsed_query);

/// The header block containing zero rows was used to initialize

/// output_format, do not output it.

/// Also do not output too much data if we're fuzzing.

if (block.rows() == 0 || (query_fuzzer_runs != 0 && processed_rows_from_blocks >= 100))

return;

/// If results are written INTO OUTFILE, we can avoid clearing progress to avoid flicker.

if (need_render_progress && tty_buf && (!select_into_file || select_into_file_and_stdout))

{

std::unique_lock lock(tty_mutex);

progress_indication.clearProgressOutput(*tty_buf, lock);

}

if (need_render_progress_table && tty_buf && (!select_into_file || select_into_file_and_stdout))

{

std::unique_lock lock(tty_mutex);

progress_table.clearTableOutput(*tty_buf, lock);

}

try

{

output_format->write(materializeBlock(

block,

!output_format->supportsSpecialSerializationKinds()));

written_first_block = true;

}

catch (const NetException &)

{

throw;

}

catch (const ErrnoException &)

{

throw;

}

catch (const Exception &)

{

//this is a bad catch. It catches too much cases unrelated to LocalFormatError

/// Catch client errors like NO_ROW_DELIMITER

throw LocalFormatError(getCurrentExceptionMessageAndPattern(print_stack_trace), getCurrentExceptionCode());

}

/// Received data block is immediately displayed to the user.

output_format->flush();

/// Restore progress bar and progress table after data block.

if (need_render_progress && tty_buf)

{

if (select_into_file && !select_into_file_and_stdout)

error_stream << "\r";

std::unique_lock lock(tty_mutex);

progress_indication.writeProgress(*tty_buf, lock);

}

if (need_render_progress_table && tty_buf && !cancelled)

{

if (!need_render_progress && select_into_file && !select_into_file_and_stdout)

error_stream << "\r";

std::unique_lock lock(tty_mutex);

progress_table.writeTable(*tty_buf, lock, progress_table_toggle_on.load(), progress_table_toggle_enabled, false);

}

}

void ClientBase::onLogData(Block & block)

{

initLogsOutputStream();

if (need_render_progress && tty_buf)

{

std::unique_lock lock(tty_mutex);

progress_indication.clearProgressOutput(*tty_buf, lock);

}

if (need_render_progress_table && tty_buf)

{

std::unique_lock lock(tty_mutex);

progress_table.clearTableOutput(*tty_buf, lock);

}

/// Logs can be unsorted, i.e. if they were combined from multiple servers (in case of distributed queries)

{

SortDescription desc;

desc.push_back(SortColumnDescription("event_time"));

desc.push_back(SortColumnDescription("event_time_microseconds"));

sortBlock(block, desc, 0, IColumn::PermutationSortStability::Stable);

}

logs_out_stream->writeLogs(block);

logs_out_stream->flush();

}

void ClientBase::onTotals(Block & block, ASTPtr parsed_query)

{

initOutputFormat(block, parsed_query);

output_format->setTotals(materializeBlock(block, !output_format->supportsSpecialSerializationKinds()));

}

void ClientBase::onExtremes(Block & block, ASTPtr parsed_query)

{

initOutputFormat(block, parsed_query);

output_format->setExtremes(materializeBlock(block, !output_format->supportsSpecialSerializationKinds()));

}

void ClientBase::onReceiveExceptionFromServer(std::unique_ptr<Exception> && e)

{

have_error = true;

server_exception = std::move(e);

resetOutput();

}

void ClientBase::onProfileInfo(const ProfileInfo & profile_info)

{

if (profile_info.hasAppliedLimit() && output_format)

output_format->setRowsBeforeLimit(profile_info.getRowsBeforeLimit());

if (profile_info.hasAppliedAggregation() && output_format)

output_format->setRowsBeforeAggregation(profile_info.getRowsBeforeAggregation());

}

void ClientBase::initOutputFormat(const Block & block, ASTPtr parsed_query)

try

{

if (!output_format)

{

/// Ignore all results when fuzzing as they can be huge.

if (query_fuzzer_runs)

{

output_format = std::make_shared<NullOutputFormat>(std::make_shared<const Block>(block));

return;

}

WriteBuffer * out_buf = nullptr;

if (!pager.empty() && !isEmbeeddedClient())

{

if (SIG_ERR == signal(SIGPIPE, SIG_IGN))

throw ErrnoException(ErrorCodes::CANNOT_SET_SIGNAL_HANDLER, "Cannot set signal handler for SIGPIPE");

if (SIG_ERR == signal(SIGQUIT, SIG_IGN))

throw ErrnoException(ErrorCodes::CANNOT_SET_SIGNAL_HANDLER, "Cannot set signal handler for SIGQUIT");

ShellCommand::Config config(pager);

config.pipe_stdin_only = true;

config.terminate_in_destructor_strategy.terminate_in_destructor = true;

config.terminate_in_destructor_strategy.termination_signal = SIGTERM;

pager_cmd = ShellCommand::execute(config);

out_buf = &pager_cmd->in;

}

else

{

out_buf = std_out.get();

}

select_into_file = false;

select_into_file_and_stdout = false;

String current_format = default_output_format;

/// The query can specify output format or output file.

if (const auto * query_with_output = dynamic_cast<const ASTQueryWithOutput *>(parsed_query.get()))

{

if (query_with_output->out_file && isEmbeeddedClient())

throw Exception(ErrorCodes::SUPPORT_IS_DISABLED, "Out files are disabled when you are running client embedded into server.");

String out_file;

if (query_with_output->out_file)

{

select_into_file = true;

const auto & out_file_node = query_with_output->out_file->as<ASTLiteral &>();

out_file = out_file_node.value.safeGet<std::string>();

std::string compression_method_string;

if (query_with_output->compression)

{

const auto & compression_method_node = query_with_output->compression->as<ASTLiteral &>();

compression_method_string = compression_method_node.value.safeGet<std::string>();

}

CompressionMethod compression_method = chooseCompressionMethod(out_file, compression_method_string);

UInt64 compression_level = 3;

if (query_with_output->compression_level)

{

const auto & compression_level_node = query_with_output->compression_level->as<ASTLiteral &>();

compression_level_node.value.tryGet<UInt64>(compression_level);

}

auto flags = O_WRONLY | O_EXCL;

auto file_exists = fs::exists(out_file);

if (file_exists && query_with_output->is_outfile_append)

flags |= O_APPEND;

else if (file_exists && query_with_output->is_outfile_truncate)

flags |= O_TRUNC;

else

flags |= O_CREAT;

chassert(out_file_buf.get() == nullptr);

out_file_buf = wrapWriteBufferWithCompressionMethod(

std::make_unique<WriteBufferFromFile>(out_file, DBMS_DEFAULT_BUFFER_SIZE, flags),

compression_method,

static_cast<int>(compression_level)

);

if (query_with_output->is_into_outfile_with_stdout)

{

select_into_file_and_stdout = true;

out_file_buf = std::make_unique<ForkWriteBuffer>(std::vector<WriteBufferPtr>{std::move(out_file_buf),

std::make_shared<WriteBufferFromFileDescriptor>(stdout_fd)});

}

// We are writing to file, so default format is the same as in non-interactive mode.

if (is_interactive && is_default_format)

current_format = "TabSeparated";

}

if (query_with_output->format_ast != nullptr)

{

if (has_vertical_output_suffix)

throw Exception(ErrorCodes::CLIENT_OUTPUT_FORMAT_SPECIFIED, "Output format already specified");

const auto & id = query_with_output->format_ast->as<ASTIdentifier &>();

current_format = id.name();

}

else if (query_with_output->out_file)

{

auto format_name = FormatFactory::instance().tryGetFormatFromFileName(out_file);

if (format_name)

current_format = *format_name;

}

}

if (has_vertical_output_suffix)

current_format = "Vertical";

bool logs_into_stdout = server_logs_file == "-";

bool extras_into_stdout = need_render_progress || logs_into_stdout;

bool select_only_into_file = select_into_file && !select_into_file_and_stdout;

if (!out_file_buf && default_output_compression_method != CompressionMethod::None)

out_file_buf = wrapWriteBufferWithCompressionMethod(out_buf, default_output_compression_method, 3, 0);

auto format_settings = getFormatSettings(client_context);

format_settings.is_writing_to_terminal = stdout_is_a_tty;

/// It is not clear how to write progress and logs

/// intermixed with data with parallel formatting.

/// It may increase code complexity significantly.

if (!extras_into_stdout || select_only_into_file)

output_format = client_context->getOutputFormatParallelIfPossible(

current_format, out_file_buf ? *out_file_buf : *out_buf, block, format_settings);

else

output_format = client_context->getOutputFormat(

current_format, out_file_buf ? *out_file_buf : *out_buf, block, format_settings);

output_format->setAutoFlush();

if ((!select_into_file || select_into_file_and_stdout)

&& stdout_is_a_tty

&& stdin_is_a_tty

&& !FormatFactory::instance().checkIfOutputFormatIsTTYFriendly(current_format))

{

stopKeystrokeInterceptorIfExists();

SCOPE_EXIT({ startKeystrokeInterceptorIfExists(); });

const auto question = fmt::format(R"(The requested output format `{}` is binary and could produce side-effects when output directly into the terminal.

If you want to output it into a file, use the "INTO OUTFILE" modifier in the query or redirect the output of the shell command.

Do you want to output it anyway? [y/N] )", current_format);

if (!ask(question, *std_in, *std_out))

output_format = std::make_shared<NullOutputFormat>(std::make_shared<const Block>(block));

*std_out << '\n';

}

}

}

catch (...)

{

if (out_file_buf)

out_file_buf->cancel();

out_file_buf.reset();

throw LocalFormatError(getCurrentExceptionMessageAndPattern(print_stack_trace), getCurrentExceptionCode());

}

void ClientBase::initLogsOutputStream()

{

if (!logs_out_stream)

{

WriteBuffer * wb = out_logs_buf.get();

bool color_logs = false;

if (!out_logs_buf)

{

if (server_logs_file.empty())

{

/// Use stderr by default

out_logs_buf = std::make_unique<AutoCanceledWriteBuffer<WriteBufferFromFileDescriptor>>(stderr_fd);

wb = out_logs_buf.get();

color_logs = stderr_is_a_tty;

}

else if (server_logs_file == "-")

{

/// Use stdout if --server_logs_file=- specified

wb = std_out.get();

color_logs = stdout_is_a_tty;

}

else

{

out_logs_buf

= std::make_unique<AutoCanceledWriteBuffer<WriteBufferFromFile>>(server_logs_file, DBMS_DEFAULT_BUFFER_SIZE, O_WRONLY | O_APPEND | O_CREAT);

wb = out_logs_buf.get();

}

}

logs_out_stream = std::make_unique<InternalTextLogs>(*wb, color_logs);

}

}

void ClientBase::adjustSettings(ContextMutablePtr context)

{

/// NOTE: Do not forget to set changed=false to avoid sending it to the server (to avoid breakage read only profiles)

/// Do not limit pretty format output in case of --pager specified or in case of stdout is not a tty.

if (!pager.empty() || !stdout_is_a_tty)

{

Settings settings = context->getSettingsCopy();

if (!context->getSettingsRef()[Setting::output_format_pretty_max_rows].changed)

{

settings[Setting::output_format_pretty_max_rows] = std::numeric_limits<UInt64>::max();

settings[Setting::output_format_pretty_max_rows].changed = false;

}

if (!context->getSettingsRef()[Setting::output_format_pretty_max_value_width].changed)

{

settings[Setting::output_format_pretty_max_value_width] = std::numeric_limits<UInt64>::max();

settings[Setting::output_format_pretty_max_value_width].changed = false;

}

context->setSettings(settings);

}

}

void ClientBase::initClientContext(ContextMutablePtr context)

{

client_context = context;

client_context->setClientName(std::string(DEFAULT_CLIENT_NAME));

client_context->setQuotaClientKey(getClientConfiguration().getString("quota_key", ""));

client_context->setQueryKindInitial();

client_context->setQueryKind(query_kind);

client_context->setQueryParameters(query_parameters);

}

bool ClientBase::isFileDescriptorSuitableForInput(int fd)

{

struct stat file_stat;

return fstat(fd, &file_stat) == 0

&& (S_ISREG(file_stat.st_mode) || S_ISLNK(file_stat.st_mode));

}

void ClientBase::setDefaultFormatsAndCompressionFromConfiguration()

{

if (getClientConfiguration().has("output-format"))

{

default_output_format = getClientConfiguration().getString("output-format");

is_default_format = false;

}

else if (getClientConfiguration().has("format"))

{

default_output_format = getClientConfiguration().getString("format");

is_default_format = false;

}

else if (getClientConfiguration().has("vertical"))

{

default_output_format = "Vertical";

is_default_format = false;

}

else if (isFileDescriptorSuitableForInput(stdout_fd))

{

std::optional<String> format_from_file_name = FormatFactory::instance().tryGetFormatFromFileDescriptor(stdout_fd);

if (format_from_file_name)

default_output_format = *format_from_file_name;

else

default_output_format = "TSV";

std::optional<String> file_name = tryGetFileNameFromFileDescriptor(stdout_fd);

if (file_name)

default_output_compression_method = chooseCompressionMethod(*file_name, "");

}

else if (is_interactive)

{

default_output_format = "PrettyCompact";

}

else

{

default_output_format = "TSV";

}

if (getClientConfiguration().has("input-format"))

{

default_input_format = getClientConfiguration().getString("input-format");

}

else if (getClientConfiguration().has("format"))

{

default_input_format = getClientConfiguration().getString("format");

}

else if (getClientConfiguration().getString("table-file", "-") != "-")

{

auto file_name = getClientConfiguration().getString("table-file");

std::optional<String> format_from_file_name = FormatFactory::instance().tryGetFormatFromFileName(file_name);

if (format_from_file_name)

default_input_format = *format_from_file_name;

else

default_input_format = "auto";

}

else

{

std::optional<String> format_from_file_name = FormatFactory::instance().tryGetFormatFromFileDescriptor(stdin_fd);

if (format_from_file_name)

default_input_format = *format_from_file_name;

else

default_input_format = "auto";

std::optional<String> file_name = tryGetFileNameFromFileDescriptor(stdin_fd);

if (file_name)

default_input_compression_method = chooseCompressionMethod(*file_name, "");

}

if (getClientConfiguration().has("insert_format_max_block_size"))

insert_format_max_block_size_rows_from_config = getClientConfiguration().getUInt64("insert_format_max_block_size");

if (getClientConfiguration().has("insert_format_max_block_size_bytes"))

insert_format_max_block_size_bytes_from_config = getClientConfiguration().getUInt64("insert_format_max_block_size_bytes");

if (getClientConfiguration().has("insert_format_min_block_size_rows"))

insert_format_min_block_size_rows_from_config = getClientConfiguration().getUInt64("insert_format_min_block_size_rows");

if (getClientConfiguration().has("insert_format_min_block_size_bytes"))

insert_format_min_block_size_bytes_from_config = getClientConfiguration().getUInt64("insert_format_min_block_size_bytes");

}

void ClientBase::initTTYBuffer(ProgressOption progress_option, ProgressOption progress_table_option)

{

if (tty_buf)

return;

if (progress_option == ProgressOption::OFF || (!is_interactive && progress_option == ProgressOption::DEFAULT))

need_render_progress = false;

if (progress_table_option == ProgressOption::OFF || (!is_interactive && progress_table_option == ProgressOption::DEFAULT))

need_render_progress_table = false;

if (!need_render_progress && !need_render_progress_table)

return;

progress_table_toggle_enabled = getClientConfiguration().getBool("enable-progress-table-toggle");

progress_table_toggle_on = !progress_table_toggle_enabled;

/// If need_render_progress and need_render_progress_table are enabled,

/// use ProgressOption that was set for the progress bar for progress table as well.

ProgressOption progress = progress_option ? progress_option : progress_table_option;

/// Output all progress bar commands to terminal at once to avoid flicker.

/// This size is usually greater than the window size.

static constexpr size_t buf_size = 1024;

/// Prefer to use an existing fd for the tty, from stdin/stdout/stderr, to allow redirecting

/// progress indication to a different terminal.

int tty_fd = -1;

if (stderr_is_a_tty || progress == ProgressOption::ERR)