Add socket injection into client

clickhouse/base/socket.cpp

@@ -213,6 +213,9 @@ const std::string & NetworkAddress::Host() const {
 }
 
 
+SocketBase::~SocketBase() = default;
+
+
 Socket::Socket(const NetworkAddress& addr)
     : handle_(SocketConnect(addr))
 {}

clickhouse/base/socket.h

@@ -57,13 +57,23 @@ class windowsErrorCategory : public std::error_category {
 
 #endif
 
-class Socket {
+
+class SocketBase {
+public:
+    virtual ~SocketBase();
+
+    virtual std::unique_ptr<InputStream> makeInputStream() const = 0;
+    virtual std::unique_ptr<OutputStream> makeOutputStream() const = 0;
+};
+
+
+class Socket : public SocketBase {
 public:
     Socket(const NetworkAddress& addr);
     Socket(Socket&& other) noexcept;
     Socket& operator=(Socket&& other) noexcept;
 
-    virtual ~Socket();
+    ~Socket() override;
 
     /// @params idle the time (in seconds) the connection needs to remain
     ///         idle before TCP starts sending keepalive probes.
@@ -75,8 +85,8 @@ class Socket {
     /// @params nodelay whether to enable TCP_NODELAY
     void SetTcpNoDelay(bool nodelay) noexcept;
 
-    virtual std::unique_ptr<InputStream> makeInputStream() const;
-    virtual std::unique_ptr<OutputStream> makeOutputStream() const;
+    std::unique_ptr<InputStream> makeInputStream() const override;
+    std::unique_ptr<OutputStream> makeOutputStream() const override;
 
 protected:
     Socket(const Socket&) = delete;

clickhouse/base/sslsocket.h

@@ -44,7 +44,7 @@ class SSLSocket : public Socket {
 public:
     explicit SSLSocket(const NetworkAddress& addr, const SSLParams & ssl_params, SSLContext& context);
     SSLSocket(SSLSocket &&) = default;
-    ~SSLSocket() = default;
+    ~SSLSocket() override = default;
 
     SSLSocket(const SSLSocket & ) = delete;
     SSLSocket& operator=(const SSLSocket & ) = delete;

clickhouse/client.cpp

@@ -84,6 +84,7 @@ std::ostream& operator<<(std::ostream& os, const ClientOptions& opt) {
 class Client::Impl {
 public:
      Impl(const ClientOptions& opts);
+     Impl(const ClientOptions& opts, std::unique_ptr<SocketBase> socket);
     ~Impl();
 
     void ExecuteQuery(Query query);
@@ -121,6 +122,8 @@ class Client::Impl {
 
     void WriteBlock(const Block& block, OutputStream* output);
 
+    void InitializeStreams(std::unique_ptr<SocketBase>&& socket);
+
 private:
     /// In case of network errors tries to reconnect to server and
     /// call fuc several times.
@@ -159,7 +162,8 @@ class Client::Impl {
     OutputStreams output_streams_;
     OutputStream* output_;
 
-    std::unique_ptr<Socket> socket_;
+    bool is_external_socket_initialization_{false};
+    std::unique_ptr<SocketBase> socket_;
 
 #if defined(WITH_OPENSSL)
     std::unique_ptr<SSLContext> ssl_context_;
@@ -193,6 +197,21 @@ Client::Impl::Impl(const ClientOptions& opts)
     }
 }
 
+Client::Impl::Impl(const ClientOptions& opts, std::unique_ptr<SocketBase> socket)
+    : options_(opts)
+    , events_(nullptr)
+    , is_external_socket_initialization_(true)
+{
+    InitializeStreams(std::move(socket));
+    if (!Handshake()) {
+        throw std::runtime_error("fail to connect to " + options_.host);
+    }
+
+    if (options_.compression_method != CompressionMethod::None) {
+        compression_ = CompressionState::Enable;
+    }
+}
+
 Client::Impl::~Impl()
 { }
 
@@ -295,6 +314,9 @@ void Client::Impl::Ping() {
 }
 
 void Client::Impl::ResetConnection() {
+    if (is_external_socket_initialization_) {
+        throw std::runtime_error("Can not reset connection, socket was provided externally");
+    }
 
     std::unique_ptr<Socket> socket;
 
@@ -336,19 +358,7 @@ void Client::Impl::ResetConnection() {
         socket->SetTcpNoDelay(options_.tcp_nodelay);
     }
 
-    OutputStreams output_streams;
-    auto socket_output = output_streams.Add(socket->makeOutputStream());
-    auto output = output_streams.AddNew<BufferedOutput>(socket_output);
-
-    InputStreams input_streams;
-    auto socket_input = input_streams.Add(socket->makeInputStream());
-    auto input = input_streams.AddNew<BufferedInput>(socket_input);
-
-    std::swap(output_streams, output_streams_);
-    std::swap(input_streams, input_streams_);
-    std::swap(socket, socket_);
-    output_ = output;
-    input_ = input;
+    InitializeStreams(std::move(socket));
 
 #if defined(WITH_OPENSSL)
     std::swap(ssl_context_, ssl_context);
@@ -705,6 +715,22 @@ void Client::Impl::SendData(const Block& block) {
     output_->Flush();
 }
 
+void Client::Impl::InitializeStreams(std::unique_ptr<SocketBase>&& socket) {
+    OutputStreams output_streams;
+    auto socket_output = output_streams.Add(socket->makeOutputStream());
+    auto output = output_streams.AddNew<BufferedOutput>(socket_output);
+
+    InputStreams input_streams;
+    auto socket_input = input_streams.Add(socket->makeInputStream());
+    auto input = input_streams.AddNew<BufferedInput>(socket_input);
+
+    std::swap(output_streams, output_streams_);
+    std::swap(input_streams, input_streams_);
+    std::swap(socket, socket_);
+    output_ = output;
+    input_ = input;
+}
+
 bool Client::Impl::SendHello() {
     WireFormat::WriteUInt64(output_, ClientCodes::Hello);
     WireFormat::WriteString(output_, std::string(DBMS_NAME) + " client");
@@ -796,6 +822,12 @@ Client::Client(const ClientOptions& opts)
 {
 }
 
+Client::Client(const ClientOptions& opts, std::unique_ptr<SocketBase> socket)
+    : options_(opts)
+    , impl_(new Impl(opts, std::move(socket)))
+{
+}
+
 Client::~Client()
 { }
 

clickhouse/client.h

@@ -163,12 +163,15 @@ struct ClientOptions {
 
 std::ostream& operator<<(std::ostream& os, const ClientOptions& options);
 
+class SocketBase;
+
 /**
  *
  */
 class Client {
 public:
      Client(const ClientOptions& opts);
+     Client(const ClientOptions& opts, std::unique_ptr<SocketBase> socket);
     ~Client();
 
     /// Intends for execute arbitrary queries.