GRPCServer.cpp

#include <Server/GRPCServer.h>
#include <limits>
#include <memory>
#include <Poco/Net/SocketAddress.h>
#if USE_GRPC

#include <Columns/ColumnString.h>
#include <Columns/ColumnsNumber.h>
#include <Common/CurrentThread.h>
#include <Common/QueryScope.h>
#include <Common/DateLUTImpl.h>
#include <Common/SettingsChanges.h>
#include <Common/setThreadName.h>
#include <Common/Stopwatch.h>
#include <Common/ThreadPool.h>
#include <Core/Settings.h>
#include <DataTypes/DataTypeFactory.h>
#include <QueryPipeline/ProfileInfo.h>
#include <Interpreters/Context.h>
#include <Interpreters/DatabaseCatalog.h>
#include <Interpreters/InternalTextLogsQueue.h>
#include <Interpreters/executeQuery.h>
#include <Interpreters/Session.h>
#include <IO/CompressionMethod.h>
#include <IO/ConcatReadBuffer.h>
#include <IO/ReadBufferFromString.h>
#include <IO/ReadHelpers.h>
#include <Parsers/parseQuery.h>
#include <Parsers/ASTIdentifier_fwd.h>
#include <Parsers/ASTInsertQuery.h>
#include <Parsers/ASTQueryWithOutput.h>
#include <Parsers/ParserQuery.h>
#include <Processors/Executors/PullingAsyncPipelineExecutor.h>
#include <Processors/Executors/PullingPipelineExecutor.h>
#include <Processors/Executors/PushingPipelineExecutor.h>
#include <Processors/Executors/CompletedPipelineExecutor.h>
#include <Processors/Executors/PipelineExecutor.h>
#include <Processors/Formats/IInputFormat.h>
#include <Processors/Formats/IOutputFormat.h>
#include <Processors/Sinks/SinkToStorage.h>
#include <Processors/Sinks/EmptySink.h>
#include <QueryPipeline/QueryPipelineBuilder.h>
#include <Server/IServer.h>
#include <Storages/IStorage.h>
#include <Poco/FileStream.h>
#include <Poco/StreamCopier.h>
#include <Poco/Util/LayeredConfiguration.h>
#include <base/range.h>
#include <Common/logger_useful.h>

#include <absl/base/log_severity.h>
#include <absl/log/globals.h>
#include <absl/log/initialize.h>
#include <absl/log/log_sink_registry.h>

#include <grpc/support/log.h>
#include <grpc++/security/server_credentials.h>
#include <grpc++/server.h>
#include <grpc++/server_builder.h>


using GRPCService = clickhouse::grpc::ClickHouse::AsyncService;
using GRPCQueryInfo = clickhouse::grpc::QueryInfo;
using GRPCResult = clickhouse::grpc::Result;
using GRPCException = clickhouse::grpc::Exception;
using GRPCProgress = clickhouse::grpc::Progress;
using GRPCObsoleteTransportCompression = clickhouse::grpc::ObsoleteTransportCompression;

namespace DB
{
namespace Setting
{
    extern const SettingsBool allow_settings_after_format_in_insert;
    extern const SettingsBool calculate_text_stack_trace;
    extern const SettingsUInt64 interactive_delay;
    extern const SettingsLogsLevel send_logs_level;
    extern const SettingsString send_logs_source_regexp;
    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 SettingsBool throw_if_no_data_to_insert;
    extern const SettingsBool use_concurrency_control;
}

namespace ErrorCodes
{
    extern const int INVALID_CONFIG_PARAMETER;
    extern const int INVALID_GRPC_QUERY_INFO;
    extern const int INVALID_SESSION_TIMEOUT;
    extern const int LOGICAL_ERROR;
    extern const int NETWORK_ERROR;
    extern const int NO_DATA_TO_INSERT;
    extern const int SUPPORT_IS_DISABLED;
    extern const int BAD_REQUEST_PARAMETER;
}

namespace
{
    /// Make grpc to pass logging messages to ClickHouse logging system.
    class GrpcLogSink : public absl::LogSink
    {
    public:
        void Send(const absl::LogEntry & entry) override
        {
            static LoggerRawPtr logger = getRawLogger("grpc");

            const auto msg = std::string(entry.text_message());
            const auto file = entry.source_filename();
            int line = entry.source_line();

            switch (entry.log_severity())
            {
                case absl::LogSeverity::kInfo:
                    LOG_INFO(logger, "{} ({}:{})", msg, file, line);
                    break;
                case absl::LogSeverity::kWarning:
                    LOG_WARNING(logger, "{} ({}:{})", msg, file, line);
                    break;
                case absl::LogSeverity::kError:
                    LOG_ERROR(logger, "{} ({}:{})", msg, file, line);
                    break;
                case absl::LogSeverity::kFatal:
                    LOG_ERROR(logger, "FATAL: {} ({}:{})", msg, file, line);
                    break;
            }
        }
    };
    GrpcLogSink grpc_log_sink;

