clickhouse_c/
async_client.rs

1//! Tokio client built on [`IolessClient`].
2
3use core::pin::Pin;
4
5use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
6use tokio::net::{TcpStream, ToSocketAddrs};
7
8use crate::alloc::Allocator;
9use crate::builder::BlockBuilder;
10use crate::client::{ClientOpts, Event, ServerInfo};
11use crate::codec::Codec;
12use crate::error::{Error, ErrorKind, Result};
13use crate::ioless::{IolessClient, Step};
14
15const DEFAULT_READ_BUF_BYTES: usize = 8 * 1024;
16
17/// Asynchronous byte stream used by [`AsyncClient`].
18///
19/// Any `AsyncRead + AsyncWrite + Unpin + Send` type implements this trait.
20pub trait AsyncTransport: AsyncRead + AsyncWrite + Unpin + Send {}
21
22impl<S: AsyncRead + AsyncWrite + Unpin + Send> AsyncTransport for S {}
23
24/// [`AsyncClient`] with transport type erased.
25///
26/// Use this alias when plaintext and TLS connections must share one type.
27/// Dynamic dispatch applies to transport reads and writes.
28pub type BoxedAsyncClient = AsyncClient<Box<dyn AsyncTransport>>;
29
30/// Asynchronous ClickHouse native protocol client.
31///
32/// Client accepts any [`AsyncTransport`]. [`connect`](Self::connect) creates a
33/// TCP connection. [`connect_tls`](Self::connect_tls) creates a rustls
34/// connection when `tls` feature is enabled. [`boxed`](Self::boxed) erases
35/// transport type.
36pub struct AsyncClient<S = TcpStream> {
37    core: IolessClient,
38    stream: S,
39    read_buf: Vec<u8>,
40}
41
42impl AsyncClient<TcpStream> {
43    /// Connects over TCP and completes Hello handshake.
44    ///
45    /// Socket uses `TCP_NODELAY`.
46    pub async fn connect<A>(
47        addr: A,
48        opts: ClientOpts,
49        codec: Option<Pin<Box<Codec>>>,
50    ) -> Result<Self>
51    where
52        A: ToSocketAddrs,
53    {
54        let sock = TcpStream::connect(addr).await?;
55        sock.set_nodelay(true).ok();
56        Self::handshake_on(sock, opts, codec).await
57    }
58}
59
60#[cfg(feature = "tls")]
61impl AsyncClient<tokio_rustls::client::TlsStream<TcpStream>> {
62    /// Connects over TCP and TLS, then completes Hello handshake.
63    ///
64    /// TLS verifies peer using `config` and `domain`. Domain is also sent as
65    /// SNI. Default configuration is available from
66    /// [`tls::default_config`](crate::tls::default_config).
67    pub async fn connect_tls<A>(
68        addr: A,
69        domain: &str,
70        opts: ClientOpts,
71        codec: Option<Pin<Box<Codec>>>,
72        config: std::sync::Arc<rustls::ClientConfig>,
73    ) -> Result<Self>
74    where
75        A: ToSocketAddrs,
76    {
77        let sock = TcpStream::connect(addr).await?;
78        sock.set_nodelay(true).ok();
79        let server_name =
80            rustls::pki_types::ServerName::try_from(domain.to_owned()).map_err(|_| {
81                Error::new(
82                    ErrorKind::Usage,
83                    format!("invalid TLS server name: {domain}"),
84                )
85            })?;
86        let tls = tokio_rustls::TlsConnector::from(config)
87            .connect(server_name, sock)
88            .await
89            .map_err(|e| Error::new(ErrorKind::Io, format!("TLS handshake: {e}")))?;
90        Self::handshake_on(tls, opts, codec).await
91    }
92}
93
94impl<S: AsyncTransport> AsyncClient<S> {
95    /// Completes Hello handshake over an existing transport.
96    pub async fn handshake_on(
97        stream: S,
98        opts: ClientOpts,
99        codec: Option<Pin<Box<Codec>>>,
100    ) -> Result<Self> {
101        let read_buf_bytes = if opts.read_buffer_bytes == 0 {
102            DEFAULT_READ_BUF_BYTES
103        } else {
104            opts.read_buffer_bytes
105        };
106        let mut client = Self {
107            core: IolessClient::new(&opts, Allocator::stdlib(), codec)?,
108            stream,
109            read_buf: vec![0; read_buf_bytes],
110        };
111        client.pump_until_ready(|core| core.handshake()).await?;
112        Ok(client)
113    }
114
115    pub async fn send_query(&mut self, sql: &str, query_id: Option<&str>) -> Result<()> {
116        self.drain_out().await?;
117        self.core.send_query(sql, query_id)?;
118        self.drain_out().await
119    }
120
121    /// Sends a Data block, or empty terminator when `builder` is `None`.
122    pub async fn send_data(&mut self, builder: Option<&BlockBuilder<'_>>) -> Result<()> {
123        self.drain_out().await?;
124        self.core.send_data(builder)?;
125        self.drain_out().await
126    }
127
128    /// Sends empty Data block that ends INSERT input.
129    pub async fn send_data_end(&mut self) -> Result<()> {
130        self.drain_out().await?;
131        self.core.send_data_end()?;
132        self.drain_out().await
133    }
134
135    /// Waits for next server event.
136    ///
137    /// Returned event owns block or exception payload.
138    pub async fn recv_event(&mut self) -> Result<Event> {
139        let mut event = None;
140        self.pump_until_ready(|core| {
141            Ok(match core.recv_event()? {
142                Step::Ready(e) => {
143                    event = Some(e);
144                    Step::Ready(())
145                }
146                Step::NeedsInput => Step::NeedsInput,
147            })
148        })
149        .await?;
150        Ok(event.expect("pump_until_ready only returns once the step stored an event"))
151    }
152
153    /// Returns server identity received during handshake.
154    pub fn server_info(&self) -> Option<ServerInfo> {
155        self.core.server_info()
156    }
157
158    /// Erases transport type without changing connection state.
159    pub fn boxed(self) -> BoxedAsyncClient
160    where
161        S: 'static,
162    {
163        AsyncClient {
164            core: self.core,
165            stream: Box::new(self.stream),
166            read_buf: self.read_buf,
167        }
168    }
169
170    /// Returns mutable access to transport-independent protocol client.
171    pub fn core(&mut self) -> &mut IolessClient {
172        &mut self.core
173    }
174
175    async fn drain_out(&mut self) -> Result<()> {
176        let mut wrote = false;
177        loop {
178            // Queue and stream are disjoint fields across await
179            let buf = self.core.pending_out();
180            if buf.is_empty() {
181                break;
182            }
183            let n = self.stream.write(buf).await?;
184            if n == 0 {
185                return Err(Error::new(ErrorKind::Io, "transport write returned zero"));
186            }
187            self.core.consume_out(n);
188            wrote = true;
189        }
190        // TLS can buffer partial record after poll_write reports progress
191        if wrote {
192            self.stream.flush().await?;
193        }
194        Ok(())
195    }
196
197    async fn pump_until_ready(
198        &mut self,
199        mut step: impl FnMut(&mut IolessClient) -> Result<Step<()>>,
200    ) -> Result<()> {
201        loop {
202            self.drain_out().await?;
203            match step(&mut self.core)? {
204                Step::Ready(()) => return self.drain_out().await,
205                Step::NeedsInput => {
206                    self.drain_out().await?;
207                    self.read_more().await?;
208                }
209            }
210        }
211    }
212
213    async fn read_more(&mut self) -> Result<()> {
214        let n = self.stream.read(&mut self.read_buf).await?;
215        if n == 0 {
216            return Err(Error::new(ErrorKind::Eof, "transport closed"));
217        }
218        self.core.submit(&self.read_buf[..n])
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    use super::{AsyncClient, Event};
225
226    #[test]
227    fn async_client_is_send() {
228        fn assert_send<T: Send>() {}
229        assert_send::<AsyncClient>();
230        assert_send::<Event>();
231    }
232}
clickhouse_c/async_client.rs

clickhouse_c/
async_client.rs
