{

UInt32 secret_key = 0;

#if USE_SSL

if (RAND_bytes(reinterpret_cast<unsigned char *>(&secret_key), sizeof(secret_key)) != 1)

throw Exception(ErrorCodes::OPENSSL_ERROR, "RAND_bytes failed: {}", getOpenSSLErrors());

#else

const int random_fd = ::open("/dev/urandom", O_RDONLY | O_CLOEXEC);

if (random_fd == -1)

throw ErrnoException(ErrorCodes::CANNOT_OPEN_FILE, "Cannot open /dev/urandom");

SCOPE_EXIT({ [[maybe_unused]] int err = ::close(random_fd); });

auto * position = reinterpret_cast<char *>(&secret_key);

size_t bytes_remaining = sizeof(secret_key);

while (bytes_remaining > 0)

{

ssize_t bytes_read = ::read(random_fd, position, bytes_remaining);

if (bytes_read == -1)

{

if (errno == EINTR)

continue;

throw ErrnoException(ErrorCodes::CANNOT_READ_ALL_DATA, "Cannot read from /dev/urandom");

}

if (bytes_read == 0)

throw Exception(ErrorCodes::CANNOT_READ_ALL_DATA, "Unexpected end of /dev/urandom");

position += bytes_read;

bytes_remaining -= bytes_read;

}

#endif

return secret_key;

{

/// Only treat the packet as a no-op when it consists of a single statement. A simple-query

/// packet may contain several `;`-separated statements; if we shortcut on the leading keyword

/// we would acknowledge the whole packet and silently skip the rest (e.g. `RESET ALL; SELECT 1`

/// or, worse, `RESET ALL; DROP TABLE t`). An interior `;` — anything other than trailing

/// whitespace after it — means there is more than one statement, so bail out and let the normal

/// multi-statement splitter handle it.

if (const size_t semicolon = query.find(';'); semicolon != String::npos)

{

for (size_t i = semicolon + 1; i < query.size(); ++i)

{

const char c = query[i];

if (c != ' ' && c != '\t' && c != '\n' && c != '\r' && c != '\f' && c != '\v')

return std::nullopt;

}

}

/// Enough to cover the longest recognized command plus its argument: a PostgreSQL identifier

/// is at most 63 bytes, and a qualified `RESET extension.setting` name consists of two of them.

/// If the normalized prefix fills this budget completely, the statement may continue beyond it,

/// so we could not verify that nothing trails the argument — treat it as not recognized.

static constexpr size_t max_prefix_len = 160;

String prefix = PostgreSQLProtocol::Messaging::CommandComplete::extractNormalizedPrefix(query, max_prefix_len);

if (prefix.size() == max_prefix_len)

return std::nullopt;

/// The multi-statement guard above rejected any statement after an interior `;`, so the only

/// `;` that can remain is a single trailing one (with trailing whitespace already collapsed).

/// Drop it so it is not mistaken for a command argument below.

while (!prefix.empty() && (prefix.back() == ';' || prefix.back() == ' '))

prefix.pop_back();

/// `extractNormalizedPrefix` already uppercased the text and collapsed runs of

/// whitespace to single spaces, so a keyword is a run of 'A'..'Z'. Take the next

/// such run, skipping a leading space; this also stops at a `;` or `*`.

size_t pos = 0;

const auto take_word = [&]() -> String

{

while (pos < prefix.size() && prefix[pos] == ' ')

++pos;

const size_t start = pos;

while (pos < prefix.size() && prefix[pos] >= 'A' && prefix[pos] <= 'Z')

++pos;

return prefix.substr(start, pos - start);

};

const auto has_more = [&]() -> bool

{

while (pos < prefix.size() && prefix[pos] == ' ')

++pos;

return pos < prefix.size();

};

/// An argument in the normalized prefix: an identifier — a letter or `_` followed by letters,

/// digits, `_` or `$` (the text is already uppercased). With `allow_dots`, a qualified name

/// such as `EXTENSION.SETTING` (for `RESET`) is a single token as well.

const auto take_identifier = [&](bool allow_dots) -> String

{

while (pos < prefix.size() && prefix[pos] == ' ')

++pos;

const size_t start = pos;

if (pos < prefix.size() && ((prefix[pos] >= 'A' && prefix[pos] <= 'Z') || prefix[pos] == '_'))

{

++pos;

while (pos < prefix.size()

&& ((prefix[pos] >= 'A' && prefix[pos] <= 'Z') || (prefix[pos] >= '0' && prefix[pos] <= '9')

|| prefix[pos] == '_' || prefix[pos] == '$' || (allow_dots && prefix[pos] == '.')))

++pos;

}

return prefix.substr(start, pos - start);

};

/// Not `const`: the early returns below move it out, and `performance-no-automatic-move`

/// (clang-tidy) rejects returning a `const` local because constness prevents the move.

String command = take_word();

/// The keyword must end at a word boundary: `RESET1FOO` must not be taken for `RESET`.

if (pos < prefix.size() && prefix[pos] != ' ')

return std::nullopt;

/// Accept the connection-cleanup commands the Skunk driver actually sends: `RESET { name | ALL }`

/// and `UNLISTEN { channel | * }`, reporting the bare keyword as the tag. `LISTEN` and `NOTIFY`

/// are deliberately NOT accepted here: unlike `UNLISTEN` (idempotent unsubscribe-all cleanup),

/// they are application-visible PostgreSQL pub/sub operations, and this handler never delivers a

/// `NotificationResponse`, so acknowledging them would turn an unsupported feature into a silent

/// false success instead of a plain error. Issue #12476 only asks for `UNLISTEN *` / `RESET ALL`.

///

/// Accept exactly one argument — an identifier (for `RESET` possibly a qualified

/// `extension.setting` name; `ALL` is itself covered as an identifier), or `*` for `UNLISTEN` —

/// and require the statement to end right after it, so that malformed variants such as a bare

/// `RESET`, `RESET foo bar` or `UNLISTEN * garbage` are not acknowledged as success but fall

/// through to the normal error path. Valid forms that no driver is known to emit — quoted

/// identifiers and multi-word variants such as `RESET SESSION AUTHORIZATION` — likewise fall

/// through, as before this change.

if (command == "UNLISTEN" || command == "RESET")

{

String arg;

if (command == "UNLISTEN" && has_more() && prefix[pos] == '*')

{

arg = "*";

++pos;

}

else

{

arg = take_identifier(/* allow_dots = */ command == "RESET");

}

if (arg.empty() || has_more())

return std::nullopt;

return command;

}

if (command == "DISCARD")

{

/// PostgreSQL accepts only `DISCARD { ALL | PLANS | SEQUENCES | TEMP | TEMPORARY }`

/// (with `TEMPORARY` normalized to `TEMP`). Reject a bare `DISCARD` or an unknown

/// subcommand such as `DISCARD FOO` instead of claiming success for a command we were

/// never asked to emulate.

String arg = take_word();

if (arg == "TEMPORARY")

arg = "TEMP";

if ((arg == "ALL" || arg == "PLANS" || arg == "SEQUENCES" || arg == "TEMP") && !has_more())

return command + " " + arg;

return std::nullopt;

}

return std::nullopt;

const Poco::Net::StreamSocket & socket_,

#if USE_SSL

const std::string & prefix_,

#endif

IServer & server_,

TCPServer & tcp_server_,

bool ssl_enabled_,

bool secure_required_,

Int32 connection_id_,

std::optional<String> default_session_user_,

VectorWithMemoryTracking<std::shared_ptr<PostgreSQLProtocol::PGAuthentication::AuthenticationMethod>> & auth_methods_,

const ProfileEvents::Event & read_event_,

const ProfileEvents::Event & write_event_)

: Poco::Net::TCPServerConnection(socket_)

#if USE_SSL

, config(server_.config())

, prefix(prefix_)

#endif

, server(server_)

, tcp_server(tcp_server_)

, ssl_enabled(ssl_enabled_)

, secure_required(secure_required_)

, connection_id(connection_id_)

, default_session_user(std::move(default_session_user_))

, read_event(read_event_)

, write_event(write_event_)

, authentication_manager(auth_methods_)

, prepared_statements_manager(std::nullopt)

{

/// `BackendKeyData` identifies every statement on this connection for cancellation.

secret_key = generateRandomUInt32();

query_id_token = generateRandomUInt32();

changeIO(socket());

#if USE_SSL

params.privateKeyFile = config.getString(prefix + Poco::Net::SSLManager::CFG_PRIV_KEY_FILE, "");

params.certificateFile = config.getString(prefix + Poco::Net::SSLManager::CFG_CERTIFICATE_FILE, params.privateKeyFile);

if (!params.privateKeyFile.empty() && !params.certificateFile.empty())

{

params.caLocation = config.getString(prefix + Poco::Net::SSLManager::CFG_CA_LOCATION, "");

if (params.caLocation.empty())

{

auto ctx = Poco::Net::SSLManager::instance().defaultServerContext();

params.caLocation = ctx->getCAPaths().caLocation;

}

params.verificationMode = Poco::Net::SSLManager::VAL_VER_MODE;

if (config.hasProperty(prefix + Poco::Net::SSLManager::CFG_VER_MODE))

{

std::string mode = config.getString(prefix + Poco::Net::SSLManager::CFG_VER_MODE);

params.verificationMode = Poco::Net::Utility::convertVerificationMode(mode);

}

params.verificationDepth = config.getInt(prefix + Poco::Net::SSLManager::CFG_VER_DEPTH, Poco::Net::SSLManager::VAL_VER_DEPTH);

params.loadDefaultCAs

= config.getBool(prefix + Poco::Net::SSLManager::CFG_ENABLE_DEFAULT_CA, Poco::Net::SSLManager::VAL_ENABLE_DEFAULT_CA);

params.cipherList = config.getString(prefix + Poco::Net::SSLManager::CFG_CIPHER_LIST, Poco::Net::SSLManager::VAL_CIPHER_LIST);

params.cipherList

= config.getString(prefix + Poco::Net::SSLManager::CFG_CYPHER_LIST, params.cipherList); // for backwards compatibility

bool require_tlsv1 = config.getBool(prefix + Poco::Net::SSLManager::CFG_REQUIRE_TLSV1, false);

bool require_tlsv1_1 = config.getBool(prefix + Poco::Net::SSLManager::CFG_REQUIRE_TLSV1_1, false);

bool require_tlsv1_2 = config.getBool(prefix + Poco::Net::SSLManager::CFG_REQUIRE_TLSV1_2, false);

if (require_tlsv1_2)

usage = Poco::Net::Context::TLSV1_2_SERVER_USE;

else if (require_tlsv1_1)

usage = Poco::Net::Context::TLSV1_1_SERVER_USE;

else if (require_tlsv1)

usage = Poco::Net::Context::TLSV1_SERVER_USE;

else

usage = Poco::Net::Context::SERVER_USE;

params.dhParamsFile = config.getString(prefix + Poco::Net::SSLManager::CFG_DH_PARAMS_FILE, "");

params.ecdhCurve = config.getString(prefix + Poco::Net::SSLManager::CFG_ECDH_CURVE, "");

std::string disabled_protocols_list = config.getString(prefix + Poco::Net::SSLManager::CFG_DISABLE_PROTOCOLS, "");

Poco::StringTokenizer dp_tok(

disabled_protocols_list, ";,", Poco::StringTokenizer::TOK_TRIM | Poco::StringTokenizer::TOK_IGNORE_EMPTY);

disabled_protocols = 0;

for (const auto & token : dp_tok)

{

if (token == "sslv2")

disabled_protocols |= Poco::Net::Context::PROTO_SSLV2;

else if (token == "sslv3")

disabled_protocols |= Poco::Net::Context::PROTO_SSLV3;

else if (token == "tlsv1")

disabled_protocols |= Poco::Net::Context::PROTO_TLSV1;

else if (token == "tlsv1_1")

disabled_protocols |= Poco::Net::Context::PROTO_TLSV1_1;

else if (token == "tlsv1_2")

disabled_protocols |= Poco::Net::Context::PROTO_TLSV1_2;

else if (token == "tlsv1_3")

disabled_protocols |= Poco::Net::Context::PROTO_TLSV1_3;

}

extended_verification = config.getBool(prefix + Poco::Net::SSLManager::CFG_EXTENDED_VERIFICATION, false);

prefer_server_ciphers = config.getBool(prefix + Poco::Net::SSLManager::CFG_PREFER_SERVER_CIPHERS, false);

}

#endif

{

in = std::make_shared<ReadBufferFromPocoSocket>(socket, read_event);

out = std::make_shared<AutoCanceledWriteBuffer<WriteBufferFromPocoSocket>>(socket, write_event);

message_transport = std::make_shared<PostgreSQLProtocol::Messaging::MessageTransport>(in.get(), out.get());

{

DB::setThreadName(ThreadName::POSTGRES_HANDLER);

session = std::make_unique<Session>(server.context(), ClientInfo::Interface::POSTGRESQL);

SCOPE_EXIT({ session.reset(); });

session->setClientConnectionId(connection_id);

/// A `CancelRequest` for this connection arrives on a different connection, so the secret has

/// to be reachable from the whole server for as long as this one is open.

server.context()->getProcessList().registerPostgreSQLCancellationKey(connection_id, secret_key, currentQueryId());

SCOPE_EXIT({ server.context()->getProcessList().unregisterPostgreSQLCancellationKey(connection_id, secret_key); });

try

{

if (!startup())

return;

/// Emit `ReadyForQuery` only at explicit protocol boundaries.

need_ready_for_query = true;

while (tcp_server.isOpen())

{

if (need_ready_for_query)

{

message_transport->send(PostgreSQLProtocol::Messaging::ReadyForQuery(), true);

need_ready_for_query = false;

}

constexpr size_t connection_check_timeout = 1; // 1 second

while (!in->poll(1000000 * connection_check_timeout))

if (!tcp_server.isOpen())

return;

PostgreSQLProtocol::Messaging::FrontMessageType message_type = message_transport->receiveMessageType();

if (!tcp_server.isOpen())

return;

/// After an extended-query error, discard through `Sync` but honor `Terminate`.

if (ignore_until_sync

&& message_type != PostgreSQLProtocol::Messaging::FrontMessageType::SYNC

&& message_type != PostgreSQLProtocol::Messaging::FrontMessageType::TERMINATE)

{

message_transport->dropMessage();

continue;

}

switch (message_type)

{

case PostgreSQLProtocol::Messaging::FrontMessageType::QUERY:

/// A simple query is a complete protocol cycle, and it also destroys the

/// unnamed prepared statement and the unnamed portal.

in_extended_query_cycle = false;

prepared_statements_manager.dropUnnamedStatementAndPortal();

processQuery();

need_ready_for_query = true;

message_transport->flush();

break;

case PostgreSQLProtocol::Messaging::FrontMessageType::TERMINATE:

LOG_DEBUG(log, "Client closed the connection");

return;

case PostgreSQLProtocol::Messaging::FrontMessageType::PARSE:

/// An extended-query cycle ends at its `Sync` or at a simple query.

in_extended_query_cycle = true;

processParseQuery();

message_transport->flush();

break;

case PostgreSQLProtocol::Messaging::FrontMessageType::BIND:

in_extended_query_cycle = true;

processBindQuery();

message_transport->flush();

break;

case PostgreSQLProtocol::Messaging::FrontMessageType::EXECUTE:

in_extended_query_cycle = true;

processExecuteQuery();

message_transport->flush();

break;

case PostgreSQLProtocol::Messaging::FrontMessageType::SYNC:

/// `Sync` ends the cycle and produces one `ReadyForQuery`.

in_extended_query_cycle = false;

processSyncQuery();

need_ready_for_query = true;

message_transport->flush();

break;

case PostgreSQLProtocol::Messaging::FrontMessageType::DESCRIBE:

in_extended_query_cycle = true;

processDescribeQuery();

message_transport->flush();

break;

case PostgreSQLProtocol::Messaging::FrontMessageType::FLUSH:

message_transport->send(

PostgreSQLProtocol::Messaging::ErrorOrNoticeResponse(

PostgreSQLProtocol::Messaging::ErrorOrNoticeResponse::ERROR,

"0A000",

"ClickHouse doesn't support extended query mechanism"),

true);

LOG_ERROR(log, "Client tried to access via extended query protocol");

message_transport->dropMessage();

recoverFromRejectedMessage();

break;

case PostgreSQLProtocol::Messaging::FrontMessageType::CLOSE:

in_extended_query_cycle = true;

processCloseQuery();

message_transport->flush();

break;

default:

message_transport->send(

PostgreSQLProtocol::Messaging::ErrorOrNoticeResponse(

PostgreSQLProtocol::Messaging::ErrorOrNoticeResponse::ERROR,

"0A000",

"Command is not supported"),

true);

LOG_ERROR(log, "Command is not supported. Command code {:d}", static_cast<Int32>(message_type));

message_transport->dropMessage();

recoverFromRejectedMessage();

}

}

}

catch (const Poco::Exception &exc)

{

log->log(exc);

}

{

/// A rejected message belonging to an extended-query cycle is recovered at that

/// cycle's `Sync`, which is where its `ReadyForQuery` comes from. Without an open

/// cycle there is no `Sync` to wait for, so the client is owed one right away.

if (in_extended_query_cycle)

ignore_until_sync = true;

else

need_ready_for_query = true;

{

Int32 payload_size = 0;

Int32 info = 0;

establishSecureConnection(payload_size, info);

if (static_cast<PostgreSQLProtocol::Messaging::FrontMessageType>(info) == PostgreSQLProtocol::Messaging::FrontMessageType::CANCEL_REQUEST)

{

LOG_DEBUG(log, "Client issued request canceling");

cancelRequest();

return false;

}

std::unique_ptr<PostgreSQLProtocol::Messaging::StartupMessage> start_up_msg = receiveStartupMessage(payload_size);

/// An empty user name means the default session user: the `default_session_user`

/// server setting, possibly overridden for this listener in the `protocols` section.

/// If the resolved name is empty too (explicitly configured to prohibit connections

/// without a user name), authentication fails on the empty user name below.

if (start_up_msg->user.empty())

start_up_msg->user = default_session_user

? *default_session_user

: String(server.context()->getServerSettings()[ServerSetting::default_session_user]);

const auto & user_name = start_up_msg->user;

if (user_name.empty())

{

auto exception = Exception(ErrorCodes::AUTHENTICATION_FAILED, "Got an empty user name from PostgreSQL startup message");

session->onAuthenticationFailure(user_name, socket().peerAddress(), exception);

message_transport->send(

PostgreSQLProtocol::Messaging::ErrorOrNoticeResponse(

PostgreSQLProtocol::Messaging::ErrorOrNoticeResponse::ERROR, "28P01", "Invalid user or password"),

true);

return false;

}

authentication_manager.authenticate(user_name, *session, *message_transport, socket().peerAddress());

try

{

session->makeSessionContext();

session->sessionContext()->setDefaultFormat("PostgreSQLWire");

if (!start_up_msg->database.empty())

{

/// `database` is a real setting, so enforce its constraints on the startup-message

/// database too, consistently with `USE`, `SET database = ...` and the HTTP

/// `?database=...` parameter.

SettingsChanges database_change;

database_change.setSetting("database", start_up_msg->database);

session->sessionContext()->checkSettingsConstraints(database_change, SettingSource::QUERY);

session->sessionContext()->setCurrentDatabase(start_up_msg->database);

}

}

catch (const Exception & exc)

{

message_transport->send(

PostgreSQLProtocol::Messaging::ErrorOrNoticeResponse(

PostgreSQLProtocol::Messaging::ErrorOrNoticeResponse::ERROR, "XX000", exc.message()),

true);

throw;

}

sendParameterStatusData(*start_up_msg);

message_transport->send(

PostgreSQLProtocol::Messaging::BackendKeyData(connection_id, secret_key), true);

LOG_DEBUG(log, "Successfully finished Startup stage");

return true;

{

bool was_secure_connection = false;

bool was_encryption_req = true;

readBinaryBigEndian(payload_size, *in);

readBinaryBigEndian(info, *in);

switch (static_cast<PostgreSQLProtocol::Messaging::FrontMessageType>(info))

{

case PostgreSQLProtocol::Messaging::FrontMessageType::SSL_REQUEST:

LOG_DEBUG(log, "Client requested SSL");

if (ssl_enabled)

{

was_secure_connection = true;

makeSecureConnectionSSL();

}

else

message_transport->send('N', true);

break;

case PostgreSQLProtocol::Messaging::FrontMessageType::GSSENC_REQUEST:

LOG_DEBUG(log, "Client requested GSSENC");

message_transport->send('N', true);

break;

default:

was_encryption_req = false;

}

if (was_encryption_req)

{

readBinaryBigEndian(payload_size, *in);

readBinaryBigEndian(info, *in);

}

if (secure_required && !was_secure_connection)

{

message_transport->send(

PostgreSQLProtocol::Messaging::ErrorOrNoticeResponse(

PostgreSQLProtocol::Messaging::ErrorOrNoticeResponse::ERROR, "XX000", "SSL connection required."),

true);

throw Exception(ErrorCodes::OPENSSL_ERROR, "SSL connection required.");

}

{

message_transport->send('S', true);

Poco::Net::Context::Ptr ctx;

if (!params.privateKeyFile.empty() && !params.certificateFile.empty())

{

ctx = Poco::Net::SSLManager::instance().getCustomServerContext(prefix);

if (!ctx)

{

ctx = new Poco::Net::Context(usage, params);

ctx->disableProtocols(disabled_protocols);

ctx->enableExtendedCertificateVerification(extended_verification);

if (prefer_server_ciphers)

ctx->preferServerCiphers();

CertificateReloader::instance().tryLoad(config, ctx->sslContext(), prefix);

ctx = Poco::Net::SSLManager::instance().setCustomServerContext(prefix, ctx);

}

}

else

{

ctx = Poco::Net::SSLManager::instance().defaultServerContext();

}

ss = std::make_shared<Poco::Net::SecureStreamSocket>(Poco::Net::SecureStreamSocket::attach(socket(), ctx));

changeIO(*ss);

{

auto & parameters = start_up_message.parameters;

if (parameters.contains("application_name"))

message_transport->send(PostgreSQLProtocol::Messaging::ParameterStatus("application_name", parameters["application_name"]));

if (parameters.contains("client_encoding"))

message_transport->send(PostgreSQLProtocol::Messaging::ParameterStatus("client_encoding", parameters["client_encoding"]));

else

message_transport->send(PostgreSQLProtocol::Messaging::ParameterStatus("client_encoding", "UTF8"));

message_transport->send(PostgreSQLProtocol::Messaging::ParameterStatus("server_version", VERSION_STRING));

message_transport->send(PostgreSQLProtocol::Messaging::ParameterStatus("server_encoding", "UTF8"));

message_transport->send(PostgreSQLProtocol::Messaging::ParameterStatus("DateStyle", "ISO"));

message_transport->flush();

{

/// The random component is a token of its own and never the secret from `BackendKeyData`:

/// `system.processes` and `system.query_log` expose query IDs verbatim, while the secret

/// authenticates `CancelRequest`. It still has to be here, because a query ID that another

/// interface can predict can be occupied to keep a PostgreSQL statement from starting.

return fmt::format("postgres:{:d}:{:d}", connection_id_, query_id_token_);

{

/// One statement, one query ID: a query ID may be held by only one query at a time across the

/// whole server, so an ID that outlived its statement would keep the next one from starting.

query_id_token = generateRandomUInt32();

const String query_id = currentQueryId();

query_context->setCurrentQueryId(query_id);

/// `CancelRequest` names the connection, so its entry has to follow the current statement.

server.context()->getProcessList().registerPostgreSQLCancellationKey(connection_id, secret_key, query_id);

{

std::unique_ptr<PostgreSQLProtocol::Messaging::CancelRequest> msg =

message_transport->receiveWithPayloadSize<PostgreSQLProtocol::Messaging::CancelRequest>(8);

/// The process ID and secret key authenticate this otherwise unauthenticated request.

/// PostgreSQL exposes no response, so report the outcome only to the log.

CancellationCode code = server.context()->getProcessList().sendCancelToPostgreSQLQuery(msg->process_id, msg->secret_key);

LOG_DEBUG(log, "Cancellation request for connection {}: {}", msg->process_id,

code == CancellationCode::CancelSent ? "sent" : "not sent");

{

/// The declared size is read from the wire before any authentication, and the message is read

/// into memory in full, so it has to be bounded. PostgreSQL uses the same limit.

static constexpr Int32 max_startup_message_size = 10000;

std::unique_ptr<PostgreSQLProtocol::Messaging::StartupMessage> message;

try

{

if (payload_size < 8 || payload_size > max_startup_message_size)

throw Exception(ErrorCodes::UNKNOWN_PACKET_FROM_CLIENT,

"Startup message declares a size of {} bytes, while it must be between 8 and {} bytes",

payload_size, max_startup_message_size);

message = message_transport->receiveWithPayloadSize<PostgreSQLProtocol::Messaging::StartupMessage>(payload_size - 8);

}

catch (const Exception &)

{

message_transport->send(

PostgreSQLProtocol::Messaging::ErrorOrNoticeResponse(

PostgreSQLProtocol::Messaging::ErrorOrNoticeResponse::ERROR, "08P01", "Can't correctly handle Startup message"),

true);

throw;

}

LOG_DEBUG(log, "Successfully received Startup message");

return message;

{

static constexpr std::string_view pg_catalog = "pg_catalog";

/// A fast path for the common case of a query that does not mention the schema at all.

if (std::search(query.begin(), query.end(), pg_catalog.begin(), pg_catalog.end(),

[](char a, char b) { return equalsCaseInsensitive(a, b); }) == query.end())

return query;

std::vector<Token> tokens;

Lexer lexer(query.data(), query.data() + query.size());

for (Token token = lexer.nextToken(); !token.isEnd(); token = lexer.nextToken())

tokens.push_back(token);

auto is_pg_catalog = [](const Token & token)

{

std::string_view text(token.begin, token.size());

return (token.type == TokenType::BareWord && equalsCaseInsensitive(text, pg_catalog))

|| (token.type == TokenType::QuotedIdentifier && text == "\"pg_catalog\"");

};

auto next_significant = [&](size_t i) -> std::optional<size_t>

{

for (size_t j = i + 1; j < tokens.size(); ++j)

if (tokens[j].isSignificant())

return j;

return std::nullopt;

};

String result;

result.reserve(query.size());

std::optional<size_t> prev_emitted_significant;

for (size_t i = 0; i < tokens.size(); ++i)

{

const Token & token = tokens[i];

if (is_pg_catalog(token)

&& (!prev_emitted_significant || tokens[*prev_emitted_significant].type != TokenType::Dot))

{

auto dot = next_significant(i);

if (dot && tokens[*dot].type == TokenType::Dot)

{

auto after_dot = next_significant(*dot);

if (after_dot

&& (tokens[*after_dot].type == TokenType::BareWord || tokens[*after_dot].type == TokenType::QuotedIdentifier))

{

/// Skip the qualifier and the dot (and anything insignificant in between).

i = *dot;

continue;

}

}

}

result.append(token.begin, token.end);

if (token.isSignificant())

prev_emitted_significant = i;

}

return result;

{

ParserCopyQuery parser_copy;

ASTPtr copy_query_parsed;

try

{

copy_query_parsed = parseQuery(parser_copy, query, 0, DBMS_DEFAULT_MAX_PARSER_DEPTH, DBMS_DEFAULT_MAX_PARSER_BACKTRACKS);

}

catch (const Exception &)

{

copy_query_parsed.reset();

}

/* The Postgres protocol for a copy query is different from simple queries such as SELECT.

if (copy_query_parsed && copy_query_parsed->as<ASTCopyQuery>()->type == ASTCopyQuery::QueryType::COPY_FROM)

{

auto * copy_query = copy_query_parsed->as<ASTCopyQuery>();

auto query_context = session->makeQueryContext();

assignStatementQueryId(query_context);

QueryScope query_scope = QueryScope::create(query_context);

String columns_to_insert;

if (!copy_query->column_names.empty())

{

for (const auto & column_name : copy_query->column_names)

columns_to_insert += fmt::format("{}, ", column_name);

columns_to_insert.pop_back();

columns_to_insert.pop_back();

columns_to_insert = "(" + columns_to_insert + ")";

}

/// The parser has already quoted each part of `table_name`.

auto [ast, io] = executeQuery(fmt::format("INSERT INTO {} {} FROM INFILE 'psql_copy'", copy_query->table_name, columns_to_insert), query_context, {}, QueryProcessingStage::Enum::Complete);

chassert(io.pipeline.pushing());

auto executor = std::make_unique<PushingPipelineExecutor>(io.pipeline);

String format;

switch (copy_query->format)

{

case ASTCopyQuery::Formats::TSV:

format = "TSV";

break;

case ASTCopyQuery::Formats::CSV:

format = "CSV";

break;

case ASTCopyQuery::Formats::Binary:

format = "RowBinary";

break;

}

const Settings & settings = query_context->getSettingsRef();

message_transport->send(PostgreSQLProtocol::Messaging::CopyInResponse(), true);

executor->start();

while (true)

{

message_transport->flush();

PostgreSQLProtocol::Messaging::FrontMessageType message_type = message_transport->receiveMessageType();

if (message_type == PostgreSQLProtocol::Messaging::FrontMessageType::COPY_DATA)

{

std::unique_ptr<PostgreSQLProtocol::Messaging::CopyInData> data_query =

message_transport->receive<PostgreSQLProtocol::Messaging::CopyInData>();

ReadBufferFromString buf(data_query->query);

auto format_ptr = FormatFactory::instance().getInput(

format,

buf,

io.pipeline.getHeader(),

query_context,

settings[Setting::max_insert_block_size],

std::nullopt,

nullptr,

nullptr,

false,

CompressionMethod::None,

false,

settings[Setting::max_insert_block_size_bytes],

settings[Setting::min_insert_block_size_rows],

settings[Setting::min_insert_block_size_bytes]);

while (true)

{

auto chunk = format_ptr->generate();

if (chunk.empty())

break;

executor->push(std::move(chunk));

}

}

else if (message_type == PostgreSQLProtocol::Messaging::FrontMessageType::COPY_COMPLETION)

{

message_transport->receive<PostgreSQLProtocol::Messaging::CopyDone>();

executor->finish();

break;

}

else

{

executor->cancel();

throw Exception(ErrorCodes::BAD_ARGUMENTS, "Received incorrect message type - expected {} or {}, got {}", PostgreSQLProtocol::Messaging::FrontMessageType::COPY_DATA, PostgreSQLProtocol::Messaging::FrontMessageType::COPY_COMPLETION, message_type);

}

}

auto command = PostgreSQLProtocol::Messaging::CommandComplete::Command::COPY;

message_transport->send(PostgreSQLProtocol::Messaging::CommandComplete(command, 0), true);

return true;

}

/* In the case of a COPY TO request, the server calculates the number of columns and then sends it to the client in CopyOutResponse.

if (copy_query_parsed && copy_query_parsed->as<ASTCopyQuery>()->type == ASTCopyQuery::QueryType::COPY_TO)

{

auto * copy_query = copy_query_parsed->as<ASTCopyQuery>();

auto query_context = session->makeQueryContext();

assignStatementQueryId(query_context);

QueryScope query_scope = QueryScope::create(query_context);

String columns_to_select = "*";

if (!copy_query->column_names.empty())

{

columns_to_select.clear();

for (const auto & column_name : copy_query->column_names)

columns_to_select += fmt::format("{}, ", column_name);

columns_to_select.pop_back();

columns_to_select.pop_back();

}

auto select_query = fmt::format("SELECT {} FROM {};", columns_to_select, copy_query->table_name);

auto [ast, io] = executeQuery(select_query, query_context, {}, QueryProcessingStage::Enum::Complete);

chassert(io.pipeline.pulling());

message_transport->send(PostgreSQLProtocol::Messaging::CopyOutResponse(static_cast<Int32>(io.pipeline.getHeader().columns())));

VectorWithMemoryTracking<char> result_buf;

WriteBufferFromVectorImpl<decltype(result_buf)> output_buffer(result_buf);

auto format_ptr = FormatFactory::instance().getOutputFormat(toString(copy_query->format), output_buffer, io.pipeline.getHeader(), query_context);

auto executor = std::make_unique<PullingPipelineExecutor>(io.pipeline);

Block block;

while (executor->pull(block))

{

output_buffer.restart(DBMS_DEFAULT_BUFFER_SIZE); // This will recreate moved vector

format_ptr->write(materializeBlock(block));

format_ptr->flush();

output_buffer.finalize();

message_transport->send(PostgreSQLProtocol::Messaging::CopyOutData(result_buf));

result_buf.clear();

}

message_transport->send(PostgreSQLProtocol::Messaging::CopyCompletionResponse(), true);

return true;

}

return false;

{

/// Output position before the currently executing statement. If a statement

/// fails when nothing has been sent for it yet, the session can be kept alive.

size_t out_bytes_before_statement = out->count();

try

{

std::unique_ptr<PostgreSQLProtocol::Messaging::Query> query =

message_transport->receive<PostgreSQLProtocol::Messaging::Query>();

if (isEmptyQuery(query->query))

{

message_transport->send(PostgreSQLProtocol::Messaging::EmptyQueryResponse());

return;

}

bool psycopg2_cond = query->query == "BEGIN" || query->query == "COMMIT"; // psycopg2 starts and ends queries with BEGIN/COMMIT commands

bool jdbc_cond = query->query.contains("SET extra_float_digits") || query->query.contains("SET application_name"); // jdbc starts with setting this parameter