    /// See also contrib/grpc/src/core/util/log.cc
    void initGRPCLogging(const Poco::Util::AbstractConfiguration & config)
    {
        static std::once_flag once_flag;
        std::call_once(once_flag, [&config]
        {
            absl::AddLogSink(&grpc_log_sink);
            absl::SetStderrThreshold(absl::LogSeverityAtLeast::kInfinity);
            absl::InitializeLog();

            absl::SetStderrThreshold(absl::LogSeverityAtLeast::kInfinity);

            const bool verbose = config.getBool("grpc.verbose_logs", false);
            static LoggerRawPtr logger = getRawLogger("grpc");

            if (verbose)
                grpc_tracer_set_enabled("all", true);

            if (logger->is(Poco::Message::PRIO_DEBUG))
            {
                absl::SetMinLogLevel(absl::LogSeverityAtLeast::kInfo);
                absl::SetVLogLevel("*grpc*/*", 2);
            }
            else if (logger->is(Poco::Message::PRIO_INFORMATION))
            {
                absl::SetMinLogLevel(absl::LogSeverityAtLeast::kInfo);
                absl::SetVLogLevel("*grpc*/*", -1);
            }
            else
            {
                absl::SetMinLogLevel(absl::LogSeverityAtLeast::kWarning);
                absl::SetVLogLevel("*grpc*/*", -1);
            }
        });
    }

    /// Gets file's contents as a string, throws an exception if failed.
    String readFile(const String & filepath)
    {
        Poco::FileInputStream ifs(filepath);
        String res;
        Poco::StreamCopier::copyToString(ifs, res);
        return res;
    }

    /// Makes credentials based on the server config.
    std::shared_ptr<grpc::ServerCredentials> makeCredentials(const Poco::Util::AbstractConfiguration & config)
    {
        if (config.getBool("grpc.enable_ssl", false))
        {
#if USE_SSL
            grpc::SslServerCredentialsOptions options;
            grpc::SslServerCredentialsOptions::PemKeyCertPair key_cert_pair;
            key_cert_pair.private_key = readFile(config.getString("grpc.ssl_key_file"));
            key_cert_pair.cert_chain = readFile(config.getString("grpc.ssl_cert_file"));
            options.pem_key_cert_pairs.emplace_back(std::move(key_cert_pair));
            if (config.getBool("grpc.ssl_require_client_auth", false))
            {
                options.client_certificate_request = GRPC_SSL_REQUEST_AND_REQUIRE_CLIENT_CERTIFICATE_AND_VERIFY;
                if (config.has("grpc.ssl_ca_cert_file"))
                    options.pem_root_certs = readFile(config.getString("grpc.ssl_ca_cert_file"));
            }
            return grpc::SslServerCredentials(options);
#else
            throw DB::Exception(DB::ErrorCodes::SUPPORT_IS_DISABLED, "Can't use SSL in grpc, because ClickHouse was built without SSL library");
#endif
        }
        return grpc::InsecureServerCredentials();
    }

    /// Transport compression makes gRPC library to compress packed Result messages before sending them through network.
    struct TransportCompression
    {
        grpc_compression_algorithm algorithm;
        grpc_compression_level level;

        /// Extracts the settings of transport compression from a query info if possible.
        static std::optional<TransportCompression> fromQueryInfo(const GRPCQueryInfo & query_info)
        {
            TransportCompression res{};
            if (!query_info.transport_compression_type().empty())
            {
                res.setAlgorithm(query_info.transport_compression_type(), ErrorCodes::INVALID_GRPC_QUERY_INFO);
                res.setLevel(query_info.transport_compression_level(), ErrorCodes::INVALID_GRPC_QUERY_INFO);
                return res;
            }

            if (query_info.has_obsolete_result_compression())
            {
                switch (query_info.obsolete_result_compression().algorithm())
                {
                    case GRPCObsoleteTransportCompression::NO_COMPRESSION: res.algorithm = GRPC_COMPRESS_NONE; break;
                    case GRPCObsoleteTransportCompression::DEFLATE: res.algorithm = GRPC_COMPRESS_DEFLATE; break;
                    case GRPCObsoleteTransportCompression::GZIP: res.algorithm = GRPC_COMPRESS_GZIP; break;
                    case GRPCObsoleteTransportCompression::STREAM_GZIP: throw Exception(ErrorCodes::INVALID_GRPC_QUERY_INFO, "STREAM_GZIP is no longer supported"); /// was flagged experimental in gRPC, removed as per v1.44
                    default: throw Exception(ErrorCodes::INVALID_GRPC_QUERY_INFO, "Unknown compression algorithm: {}", GRPCObsoleteTransportCompression::CompressionAlgorithm_Name(query_info.obsolete_result_compression().algorithm()));
                }

                switch (query_info.obsolete_result_compression().level())
                {
                    case GRPCObsoleteTransportCompression::COMPRESSION_NONE: res.level = GRPC_COMPRESS_LEVEL_NONE; break;
                    case GRPCObsoleteTransportCompression::COMPRESSION_LOW: res.level = GRPC_COMPRESS_LEVEL_LOW; break;
                    case GRPCObsoleteTransportCompression::COMPRESSION_MEDIUM: res.level = GRPC_COMPRESS_LEVEL_MED; break;
                    case GRPCObsoleteTransportCompression::COMPRESSION_HIGH: res.level = GRPC_COMPRESS_LEVEL_HIGH; break;
                    default: throw Exception(ErrorCodes::INVALID_GRPC_QUERY_INFO, "Unknown compression level: {}", GRPCObsoleteTransportCompression::CompressionLevel_Name(query_info.obsolete_result_compression().level()));
                }
                return res;
            }

            return std::nullopt;
        }

