netget/server/clickhouse/
mod.rs

1//! ClickHouse native TCP server. Rust owns the handshake, revision negotiation, packet and
2//! block encoding, compression, the INSERT exchange and every bound; the handler answers each
3//! query and decides each INSERT's rows.
4pub mod actions;
5pub mod wire;
6
7use crate::llm::action_helper::call_llm;
8use crate::llm::actions::protocol_trait::ActionResult;
9use crate::logging::emit::Log;
10use crate::protocol::{Event, EventType, SpawnContext};
11use crate::server::connection::ConnectionId;
12use crate::state::server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo};
13use anyhow::{Context, Result};
14use serde_json::{json, Map, Value};
15use std::{net::SocketAddr, time::Duration};
16use tokio::{
17    io::{AsyncWriteExt, BufReader, ReadHalf, WriteHalf},
18    net::{TcpListener, TcpStream},
19};
20use wire::{client_packet, server_packet, Block, Source};
21
22pub const IDLE_TIMEOUT: Duration = Duration::from_secs(300);
23pub const MAX_IDLE_TIMEOUT_SECS: u64 = 86400;
24const WRITE_TIMEOUT: Duration = Duration::from_secs(30);
25const AUTHENTICATION_FAILED: i32 = 516;
26const UNKNOWN_EXCEPTION: i32 = 1002;
27const TOO_MANY_SIMULTANEOUS_QUERIES: i32 = 202;
28const UNKNOWN_PACKET_FROM_CLIENT: i32 = 101;
29
30#[derive(Clone)]
31struct Config {
32    credentials: Option<(String, String)>,
33    idle: Duration,
34}
35
36fn config(ctx: &SpawnContext) -> Result<Config> {
37    let params = ctx.startup_params.as_ref();
38    let get = |name: &str| -> Result<Option<String>> {
39        Ok(params
40            .map(|p| p.get_optional_string(name))
41            .transpose()?
42            .flatten())
43    };
44    let credentials = match (get("user")?, get("password")?) {
45        (Some(u), Some(p)) => Some((u, p)),
46        (None, None) => None,
47        _ => anyhow::bail!("user and password are configured together"),
48    };
49    let secs = params
50        .map(|p| p.get_optional_u64("idle_timeout_secs"))
51        .transpose()?
52        .flatten()
53        .unwrap_or(IDLE_TIMEOUT.as_secs());
54    anyhow::ensure!(
55        (1..=MAX_IDLE_TIMEOUT_SECS).contains(&secs),
56        "idle_timeout_secs must be between 1 and {MAX_IDLE_TIMEOUT_SECS}"
57    );
58    Ok(Config {
59        credentials,
60        idle: Duration::from_secs(secs),
61    })
62}
63
64pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
65    let cfg = config(&ctx)?;
66    let listener = TcpListener::bind(ctx.legacy_listen_addr()).await?;
67    let addr = listener.local_addr()?;
68    Log::new(Some(&ctx.status_tx)).info(format!("ClickHouse native server listening on {addr}"));
69    let state = ctx.state.clone();
70    let server_id = ctx.server_id;
71    let accept = tokio::spawn(async move {
72        let limiter = crate::server::accept_bounded::ConnectionLimiter::new(
73            crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS,
74        );
75        loop {
76            let (socket, peer, permit) = match crate::server::accept_bounded::accept_bounded(
77                &listener,
78                &limiter,
79                b"",
80                "ClickHouse",
81                Some(&ctx.status_tx),
82            )
83            .await
84            {
85                Ok(v) => v,
86                Err(_) => break,
87            };
88            let id = ConnectionId::new(ctx.state.get_next_unified_id().await);
89            let now = crate::utils::clock::Instant::now();
90            ctx.state
91                .add_connection_to_server(
92                    server_id,
93                    ConnectionState {
94                        id,
95                        remote_addr: peer,
96                        local_addr: addr,
97                        bytes_sent: 0,
98                        bytes_received: 0,
99                        packets_sent: 0,
100                        packets_received: 0,
101                        last_activity: now,
102                        status: ConnectionStatus::Active,
103                        status_changed_at: now,
104                        protocol_info: ProtocolConnectionInfo::empty(),
105                    },
106                )
107                .await;
108            let child = ctx.clone();
109            let cfg = cfg.clone();
110            ctx.state
111                .spawn_server_task(server_id, async move {
112                    let _permit = permit;
113                    if let Err(e) = session(&child, id, socket, peer, &cfg).await {
114                        Log::new(Some(&child.status_tx))
115                            .warn(format!("ClickHouse connection {id} ended: {e}"));
116                    }
117                    child
118                        .state
119                        .update_connection_status(server_id, id, ConnectionStatus::Closed)
120                        .await;
121                    let _ = child.status_tx.send("__UPDATE_UI__".into());
122                })
123                .await;
124        }
125    });
126    state.register_server_task(server_id, accept).await;
127    Ok(addr)
128}
129
130fn outcome(ctx: &SpawnContext, id: ConnectionId, operation: &str, decision: &str) {
131    let summary = format!("ClickHouse connection {id} operation={operation} decision={decision}");
132    let log = Log::new(Some(&ctx.status_tx));
133    if decision.starts_with("fail_closed_") || decision == "model_silent" {
134        log.error(summary);
135    } else {
136        log.info(summary);
137    }
138}
139
140/// The handler's answer, or the exception packet to send instead (already logged).
141async fn decide(
142    ctx: &SpawnContext,
143    id: ConnectionId,
144    event_type: &'static EventType,
145    data: Value,
146    allowed: &[&str],
147) -> std::result::Result<(String, Value), Vec<u8>> {
148    let op = event_type.id.clone();
149    let generic = || {
150        wire::exception(
151            UNKNOWN_EXCEPTION,
152            crate::utils::WireFailure::Unavailable.prefixed_text(),
153        )
154    };
155    let result = match call_llm(
156        &ctx.llm_client,
157        &ctx.state,
158        ctx.server_id,
159        Some(id),
160        &event_type_event(event_type, data),
161        &actions::ClickhouseProtocol,
162    )
163    .await
164    {
165        Ok(r) => r,
166        Err(error) => {
167            outcome(ctx, id, &op, "fail_closed_llm_error");
168            let code = if crate::llm::is_overload_error(&error) {
169                TOO_MANY_SIMULTANEOUS_QUERIES
170            } else {
171                UNKNOWN_EXCEPTION
172            };
173            return Err(wire::exception(
174                code,
175                crate::utils::wire_failure::prefixed_wire_failure_text(&error),
176            ));
177        }
178    };
179    if !result.failures.is_empty() {
180        outcome(ctx, id, &op, "fail_closed_invalid_reply");
181        return Err(generic());
182    }
183    let mut found = Vec::new();
184    let mut pending = result.protocol_results;
185    while let Some(r) = pending.pop() {
186        match r {
187            ActionResult::Custom { name, data } if name.starts_with("clickhouse_") => {
188                found.push((name, data))
189            }
190            ActionResult::Multiple(items) => pending.extend(items),
191            _ => {}
192        }
193    }
194    match found.len() {
195        0 => {
196            outcome(ctx, id, &op, "model_silent");
197            Err(generic())
198        }
199        1 => {
200            let (name, data) = found.remove(0);
201            if name == "clickhouse_exception" {
202                outcome(ctx, id, &op, "model_reject");
203                let code = data["code"]
204                    .as_i64()
205                    .unwrap_or(i64::from(UNKNOWN_EXCEPTION)) as i32;
206                Err(wire::exception(
207                    code,
208                    data["message"].as_str().unwrap_or_default(),
209                ))
210            } else if allowed.contains(&name.as_str()) {
211                outcome(ctx, id, &op, "model_answer");
212                Ok((name, data))
213            } else {
214                outcome(ctx, id, &op, "fail_closed_invalid_reply");
215                Err(generic())
216            }
217        }
218        _ => {
219            outcome(ctx, id, &op, "fail_closed_invalid_reply");
220            Err(generic())
221        }
222    }
223}
224
225fn event_type_event(event_type: &'static EventType, data: Value) -> Event {
226    Event::new(event_type, data)
227}
228
229struct Conn<'a> {
230    ctx: &'a SpawnContext,
231    id: ConnectionId,
232    reader: BufReader<ReadHalf<TcpStream>>,
233    writer: WriteHalf<TcpStream>,
234}
235
236impl Conn<'_> {
237    async fn send(&mut self, bytes: &[u8]) -> Result<()> {
238        tokio::time::timeout(WRITE_TIMEOUT, self.writer.write_all(bytes))
239            .await
240            .context("ClickHouse write deadline")??;
241        self.ctx
242            .state
243            .update_connection_stats(
244                self.ctx.server_id,
245                self.id,
246                None,
247                Some(bytes.len() as u64),
248                None,
249                Some(1),
250            )
251            .await;
252        Ok(())
253    }
254
255    /// Read a client Data packet (after its type): the table name, then the block.
256    async fn read_data(&mut self, compressed: bool) -> Result<Block> {
257        let mut plain = Source::plain(&mut self.reader);
258        let table = plain.string().await?;
259        anyhow::ensure!(table.is_empty(), "external tables are not supported");
260        let mut source = Source::new(&mut self.reader, compressed);
261        let block = Block::decode(&mut source).await?;
262        anyhow::ensure!(
263            source.drained(),
264            "a compressed frame carried bytes past its block"
265        );
266        Ok(block)
267    }
268
269    async fn expect_data(&mut self, compressed: bool, idle: Duration) -> Result<Block> {
270        let kind = wire::read_packet_type(&mut self.reader, idle)
271            .await?
272            .context("client closed before sending its data")?;
273        anyhow::ensure!(
274            kind == client_packet::DATA,
275            "expected a Data packet, got {kind}"
276        );
277        tokio::time::timeout(wire::IO_TIMEOUT, self.read_data(compressed))
278            .await
279            .context("ClickHouse block deadline")?
280    }
281}
282
283fn same_secret(a: &[u8], b: &[u8]) -> bool {
284    a.len() == b.len() && a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
285}
286
287async fn session(
288    ctx: &SpawnContext,
289    id: ConnectionId,
290    socket: TcpStream,
291    peer: SocketAddr,
292    cfg: &Config,
293) -> Result<()> {
294    let (reader, writer) = tokio::io::split(socket);
295    let mut c = Conn {
296        ctx,
297        id,
298        reader: BufReader::new(reader),
299        writer,
300    };
301    let Some(kind) = wire::read_packet_type(&mut c.reader, cfg.idle).await? else {
302        return Ok(());
303    };
304    anyhow::ensure!(
305        kind == client_packet::HELLO,
306        "first packet {kind} is not a hello"
307    );
308    let (client_name, client_revision, database, user, password) =
309        tokio::time::timeout(wire::IO_TIMEOUT, async {
310            let mut s = Source::plain(&mut c.reader);
311            let name = s.string().await?;
312            s.varuint().await?;
313            s.varuint().await?;
314            let revision = s.varuint().await?;
315            anyhow::Ok((
316                name,
317                revision,
318                s.string().await?,
319                s.string().await?,
320                s.raw_string().await?,
321            ))
322        })
323        .await
324        .context("ClickHouse hello deadline")??;
325    let revision = client_revision.min(wire::REVISION);
326    if let Some((u, p)) = &cfg.credentials {
327        if user != *u || !same_secret(&password, p.as_bytes()) {
328            c.send(&wire::exception(AUTHENTICATION_FAILED, &format!("{user}: Authentication failed: password is incorrect, or there is no user with such name."))).await?;
329            return Ok(());
330        }
331    }
332    let mut hello = Vec::new();
333    wire::put_varuint(&mut hello, server_packet::HELLO);
334    wire::put_string(&mut hello, b"ClickHouse");
335    wire::put_varuint(&mut hello, wire::VERSION_MAJOR);
336    wire::put_varuint(&mut hello, wire::VERSION_MINOR);
337    wire::put_varuint(&mut hello, wire::REVISION);
338    if revision >= 54058 {
339        wire::put_string(&mut hello, b"UTC");
340    }
341    if revision >= 54372 {
342        wire::put_string(&mut hello, b"netget");
343    }
344    if revision >= 54401 {
345        wire::put_varuint(&mut hello, wire::VERSION_PATCH);
346    }
347    c.send(&hello).await?;
348    Log::new(Some(&ctx.status_tx)).info(format!(
349        "ClickHouse connection {id}: {client_name} as {user} (revision {revision})"
350    ));
351    loop {
352        let Some(kind) = wire::read_packet_type(&mut c.reader, cfg.idle).await? else {
353            return Ok(());
354        };
355        match kind {
356            client_packet::PING => {
357                let mut pong = Vec::new();
358                wire::put_varuint(&mut pong, server_packet::PONG);
359                c.send(&pong).await?;
360            }
361            client_packet::CANCEL => {}
362            client_packet::QUERY => {
363                let q = tokio::time::timeout(wire::IO_TIMEOUT, async {
364                    let mut s = Source::plain(&mut c.reader);
365                    wire::read_query(&mut s, revision).await
366                })
367                .await
368                .context("ClickHouse query deadline")??;
369                ctx.state
370                    .update_connection_stats(
371                        ctx.server_id,
372                        id,
373                        Some(q.query.len() as u64),
374                        None,
375                        Some(1),
376                        None,
377                    )
378                    .await;
379                // The external tables end with an empty block, which always follows a query.
380                let external = c.expect_data(q.compression, wire::IO_TIMEOUT).await?;
381                if !external.is_empty() {
382                    c.send(&wire::exception(
383                        UNKNOWN_PACKET_FROM_CLIENT,
384                        "external tables are not supported",
385                    ))
386                    .await?;
387                    continue;
388                }
389                query(&mut c, &q, &database, &user, peer, cfg).await?;
390            }
391            other => {
392                c.send(&wire::exception(
393                    UNKNOWN_PACKET_FROM_CLIENT,
394                    &format!("Unknown packet {other} from client"),
395                ))
396                .await?;
397                return Ok(());
398            }
399        }
400    }
401}
402
403async fn query(
404    c: &mut Conn<'_>,
405    q: &wire::Query,
406    database: &str,
407    user: &str,
408    peer: SocketAddr,
409    cfg: &Config,
410) -> Result<()> {
411    let (ctx, id) = (c.ctx, c.id);
412    let settings: Map<String, Value> = q
413        .settings
414        .iter()
415        .map(|(k, v)| (k.clone(), json!(v)))
416        .collect();
417    let answer = decide(
418        ctx,
419        id,
420        &actions::QUERY_EVENT,
421        json!({"query": q.query, "query_id": q.query_id, "database": database, "user": user, "settings": settings, "remote_addr": peer.to_string()}),
422        &["clickhouse_result", "clickhouse_ok", "clickhouse_insert"],
423    )
424    .await;
425    let (name, data) = match answer {
426        Ok(a) => a,
427        Err(packet) => return c.send(&packet).await,
428    };
429    let mut eos = Vec::new();
430    wire::put_varuint(&mut eos, server_packet::END_OF_STREAM);
431    match name.as_str() {
432        "clickhouse_result" => {
433            let block = Block::from_json(&data["columns"], &data["rows"])?;
434            let header = Block {
435                columns: block.columns.clone(),
436                rows: Vec::new(),
437            };
438            let mut out = wire::data_packet(server_packet::DATA, &header, q.compression)?;
439            if !block.rows.is_empty() {
440                out.extend(wire::data_packet(
441                    server_packet::DATA,
442                    &block,
443                    q.compression,
444                )?);
445            }
446            out.extend(wire::progress(block.rows.len() as u64, 0, 0));
447            out.extend(eos);
448            c.send(&out).await
449        }
450        "clickhouse_ok" => c.send(&eos).await,
451        _ => {
452            // INSERT: send the header, read row blocks until the empty one, then ask.
453            let header = Block::from_json(&data["columns"], &Value::Null)?;
454            c.send(&wire::data_packet(
455                server_packet::DATA,
456                &header,
457                q.compression,
458            )?)
459            .await?;
460            let mut rows: Vec<Value> = Vec::new();
461            let mut columns = header.columns_json();
462            loop {
463                let block = c.expect_data(q.compression, cfg.idle).await?;
464                if block.is_empty() {
465                    break;
466                }
467                anyhow::ensure!(
468                    rows.len() + block.rows.len() <= wire::MAX_ROWS,
469                    "INSERT of more than {} rows",
470                    wire::MAX_ROWS
471                );
472                columns = block.columns_json();
473                rows.extend(block.rows.into_iter().map(Value::Array));
474            }
475            let decided = decide(
476                ctx,
477                id,
478                &actions::INSERT_DATA_EVENT,
479                json!({"query": q.query, "columns": columns, "rows": rows, "user": user}),
480                &["clickhouse_ok"],
481            )
482            .await;
483            match decided {
484                Ok(_) => {
485                    let mut out = wire::progress(0, 0, rows.len() as u64);
486                    out.extend(eos);
487                    c.send(&out).await
488                }
489                Err(packet) => c.send(&packet).await,
490            }
491        }
492    }
493}
netget/server/clickhouse/mod.rs

netget/server/clickhouse/
mod.rs