        /// Extracts the settings of transport compression from the server configuration.
        static TransportCompression fromConfiguration(const Poco::Util::AbstractConfiguration & config)
        {
            TransportCompression res{};
            if (config.has("grpc.transport_compression_type"))
            {
                res.setAlgorithm(config.getString("grpc.transport_compression_type"), ErrorCodes::INVALID_CONFIG_PARAMETER);
                res.setLevel(config.getInt("grpc.transport_compression_level", 0), ErrorCodes::INVALID_CONFIG_PARAMETER);
            }
            else
            {
                res.setAlgorithm(config.getString("grpc.compression", "none"), ErrorCodes::INVALID_CONFIG_PARAMETER);
                res.setLevel(config.getString("grpc.compression_level", "none"), ErrorCodes::INVALID_CONFIG_PARAMETER);
            }
            return res;
        }

    private:
        void setAlgorithm(const String & str, int error_code)
        {
            if (str == "none")
                algorithm = GRPC_COMPRESS_NONE;
            else if (str == "deflate")
                algorithm = GRPC_COMPRESS_DEFLATE;
            else if (str == "gzip")
                algorithm = GRPC_COMPRESS_GZIP;
            else if (str == "stream_gzip")
                throw Exception(ErrorCodes::INVALID_GRPC_QUERY_INFO, "STREAM_GZIP is no longer supported"); /// was flagged experimental in gRPC, removed as per v1.44
            else
                throw Exception(error_code, "Unknown compression algorithm: '{}'", str);
        }

        void setLevel(const String & str, int error_code)
        {
            if (str == "none")
                level = GRPC_COMPRESS_LEVEL_NONE;
            else if (str == "low")
                level = GRPC_COMPRESS_LEVEL_LOW;
            else if (str == "medium")
                level = GRPC_COMPRESS_LEVEL_MED;
            else if (str == "high")
                level = GRPC_COMPRESS_LEVEL_HIGH;
            else
                throw Exception(error_code, "Unknown compression level: '{}'", str);
        }

        void setLevel(int level_, int error_code)
        {
            if (0 <= level_ && level_ < GRPC_COMPRESS_LEVEL_COUNT)
                level = static_cast<grpc_compression_level>(level_);
            else
                throw Exception(error_code, "Compression level {} is out of range 0..{}", level_, GRPC_COMPRESS_LEVEL_COUNT - 1);
        }
    };

    /// Gets session's timeout from query info or from the server config.
    std::chrono::steady_clock::duration getSessionTimeout(const GRPCQueryInfo & query_info, const Poco::Util::AbstractConfiguration & config)
    {
        auto session_timeout = query_info.session_timeout();
        if (session_timeout)
        {
            auto max_session_timeout = config.getUInt("max_session_timeout", 3600);
            if (session_timeout > max_session_timeout)
                throw Exception(ErrorCodes::INVALID_SESSION_TIMEOUT, "Session timeout '{}' is larger than max_session_timeout: {}. "
                    "Maximum session timeout could be modified in configuration file.",
                    std::to_string(session_timeout), std::to_string(max_session_timeout));
        }
        else
            session_timeout = config.getInt("default_session_timeout", 60);
        return std::chrono::seconds(session_timeout);
    }

    /// Generates a description of a query by a specified query info.
    /// This description is used for logging only.
    String getQueryDescription(const GRPCQueryInfo & query_info)
    {
        String str;
        if (!query_info.query().empty())
        {
            std::string_view query = query_info.query();
            constexpr size_t max_query_length_to_log = 64;
            if (query.length() > max_query_length_to_log)
                query.remove_suffix(query.length() - max_query_length_to_log);
            if (size_t format_pos = query.find(" FORMAT "); format_pos != String::npos)
                query.remove_suffix(query.length() - format_pos - strlen(" FORMAT "));
            str.append("\"").append(query);
            if (query != query_info.query())
                str.append("...");
            str.append("\"");
        }
        if (!query_info.query_id().empty())
            str.append(str.empty() ? "" : ", ").append("query_id: ").append(query_info.query_id());
        if (!query_info.input_data().empty())
            str.append(str.empty() ? "" : ", ").append("input_data: ").append(std::to_string(query_info.input_data().size())).append(" bytes");
        if (query_info.external_tables_size())
            str.append(str.empty() ? "" : ", ").append("external tables: ").append(std::to_string(query_info.external_tables_size()));
        return str;
    }

    /// Generates a description of a result.
    /// This description is used for logging only.
    String getResultDescription(const GRPCResult & result)
    {
        String str;
        if (!result.output().empty())
            str.append("output: ").append(std::to_string(result.output().size())).append(" bytes");
        if (!result.totals().empty())
            str.append(str.empty() ? "" : ", ").append("totals");
        if (!result.extremes().empty())
            str.append(str.empty() ? "" : ", ").append("extremes");
        if (result.has_progress())
            str.append(str.empty() ? "" : ", ").append("progress");
        if (result.logs_size())
            str.append(str.empty() ? "" : ", ").append("logs: ").append(std::to_string(result.logs_size())).append(" entries");
        if (result.cancelled())
            str.append(str.empty() ? "" : ", ").append("cancelled");
        if (result.has_exception())
            str.append(str.empty() ? "" : ", ").append("exception");
        return str;
    }

    using CompletionCallback = std::function<void(bool)>;

    /// Requests a connection and provides low-level interface for reading and writing.
    class BaseResponder
    {
    public:
        virtual ~BaseResponder() = default;

        virtual void start(GRPCService & grpc_service,
                           grpc::ServerCompletionQueue & new_call_queue,
                           grpc::ServerCompletionQueue & notification_queue,
                           const CompletionCallback & callback) = 0;

        virtual void read(GRPCQueryInfo & query_info_, const CompletionCallback & callback) = 0;
        virtual void write(const GRPCResult & result, const CompletionCallback & callback) = 0;
        virtual void writeAndFinish(const GRPCResult & result, const grpc::Status & status, const CompletionCallback & callback) = 0;

        Poco::Net::SocketAddress getClientAddress() const
        {
            /// Returns a string like ipv4:127.0.0.1:55930 or ipv6:%5B::1%5D:55930
            String uri_encoded_peer = grpc_context.peer();

            constexpr const std::string_view ipv4_prefix = "ipv4:";
            constexpr const std::string_view ipv6_prefix = "ipv6:";

            bool ipv4 = uri_encoded_peer.starts_with(ipv4_prefix);
            bool ipv6 = uri_encoded_peer.starts_with(ipv6_prefix);

            if (!ipv4 && !ipv6)
                throw Exception(ErrorCodes::LOGICAL_ERROR, "Expected ipv4 or ipv6 protocol in peer address, got {}", uri_encoded_peer);

            auto prefix = ipv4 ? ipv4_prefix : ipv6_prefix;
            auto family = ipv4 ? Poco::Net::AddressFamily::Family::IPv4 : Poco::Net::AddressFamily::Family::IPv6;

            uri_encoded_peer= uri_encoded_peer.substr(prefix.length());

            String peer;
            Poco::URI::decode(uri_encoded_peer, peer);

            return Poco::Net::SocketAddress{family, peer};
        }

        std::optional<String> getClientHeader(const String & key) const
        {
            const auto & client_metadata = grpc_context.client_metadata();
            auto it = client_metadata.find(key);
            if (it != client_metadata.end())
                return String{it->second.data(), it->second.size()};
            return std::nullopt;
        }

        void setTransportCompression(const TransportCompression & transport_compression)
        {
            grpc_context.set_compression_algorithm(transport_compression.algorithm);
            grpc_context.set_compression_level(transport_compression.level);
        }

    protected:
        CompletionCallback * getCallbackPtr(const CompletionCallback & callback)
        {
            /// It would be better to pass callbacks to gRPC calls.
            /// However gRPC calls can be tagged with `void *` tags only.
            /// The map `callbacks` here is used to keep callbacks until they're called.
            std::lock_guard lock{mutex};
            size_t callback_id = next_callback_id++;
            auto & callback_in_map = callbacks[callback_id];
            callback_in_map = [this, callback, callback_id](bool ok)
            {
                CompletionCallback callback_to_call;
                {
                    std::lock_guard lock2{mutex};
                    callback_to_call = callback;
                    callbacks.erase(callback_id);
                }
                callback_to_call(ok);
            };
            return &callback_in_map;
        }

        grpc::ServerContext grpc_context;

    private:
        grpc::ServerAsyncReaderWriter<GRPCResult, GRPCQueryInfo> reader_writer{&grpc_context};
        std::unordered_map<size_t, CompletionCallback> callbacks;
        size_t next_callback_id = 0;
        std::mutex mutex;
    };

    enum CallType
    {
        CALL_SIMPLE,             /// ExecuteQuery() call
        CALL_WITH_STREAM_INPUT,  /// ExecuteQueryWithStreamInput() call
        CALL_WITH_STREAM_OUTPUT, /// ExecuteQueryWithStreamOutput() call
        CALL_WITH_STREAM_IO,     /// ExecuteQueryWithStreamIO() call
        CALL_MAX,
    };

    const char * getCallName(CallType call_type)
    {
        switch (call_type)
        {
            case CALL_SIMPLE: return "ExecuteQuery()";
            case CALL_WITH_STREAM_INPUT: return "ExecuteQueryWithStreamInput()";
            case CALL_WITH_STREAM_OUTPUT: return "ExecuteQueryWithStreamOutput()";
            case CALL_WITH_STREAM_IO: return "ExecuteQueryWithStreamIO()";
            case CALL_MAX: break;
        }
        UNREACHABLE();
    }

    bool isInputStreaming(CallType call_type)
    {
        return (call_type == CALL_WITH_STREAM_INPUT) || (call_type == CALL_WITH_STREAM_IO);
    }

    bool isOutputStreaming(CallType call_type)
    {
        return (call_type == CALL_WITH_STREAM_OUTPUT) || (call_type == CALL_WITH_STREAM_IO);
    }

    template <enum CallType call_type>
    class Responder;

    template<>
    class Responder<CALL_SIMPLE> : public BaseResponder
    {
    public:
        void start(GRPCService & grpc_service,
                  grpc::ServerCompletionQueue & new_call_queue,
                  grpc::ServerCompletionQueue & notification_queue,
                  const CompletionCallback & callback) override
        {
            grpc_service.RequestExecuteQuery(&grpc_context, &query_info.emplace(), &response_writer, &new_call_queue, &notification_queue, getCallbackPtr(callback));
        }

        void read(GRPCQueryInfo & query_info_, const CompletionCallback & callback) override
        {
            if (!query_info.has_value())
            {
                callback(false);
                return;
            }

            query_info_ = std::move(query_info).value();
            query_info.reset();
            callback(true);
        }

        void write(const GRPCResult &, const CompletionCallback &) override
        {
            throw Exception(ErrorCodes::LOGICAL_ERROR, "Responder<CALL_SIMPLE>::write() should not be called");
        }

        void writeAndFinish(const GRPCResult & result, const grpc::Status & status, const CompletionCallback & callback) override
        {
            response_writer.Finish(result, status, getCallbackPtr(callback));
        }

    private:
        grpc::ServerAsyncResponseWriter<GRPCResult> response_writer{&grpc_context};
        std::optional<GRPCQueryInfo> query_info;
    };

    template<>
    class Responder<CALL_WITH_STREAM_INPUT> : public BaseResponder
    {
    public:
        void start(GRPCService & grpc_service,
                  grpc::ServerCompletionQueue & new_call_queue,
                  grpc::ServerCompletionQueue & notification_queue,
                  const CompletionCallback & callback) override
        {
            grpc_service.RequestExecuteQueryWithStreamInput(&grpc_context, &reader, &new_call_queue, &notification_queue, getCallbackPtr(callback));
        }

        void read(GRPCQueryInfo & query_info_, const CompletionCallback & callback) override
        {
            reader.Read(&query_info_, getCallbackPtr(callback));
        }

        void write(const GRPCResult &, const CompletionCallback &) override
        {
            throw Exception(ErrorCodes::LOGICAL_ERROR, "Responder<CALL_WITH_STREAM_INPUT>::write() should not be called");
        }

        void writeAndFinish(const GRPCResult & result, const grpc::Status & status, const CompletionCallback & callback) override
        {
            reader.Finish(result, status, getCallbackPtr(callback));
        }

    private:
        grpc::ServerAsyncReader<GRPCResult, GRPCQueryInfo> reader{&grpc_context};
    };

    template<>
    class Responder<CALL_WITH_STREAM_OUTPUT> : public BaseResponder
    {
    public:
        void start(GRPCService & grpc_service,
                  grpc::ServerCompletionQueue & new_call_queue,
                  grpc::ServerCompletionQueue & notification_queue,
                  const CompletionCallback & callback) override
        {
            grpc_service.RequestExecuteQueryWithStreamOutput(&grpc_context, &query_info.emplace(), &writer, &new_call_queue, &notification_queue, getCallbackPtr(callback));
        }

        void read(GRPCQueryInfo & query_info_, const CompletionCallback & callback) override
        {
            if (!query_info.has_value())
            {
                callback(false);
                return;
            }

            query_info_ = std::move(query_info).value();
            query_info.reset();
            callback(true);
        }

        void write(const GRPCResult & result, const CompletionCallback & callback) override
        {
            writer.Write(result, getCallbackPtr(callback));
        }

        void writeAndFinish(const GRPCResult & result, const grpc::Status & status, const CompletionCallback & callback) override
        {
            writer.WriteAndFinish(result, {}, status, getCallbackPtr(callback));
        }

    private:
        grpc::ServerAsyncWriter<GRPCResult> writer{&grpc_context};
        std::optional<GRPCQueryInfo> query_info;
    };

    template<>
    class Responder<CALL_WITH_STREAM_IO> : public BaseResponder
    {
    public:
        void start(GRPCService & grpc_service,
                  grpc::ServerCompletionQueue & new_call_queue,
                  grpc::ServerCompletionQueue & notification_queue,
                  const CompletionCallback & callback) override
        {
            grpc_service.RequestExecuteQueryWithStreamIO(&grpc_context, &reader_writer, &new_call_queue, &notification_queue, getCallbackPtr(callback));
        }

        void read(GRPCQueryInfo & query_info_, const CompletionCallback & callback) override
        {
            reader_writer.Read(&query_info_, getCallbackPtr(callback));
        }

        void write(const GRPCResult & result, const CompletionCallback & callback) override
        {
            reader_writer.Write(result, getCallbackPtr(callback));
        }

        void writeAndFinish(const GRPCResult & result, const grpc::Status & status, const CompletionCallback & callback) override
        {
            reader_writer.WriteAndFinish(result, {}, status, getCallbackPtr(callback));
        }

    private:
        grpc::ServerAsyncReaderWriter<GRPCResult, GRPCQueryInfo> reader_writer{&grpc_context};
    };

    std::unique_ptr<BaseResponder> makeResponder(CallType call_type)
    {
        switch (call_type)
        {
            case CALL_SIMPLE: return std::make_unique<Responder<CALL_SIMPLE>>();
            case CALL_WITH_STREAM_INPUT: return std::make_unique<Responder<CALL_WITH_STREAM_INPUT>>();
            case CALL_WITH_STREAM_OUTPUT: return std::make_unique<Responder<CALL_WITH_STREAM_OUTPUT>>();
            case CALL_WITH_STREAM_IO: return std::make_unique<Responder<CALL_WITH_STREAM_IO>>();
            case CALL_MAX: break;
        }
        UNREACHABLE();
    }


    /// Implementation of ReadBuffer, which just calls a callback.
    class ReadBufferFromCallback : public ReadBuffer
    {
    public:
        explicit ReadBufferFromCallback(const std::function<std::pair<const void *, size_t>(void)> & callback_)
            : ReadBuffer(nullptr, 0), callback(callback_) {}

    private:
        bool nextImpl() override
        {
            const void * new_pos = nullptr;
            size_t new_size = 0;
            std::tie(new_pos, new_size) = callback();
            if (!new_size)
                return false;
            BufferBase::set(static_cast<BufferBase::Position>(const_cast<char *>(static_cast<const char *>(new_pos))), new_size, 0);
            return true;
        }

        std::function<std::pair<const void *, size_t>(void)> callback;
    };


    /// A boolean state protected by mutex able to wait until other thread sets it to a specific value.
    class BoolState
    {
    public:
        explicit BoolState(bool initial_value) : value(initial_value) {}

        bool get() const
        {
            std::lock_guard lock{mutex};
            return value;
        }

        void set(bool new_value)
        {
            std::lock_guard lock{mutex};
            if (value == new_value)
                return;
            value = new_value;
            changed.notify_all();
        }

        void wait(bool wanted_value) const
        {
            std::unique_lock lock{mutex};
            changed.wait(lock, [this, wanted_value]() { return value == wanted_value; });
        }

    private:
        bool value;
        mutable std::mutex mutex;
        mutable std::condition_variable changed;
    };


    /// Handles a connection after a responder is started (i.e. after getting a new call).
    class Call // NOLINT(clang-analyzer-optin.performance.Padding)
    {
    public:
        Call(CallType call_type_, std::unique_ptr<BaseResponder> responder_, IServer & iserver_, LoggerRawPtr log_);
        ~Call();

        void start(const std::function<void(void)> & on_finish_call_callback);

    private:
        void run();

        void receiveQuery();
        void executeQuery();

        void processInput();
        void initializePipeline(const Block & header);
        void createExternalTables();

        void generateOutput();

        void finishQuery();
        void onException(const Exception & exception);
        void onFatalError();
        void releaseQueryIDAndSessionID();
        void close();

        void readQueryInfo();
        void throwIfFailedToReadQueryInfo();
        bool isQueryCancelled();

        void addQueryDetailsToResult();
        void addOutputFormatToResult();
        void addOutputColumnsNamesAndTypesToResult(const Block & headers);
        void addProgressToResult();
        void addTotalsToResult(const Block & totals);
        void addExtremesToResult(const Block & extremes);
        void addProfileInfoToResult(const ProfileInfo & info);
        void addLogsToResult();
        void sendResult();
        void throwIfFailedToSendResult();
        void sendException(const Exception & exception);

        const CallType call_type;
        std::unique_ptr<BaseResponder> responder;
        IServer & iserver;
        LoggerRawPtr log = nullptr;

        std::optional<Session> session;
        ContextMutablePtr query_context;
        std::optional<QueryScope> query_scope;
        OpenTelemetry::TracingContextHolderPtr thread_trace_context;
        String query_text;
        ASTPtr ast;
        ASTInsertQuery * insert_query = nullptr;
        String input_format;
        String input_data_delimiter;
        CompressionMethod input_compression_method = CompressionMethod::None;
        PODArray<char> output;
        String output_format;
        bool send_output_columns_names_and_types = false;
        CompressionMethod output_compression_method = CompressionMethod::None;
        int output_compression_level = 0;

        uint64_t interactive_delay = 100000;
        bool send_exception_with_stacktrace = true;
        bool input_function_is_used = false;

        BlockIO io;
        Progress progress;
        InternalTextLogsQueuePtr logs_queue;

        GRPCQueryInfo query_info; /// We reuse the same messages multiple times.
        GRPCResult result;

        bool initial_query_info_read = false;
        bool finalize = false;
        bool responder_finished = false;
        bool cancelled = false;

        std::unique_ptr<ReadBuffer> read_buffer;
        std::unique_ptr<WriteBuffer> write_buffer;
        WriteBufferFromVector<PODArray<char>> * nested_write_buffer = nullptr;
        WriteBuffer * compressing_write_buffer = nullptr;
        std::unique_ptr<QueryPipeline> pipeline;
        std::unique_ptr<PullingPipelineExecutor> pipeline_executor;
        std::shared_ptr<IOutputFormat> output_format_processor;
        bool need_input_data_from_insert_query = true;
        bool need_input_data_from_query_info = true;
        bool need_input_data_delimiter = false;

        Stopwatch query_time;
        UInt64 waited_for_client_reading = 0;
        UInt64 waited_for_client_writing = 0;

        /// The following fields are accessed both from call_thread and queue_thread.
        BoolState reading_query_info{false};
        std::atomic<bool> failed_to_read_query_info = false;
        GRPCQueryInfo next_query_info_while_reading;
        std::atomic<bool> want_to_cancel = false;
        std::atomic<bool> check_query_info_contains_cancel_only = false;
        BoolState sending_result{false};
        std::atomic<bool> failed_to_send_result = false;

        ThreadFromGlobalPool call_thread;
    };

    Call::Call(CallType call_type_, std::unique_ptr<BaseResponder> responder_, IServer & iserver_, LoggerRawPtr log_)
        : call_type(call_type_), responder(std::move(responder_)), iserver(iserver_), log(log_)
    {
    }

    Call::~Call()
    {
        if (call_thread.joinable())
            call_thread.join();
    }

    void Call::start(const std::function<void(void)> & on_finish_call_callback)
    {
        auto runner_function = [this, on_finish_call_callback]
        {
            try
            {
                run();
            }
            catch (...)
            {
                tryLogCurrentException("GRPCServer");
            }
            on_finish_call_callback();
        };
        call_thread = ThreadFromGlobalPool(runner_function);
    }

    void Call::run()
    {
        try
        {
            DB::setThreadName(ThreadName::GRPC_SERVER_CALL);
            receiveQuery();
            executeQuery();
            processInput();
            generateOutput();
            finishQuery();
        }
        catch (Exception & exception)
        {
            onException(exception);
        }
        catch (Poco::Exception & exception)
        {
            onException(Exception{Exception::CreateFromPocoTag{}, exception});
        }
        catch (std::exception & exception)
        {
            onException(Exception{Exception::CreateFromSTDTag{}, exception});
        }
    }

    void Call::receiveQuery()
    {
        LOG_INFO(log, "Handling call {}", getCallName(call_type));

        readQueryInfo();

        if (query_info.cancel())
            throw Exception(ErrorCodes::INVALID_GRPC_QUERY_INFO, "Initial query info cannot set the 'cancel' field");

        LOG_DEBUG(log, "Received initial QueryInfo: {}", getQueryDescription(query_info));
    }

    void Call::executeQuery()
    {
        /// Retrieve user credentials.
        std::string user = query_info.user_name();
        std::string password = query_info.password();
        std::string quota_key = query_info.quota();
        Poco::Net::SocketAddress user_address = responder->getClientAddress();

        if (user.empty())
        {
            user = "default";
            password = "";
        }

        /// Authentication.
        session.emplace(iserver.context(), ClientInfo::Interface::GRPC);
        session->authenticate(user, password, user_address);
        session->setQuotaClientKey(quota_key);

        ClientInfo client_info = session->getClientInfo();

        /// Parse the OpenTelemetry traceparent header.
        auto traceparent = responder->getClientHeader("traceparent");
        if (traceparent)
        {
            String error;
            if (!client_info.client_trace_context.parseTraceparentHeader(traceparent.value(), error))
            {
                throw Exception(ErrorCodes::BAD_REQUEST_PARAMETER,
                    "Failed to parse OpenTelemetry traceparent header '{}': {}",
                    traceparent.value(), error);
            }
            auto tracestate = responder->getClientHeader("tracestate");
            client_info.client_trace_context.tracestate = tracestate.value_or("");
        }

        /// The user could specify session identifier and session timeout.
        /// It allows to modify settings, create temporary tables and reuse them in subsequent requests.
        if (!query_info.session_id().empty())
        {
            session->makeSessionContext(
                query_info.session_id(), getSessionTimeout(query_info, iserver.config()), query_info.session_check());
        }

        query_context = session->makeQueryContext(std::move(client_info));

        /// Prepare settings.
        SettingsChanges settings_changes;
        for (const auto & [key, value] : query_info.settings())
        {
            settings_changes.push_back({key, value});
        }
        query_context->checkSettingsConstraints(settings_changes, SettingSource::QUERY);
        query_context->applySettingsChanges(settings_changes);

        query_context->setCurrentQueryId(query_info.query_id());
        query_scope = QueryScope::create(query_context, /* fatal_error_callback */ [this]{ onFatalError(); });

        /// Set up tracing context for this query on current thread
        thread_trace_context = std::make_unique<OpenTelemetry::TracingContextHolder>("GRPCServer",
            query_context->getClientInfo().client_trace_context,
            query_context->getSettingsRef(),
            query_context->getOpenTelemetrySpanLog());
        thread_trace_context->root_span.kind = OpenTelemetry::SpanKind::SERVER;

        /// Prepare for sending exceptions and logs.
        const Settings & settings = query_context->getSettingsRef();
        send_exception_with_stacktrace = settings[Setting::calculate_text_stack_trace];
        const auto client_logs_level = settings[Setting::send_logs_level];
        if (client_logs_level != LogsLevel::none)
        {
            logs_queue = std::make_shared<InternalTextLogsQueue>();
            logs_queue->max_priority = Poco::Logger::parseLevel(client_logs_level.toString());
            logs_queue->setSourceRegexp(settings[Setting::send_logs_source_regexp]);
            CurrentThread::attachInternalTextLogsQueue(logs_queue, client_logs_level);
        }

        /// Set the current database if specified.
        if (!query_info.database().empty())
            query_context->setCurrentDatabase(query_info.database());

        /// Apply transport compression for this call.
        if (auto transport_compression = TransportCompression::fromQueryInfo(query_info))
            responder->setTransportCompression(*transport_compression);

        /// The interactive delay will be used to show progress.
        interactive_delay = settings[Setting::interactive_delay];
        query_context->setProgressCallback([this](const Progress & value) { return progress.incrementPiecewiseAtomically(value); });

        /// Parse the query.
        query_text = std::move(*(query_info.mutable_query()));
        const char * begin = query_text.data();
        const char * end = begin + query_text.size();
        ParserQuery parser(end, settings[Setting::allow_settings_after_format_in_insert]);
        ast = parseQuery(parser, begin, end, "", settings[Setting::max_query_size], settings[Setting::max_parser_depth], settings[Setting::max_parser_backtracks]);

        /// Choose input format.
        insert_query = ast->as<ASTInsertQuery>();
        if (insert_query)
        {
            input_format = insert_query->format;
            if (input_format.empty())
                input_format = "Values";
        }

        input_data_delimiter = query_info.input_data_delimiter();

        /// Choose output format.
        query_context->setDefaultFormat(query_info.output_format());
        if (const auto * ast_query_with_output = dynamic_cast<const ASTQueryWithOutput *>(ast.get());
            ast_query_with_output && ast_query_with_output->format_ast)
        {
            output_format = getIdentifierName(ast_query_with_output->format_ast);
        }
        if (output_format.empty())
            output_format = query_context->getDefaultFormat();

        send_output_columns_names_and_types = query_info.send_output_columns();

        /// Choose compression.
        String input_compression_method_str = query_info.input_compression_type();
        if (input_compression_method_str.empty())
            input_compression_method_str = query_info.obsolete_compression_type();
        input_compression_method = chooseCompressionMethod("", input_compression_method_str);

        String output_compression_method_str = query_info.output_compression_type();