From ecd96bd463a1dc2e5f8c62b5dd6cc1afb8f93baa Mon Sep 17 00:00:00 2001 From: Derek Date: Tue, 10 Mar 2026 12:34:04 +1100 Subject: [PATCH 1/6] feat(native): add Bool type and sparse serialization support - Map 'Bool' type string to UInt8 in the column type parser (Bool is a UInt8 alias on the wire: true=1, false=0) - Implement sparse (custom) column serialization in the data block reader: read offset groups to locate non-default values, read only those values, then reconstruct the full column with zero/empty/null defaults at the remaining positions - Add native_bool_type integration test covering SELECT of true/false values; also exercises the sparse code path on the cluster where ClickHouse sends Bool columns with custom_ser=1 --- src/native/columns.rs | 1457 +++++++++++++++++++++++++++++++++++++++++ src/native/reader.rs | 490 ++++++++++++++ tests/it/native.rs | 1192 +++++++++++++++++++++++++++++++++ 3 files changed, 3139 insertions(+) create mode 100644 src/native/columns.rs create mode 100644 src/native/reader.rs create mode 100644 tests/it/native.rs diff --git a/src/native/columns.rs b/src/native/columns.rs new file mode 100644 index 00000000..99be6412 --- /dev/null +++ b/src/native/columns.rs @@ -0,0 +1,1457 @@ +//! Native binary column type system and data reader. +//! +//! Reads ClickHouse native binary column data and re-serializes it as +//! RowBinary so the existing `rowbinary::deserialize_row` machinery can consume it. +//! +//! RowBinary and native binary formats are identical for scalar types — the +//! only difference is layout (columnar vs row-oriented). Nullable is the only +//! type that differs structurally. + +use tokio::io::AsyncReadExt; + +use crate::error::{Error, Result}; +use crate::native::io::ClickHouseRead; + +/// Supported ClickHouse column types for native transport. +#[derive(Debug, Clone)] +pub(crate) enum ColumnType { + UInt8, + UInt16, + UInt32, + UInt64, + Int8, + Int16, + Int32, + Int64, + Int128, + UInt128, + Int256, + UInt256, + Float32, + Float64, + /// BFloat16 — 16-bit brain float, 2 bytes on wire. + BFloat16, + /// Decimal32/64/128/256 — wire format identical to Int32/64/128/256 (raw LE bytes). + Decimal32, + Decimal64, + Decimal128, + Decimal256, + String, + FixedString(usize), + Uuid, + /// IPv4 — stored as 4-byte little-endian UInt32. + IPv4, + /// IPv6 — stored as 16 bytes. + IPv6, + Date, + Date32, + DateTime, + DateTime64, + /// Time — stored as UInt32 (seconds since midnight). + Time, + /// Time64 — stored as Int64 (ticks since midnight at given precision). + Time64, + Nullable(Box), + LowCardinality(Box), + /// Enum8/Enum16 — wire-compatible with UInt8/UInt16 respectively. + Enum8, + Enum16, + /// SimpleAggregateFunction(func, T) — wire-compatible with inner type T. + SimpleAggregateFunction(Box), + /// Array(T) — n cumulative u64 offsets, then all values packed as T column. + Array(Box), + /// Tuple(T1, T2, ...) — each field stored as a separate columnar block. + Tuple(Vec), + /// Map(K, V) — n cumulative u64 offsets, then K column, then V column. + Map(Box, Box), + /// JSON (legacy Object('json')) — wire format is a length-prefixed String. + Json, + /// Point — pair of Float64 (16 bytes), ClickHouse geo type. + Point, + /// Variant(T1, T2, ...) — discriminated union (ClickHouse 24.x+). + /// Wire prefix: u64 version (=0). + /// Wire data: u8[n] discriminators (255=NULL, 0..k-1 = type index in definition order), + /// then per-variant sub-columns in definition order. + Variant(Vec), + /// New JSON type (ClickHouse 24.x+). Complex path-based columnar format. + /// Wire prefix: u64 JSON version (1=string, 2=object-v2, 3=object-v3). + /// Wire data (v2): per-path Dynamic v2 headers + discriminators + values + n×u64 shared data. + NewJson, + /// Standalone Dynamic type (ClickHouse 24.x+). + /// Wire prefix: u64 version (1=deprecated, 2=intermediate, 3=flat). + /// Wire data: discriminators + per-type column data. + Dynamic, +} + +impl ColumnType { + /// Parse a ClickHouse type name into `ColumnType`. + /// + /// Returns `None` for unsupported types. + pub(crate) fn parse(type_str: &str) -> Option { + let type_str = type_str.trim(); + + if let Some(inner) = strip_outer(type_str, "Nullable") { + return Self::parse(inner).map(|t| Self::Nullable(Box::new(t))); + } + + if let Some(inner) = strip_outer(type_str, "LowCardinality") { + return Self::parse(inner).map(|t| Self::LowCardinality(Box::new(t))); + } + + if let Some(n_str) = strip_outer(type_str, "FixedString") { + return n_str.parse::().ok().map(Self::FixedString); + } + + if type_str.starts_with("DateTime64(") { + return Some(Self::DateTime64); + } + if type_str.starts_with("DateTime(") { + return Some(Self::DateTime); + } + if type_str.starts_with("Time64(") { + return Some(Self::Time64); + } + + // Decimal variants — scale is not needed for wire reading (raw LE bytes). + if type_str.starts_with("Decimal32(") { + return Some(Self::Decimal32); + } + if type_str.starts_with("Decimal64(") { + return Some(Self::Decimal64); + } + if type_str.starts_with("Decimal128(") { + return Some(Self::Decimal128); + } + if type_str.starts_with("Decimal256(") { + return Some(Self::Decimal256); + } + // Generic Decimal(precision, scale) — map to Decimal32/64/128/256 by precision. + if let Some(args_str) = strip_outer(type_str, "Decimal") { + let args = split_type_args(args_str); + if args.len() == 2 { + if let Ok(precision) = args[0].trim().parse::() { + return Some(if precision <= 9 { + Self::Decimal32 + } else if precision <= 18 { + Self::Decimal64 + } else if precision <= 38 { + Self::Decimal128 + } else { + Self::Decimal256 + }); + } + } + return None; + } + + // Enum8(...) / Enum16(...) — wire format = UInt8/UInt16 + if type_str.starts_with("Enum8(") { + return Some(Self::Enum8); + } + if type_str.starts_with("Enum16(") { + return Some(Self::Enum16); + } + + // Array(T) + if let Some(inner_str) = strip_outer(type_str, "Array") { + return Self::parse(inner_str).map(|t| Self::Array(Box::new(t))); + } + + // Tuple(T1, T2, ...) + if let Some(args_str) = strip_outer(type_str, "Tuple") { + let arg_strings = split_type_args(args_str); + let fields: Vec = + arg_strings.iter().filter_map(|s| Self::parse(s)).collect(); + // Only accept if all fields parsed successfully. + if !fields.is_empty() && fields.len() == arg_strings.len() { + return Some(Self::Tuple(fields)); + } + return None; + } + + // Map(K, V) + if let Some(args_str) = strip_outer(type_str, "Map") { + let args = split_type_args(args_str); + if args.len() == 2 { + let k = Self::parse(args[0])?; + let v = Self::parse(args[1])?; + return Some(Self::Map(Box::new(k), Box::new(v))); + } + return None; + } + + // SimpleAggregateFunction(func, T) — strip wrapper, read as T + if type_str.starts_with("SimpleAggregateFunction(") { + if let Some(rest) = type_str.strip_prefix("SimpleAggregateFunction(") { + // Find first ", " at depth 0 to split function name from type + if let Some(comma_pos) = find_first_comma_at_depth0(rest) { + let inner_str = rest[comma_pos + 1..].trim(); + let inner_str = inner_str.strip_suffix(')').unwrap_or(inner_str); + if let Some(inner) = Self::parse(inner_str) { + return Some(Self::SimpleAggregateFunction(Box::new(inner))); + } + } + } + return None; + } + + // Variant(T1, T2, ...) — discriminated union + if let Some(args_str) = strip_outer(type_str, "Variant") { + let arg_strings = split_type_args(args_str); + let fields: Vec = + arg_strings.iter().filter_map(|s| Self::parse(s)).collect(); + if !fields.is_empty() && fields.len() == arg_strings.len() { + return Some(Self::Variant(fields)); + } + return None; + } + + // Dynamic(N) — with optional max_types param + if type_str.starts_with("Dynamic(") { + return Some(Self::Dynamic); + } + + match type_str { + // Bool is an alias for UInt8 (true=1, false=0) on the wire. + "Bool" => Some(Self::UInt8), + "UInt8" => Some(Self::UInt8), + "UInt16" => Some(Self::UInt16), + "UInt32" => Some(Self::UInt32), + "UInt64" => Some(Self::UInt64), + "Int8" => Some(Self::Int8), + "Int16" => Some(Self::Int16), + "Int32" => Some(Self::Int32), + "Int64" => Some(Self::Int64), + "Int128" => Some(Self::Int128), + "UInt128" => Some(Self::UInt128), + "Int256" => Some(Self::Int256), + "UInt256" => Some(Self::UInt256), + "Float32" => Some(Self::Float32), + "Float64" => Some(Self::Float64), + "BFloat16" => Some(Self::BFloat16), + "String" => Some(Self::String), + "UUID" => Some(Self::Uuid), + "IPv4" => Some(Self::IPv4), + "IPv6" => Some(Self::IPv6), + "Date" => Some(Self::Date), + "Date32" => Some(Self::Date32), + "DateTime" => Some(Self::DateTime), + "Time" => Some(Self::Time), + // New JSON type (ClickHouse 24.x+) — path-based columnar format. + "JSON" => Some(Self::NewJson), + "Dynamic" => Some(Self::Dynamic), + // Legacy Object('json') — stored as a plain String on the wire. + "Object('json')" => Some(Self::Json), + // Geo types + "Point" => Some(Self::Point), + _ => None, + } + } + + /// Fixed wire-format size in bytes; `None` for variable-length types. + pub(crate) fn fixed_size(&self) -> Option { + match self { + Self::UInt8 | Self::Int8 | Self::Enum8 => Some(1), + Self::BFloat16 | Self::UInt16 | Self::Int16 | Self::Date | Self::Enum16 => Some(2), + Self::UInt32 + | Self::Int32 + | Self::Float32 + | Self::DateTime + | Self::Date32 + | Self::Decimal32 + | Self::IPv4 + | Self::Time => Some(4), + Self::UInt64 + | Self::Int64 + | Self::Float64 + | Self::DateTime64 + | Self::Decimal64 + | Self::Time64 => Some(8), + Self::Int128 | Self::UInt128 | Self::Uuid | Self::IPv6 | Self::Decimal128 => Some(16), + Self::Int256 | Self::UInt256 | Self::Decimal256 => Some(32), + // Point = 2 × Float64 + Self::Point => Some(16), + Self::FixedString(n) => Some(*n), + Self::String + | Self::Json + | Self::Nullable(_) + | Self::LowCardinality(_) + | Self::SimpleAggregateFunction(_) + | Self::Array(_) + | Self::Tuple(_) + | Self::Map(_, _) + | Self::Variant(_) + | Self::NewJson + | Self::Dynamic => None, + } + } +} + +// Strip "TypeName(" prefix and ")" suffix, returning the contents. +fn strip_outer<'a>(s: &'a str, name: &str) -> Option<&'a str> { + let prefix = format!("{name}("); + s.strip_prefix(prefix.as_str())?.strip_suffix(')') +} + +/// Split a comma-separated type argument list respecting parentheses depth. +/// +/// `"String, UInt64"` → `["String", "UInt64"]` +/// `"Array(String), UInt64"` → `["Array(String)", "UInt64"]` +fn split_type_args(s: &str) -> Vec<&str> { + let mut result = Vec::new(); + let mut depth = 0usize; + let mut start = 0; + for (i, c) in s.char_indices() { + match c { + '(' => depth += 1, + ')' => depth = depth.saturating_sub(1), + ',' if depth == 0 => { + result.push(s[start..i].trim()); + start = i + 1; + } + _ => {} + } + } + let tail = s[start..].trim(); + if !tail.is_empty() { + result.push(tail); + } + result +} + +/// Find the byte offset of the first ',' at parentheses depth 0. +fn find_first_comma_at_depth0(s: &str) -> Option { + let mut depth = 0usize; + for (i, c) in s.char_indices() { + match c { + '(' => depth += 1, + ')' => depth = depth.saturating_sub(1), + ',' if depth == 0 => return Some(i), + _ => {} + } + } + None +} + +/// Per-row RowBinary bytes for a single column's values. +/// +/// Each element is the RowBinary-encoded bytes for that row's field value. +pub(crate) type ColumnData = Vec>; + +/// Read all `num_rows` values for `col_type` from the native binary stream. +/// +/// Returns per-row RowBinary bytes ready for concatenation with other column data. +/// +/// Uses `Box::pin` internally for the recursive async cases (Nullable, LowCardinality). +pub(crate) fn read_column<'a, R: ClickHouseRead + 'a>( + reader: &'a mut R, + col_type: &'a ColumnType, + num_rows: u64, +) -> std::pin::Pin> + Send + 'a>> { + Box::pin(async move { + let n = num_rows as usize; + + match col_type { + ColumnType::UInt8 + | ColumnType::Int8 + | ColumnType::Enum8 + | ColumnType::BFloat16 + | ColumnType::UInt16 + | ColumnType::Int16 + | ColumnType::Enum16 + | ColumnType::UInt32 + | ColumnType::Int32 + | ColumnType::Float32 + | ColumnType::Date + | ColumnType::Date32 + | ColumnType::DateTime + | ColumnType::Decimal32 + | ColumnType::IPv4 + | ColumnType::Time + | ColumnType::UInt64 + | ColumnType::Int64 + | ColumnType::Float64 + | ColumnType::DateTime64 + | ColumnType::Decimal64 + | ColumnType::Time64 + | ColumnType::Uuid + | ColumnType::Int128 + | ColumnType::UInt128 + | ColumnType::IPv6 + | ColumnType::Decimal128 + | ColumnType::Int256 + | ColumnType::UInt256 + | ColumnType::Decimal256 + | ColumnType::Point => { + let size = col_type.fixed_size().expect("size is known for fixed type"); + read_fixed_column(reader, n, size).await + } + + ColumnType::String | ColumnType::Json => read_string_column(reader, n).await, + ColumnType::FixedString(size) => read_fixed_string_column(reader, n, *size).await, + ColumnType::Nullable(inner) => read_nullable_column(reader, n, inner).await, + ColumnType::LowCardinality(inner) => { + read_low_cardinality_column(reader, n, inner).await + } + ColumnType::SimpleAggregateFunction(inner) => { + read_column(reader, inner, num_rows).await + } + ColumnType::Array(inner) => read_array_column(reader, n, inner).await, + ColumnType::Tuple(fields) => read_tuple_column(reader, n, fields).await, + ColumnType::Map(key_type, val_type) => { + read_map_column(reader, n, key_type, val_type).await + } + ColumnType::Variant(variant_types) => { + read_variant_column(reader, n, variant_types).await + } + ColumnType::NewJson => read_json_column(reader, n).await, + ColumnType::Dynamic => read_dynamic_column(reader, n).await, + } + }) +} + +async fn read_fixed_column( + reader: &mut R, + n: usize, + size: usize, +) -> Result { + let mut result = Vec::with_capacity(n); + for _ in 0..n { + let mut buf = vec![0u8; size]; + reader.read_exact(&mut buf).await?; + result.push(buf); + } + Ok(result) +} + +async fn read_fixed_string_column( + reader: &mut R, + n: usize, + size: usize, +) -> Result { + let mut result = Vec::with_capacity(n); + for _ in 0..n { + let mut buf = vec![0u8; size]; + reader.read_exact(&mut buf).await?; + // Re-encode as RowBinary String: varint(len) + bytes + let mut row = Vec::with_capacity(size + 9); + write_var_uint(size as u64, &mut row); + row.extend_from_slice(&buf); + result.push(row); + } + Ok(result) +} + +async fn read_string_column(reader: &mut R, n: usize) -> Result { + let mut result = Vec::with_capacity(n); + for _ in 0..n { + // read_string() returns raw bytes (varint length already consumed) + let s = reader.read_string().await?; + // Re-encode as RowBinary: varint(len) + bytes + let mut row = Vec::with_capacity(s.len() + 9); + write_var_uint(s.len() as u64, &mut row); + row.extend_from_slice(&s); + result.push(row); + } + Ok(result) +} + +async fn read_nullable_column( + reader: &mut R, + n: usize, + inner: &ColumnType, +) -> Result { + // Native: N null-flags (1 byte each: 1=null, 0=has-value) then N values + let mut null_flags = vec![0u8; n]; + reader.read_exact(&mut null_flags).await?; + + // All N values are always present (native sends placeholder for nulls too) + let inner_data = read_column(reader, inner, n as u64).await?; + + let mut result = Vec::with_capacity(n); + for (flag, value) in null_flags.into_iter().zip(inner_data.into_iter()) { + if flag != 0 { + // NULL — RowBinary: 1 byte = 1 + result.push(vec![1u8]); + } else { + // Not null — RowBinary: 0 byte then value + let mut row = Vec::with_capacity(1 + value.len()); + row.push(0u8); + row.extend_from_slice(&value); + result.push(row); + } + } + Ok(result) +} + +/// LowCardinality column reader. +/// +/// Wire format uses fixed uint64 (little-endian) for sizes, not varint. +/// +/// ```text +/// u64 state_and_type +/// bits 0-1: index size (0=U8, 1=U16, 2=U32, 3=U64) +/// bit 8: has global dictionary +/// bit 9: has additional keys (new rows not in global dict) +/// if bit 8 set: +/// u64 global_dict_size +/// global_dict_size × inner_type values +/// if bit 9 set: +/// u64 additional_keys_size +/// additional_keys_size × inner_type values +/// u64 num_indices (must equal num_rows) +/// num_indices × index_bytes (indices into combined dict) +/// ``` +async fn read_low_cardinality_column( + reader: &mut R, + n: usize, + inner: &ColumnType, +) -> Result { + use tokio::io::AsyncReadExt as _; + + // Wire format starts with a serialization version u64 (= 1). + let _version = reader.read_u64_le().await?; + + let state = reader.read_u64_le().await?; + let index_type = (state & 0x03) as u8; + // Bit 8: NEED_GLOBAL_DICTIONARY — server sends a shared global dict + let has_global_dict = (state & 0x100) != 0; + // Bit 9: HAS_ADDITIONAL_KEYS — server sends per-block additional keys + let has_additional_keys = (state & 0x200) != 0; + + // Indices 0.. reference additional_keys first, then global_dict. + // Build combined dict in that order. + let mut additional: ColumnData = Vec::new(); + let mut global: ColumnData = Vec::new(); + + if has_global_dict { + let sz = reader.read_u64_le().await?; + global = read_column(reader, inner, sz).await?; + } + + if has_additional_keys { + let sz = reader.read_u64_le().await?; + additional = read_column(reader, inner, sz).await?; + } + + // Combined dict: additional_keys first (indices 0..additional.len()), + // then global_dict (indices additional.len()..). + // For the common case (only additional_keys, no global dict), dict = additional. + let mut dict: ColumnData = additional; + dict.extend(global); + + // If neither flag is set the entire dict is sent as a single section + // (older / simpler LowCardinality without shared dictionaries). + if !has_global_dict && !has_additional_keys { + let sz = reader.read_u64_le().await?; + dict.extend(read_column(reader, inner, sz).await?); + } + + let num_indices = reader.read_u64_le().await?; + if num_indices != n as u64 { + return Err(Error::BadResponse(format!( + "native protocol: LowCardinality index count {num_indices} != row count {n}" + ))); + } + + let index_bytes = match index_type { + 0 => 1usize, + 1 => 2, + 2 => 4, + 3 => 8, + other => { + return Err(Error::BadResponse(format!( + "native protocol: unknown LowCardinality index type {other}" + ))); + } + }; + + let dict_size = dict.len(); + let mut result = Vec::with_capacity(n); + for _ in 0..n { + let idx = read_index(reader, index_bytes).await? as usize; + let value = dict.get(idx).ok_or_else(|| { + Error::BadResponse(format!( + "native protocol: LowCardinality index {idx} out of range (dict size {dict_size})" + )) + })?; + result.push(value.clone()); + } + Ok(result) +} + +/// Array(T) column reader. +/// +/// Native wire format: +/// ```text +/// n × u64 cumulative end-offsets (last value = total element count) +/// total_elements × T values packed as a regular T column +/// ``` +/// Output RowBinary per row: varuint(count) + count × T_rowbinary +async fn read_array_column( + reader: &mut R, + n: usize, + inner: &ColumnType, +) -> Result { + use tokio::io::AsyncReadExt as _; + + let mut offsets = Vec::with_capacity(n); + for _ in 0..n { + offsets.push(reader.read_u64_le().await?); + } + + let total = offsets.last().copied().unwrap_or(0); + let all_values = read_column(reader, inner, total).await?; + + let mut result = Vec::with_capacity(n); + let mut prev = 0usize; + for &end in &offsets { + let end = end as usize; + let count = end - prev; + let mut row = Vec::new(); + write_var_uint(count as u64, &mut row); + for v in &all_values[prev..end] { + row.extend_from_slice(v); + } + result.push(row); + prev = end; + } + Ok(result) +} + +/// Tuple(T1, T2, ...) column reader. +/// +/// Native wire format: each field is its own complete columnar block in field order. +/// Output RowBinary per row: T1_bytes + T2_bytes + ... (simple concatenation). +async fn read_tuple_column( + reader: &mut R, + n: usize, + fields: &[ColumnType], +) -> Result { + let mut rows = vec![Vec::new(); n]; + for field_type in fields { + let field_data = read_column(reader, field_type, n as u64).await?; + for (row, cell) in rows.iter_mut().zip(field_data.into_iter()) { + row.extend_from_slice(&cell); + } + } + Ok(rows) +} + +/// Map(K, V) column reader. +/// +/// Native wire format: +/// ```text +/// n × u64 cumulative end-offsets +/// total_entries × K key column +/// total_entries × V value column +/// ``` +/// Output RowBinary per row: varuint(count) + count × (K_bytes + V_bytes) +async fn read_map_column( + reader: &mut R, + n: usize, + key_type: &ColumnType, + val_type: &ColumnType, +) -> Result { + use tokio::io::AsyncReadExt as _; + + let mut offsets = Vec::with_capacity(n); + for _ in 0..n { + offsets.push(reader.read_u64_le().await?); + } + + let total = offsets.last().copied().unwrap_or(0); + let keys = read_column(reader, key_type, total).await?; + let vals = read_column(reader, val_type, total).await?; + + let mut result = Vec::with_capacity(n); + let mut prev = 0usize; + for &end in &offsets { + let end = end as usize; + let count = end - prev; + let mut row = Vec::new(); + write_var_uint(count as u64, &mut row); + for i in prev..end { + row.extend_from_slice(&keys[i]); + row.extend_from_slice(&vals[i]); + } + result.push(row); + prev = end; + } + Ok(result) +} + +/// Variant(T1, T2, ...) column reader. +/// +/// Wire format: +/// ```text +/// u64 version (= 0) +/// n × u8 discriminators (255 = NULL, 0..k-1 = type index in definition order) +/// for each variant type Ti in order: +/// [rows where discriminator == i, in original row order] +/// ``` +/// Output: per-row JSON string encoded as RowBinary String. +async fn read_variant_column( + reader: &mut R, + n: usize, + variant_types: &[ColumnType], +) -> Result { + use tokio::io::AsyncReadExt as _; + + // Wire prefix: u64 version = 0 + let _version = reader.read_u64_le().await?; + + let mut discriminators = vec![0u8; n]; + reader.read_exact(&mut discriminators).await?; + + let k = variant_types.len(); + let mut type_counts = vec![0u64; k]; + for &d in &discriminators { + if (d as usize) < k { + type_counts[d as usize] += 1; + } + } + + let mut type_values: Vec = Vec::with_capacity(k); + for (i, col_type) in variant_types.iter().enumerate() { + type_values.push(read_column(reader, col_type, type_counts[i]).await?); + } + + let mut type_cursors = vec![0usize; k]; + let mut result = Vec::with_capacity(n); + for &d in &discriminators { + let json_bytes: Vec = if d == 255 || (d as usize) >= k { + b"null".to_vec() + } else { + let idx = d as usize; + let cursor = type_cursors[idx]; + type_cursors[idx] += 1; + rowbinary_to_json(&type_values[idx][cursor], &variant_types[idx]) + }; + let mut row = Vec::with_capacity(json_bytes.len() + 9); + write_var_uint(json_bytes.len() as u64, &mut row); + row.extend_from_slice(&json_bytes); + result.push(row); + } + Ok(result) +} + +/// New JSON column (ClickHouse 24.x+) reader. +/// +/// Dispatches based on the wire serialization version: +/// - `1`: each row is a plain JSON string (String column format) +/// - `2`: path-based object format with Dynamic v1/v2 sub-columns + shared data +/// - `3`: path-based object format with Dynamic v3 sub-columns (no shared data) +async fn read_json_column(reader: &mut R, n: usize) -> Result { + use tokio::io::AsyncReadExt as _; + + let version = reader.read_u64_le().await?; + match version { + 1 => read_string_column(reader, n).await, + 2 => read_json_object_v2_column(reader, n).await, + 3 => read_json_object_v3_column(reader, n).await, + _ => Err(Error::BadResponse(format!( + "native protocol: unsupported JSON serialization version: {version}" + ))), + } +} + +/// JSON v2 object column reader. +/// +/// Wire format (after the u64 version=2 already consumed): +/// ```text +/// varuint numDynamicPaths +/// String[] pathNames (sorted alphabetically) +/// for each path: +/// u64 dynVersion (1 or 2) +/// [if dynVersion==1: varuint maxTypes] +/// varuint numTypes (server types, excluding SharedVariant) +/// String[] typeNames +/// u64 variantVersion (= 0) +/// for each path: +/// u8[n] discriminators (index in sorted(typeNames+"SharedVariant"), 255=NULL) +/// for each type in sorted order: column data +/// n × u64 shared data (discard) +/// ``` +async fn read_json_object_v2_column( + reader: &mut R, + n: usize, +) -> Result { + use tokio::io::AsyncReadExt as _; + + let num_paths = reader.read_var_uint().await? as usize; + let mut path_names: Vec = Vec::with_capacity(num_paths); + for _ in 0..num_paths { + path_names.push(reader.read_utf8_string().await?); + } + + // Read Dynamic v1/v2 header for each path. + // sorted_types[p] = Vec<(type_name, col_type)> in sorted order (SharedVariant included). + let mut path_sorted_types: Vec> = Vec::with_capacity(num_paths); + for path_name in &path_names { + let dyn_version = reader.read_u64_le().await?; + if dyn_version == 1 { + // v1 has an extra maxTypes field before numTypes + let _max_types = reader.read_var_uint().await?; + } else if dyn_version != 2 { + return Err(Error::BadResponse(format!( + "native protocol: unexpected Dynamic version {dyn_version} in JSON v2 path \"{path_name}\"" + ))); + } + + let num_types = reader.read_var_uint().await? as usize; + let mut type_names: Vec = Vec::with_capacity(num_types + 1); + for _ in 0..num_types { + type_names.push(reader.read_utf8_string().await?); + } + // SharedVariant is implicit — add and sort to get the discriminator indices. + type_names.push("SharedVariant".to_string()); + type_names.sort(); + + let _variant_version = reader.read_u64_le().await?; + + let types: Vec<(String, ColumnType)> = type_names + .into_iter() + .map(|name| { + let ct = if name == "SharedVariant" { + ColumnType::String + } else { + ColumnType::parse(&name).unwrap_or(ColumnType::String) + }; + (name, ct) + }) + .collect(); + + path_sorted_types.push(types); + } + + // Read data for each path: discriminators then per-type column values. + let mut path_discriminators: Vec> = Vec::with_capacity(num_paths); + let mut path_values: Vec> = Vec::with_capacity(num_paths); + + for types in &path_sorted_types { + let k = types.len(); + + let mut discriminators = vec![0u8; n]; + reader.read_exact(&mut discriminators).await?; + + let mut type_counts = vec![0u64; k]; + for &d in &discriminators { + if (d as usize) < k { + type_counts[d as usize] += 1; + } + } + + let mut col_values: Vec = Vec::with_capacity(k); + for (i, (_, col_type)) in types.iter().enumerate() { + col_values.push(read_column(reader, col_type, type_counts[i]).await?); + } + + path_discriminators.push(discriminators); + path_values.push(col_values); + } + + // Discard shared data: n × u64 (one u64 per row, unused by us). + for _ in 0..n { + let _ = reader.read_u64_le().await?; + } + + // Build per-row JSON objects by reassembling path values. + let mut path_cursors: Vec> = path_sorted_types + .iter() + .map(|types| vec![0usize; types.len()]) + .collect(); + + let mut result = Vec::with_capacity(n); + for row_i in 0..n { + let mut json = b"{".to_vec(); + let mut first = true; + + for (path_idx, path_name) in path_names.iter().enumerate() { + let disc = path_discriminators[path_idx][row_i] as usize; + let k = path_sorted_types[path_idx].len(); + + if disc == 255 || disc >= k { + // Absent / NULL — omit key from output. + continue; + } + + let cursor = path_cursors[path_idx][disc]; + path_cursors[path_idx][disc] += 1; + + let (type_name, col_type) = &path_sorted_types[path_idx][disc]; + if type_name == "SharedVariant" { + // SharedVariant stores overflow values in an opaque binary format; skip. + continue; + } + + if !first { + json.push(b','); + } + first = false; + + json.extend_from_slice(&json_quote_bytes(path_name.as_bytes())); + json.push(b':'); + + let cell = &path_values[path_idx][disc][cursor]; + json.extend_from_slice(&rowbinary_to_json(cell, col_type)); + } + + json.push(b'}'); + + let mut row = Vec::with_capacity(json.len() + 9); + write_var_uint(json.len() as u64, &mut row); + row.extend_from_slice(&json); + result.push(row); + } + + Ok(result) +} + +/// JSON v3 object column reader (new flat format, ClickHouse 25.6+). +/// +/// Wire format (after the u64 version=3 already consumed): +/// ```text +/// varuint numDynamicPaths +/// String[] pathNames +/// for each path: +/// varuint numTypes +/// String[] typeNames +/// for each path: +/// discriminators (u8/u16/u32/u64 depending on numTypes+1) +/// for each type in order: column data +/// ``` +/// No shared data section in v3. +async fn read_json_object_v3_column( + reader: &mut R, + n: usize, +) -> Result { + let num_paths = reader.read_var_uint().await? as usize; + let mut path_names: Vec = Vec::with_capacity(num_paths); + for _ in 0..num_paths { + path_names.push(reader.read_utf8_string().await?); + } + + let mut path_col_types: Vec> = Vec::with_capacity(num_paths); + let mut path_total_types: Vec = Vec::with_capacity(num_paths); + + for _ in 0..num_paths { + let num_types = reader.read_var_uint().await? as usize; + let mut col_types: Vec = Vec::with_capacity(num_types); + for _ in 0..num_types { + let name = reader.read_utf8_string().await?; + col_types.push(ColumnType::parse(&name).unwrap_or(ColumnType::String)); + } + path_total_types.push(num_types); + path_col_types.push(col_types); + } + + let mut path_discriminators: Vec> = Vec::with_capacity(num_paths); + let mut path_values: Vec> = Vec::with_capacity(num_paths); + + for (p, col_types) in path_col_types.iter().enumerate() { + let total_types = path_total_types[p]; + // NULL discriminator = total_types; discriminator range = [0, total_types]. + let disc_size = if total_types <= 254 { 1usize } + else if total_types <= 65535 { 2 } + else if total_types <= u32::MAX as usize { 4 } + else { 8 }; + + let mut discriminators: Vec = Vec::with_capacity(n); + for _ in 0..n { + discriminators.push(read_index(reader, disc_size).await? as usize); + } + + let mut type_counts = vec![0u64; total_types]; + for &d in &discriminators { + if d < total_types { + type_counts[d] += 1; + } + } + + let mut col_values: Vec = Vec::with_capacity(total_types); + for (i, col_type) in col_types.iter().enumerate() { + col_values.push(read_column(reader, col_type, type_counts[i]).await?); + } + + path_discriminators.push(discriminators); + path_values.push(col_values); + } + + let mut path_cursors: Vec> = path_col_types + .iter() + .map(|types| vec![0usize; types.len()]) + .collect(); + + let mut result = Vec::with_capacity(n); + for row_i in 0..n { + let mut json = b"{".to_vec(); + let mut first = true; + + for (path_idx, path_name) in path_names.iter().enumerate() { + let disc = path_discriminators[path_idx][row_i]; + let total_types = path_total_types[path_idx]; + + if disc == total_types || disc > total_types { + // NULL — omit key. + continue; + } + + let cursor = path_cursors[path_idx][disc]; + path_cursors[path_idx][disc] += 1; + + if !first { + json.push(b','); + } + first = false; + + json.extend_from_slice(&json_quote_bytes(path_name.as_bytes())); + json.push(b':'); + + let cell = &path_values[path_idx][disc][cursor]; + let col_type = &path_col_types[path_idx][disc]; + json.extend_from_slice(&rowbinary_to_json(cell, col_type)); + } + + json.push(b'}'); + + let mut row = Vec::with_capacity(json.len() + 9); + write_var_uint(json.len() as u64, &mut row); + row.extend_from_slice(&json); + result.push(row); + } + + Ok(result) +} + +/// Standalone Dynamic column (ClickHouse 24.x+) reader. +/// +/// Dispatches based on the wire serialization version prefix: +/// - `1`: deprecated format (maxTypes + totalTypes + sorted types + SharedVariant + variantVersion) +/// - `2`: intermediate format (totalTypes + sorted types + SharedVariant + variantVersion) +/// - `3`: flat format (totalTypes + types, NULL = totalTypes, no SharedVariant) +async fn read_dynamic_column(reader: &mut R, n: usize) -> Result { + use tokio::io::AsyncReadExt as _; + + let version = reader.read_u64_le().await?; + match version { + 1 => read_dynamic_v1v2_column(reader, n, true).await, + 2 => read_dynamic_v1v2_column(reader, n, false).await, + 3 => read_dynamic_v3_column(reader, n).await, + _ => Err(Error::BadResponse(format!( + "native protocol: unsupported Dynamic serialization version: {version}" + ))), + } +} + +/// Dynamic v1/v2 column reader. +/// +/// v1 has an extra `maxTypes` varuint before `totalTypes`; v2 does not. +/// Both add "SharedVariant" to the type list and sort alphabetically. +/// NULL discriminator = 255. +async fn read_dynamic_v1v2_column( + reader: &mut R, + n: usize, + has_max_types: bool, +) -> Result { + use tokio::io::AsyncReadExt as _; + + if has_max_types { + let _max_types = reader.read_var_uint().await?; + } + + let total_types = reader.read_var_uint().await? as usize; + let mut type_names: Vec = Vec::with_capacity(total_types + 1); + for _ in 0..total_types { + type_names.push(reader.read_utf8_string().await?); + } + type_names.push("SharedVariant".to_string()); + type_names.sort(); + + let _variant_version = reader.read_u64_le().await?; + + let col_types: Vec = type_names + .iter() + .map(|name| { + if name == "SharedVariant" { + ColumnType::String + } else { + ColumnType::parse(name).unwrap_or(ColumnType::String) + } + }) + .collect(); + + let k = col_types.len(); + let mut discriminators = vec![0u8; n]; + reader.read_exact(&mut discriminators).await?; + + let mut type_counts = vec![0u64; k]; + for &d in &discriminators { + if (d as usize) < k { + type_counts[d as usize] += 1; + } + } + + let mut type_values: Vec = Vec::with_capacity(k); + for (i, col_type) in col_types.iter().enumerate() { + type_values.push(read_column(reader, col_type, type_counts[i]).await?); + } + + let mut type_cursors = vec![0usize; k]; + let mut result = Vec::with_capacity(n); + for &d in &discriminators { + let json_bytes: Vec = if d == 255 || (d as usize) >= k { + b"null".to_vec() + } else { + let idx = d as usize; + let cursor = type_cursors[idx]; + type_cursors[idx] += 1; + if type_names[idx] == "SharedVariant" { + b"null".to_vec() + } else { + rowbinary_to_json(&type_values[idx][cursor], &col_types[idx]) + } + }; + let mut row = Vec::with_capacity(json_bytes.len() + 9); + write_var_uint(json_bytes.len() as u64, &mut row); + row.extend_from_slice(&json_bytes); + result.push(row); + } + Ok(result) +} + +/// Dynamic v3 column reader (new flat format, ClickHouse 25.6+). +/// +/// No SharedVariant; NULL discriminator = totalTypes. +/// Discriminator width scales with totalTypes: u8/u16/u32/u64. +async fn read_dynamic_v3_column(reader: &mut R, n: usize) -> Result { + let total_types = reader.read_var_uint().await? as usize; + let mut type_names: Vec = Vec::with_capacity(total_types); + let mut col_types: Vec = Vec::with_capacity(total_types); + for _ in 0..total_types { + let name = reader.read_utf8_string().await?; + col_types.push(ColumnType::parse(&name).unwrap_or(ColumnType::String)); + type_names.push(name); + } + + let disc_size = if total_types <= 254 { 1usize } + else if total_types <= 65535 { 2 } + else if total_types <= u32::MAX as usize { 4 } + else { 8 }; + let null_disc = total_types; + + let mut discriminators: Vec = Vec::with_capacity(n); + for _ in 0..n { + discriminators.push(read_index(reader, disc_size).await? as usize); + } + + let mut type_counts = vec![0u64; total_types]; + for &d in &discriminators { + if d < total_types { + type_counts[d] += 1; + } + } + + let mut type_values: Vec = Vec::with_capacity(total_types); + for (i, col_type) in col_types.iter().enumerate() { + type_values.push(read_column(reader, col_type, type_counts[i]).await?); + } + + let mut type_cursors = vec![0usize; total_types]; + let mut result = Vec::with_capacity(n); + for &d in &discriminators { + let json_bytes: Vec = if d == null_disc || d > total_types { + b"null".to_vec() + } else { + let cursor = type_cursors[d]; + type_cursors[d] += 1; + rowbinary_to_json(&type_values[d][cursor], &col_types[d]) + }; + let mut row = Vec::with_capacity(json_bytes.len() + 9); + write_var_uint(json_bytes.len() as u64, &mut row); + row.extend_from_slice(&json_bytes); + result.push(row); + } + Ok(result) +} + +/// Convert a RowBinary-encoded value for `col_type` into JSON bytes. +/// +/// Returns `b"null"` on any parse error rather than propagating — callers should +/// treat this as a best-effort JSON representation for use in Dynamic/Variant columns. +fn rowbinary_to_json(bytes: &[u8], col_type: &ColumnType) -> Vec { + match rowbinary_to_json_inner(bytes, col_type) { + Ok((json, _)) => json, + Err(()) => b"null".to_vec(), + } +} + +/// Inner parser: returns `(json_bytes, bytes_consumed)` or `Err(())` on underflow. +#[allow(clippy::too_many_lines)] +fn rowbinary_to_json_inner(bytes: &[u8], col_type: &ColumnType) -> Result<(Vec, usize), ()> { + macro_rules! fixed { + ($n:expr, $t:ty, $fmt:expr) => {{ + if bytes.len() < $n { return Err(()); } + let v = <$t>::from_le_bytes(bytes[..$n].try_into().unwrap()); + (format!($fmt, v).into_bytes(), $n) + }}; + } + + Ok(match col_type { + ColumnType::UInt8 => fixed!(1, u8, "{}"), + ColumnType::UInt16 => fixed!(2, u16, "{}"), + ColumnType::UInt32 | ColumnType::IPv4 | ColumnType::Time => fixed!(4, u32, "{}"), + ColumnType::UInt64 => fixed!(8, u64, "{}"), + ColumnType::Int8 => { + if bytes.is_empty() { return Err(()); } + ((bytes[0] as i8).to_string().into_bytes(), 1) + } + ColumnType::Int16 => fixed!(2, i16, "{}"), + ColumnType::Int32 | ColumnType::Decimal32 => fixed!(4, i32, "{}"), + ColumnType::Date32 => fixed!(4, i32, "{}"), + ColumnType::Int64 | ColumnType::Time64 | ColumnType::Decimal64 => fixed!(8, i64, "{}"), + ColumnType::Int128 | ColumnType::Decimal128 => { + if bytes.len() < 16 { return Err(()); } + let v = i128::from_le_bytes(bytes[..16].try_into().unwrap()); + (v.to_string().into_bytes(), 16) + } + ColumnType::UInt128 => { + if bytes.len() < 16 { return Err(()); } + let v = u128::from_le_bytes(bytes[..16].try_into().unwrap()); + (v.to_string().into_bytes(), 16) + } + ColumnType::Int256 | ColumnType::UInt256 | ColumnType::Decimal256 => { + // 32-byte big integer — emit as hex string for safety + if bytes.len() < 32 { return Err(()); } + let hex: String = bytes[..32].iter().rev().map(|b| format!("{b:02x}")).collect(); + (format!("\"{hex}\"").into_bytes(), 32) + } + ColumnType::Float32 => { + if bytes.len() < 4 { return Err(()); } + let v = f32::from_le_bytes(bytes[..4].try_into().unwrap()); + (format_float_json(v as f64).into_bytes(), 4) + } + ColumnType::Float64 => { + if bytes.len() < 8 { return Err(()); } + let v = f64::from_le_bytes(bytes[..8].try_into().unwrap()); + (format_float_json(v).into_bytes(), 8) + } + ColumnType::BFloat16 => { + // BFloat16 is u16 mantissa — convert via f32 + if bytes.len() < 2 { return Err(()); } + let raw = u16::from_le_bytes([bytes[0], bytes[1]]); + let v = f32::from_bits((raw as u32) << 16); + (format_float_json(v as f64).into_bytes(), 2) + } + ColumnType::Date => { + if bytes.len() < 2 { return Err(()); } + let days = u16::from_le_bytes([bytes[0], bytes[1]]) as u32; + (format!("\"{days}\"").into_bytes(), 2) + } + ColumnType::DateTime | ColumnType::DateTime64 => { + let size = col_type.fixed_size().unwrap_or(4); + if bytes.len() < size { return Err(()); } + let v: u64 = match size { + 4 => u32::from_le_bytes(bytes[..4].try_into().unwrap()) as u64, + 8 => u64::from_le_bytes(bytes[..8].try_into().unwrap()), + _ => return Err(()), + }; + (format!("{v}").into_bytes(), size) + } + ColumnType::Uuid => { + if bytes.len() < 16 { return Err(()); } + // UUID is stored as two u64s in big-endian byte order within ClickHouse + let hi = u64::from_be_bytes(bytes[..8].try_into().unwrap()); + let lo = u64::from_be_bytes(bytes[8..16].try_into().unwrap()); + let s = format!( + "\"{:08x}-{:04x}-{:04x}-{:04x}-{:012x}\"", + (hi >> 32) as u32, + (hi >> 16) as u16, + hi as u16, + (lo >> 48) as u16, + lo & 0x0000_ffff_ffff_ffff + ); + (s.into_bytes(), 16) + } + ColumnType::IPv6 => { + if bytes.len() < 16 { return Err(()); } + let hex: String = bytes[..16].chunks(2).map(|c| format!("{:02x}{:02x}", c[0], c[1])).collect::>().join(":"); + (format!("\"[{hex}]\"").into_bytes(), 16) + } + ColumnType::Point => { + // 2 × f64 LE + if bytes.len() < 16 { return Err(()); } + let x = f64::from_le_bytes(bytes[..8].try_into().unwrap()); + let y = f64::from_le_bytes(bytes[8..16].try_into().unwrap()); + (format!("[{},{}]", format_float_json(x), format_float_json(y)).into_bytes(), 16) + } + ColumnType::Enum8 => { + if bytes.is_empty() { return Err(()); } + ((bytes[0] as i8).to_string().into_bytes(), 1) + } + ColumnType::Enum16 => fixed!(2, i16, "{}"), + // String types: RowBinary format = varuint(len) + bytes + ColumnType::String + | ColumnType::FixedString(_) + | ColumnType::Json => { + let (len, hdr) = read_var_uint_from_slice(bytes).ok_or(())?; + let len = len as usize; + let end = hdr + len; + if bytes.len() < end { return Err(()); } + (json_quote_bytes(&bytes[hdr..end]), end) + } + ColumnType::Nullable(inner) => { + if bytes.is_empty() { return Err(()); } + if bytes[0] != 0 { + (b"null".to_vec(), 1) + } else { + let (json, consumed) = rowbinary_to_json_inner(&bytes[1..], inner)?; + (json, 1 + consumed) + } + } + ColumnType::LowCardinality(inner) => { + // After LowCardinality expansion, individual cells are the inner type's bytes + rowbinary_to_json_inner(bytes, inner)? + } + ColumnType::SimpleAggregateFunction(inner) => { + rowbinary_to_json_inner(bytes, inner)? + } + ColumnType::Array(inner) => { + let (count, hdr) = read_var_uint_from_slice(bytes).ok_or(())?; + let mut pos = hdr; + let mut json = b"[".to_vec(); + for i in 0..count { + if i > 0 { json.push(b','); } + let (elem, consumed) = rowbinary_to_json_inner(&bytes[pos..], inner)?; + json.extend_from_slice(&elem); + pos += consumed; + } + json.push(b']'); + (json, pos) + } + ColumnType::Tuple(fields) => { + let mut pos = 0; + let mut json = b"[".to_vec(); + for (i, field_type) in fields.iter().enumerate() { + if i > 0 { json.push(b','); } + let (elem, consumed) = rowbinary_to_json_inner(&bytes[pos..], field_type)?; + json.extend_from_slice(&elem); + pos += consumed; + } + json.push(b']'); + (json, pos) + } + ColumnType::Map(key_type, val_type) => { + let (count, hdr) = read_var_uint_from_slice(bytes).ok_or(())?; + let mut pos = hdr; + let mut json = b"{".to_vec(); + for i in 0..count { + if i > 0 { json.push(b','); } + let (k, kc) = rowbinary_to_json_inner(&bytes[pos..], key_type)?; + pos += kc; + json.extend_from_slice(&k); + json.push(b':'); + let (v, vc) = rowbinary_to_json_inner(&bytes[pos..], val_type)?; + pos += vc; + json.extend_from_slice(&v); + } + json.push(b'}'); + (json, pos) + } + // Dynamic/Variant/NewJson cells are already JSON strings (varuint + bytes) + ColumnType::Dynamic | ColumnType::NewJson | ColumnType::Variant(_) => { + let (len, hdr) = read_var_uint_from_slice(bytes).ok_or(())?; + let len = len as usize; + let end = hdr + len; + if bytes.len() < end { return Err(()); } + (bytes[hdr..end].to_vec(), end) + } + }) +} + +/// Format a float for JSON: avoids NaN/Infinity (not valid JSON), uses finite repr. +fn format_float_json(v: f64) -> String { + if v.is_nan() || v.is_infinite() { + "null".to_string() + } else { + // Use Rust's default float formatting (no trailing zeros) + format!("{v}") + } +} + +/// JSON-quote raw bytes as a UTF-8 string (or escaped if not valid UTF-8). +fn json_quote_bytes(bytes: &[u8]) -> Vec { + let mut out = vec![b'"']; + for &b in bytes { + match b { + b'"' => { out.push(b'\\'); out.push(b'"'); } + b'\\' => { out.push(b'\\'); out.push(b'\\'); } + b'\n' => { out.push(b'\\'); out.push(b'n'); } + b'\r' => { out.push(b'\\'); out.push(b'r'); } + b'\t' => { out.push(b'\\'); out.push(b't'); } + 0x00..=0x1f => { + // Control character — escape as \uXXXX + out.extend_from_slice(format!("\\u{b:04x}").as_bytes()); + } + _ => out.push(b), + } + } + out.push(b'"'); + out +} + +/// Read a varuint (LEB128) from a byte slice. Returns `(value, bytes_consumed)`. +fn read_var_uint_from_slice(bytes: &[u8]) -> Option<(u64, usize)> { + let mut value = 0u64; + let mut shift = 0u32; + for (i, &b) in bytes.iter().enumerate() { + value |= ((b & 0x7f) as u64) << shift; + if b & 0x80 == 0 { + return Some((value, i + 1)); + } + shift += 7; + if shift >= 63 { return None; } // overflow guard + } + None // ran out of bytes +} + +async fn read_index(reader: &mut R, bytes: usize) -> Result { + Ok(match bytes { + 1 => u64::from(reader.read_u8().await?), + 2 => u64::from(reader.read_u16_le().await?), + 4 => u64::from(reader.read_u32_le().await?), + 8 => reader.read_u64_le().await?, + _ => unreachable!(), + }) +} + +/// Write a LEB128 varint (ClickHouse 63-bit variant) into a buffer. +pub(crate) fn write_var_uint(mut value: u64, buf: &mut Vec) { + loop { + let byte = (value & 0x7F) as u8; + value >>= 7; + if value == 0 { + buf.push(byte); + break; + } + buf.push(byte | 0x80); + } +} + +/// Transpose columnar data into row-oriented RowBinary bytes. +/// +/// `column_data` contains one `ColumnData` per column. +/// Returns one `Vec` per row, suitable for `rowbinary::deserialize_row()`. +pub(crate) fn transpose_to_rowbinary( + column_data: Vec, + num_rows: u64, +) -> Vec> { + let n = num_rows as usize; + let mut rows = vec![Vec::new(); n]; + for col in column_data { + for (row_idx, cell) in col.into_iter().enumerate().take(n) { + rows[row_idx].extend_from_slice(&cell); + } + } + rows +} diff --git a/src/native/reader.rs b/src/native/reader.rs new file mode 100644 index 00000000..0c4ec39c --- /dev/null +++ b/src/native/reader.rs @@ -0,0 +1,490 @@ +//! Packet reader for ClickHouse native protocol. +//! +//! Reads and dispatches server packets: hello, exception, progress, +//! profile info, data blocks, and end-of-stream markers. + +use std::str::FromStr; + +use tokio::io::AsyncReadExt; + +use crate::error::{Error, Result}; +use crate::native::block_info::BlockInfo; +use crate::native::columns::{self, ColumnData, ColumnType, transpose_to_rowbinary, write_var_uint}; +use crate::native::sparse::{SparseDeserializeState, read_sparse_offsets}; +use crate::native::compression::decompress_data; +use crate::native::error_codes::{self, ServerError}; +use crate::native::io::ClickHouseRead; +use crate::native::protocol::DBMS_MIN_PROTOCOL_VERSION_WITH_CUSTOM_SERIALIZATION; +use crate::native::protocol::{ + ChunkedProtocolMode, NativeCompressionMethod, ProfileInfo, Progress, ServerException, + ServerHello, ServerPacketId, TableColumns, + DBMS_MIN_PROTOCOL_VERSION_WITH_CHUNKED_PACKETS, + DBMS_MIN_PROTOCOL_VERSION_WITH_PASSWORD_COMPLEXITY_RULES, + DBMS_MIN_PROTOCOL_VERSION_WITH_SERVER_QUERY_TIME_IN_PROGRESS, + DBMS_MIN_PROTOCOL_VERSION_WITH_TOTAL_BYTES_IN_PROGRESS, + DBMS_MIN_REVISION_WITH_CLIENT_WRITE_INFO, DBMS_MIN_REVISION_WITH_INTERSERVER_SECRET_V2, + DBMS_MIN_REVISION_WITH_QUERY_PLAN_SERIALIZATION, + DBMS_MIN_REVISION_WITH_ROWS_BEFORE_AGGREGATION, DBMS_MIN_REVISION_WITH_SERVER_DISPLAY_NAME, + DBMS_MIN_REVISION_WITH_SERVER_LOGS, DBMS_MIN_REVISION_WITH_SERVER_SETTINGS, + DBMS_MIN_REVISION_WITH_SERVER_TIMEZONE, DBMS_MIN_REVISION_WITH_VERSION_PATCH, + DBMS_MIN_REVISION_WITH_VERSIONED_CLUSTER_FUNCTION_PROTOCOL, + DBMS_MIN_REVISION_WITH_VERSIONED_PARALLEL_REPLICAS_PROTOCOL, +}; + +/// Server packet after dispatch. +#[derive(Debug)] +#[allow(unused)] +pub(crate) enum ServerPacket { + Hello(ServerHello), + Data(DataBlock), + Exception(ServerError), + Progress(Progress), + Pong, + EndOfStream, + ProfileInfo(ProfileInfo), + TableColumns(TableColumns), +} + +/// A fully-read data block from the server. +#[derive(Debug)] +pub(crate) struct DataBlock { + pub(crate) info: BlockInfo, + pub(crate) num_columns: u64, + pub(crate) num_rows: u64, + /// Column name + type. + pub(crate) column_headers: Vec, + /// Row-oriented RowBinary bytes, one `Vec` per row. + /// + /// Each element is the complete RowBinary bytes for one row, ready for + /// `rowbinary::deserialize_row()`. Empty if `num_rows == 0`. + pub(crate) row_data: Vec>, +} + +/// Column name + type string from a data block header. +#[derive(Debug, Clone)] +pub(crate) struct ColumnHeader { + pub(crate) name: String, + pub(crate) type_name: String, +} + +/// Read server hello response. +pub(crate) async fn read_hello( + reader: &mut R, + client_revision: u64, + chunked_modes: (ChunkedProtocolMode, ChunkedProtocolMode), +) -> Result { + let packet_id = ServerPacketId::from_u64(reader.read_var_uint().await?)?; + match packet_id { + ServerPacketId::Hello => { + read_hello_body(reader, client_revision, chunked_modes).await + } + ServerPacketId::Exception => { + let exc = read_exception(reader).await?; + Err(Error::BadResponse(format!( + "server exception during hello: {}: {}", + exc.name, exc.message + ))) + } + other => Err(Error::BadResponse(format!( + "native protocol: expected hello, got {other:?}" + ))), + } +} + +/// Read the body of a server hello packet (after packet ID). +async fn read_hello_body( + reader: &mut R, + client_revision: u64, + chunked_modes: (ChunkedProtocolMode, ChunkedProtocolMode), +) -> Result { + let server_name = reader.read_utf8_string().await?; + let major = reader.read_var_uint().await?; + let minor = reader.read_var_uint().await?; + let server_revision = reader.read_var_uint().await?; + let revision = std::cmp::min(server_revision, client_revision); + + if revision >= DBMS_MIN_REVISION_WITH_VERSIONED_PARALLEL_REPLICAS_PROTOCOL { + let _ = reader.read_var_uint().await?; + } + + let timezone = if revision >= DBMS_MIN_REVISION_WITH_SERVER_TIMEZONE { + Some(reader.read_utf8_string().await?) + } else { + None + }; + + let display_name = if revision >= DBMS_MIN_REVISION_WITH_SERVER_DISPLAY_NAME { + Some(reader.read_utf8_string().await?) + } else { + None + }; + + let patch = if revision >= DBMS_MIN_REVISION_WITH_VERSION_PATCH { + reader.read_var_uint().await? + } else { + revision + }; + + let (chunked_send, chunked_recv) = + if revision >= DBMS_MIN_PROTOCOL_VERSION_WITH_CHUNKED_PACKETS { + let srv_send = ChunkedProtocolMode::from_str( + &String::from_utf8_lossy(&reader.read_string().await?), + ) + .unwrap_or_default(); + let srv_recv = ChunkedProtocolMode::from_str( + &String::from_utf8_lossy(&reader.read_string().await?), + ) + .unwrap_or_default(); + + ( + ChunkedProtocolMode::negotiate(srv_send, chunked_modes.0, "send")?, + ChunkedProtocolMode::negotiate(srv_recv, chunked_modes.1, "recv")?, + ) + } else { + ( + ChunkedProtocolMode::default(), + ChunkedProtocolMode::default(), + ) + }; + + if revision >= DBMS_MIN_PROTOCOL_VERSION_WITH_PASSWORD_COMPLEXITY_RULES { + let rules_size = reader.read_var_uint().await?; + for _ in 0..rules_size { + drop(reader.read_utf8_string().await?); + drop(reader.read_utf8_string().await?); + } + } + + if revision >= DBMS_MIN_REVISION_WITH_INTERSERVER_SECRET_V2 { + let _ = reader.read_u64_le().await?; + } + + // Skip server settings + if revision >= DBMS_MIN_REVISION_WITH_SERVER_SETTINGS { + skip_settings(reader).await?; + } + + if revision >= DBMS_MIN_REVISION_WITH_QUERY_PLAN_SERIALIZATION { + let _ = reader.read_var_uint().await?; + } + + if revision >= DBMS_MIN_REVISION_WITH_VERSIONED_CLUSTER_FUNCTION_PROTOCOL { + let _ = reader.read_var_uint().await?; + } + + Ok(ServerHello { + server_name, + version: (major, minor, patch), + revision_version: revision, + timezone, + display_name, + chunked_send, + chunked_recv, + }) +} + +/// Skip settings key/value pairs from the wire. +async fn skip_settings(reader: &mut R) -> Result<()> { + loop { + let name = reader.read_utf8_string().await?; + if name.is_empty() { + break; + } + // Each setting: flag (varuint) + value_string + let _is_important = reader.read_var_uint().await?; + let _value = reader.read_string().await?; + } + Ok(()) +} + +/// Read a server exception from the wire. +pub(crate) async fn read_exception( + reader: &mut R, +) -> Result { + let code = reader.read_i32_le().await?; + let name = reader.read_utf8_string().await?; + let message = + String::from_utf8_lossy(&reader.read_string().await?).to_string(); + let stack_trace = reader.read_utf8_string().await?; + let has_nested = reader.read_u8().await? != 0; + + Ok(ServerException { + code, + name, + message, + stack_trace, + has_nested, + }) +} + +/// Read progress from the wire. +pub(crate) async fn read_progress( + reader: &mut R, + revision: u64, +) -> Result { + let read_rows = reader.read_var_uint().await?; + let read_bytes = reader.read_var_uint().await?; + + let total_rows_to_read = if revision >= DBMS_MIN_REVISION_WITH_SERVER_LOGS { + reader.read_var_uint().await? + } else { + 0 + }; + + let total_bytes_to_read = + if revision >= DBMS_MIN_PROTOCOL_VERSION_WITH_TOTAL_BYTES_IN_PROGRESS { + Some(reader.read_var_uint().await?) + } else { + None + }; + + let written = if revision >= DBMS_MIN_REVISION_WITH_CLIENT_WRITE_INFO { + Some(( + reader.read_var_uint().await?, + reader.read_var_uint().await?, + )) + } else { + None + }; + + let elapsed_ns = + if revision >= DBMS_MIN_PROTOCOL_VERSION_WITH_SERVER_QUERY_TIME_IN_PROGRESS { + Some(reader.read_var_uint().await?) + } else { + None + }; + + Ok(Progress { + read_rows, + read_bytes, + total_rows_to_read, + total_bytes_to_read, + written_rows: written.map(|w| w.0), + written_bytes: written.map(|w| w.1), + elapsed_ns, + }) +} + +/// Read profile info from the wire. +pub(crate) async fn read_profile_info( + reader: &mut R, + revision: u64, +) -> Result { + let rows = reader.read_var_uint().await?; + let blocks = reader.read_var_uint().await?; + let bytes = reader.read_var_uint().await?; + let applied_limit = reader.read_u8().await? != 0; + let rows_before_limit = reader.read_var_uint().await?; + let calculated_rows_before_limit = reader.read_u8().await? != 0; + + let (applied_aggregation, rows_before_aggregation) = + if revision >= DBMS_MIN_REVISION_WITH_ROWS_BEFORE_AGGREGATION { + (reader.read_u8().await? != 0, reader.read_var_uint().await?) + } else { + (false, 0) + }; + + Ok(ProfileInfo { + rows, + blocks, + bytes, + applied_limit, + rows_before_limit, + calculated_rows_before_limit, + applied_aggregation, + rows_before_aggregation, + }) +} + +/// Read table columns packet from the wire. +pub(crate) async fn read_table_columns( + reader: &mut R, +) -> Result { + Ok(TableColumns { + name: reader.read_utf8_string().await?, + description: reader.read_utf8_string().await?, + }) +} + +/// Read and dispatch a single server packet. +/// +/// Log and ProfileEvents blocks are consumed and skipped; the next +/// packet is returned instead (loop, not recursion). +pub(crate) async fn read_packet( + reader: &mut R, + revision: u64, + compression: NativeCompressionMethod, +) -> Result { + loop { + let packet_id = ServerPacketId::from_u64(reader.read_var_uint().await?)?; + + match packet_id { + ServerPacketId::Data | ServerPacketId::Totals | ServerPacketId::Extremes => { + return read_data_packet(reader, revision, compression).await; + } + ServerPacketId::Exception => { + let exc = read_exception(reader).await?; + return Ok(ServerPacket::Exception( + error_codes::map_exception_to_error(exc), + )); + } + ServerPacketId::Progress => { + return read_progress(reader, revision) + .await + .map(ServerPacket::Progress); + } + ServerPacketId::Pong => return Ok(ServerPacket::Pong), + ServerPacketId::EndOfStream => return Ok(ServerPacket::EndOfStream), + ServerPacketId::ProfileInfo => { + return read_profile_info(reader, revision) + .await + .map(ServerPacket::ProfileInfo); + } + ServerPacketId::TableColumns => { + return read_table_columns(reader) + .await + .map(ServerPacket::TableColumns); + } + ServerPacketId::Log | ServerPacketId::ProfileEvents => { + // ClickHouse sends Log and ProfileEvents blocks through the + // UNCOMPRESSED stream even when write_compression=1, so we + // always read them raw (None), never try to decompress. + read_data_packet(reader, revision, NativeCompressionMethod::None).await?; + } + other => { + return Err(Error::BadResponse(format!( + "native protocol: unhandled server packet: {other:?}" + ))); + } + } + } +} + +/// Read a full data block from the stream. +async fn read_data_packet( + reader: &mut R, + revision: u64, + compression: NativeCompressionMethod, +) -> Result { + // Temp table name (empty for normal queries) + let _table_name = reader.read_string().await?; + + match compression { + NativeCompressionMethod::None => read_data_block(reader, revision).await, + _ => { + let decompressed = decompress_data(reader, compression).await?; + let mut cursor = std::io::Cursor::new(decompressed); + read_data_block(&mut cursor, revision).await + } + } +} + +/// Read a data block from already-decompressed bytes. +async fn read_data_block(reader: &mut R, revision: u64) -> Result { + let info = BlockInfo::read_async(reader).await?; + let num_columns = reader.read_var_uint().await?; + let num_rows = reader.read_var_uint().await?; + + let has_custom_serialization = + revision >= DBMS_MIN_PROTOCOL_VERSION_WITH_CUSTOM_SERIALIZATION; + + let mut column_headers = Vec::with_capacity(num_columns as usize); + let mut column_data: Vec = Vec::with_capacity(num_columns as usize); + + for _ in 0..num_columns { + let name = reader.read_utf8_string().await?; + let type_name = reader.read_utf8_string().await?; + + // Newer servers send a custom serialization flag per column. + // 0 = normal, 1 = sparse (only non-default values are stored with offset groups). + let is_sparse = if has_custom_serialization { + reader.read_u8().await? != 0 + } else { + false + }; + + let col_type = ColumnType::parse(&type_name).ok_or_else(|| { + Error::BadResponse(format!( + "native protocol: unsupported column type '{type_name}' for column '{name}'" + )) + })?; + + let data = if num_rows == 0 { + Vec::new() + } else if is_sparse { + read_sparse_column(reader, &col_type, num_rows as usize).await? + } else { + columns::read_column(reader, &col_type, num_rows).await? + }; + + column_headers.push(ColumnHeader { name, type_name }); + column_data.push(data); + } + + let row_data = if num_rows == 0 { + Vec::new() + } else { + transpose_to_rowbinary(column_data, num_rows) + }; + + Ok(ServerPacket::Data(DataBlock { + info, + num_columns, + num_rows, + column_headers, + row_data, + })) +} + +/// Read a sparsely-serialized column. +/// +/// Sparse format: offset groups (varuint, final has END_OF_GRANULE_FLAG) identify +/// positions of non-default values. Only those values follow in the stream. +/// All other row positions get the type's default (zero/empty/null). +async fn read_sparse_column( + reader: &mut R, + col_type: &ColumnType, + num_rows: usize, +) -> Result { + let mut state = SparseDeserializeState::default(); + let non_default_positions = read_sparse_offsets(reader, num_rows, &mut state).await?; + + let non_default_count = non_default_positions.len(); + let non_default_data = if non_default_count > 0 { + columns::read_column(reader, col_type, non_default_count as u64).await? + } else { + Vec::new() + }; + + let default = sparse_default_bytes(col_type); + let mut result = vec![default; num_rows]; + for (i, pos) in non_default_positions.into_iter().enumerate() { + if pos < num_rows { + result[pos] = non_default_data[i].clone(); + } + } + Ok(result) +} + +/// Return the RowBinary-encoded default value for a type in sparse context. +/// +/// The sparse default is the column's "zero" value: 0 for numerics, empty for +/// strings, NULL for Nullable. +fn sparse_default_bytes(col_type: &ColumnType) -> Vec { + if let Some(size) = col_type.fixed_size() { + return vec![0u8; size]; + } + match col_type { + ColumnType::String | ColumnType::Json => vec![0u8], // varuint(0) = empty string + ColumnType::FixedString(n) => { + let mut v = Vec::with_capacity(*n + 9); + write_var_uint(*n as u64, &mut v); + v.extend(std::iter::repeat(0u8).take(*n)); + v + } + ColumnType::Nullable(_) => vec![0x01], // NULL + // Complex types (Array, Tuple, Map, etc.) are unlikely to be sparse, but + // return an empty vec as a safe fallback. + _ => vec![], + } +} diff --git a/tests/it/native.rs b/tests/it/native.rs new file mode 100644 index 00000000..1ffedbbb --- /dev/null +++ b/tests/it/native.rs @@ -0,0 +1,1192 @@ +//! Integration tests for the native TCP transport. +//! +//! These tests require a running ClickHouse server on port 9000. +//! Run with: `cargo test --test it --features native-transport -- native::` + +#![cfg(feature = "native-transport")] + +use clickhouse::native::NativeClient; +use clickhouse::Row; +use serde::{Deserialize, Serialize}; + +fn get_native_client() -> NativeClient { + let host = std::env::var("CLICKHOUSE_HOST").unwrap_or_else(|_| "localhost".into()); + let port = std::env::var("CLICKHOUSE_NATIVE_PORT").unwrap_or_else(|_| "9000".into()); + let user = std::env::var("CLICKHOUSE_USER").unwrap_or_else(|_| "default".into()); + let password = std::env::var("CLICKHOUSE_PASSWORD").unwrap_or_else(|_| "".into()); + + let client = NativeClient::default() + .with_addr(format!("{host}:{port}")) + .with_database("default") + .with_user(user) + .with_password(password); + + // On a replicated cluster, write to a quorum of replicas before returning + // and ensure SELECT only reads quorum-committed data. This gives + // read-after-write consistency without pinning connections to a single node. + // + // CLICKHOUSE_INSERT_QUORUM: number of replicas that must acknowledge each + // INSERT — set to the replica count for your cluster (default: 2). + if std::env::var("CLICKHOUSE_CLUSTER").is_ok() { + let quorum = std::env::var("CLICKHOUSE_INSERT_QUORUM") + .unwrap_or_else(|_| "2".into()); + client + .with_setting("insert_quorum", quorum) + .with_setting("insert_quorum_timeout", "30000") + .with_setting("select_sequential_consistency", "1") + } else { + client + } +} + +/// Create a unique test database for isolation (mirrors `prepare_database!` for HTTP tests). +/// +/// When `CLICKHOUSE_CLUSTER` is set, databases are created ON CLUSTER so all +/// nodes see the database immediately — required for multi-node setups. +async fn prepare_native_database(test_name: &str) -> NativeClient { + let client = get_native_client(); + let db = format!("chrs_native_{test_name}"); + let cluster = std::env::var("CLICKHOUSE_CLUSTER").ok(); + + let drop_sql = match &cluster { + Some(c) => format!("DROP DATABASE IF EXISTS {db} ON CLUSTER {c}"), + None => format!("DROP DATABASE IF EXISTS {db}"), + }; + let create_sql = match &cluster { + Some(c) => format!("CREATE DATABASE {db} ON CLUSTER {c}"), + None => format!("CREATE DATABASE {db}"), + }; + + client + .query(&drop_sql) + .execute() + .await + .unwrap_or_else(|e| panic!("drop db {db}: {e}")); + + client + .query(&create_sql) + .execute() + .await + .unwrap_or_else(|e| panic!("create db {db}: {e}")); + + client.with_database(db) +} + +/// Returns the ON CLUSTER clause if `CLICKHOUSE_CLUSTER` is set, otherwise empty. +fn on_cluster() -> String { + std::env::var("CLICKHOUSE_CLUSTER") + .map(|c| format!(" ON CLUSTER '{c}'")) + .unwrap_or_default() +} + +/// Returns the table engine clause for tests. +/// +/// Local Docker: `ENGINE = Memory` — fast, no persistence needed. +/// External cluster: `ReplicatedMergeTree` with a per-table `{uuid}` ZK path so +/// each CREATE TABLE gets a unique ZooKeeper node (no stale-replica conflicts on +/// re-runs). The `{uuid}` macro is substituted by ClickHouse at CREATE time. +fn test_engine(order_by: &str) -> String { + if std::env::var("CLICKHOUSE_CLUSTER").is_ok() { + format!( + "ENGINE = ReplicatedMergeTree(\ + '/clickhouse/tables/{{database}}/{{table}}/{{uuid}}', '{{replica}}'\ + ) ORDER BY {order_by}" + ) + } else { + "ENGINE = Memory".to_string() + } +} + +#[tokio::test] +async fn native_ping() { + let client = get_native_client(); + client.ping().await.expect("ping failed"); +} + +#[tokio::test] +async fn native_ddl() { + let client = prepare_native_database("ddl").await; + + client + .query(&format!("CREATE TABLE t{} (n UInt32) {}", on_cluster(), test_engine("tuple()"))) + .execute() + .await + .expect("CREATE TABLE failed"); + + client + .query("DROP TABLE t") + .execute() + .await + .expect("DROP TABLE failed"); +} + +#[tokio::test] +async fn native_scalar_types() { + let client = prepare_native_database("scalar").await; + + client + .query(&format!( + "CREATE TABLE t{} ( + u8 UInt8, u16 UInt16, u32 UInt32, u64 UInt64, + i8 Int8, i16 Int16, i32 Int32, i64 Int64, + f32 Float32, f64 Float64 + ) {}", + on_cluster(), test_engine("tuple()") + )) + .execute() + .await + .expect("CREATE failed"); + + client + .query( + "INSERT INTO t VALUES \ + (1, 2, 3, 4, -1, -2, -3, -4, 1.5, 2.5)", + ) + .execute() + .await + .expect("INSERT failed"); + + #[derive(Debug, Row, Deserialize, PartialEq)] + struct ScalarRow { + u8: u8, + u16: u16, + u32: u32, + u64: u64, + i8: i8, + i16: i16, + i32: i32, + i64: i64, + f32: f32, + f64: f64, + } + + let mut cursor = client + .query("SELECT * FROM t") + .fetch::() + .expect("fetch failed"); + + let row = cursor.next().await.expect("no error").expect("no row"); + assert_eq!(row.u8, 1); + assert_eq!(row.u16, 2); + assert_eq!(row.u32, 3); + assert_eq!(row.u64, 4); + assert_eq!(row.i8, -1); + assert_eq!(row.i16, -2); + assert_eq!(row.i32, -3); + assert_eq!(row.i64, -4); + assert!((row.f32 - 1.5f32).abs() < f32::EPSILON); + assert!((row.f64 - 2.5f64).abs() < f64::EPSILON); +} + +#[tokio::test] +async fn native_string_types() { + let client = prepare_native_database("strings").await; + + client + .query(&format!("CREATE TABLE t{} (s String, fs FixedString(4)) {}", on_cluster(), test_engine("tuple()"))) + .execute() + .await + .expect("CREATE failed"); + + client + .query("INSERT INTO t VALUES ('hello', 'abcd')") + .execute() + .await + .expect("INSERT failed"); + + #[derive(Debug, Row, Deserialize, PartialEq)] + struct StringRow { + s: String, + fs: String, + } + + let rows = client + .query("SELECT s, fs FROM t") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].s, "hello"); + assert_eq!(rows[0].fs, "abcd"); +} + +#[tokio::test] +async fn native_nullable() { + let client = prepare_native_database("nullable").await; + + client + .query(&format!("CREATE TABLE t{} (n Nullable(UInt32)) {}", on_cluster(), test_engine("tuple()"))) + .execute() + .await + .expect("CREATE failed"); + + client + .query("INSERT INTO t VALUES (1), (NULL), (3)") + .execute() + .await + .expect("INSERT failed"); + + #[derive(Debug, Row, Deserialize, PartialEq)] + struct NullableRow { + n: Option, + } + + let rows = client + .query("SELECT n FROM t ORDER BY n ASC NULLS LAST") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 3); + assert_eq!(rows[0].n, Some(1)); + assert_eq!(rows[1].n, Some(3)); + assert_eq!(rows[2].n, None); +} + +#[tokio::test] +async fn native_low_cardinality_string() { + let client = prepare_native_database("lowcard").await; + + client + .query( + &format!("CREATE TABLE t{} (id UInt32, tag LowCardinality(String)) {}", on_cluster(), test_engine("tuple()")), + ) + .execute() + .await + .expect("CREATE failed"); + + client + .query("INSERT INTO t VALUES (1, 'foo'), (2, 'bar'), (3, 'foo')") + .execute() + .await + .expect("INSERT failed"); + + #[derive(Debug, Row, Deserialize, PartialEq)] + struct LCRow { + id: u32, + tag: String, + } + + let rows = client + .query("SELECT id, tag FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 3); + assert_eq!(rows[0], LCRow { id: 1, tag: "foo".into() }); + assert_eq!(rows[1], LCRow { id: 2, tag: "bar".into() }); + assert_eq!(rows[2], LCRow { id: 3, tag: "foo".into() }); +} + +#[tokio::test] +async fn native_multiple_blocks() { + let client = prepare_native_database("multiblock").await; + + client + .query(&format!("CREATE TABLE t{} (n UInt64) {}", on_cluster(), test_engine("n"))) + .execute() + .await + .expect("CREATE failed"); + + // Insert 10,000 rows — server will send multiple blocks + client + .query( + "INSERT INTO t SELECT number FROM system.numbers LIMIT 10000", + ) + .execute() + .await + .expect("INSERT failed"); + + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one::() + .await + .expect("count failed"); + + assert_eq!(count, 10_000); +} + +#[tokio::test] +async fn native_array_type() { + let client = prepare_native_database("array").await; + + client + .query(&format!("CREATE TABLE t{} (id UInt32, tags Array(String)) {}", on_cluster(), test_engine("tuple()"))) + .execute() + .await + .expect("CREATE failed"); + + client + .query("INSERT INTO t VALUES (1, ['a', 'b', 'c']), (2, []), (3, ['x'])") + .execute() + .await + .expect("INSERT failed"); + + #[derive(Debug, Row, Deserialize, PartialEq)] + struct ArrayRow { + id: u32, + tags: Vec, + } + + let rows = client + .query("SELECT id, tags FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 3); + assert_eq!(rows[0], ArrayRow { id: 1, tags: vec!["a".into(), "b".into(), "c".into()] }); + assert_eq!(rows[1], ArrayRow { id: 2, tags: vec![] }); + assert_eq!(rows[2], ArrayRow { id: 3, tags: vec!["x".into()] }); +} + +#[tokio::test] +async fn native_tuple_type() { + let client = prepare_native_database("tuple").await; + + client + .query(&format!("CREATE TABLE t{} (id UInt32, pair Tuple(String, UInt32)) {}", on_cluster(), test_engine("tuple()"))) + .execute() + .await + .expect("CREATE failed"); + + client + .query("INSERT INTO t VALUES (1, ('hello', 42)), (2, ('world', 7))") + .execute() + .await + .expect("INSERT failed"); + + #[derive(Debug, Row, Deserialize, PartialEq)] + struct TupleRow { + id: u32, + pair: (String, u32), + } + + let rows = client + .query("SELECT id, pair FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 2); + assert_eq!(rows[0], TupleRow { id: 1, pair: ("hello".into(), 42) }); + assert_eq!(rows[1], TupleRow { id: 2, pair: ("world".into(), 7) }); +} + +#[tokio::test] +async fn native_map_type() { + let client = prepare_native_database("map").await; + + client + .query( + &format!("CREATE TABLE t{} (id UInt32, attrs Map(String, UInt32)) {}", on_cluster(), test_engine("tuple()")), + ) + .execute() + .await + .expect("CREATE failed"); + + client + .query( + "INSERT INTO t VALUES (1, {'age': 30, 'score': 100}), (2, {})", + ) + .execute() + .await + .expect("INSERT failed"); + + // Maps deserialize as JSON strings from the native transport + #[derive(Debug, Row, Deserialize)] + struct MapRow { + id: u32, + attrs: String, + } + + let rows = client + .query("SELECT id, CAST(attrs, 'String') AS attrs FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].id, 1); + assert_eq!(rows[1].id, 2); + assert_eq!(rows[1].attrs, "{}"); +} + +#[tokio::test] +async fn native_json_legacy() { + let client = prepare_native_database("json_legacy").await; + + // Object('json') is the legacy JSON type — stored as String on the wire. + client + .query( + &format!("CREATE TABLE t{} (id UInt32, data Object('json')) {} \ + SETTINGS allow_experimental_object_type = 1", + on_cluster(), test_engine("tuple()")), + ) + .execute() + .await + .unwrap_or_else(|e| { + // Legacy Object type may not be available on all server versions — skip + eprintln!("SKIP native_json_legacy: {e}"); + }); + + // Insert and query are separate — if CREATE failed, just verify we skip cleanly + let rows_result = client + .query( + "SELECT id, CAST(data, 'String') AS data FROM t ORDER BY id ASC", + ) + .fetch_all::<(u32, String)>() + .await; + + // If table doesn't exist (CREATE failed), just ensure we don't panic + match rows_result { + Ok(rows) => { + // If data was inserted, verify round-trip + assert!(rows.len() <= 10, "unexpected row count"); + } + Err(e) => { + eprintln!("SKIP native_json_legacy query: {e}"); + } + } +} + +#[tokio::test] +async fn native_variant_type() { + let client = prepare_native_database("variant").await; + + // Variant type requires ClickHouse 24.x+ with allow_experimental_variant_type + let create_result = client + .query( + &format!("CREATE TABLE t{} (id UInt32, val Variant(String, UInt64)) {} \ + SETTINGS allow_experimental_variant_type = 1", + on_cluster(), test_engine("tuple()")), + ) + .execute() + .await; + + if let Err(e) = create_result { + eprintln!("SKIP native_variant_type (server may not support Variant): {e}"); + return; + } + + client + .query( + "INSERT INTO t VALUES \ + (1, 'hello'::Variant(String, UInt64)), \ + (2, 42::Variant(String, UInt64)), \ + (3, NULL)", + ) + .execute() + .await + .expect("INSERT failed"); + + // Variant cells come back as JSON strings + #[derive(Debug, Row, Deserialize)] + struct VariantRow { + id: u32, + val: String, + } + + let rows = client + .query("SELECT id, val FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 3); + assert_eq!(rows[0].id, 1); + assert_eq!(rows[1].id, 2); + assert_eq!(rows[2].id, 3); + // Values are JSON strings: "hello", 42, null + assert!(rows[0].val.contains("hello") || rows[0].val == "\"hello\"", + "unexpected: {:?}", rows[0].val); + assert!(rows[1].val == "42" || rows[1].val.contains("42"), + "unexpected: {:?}", rows[1].val); + assert_eq!(rows[2].val, "null"); +} + +#[tokio::test] +async fn native_json_new_type() { + let client = prepare_native_database("json_new").await; + + // New JSON type (ClickHouse 24.x+) + let create_result = client + .query( + &format!("CREATE TABLE t{} (id UInt32, data JSON) {} \ + SETTINGS allow_experimental_json_type = 1", + on_cluster(), test_engine("tuple()")), + ) + .execute() + .await; + + if let Err(e) = create_result { + eprintln!("SKIP native_json_new_type (server may not support JSON type): {e}"); + return; + } + + client + .query( + "INSERT INTO t VALUES \ + (1, '{\"name\": \"Alice\", \"age\": 30}'), \ + (2, '{\"name\": \"Bob\", \"score\": 95.5}')", + ) + .execute() + .await + .expect("INSERT failed"); + + // JSON columns come back as JSON strings via Dynamic wire format + #[derive(Debug, Row, Deserialize)] + struct JsonRow { + id: u32, + data: String, + } + + let rows = client + .query("SELECT id, data FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].id, 1); + assert_eq!(rows[1].id, 2); + // Data should be non-empty JSON representations + assert!(!rows[0].data.is_empty(), "JSON data should not be empty"); + assert!(!rows[1].data.is_empty(), "JSON data should not be empty"); +} + +#[tokio::test] +async fn native_ip_types() { + use std::net::{Ipv4Addr, Ipv6Addr}; + + let client = prepare_native_database("ip").await; + + client + .query( + &format!("CREATE TABLE t{} (id UInt32, v4 IPv4, v6 IPv6) {}", on_cluster(), test_engine("tuple()")), + ) + .execute() + .await + .expect("CREATE failed"); + + client + .query( + "INSERT INTO t VALUES \ + (1, '192.168.1.1', '::1'), \ + (2, '10.0.0.1', '2001:db8::1')", + ) + .execute() + .await + .expect("INSERT failed"); + + // IPv4 deserializes as u32 (raw LE bytes); IPv6 as [u8; 16] + #[derive(Debug, Row, Deserialize, PartialEq)] + struct IpRow { + id: u32, + // ClickHouse IPv4 = u32 in RowBinary; use serde helper or raw u32 + #[serde(with = "clickhouse::serde::ipv4")] + v4: Ipv4Addr, + // IPv6 = 16-byte array + v6: [u8; 16], + } + + let rows = client + .query("SELECT id, v4, v6 FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].id, 1); + assert_eq!(rows[0].v4, Ipv4Addr::new(192, 168, 1, 1)); + assert_eq!(rows[1].v4, Ipv4Addr::new(10, 0, 0, 1)); +} + +#[tokio::test] +async fn native_decimal_type() { + let client = prepare_native_database("decimal").await; + + client + .query( + &format!("CREATE TABLE t{} (id UInt32, price Decimal64(2)) {}", on_cluster(), test_engine("tuple()")), + ) + .execute() + .await + .expect("CREATE failed"); + + client + .query("INSERT INTO t VALUES (1, 12.34), (2, 99.99), (3, 0.01)") + .execute() + .await + .expect("INSERT failed"); + + // Decimal64 is stored as i64 (scaled integer) — maps to i64 in Rust + #[derive(Debug, Row, Deserialize)] + struct DecimalRow { + id: u32, + price: i64, // raw scaled integer: 1234, 9999, 1 + } + + let rows = client + .query("SELECT id, price FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 3); + assert_eq!(rows[0].price, 1234); // 12.34 × 100 + assert_eq!(rows[1].price, 9999); // 99.99 × 100 + assert_eq!(rows[2].price, 1); // 0.01 × 100 +} + +#[tokio::test] +async fn native_empty_result() { + let client = prepare_native_database("empty").await; + + client + .query(&format!("CREATE TABLE t{} (n UInt32) {}", on_cluster(), test_engine("tuple()"))) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Deserialize)] + struct Row { + n: u32, + } + + let rows = client + .query("SELECT n FROM t") + .fetch_all::() + .await + .expect("fetch failed"); + + assert!(rows.is_empty()); +} + +#[tokio::test] +async fn native_bool_type() { + let client = prepare_native_database("bool").await; + + client + .query(&format!( + "CREATE TABLE t{} (a Bool, b Bool) {}", + on_cluster(), + test_engine("tuple()") + )) + .execute() + .await + .expect("CREATE failed"); + + client + .query("INSERT INTO t VALUES (true, false)") + .execute() + .await + .expect("INSERT failed"); + + #[derive(Debug, Row, Deserialize, PartialEq)] + struct BoolRow { + a: bool, + b: bool, + } + + let rows = client + .query("SELECT a, b FROM t") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0], BoolRow { a: true, b: false }); +} + +#[tokio::test] +async fn native_insert_scalars() { + let client = prepare_native_database("insert_scalars").await; + + client + .query(&format!( + "CREATE TABLE t{} ( + id UInt32, + val Int64, + f Float64 + ) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Serialize, Deserialize, PartialEq)] + struct ScalarRow { + id: u32, + val: i64, + f: f64, + } + + let mut insert = client.insert::("t"); + insert + .write(&ScalarRow { id: 1, val: -100, f: 3.14 }) + .await + .expect("write 1 failed"); + insert + .write(&ScalarRow { id: 2, val: 200, f: 2.718 }) + .await + .expect("write 2 failed"); + insert.end().await.expect("end failed"); + + let rows = client + .query("SELECT id, val, f FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].id, 1); + assert_eq!(rows[0].val, -100); + assert!((rows[0].f - 3.14f64).abs() < 1e-10); + assert_eq!(rows[1].id, 2); + assert_eq!(rows[1].val, 200); + assert!((rows[1].f - 2.718f64).abs() < 1e-10); +} + +#[tokio::test] +async fn native_insert_strings() { + let client = prepare_native_database("insert_strings").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, name String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Serialize, Deserialize, PartialEq)] + struct StringRow { + id: u32, + name: String, + } + + let mut insert = client.insert::("t"); + insert + .write(&StringRow { id: 1, name: "Alice".into() }) + .await + .expect("write failed"); + insert + .write(&StringRow { id: 2, name: "Bob".into() }) + .await + .expect("write failed"); + insert.end().await.expect("end failed"); + + let rows = client + .query("SELECT id, name FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 2); + assert_eq!(rows[0], StringRow { id: 1, name: "Alice".into() }); + assert_eq!(rows[1], StringRow { id: 2, name: "Bob".into() }); +} + +#[tokio::test] +async fn native_insert_nullable() { + let client = prepare_native_database("insert_nullable").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, val Nullable(Int32)) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Serialize, Deserialize, PartialEq)] + struct NullRow { + id: u32, + val: Option, + } + + let mut insert = client.insert::("t"); + insert + .write(&NullRow { id: 1, val: Some(42) }) + .await + .expect("write 1 failed"); + insert + .write(&NullRow { id: 2, val: None }) + .await + .expect("write 2 failed"); + insert + .write(&NullRow { id: 3, val: Some(-7) }) + .await + .expect("write 3 failed"); + insert.end().await.expect("end failed"); + + let rows = client + .query("SELECT id, val FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 3); + assert_eq!(rows[0], NullRow { id: 1, val: Some(42) }); + assert_eq!(rows[1], NullRow { id: 2, val: None }); + assert_eq!(rows[2], NullRow { id: 3, val: Some(-7) }); +} + +#[tokio::test] +async fn native_insert_empty() { + // Calling end() without any writes should not error. + let client = prepare_native_database("insert_empty").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Serialize, Deserialize)] + struct EmptyRow { + id: u32, + } + + let insert = client.insert::("t"); + insert.end().await.expect("empty end failed"); + + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one::() + .await + .expect("count failed"); + + assert_eq!(count, 0); +} + +#[tokio::test] +async fn native_inserter_basic() { + let client = prepare_native_database("inserter_basic").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, val String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Serialize, Deserialize, PartialEq)] + struct TestRow { + id: u32, + val: String, + } + + let mut inserter = client + .inserter::("t") + .with_max_rows(10_000); + + for i in 0u32..100 { + inserter + .write(&TestRow { id: i, val: format!("item_{i}") }) + .await + .expect("write failed"); + } + inserter.commit().await.expect("commit failed"); + inserter.end().await.expect("end failed"); + + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one::() + .await + .expect("count failed"); + + assert_eq!(count, 100); +} + +#[tokio::test] +async fn native_schema_cache() { + let client = prepare_native_database("schema_cache").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, name String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + // Cache is empty before any INSERT. + assert!(client.cached_schema("t").is_none()); + + // After an INSERT, the cache is populated from the server's column headers. + #[derive(Debug, Row, Serialize)] + struct TestRow { + id: u32, + name: String, + } + + let mut insert = client.insert::("t"); + insert + .write(&TestRow { id: 1, name: "x".into() }) + .await + .expect("write failed"); + insert.end().await.expect("end failed"); + + let schema = client.cached_schema("t").expect("schema should be cached after INSERT"); + // Server returns the actual column types; just verify names are present. + let names: Vec<&str> = schema.iter().map(|(n, _)| n.as_str()).collect(); + assert!(names.contains(&"id"), "expected 'id' in schema"); + assert!(names.contains(&"name"), "expected 'name' in schema"); + + // fetch_schema should also populate the cache. + client.clear_cached_schema("t"); + let fetched = client.fetch_schema("t").await.expect("fetch_schema failed"); + assert!(!fetched.is_empty()); +} + +#[tokio::test] +async fn native_insert_array() { + let client = prepare_native_database("insert_array").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, tags Array(String)) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Serialize, Deserialize, PartialEq)] + struct ArrayRow { + id: u32, + tags: Vec, + } + + let mut insert = client.insert::("t"); + insert + .write(&ArrayRow { id: 1, tags: vec!["alpha".into(), "beta".into()] }) + .await + .expect("write 1 failed"); + insert + .write(&ArrayRow { id: 2, tags: vec![] }) + .await + .expect("write 2 (empty array) failed"); + insert + .write(&ArrayRow { id: 3, tags: vec!["gamma".into()] }) + .await + .expect("write 3 failed"); + insert.end().await.expect("end failed"); + + let rows = client + .query("SELECT id, tags FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 3); + assert_eq!(rows[0], ArrayRow { id: 1, tags: vec!["alpha".into(), "beta".into()] }); + assert_eq!(rows[1], ArrayRow { id: 2, tags: vec![] }); + assert_eq!(rows[2], ArrayRow { id: 3, tags: vec!["gamma".into()] }); +} + +#[tokio::test] +async fn native_insert_nested_array() { + let client = prepare_native_database("insert_nested_array").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, vals Array(UInt32)) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Serialize, Deserialize, PartialEq)] + struct NumArrayRow { + id: u32, + vals: Vec, + } + + let mut insert = client.insert::("t"); + insert + .write(&NumArrayRow { id: 1, vals: vec![10, 20, 30] }) + .await + .expect("write failed"); + insert + .write(&NumArrayRow { id: 2, vals: vec![1] }) + .await + .expect("write failed"); + insert.end().await.expect("end failed"); + + let rows = client + .query("SELECT id, vals FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].vals, vec![10u32, 20, 30]); + assert_eq!(rows[1].vals, vec![1u32]); +} + +#[tokio::test] +async fn native_insert_tuple() { + let client = prepare_native_database("insert_tuple").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, pair Tuple(String, UInt32)) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Serialize, Deserialize, PartialEq)] + struct TupleRow { + id: u32, + pair: (String, u32), + } + + let mut insert = client.insert::("t"); + insert + .write(&TupleRow { id: 1, pair: ("hello".into(), 42) }) + .await + .expect("write 1 failed"); + insert + .write(&TupleRow { id: 2, pair: ("world".into(), 7) }) + .await + .expect("write 2 failed"); + insert.end().await.expect("end failed"); + + let rows = client + .query("SELECT id, pair FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 2); + assert_eq!(rows[0], TupleRow { id: 1, pair: ("hello".into(), 42) }); + assert_eq!(rows[1], TupleRow { id: 2, pair: ("world".into(), 7) }); +} + +#[tokio::test] +async fn native_insert_map() { + let client = prepare_native_database("insert_map").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, counts Map(String, UInt32)) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + // ClickHouse Map serializes in RowBinary as varuint(n) + [k1, v1, k2, v2, ...] + // HashMap does this via serde map serialization. + use std::collections::HashMap; + + #[derive(Debug, Row, Serialize, Deserialize)] + struct MapRow { + id: u32, + counts: HashMap, + } + + let mut insert = client.insert::("t"); + let mut m1 = HashMap::new(); + m1.insert("a".to_string(), 1u32); + m1.insert("b".to_string(), 2u32); + insert.write(&MapRow { id: 1, counts: m1 }).await.expect("write 1 failed"); + + let m2 = HashMap::new(); + insert.write(&MapRow { id: 2, counts: m2 }).await.expect("write 2 (empty map) failed"); + + insert.end().await.expect("end failed"); + + // Read back: verify map size and spot-check a value + #[derive(Debug, Row, Deserialize)] + struct ReadRow { + id: u32, + n: u64, // number of entries + val_a: u32, // counts['a'] for row 1 + } + + let rows = client + .query("SELECT id, length(counts) AS n, counts['a'] AS val_a FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].id, 1); + assert_eq!(rows[0].n, 2); // two entries + assert_eq!(rows[0].val_a, 1); // counts['a'] == 1 + assert_eq!(rows[1].id, 2); + assert_eq!(rows[1].n, 0); // empty map +} + +#[tokio::test] +async fn native_insert_lz4() { + // Verify that INSERT works correctly with LZ4 compression enabled. + let client = prepare_native_database("insert_lz4").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, name String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + // Build a client with LZ4 compression pointing at the same database. + let lz4_client = get_native_client() + .with_lz4() + .with_database("chrs_native_insert_lz4"); + + #[derive(Debug, Row, Serialize, Deserialize)] + struct Row { + id: u32, + name: String, + } + + let mut insert = lz4_client.insert::("t"); + insert.write(&Row { id: 1, name: "alice".to_string() }).await.expect("write 1 failed"); + insert.write(&Row { id: 2, name: "bob".to_string() }).await.expect("write 2 failed"); + insert.end().await.expect("end failed"); + + // Read back without compression to confirm data integrity. + let rows = client + .query("SELECT id, name FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].id, 1); + assert_eq!(rows[0].name, "alice"); + assert_eq!(rows[1].id, 2); + assert_eq!(rows[1].name, "bob"); +} From a942a2adad5ce1cfa835cffaa9020ce3637c5565 Mon Sep 17 00:00:00 2001 From: Derek Date: Tue, 10 Mar 2026 12:58:28 +1100 Subject: [PATCH 2/6] =?UTF-8?q?feat(native):=20comprehensive=20type=20cove?= =?UTF-8?q?rage=20=E2=80=94=20extended=20scalars,=20DateTime,=20Decimal,?= =?UTF-8?q?=20Enum,=20UUID,=20BFloat16,=20Point,=20Time?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds tests and bug fixes for all previously untested types: - UInt128/Int128 (u128/i128), UInt256/Int256 ([u8;32] raw LE) - BFloat16 (u16 raw bits; 1.0 = 0x3F80, 2.0 = 0x4000) - UUID ([u8;16], ClickHouse stores as two LE uint64s: low-value at byte 8) - Enum8/Enum16 (wire-compatible with Int8/Int16; read as i8/i16) - Date/Date32/DateTime/DateTime64 (3, 6, 9 precision) with epoch-relative int values - Time/Time64(3) — seconds/milliseconds since midnight - Decimal32/Decimal128/Decimal256 (all sizes; Decimal64 was already tested) - Point (geo) — fix: Point is Tuple(Float64,Float64) in columnar format, not a flat 16-byte blob; all x-values come before all y-values in the native stream, so we must read two separate Float64 sub-columns and interleave per-row --- src/native/columns.rs | 22 ++- tests/it/native.rs | 389 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 406 insertions(+), 5 deletions(-) diff --git a/src/native/columns.rs b/src/native/columns.rs index 99be6412..482408b5 100644 --- a/src/native/columns.rs +++ b/src/native/columns.rs @@ -269,8 +269,6 @@ impl ColumnType { | Self::Time64 => Some(8), Self::Int128 | Self::UInt128 | Self::Uuid | Self::IPv6 | Self::Decimal128 => Some(16), Self::Int256 | Self::UInt256 | Self::Decimal256 => Some(32), - // Point = 2 × Float64 - Self::Point => Some(16), Self::FixedString(n) => Some(*n), Self::String | Self::Json @@ -282,7 +280,9 @@ impl ColumnType { | Self::Map(_, _) | Self::Variant(_) | Self::NewJson - | Self::Dynamic => None, + | Self::Dynamic + // Point is Tuple(Float64, Float64) in columnar format — not a flat 16-byte blob. + | Self::Point => None, } } } @@ -381,12 +381,24 @@ pub(crate) fn read_column<'a, R: ClickHouseRead + 'a>( | ColumnType::Decimal128 | ColumnType::Int256 | ColumnType::UInt256 - | ColumnType::Decimal256 - | ColumnType::Point => { + | ColumnType::Decimal256 => { let size = col_type.fixed_size().expect("size is known for fixed type"); read_fixed_column(reader, n, size).await } + // Point = Tuple(Float64, Float64) in the native columnar format: + // all N x-values come first, then all N y-values. + // Transpose here so each row becomes the 16 raw bytes [f64(x) || f64(y)]. + ColumnType::Point => { + let x_col = read_fixed_column(reader, n, 8).await?; + let y_col = read_fixed_column(reader, n, 8).await?; + Ok(x_col + .into_iter() + .zip(y_col) + .map(|(mut x, y)| { x.extend_from_slice(&y); x }) + .collect()) + } + ColumnType::String | ColumnType::Json => read_string_column(reader, n).await, ColumnType::FixedString(size) => read_fixed_string_column(reader, n, *size).await, ColumnType::Nullable(inner) => read_nullable_column(reader, n, inner).await, diff --git a/tests/it/native.rs b/tests/it/native.rs index 1ffedbbb..57c1ca8b 100644 --- a/tests/it/native.rs +++ b/tests/it/native.rs @@ -701,6 +701,395 @@ async fn native_bool_type() { assert_eq!(rows[0], BoolRow { a: true, b: false }); } +// --------------------------------------------------------------------------- +// Extended integer types: UInt128 / Int128 / UInt256 / Int256 +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn native_extended_int_types() { + let client = prepare_native_database("ext_int").await; + + client + .query(&format!( + "CREATE TABLE t{} ( + u128 UInt128, i128 Int128, + u256 UInt256, i256 Int256 + ) {}", + on_cluster(), + test_engine("tuple()") + )) + .execute() + .await + .expect("CREATE failed"); + + client + .query("INSERT INTO t VALUES (42, -42, 42, -42)") + .execute() + .await + .expect("INSERT failed"); + + // 256-bit types have no native Rust equivalent — read as raw 32-byte LE arrays. + #[derive(Debug, Row, Deserialize)] + struct ExtIntRow { + u128: u128, + i128: i128, + u256: [u8; 32], + i256: [u8; 32], + } + + let rows = client + .query("SELECT u128, i128, u256, i256 FROM t") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].u128, 42u128); + assert_eq!(rows[0].i128, -42i128); + // UInt256(42): first byte = 42, rest zero (LE) + assert_eq!(rows[0].u256[0], 42); + assert!(rows[0].u256[1..].iter().all(|&b| b == 0)); + // Int256(-42): two's complement 32-byte LE — last bytes all 0xFF + assert_eq!(rows[0].i256[31], 0xFF); +} + +// --------------------------------------------------------------------------- +// BFloat16 (brain float, 2-byte) and UUID +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn native_bfloat16_uuid() { + let client = prepare_native_database("bf16_uuid").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, bf BFloat16, uuid UUID) {}", + on_cluster(), + test_engine("tuple()") + )) + .execute() + .await + .expect("CREATE failed"); + + client + .query( + "INSERT INTO t VALUES \ + (1, 1.0, '00000000-0000-0000-0000-000000000001'), \ + (2, 2.0, 'ffffffff-ffff-ffff-ffff-ffffffffffff')", + ) + .execute() + .await + .expect("INSERT failed"); + + // BFloat16 = 2 raw bytes; read as u16 (raw bit pattern). + // UUID = 16 bytes. + #[derive(Debug, Row, Deserialize)] + struct Bf16UuidRow { + id: u32, + bf: u16, // raw BFloat16 bits + uuid: [u8; 16], + } + + let rows = client + .query("SELECT id, bf, uuid FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 2); + // BFloat16(1.0) = 0x3F80 = 16256 + assert_eq!(rows[0].bf, 0x3F80u16); + // BFloat16(2.0) = 0x4000 = 16384 + assert_eq!(rows[1].bf, 0x4000u16); + // ClickHouse stores UUID as two LE uint64s: high 8 bytes then low 8 bytes. + // UUID 00000000-0000-0000-0000-000000000001: + // high u64 = 0 → bytes [0..8] all zero + // low u64 = 1 → bytes [8..16] = [1, 0, 0, 0, 0, 0, 0, 0] (LE) + assert_eq!(rows[0].uuid[8], 1); + assert!(rows[0].uuid[..8].iter().all(|&b| b == 0)); + assert!(rows[0].uuid[9..].iter().all(|&b| b == 0)); + // UUID all-0xff + assert!(rows[1].uuid.iter().all(|&b| b == 0xFF)); +} + +// --------------------------------------------------------------------------- +// Enum8 and Enum16 +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn native_enum_types() { + let client = prepare_native_database("enums").await; + + client + .query(&format!( + "CREATE TABLE t{} ( + id UInt32, + e8 Enum8('low' = 1, 'med' = 2, 'high' = 3), + e16 Enum16('pending' = 100, 'active' = 200, 'closed' = 300) + ) {}", + on_cluster(), + test_engine("tuple()") + )) + .execute() + .await + .expect("CREATE failed"); + + client + .query( + "INSERT INTO t VALUES \ + (1, 'low', 'pending'), \ + (2, 'high', 'active'), \ + (3, 'med', 'closed')", + ) + .execute() + .await + .expect("INSERT failed"); + + // Enum8/16 are wire-compatible with Int8/Int16 — deserialize as raw integer discriminant. + #[derive(Debug, Row, Deserialize)] + struct EnumRow { + id: u32, + e8: i8, + e16: i16, + } + + let rows = client + .query("SELECT id, e8, e16 FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 3); + assert_eq!(rows[0].e8, 1); // 'low' = 1 + assert_eq!(rows[0].e16, 100); // 'pending' = 100 + assert_eq!(rows[1].e8, 3); // 'high' = 3 + assert_eq!(rows[1].e16, 200); // 'active' = 200 + assert_eq!(rows[2].e8, 2); // 'med' = 2 + assert_eq!(rows[2].e16, 300); // 'closed' = 300 +} + +// --------------------------------------------------------------------------- +// Date, Date32, DateTime, DateTime64 (multiple precisions) +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn native_datetime_all() { + let client = prepare_native_database("datetimes").await; + + client + .query(&format!( + "CREATE TABLE t{} ( + id UInt32, + d Date, + d32 Date32, + dt DateTime, + dt3 DateTime64(3), + dt6 DateTime64(6), + dt9 DateTime64(9) + ) {}", + on_cluster(), + test_engine("tuple()") + )) + .execute() + .await + .expect("CREATE failed"); + + client + .query( + "INSERT INTO t VALUES \ + (1, '1970-01-01', '1970-01-01', '1970-01-01 00:00:01', \ + '1970-01-01 00:00:00.001', '1970-01-01 00:00:00.000001', \ + '1970-01-01 00:00:00.000000001')", + ) + .execute() + .await + .expect("INSERT failed"); + + // Date = u16 (days since 1970-01-01), DateTime = u32 (unix seconds), + // DateTime64(N) = i64 (scaled: ×10^N from epoch). + #[derive(Debug, Row, Deserialize)] + struct DtRow { + id: u32, + d: u16, + d32: i32, + dt: u32, + dt3: i64, + dt6: i64, + dt9: i64, + } + + let rows = client + .query("SELECT id, d, d32, dt, dt3, dt6, dt9 FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].d, 0); // 1970-01-01 = day 0 + assert_eq!(rows[0].d32, 0); + assert_eq!(rows[0].dt, 1); // 1 second past epoch + assert_eq!(rows[0].dt3, 1); // 1 millisecond + assert_eq!(rows[0].dt6, 1); // 1 microsecond + assert_eq!(rows[0].dt9, 1); // 1 nanosecond +} + +// --------------------------------------------------------------------------- +// Decimal32 / Decimal128 / Decimal256 (Decimal64 already in native_decimal_type) +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn native_decimal_all_sizes() { + let client = prepare_native_database("decimals").await; + + client + .query(&format!( + "CREATE TABLE t{} ( + id UInt32, + d32 Decimal32(2), + d128 Decimal128(4), + d256 Decimal256(6) + ) {}", + on_cluster(), + test_engine("tuple()") + )) + .execute() + .await + .expect("CREATE failed"); + + client + .query("INSERT INTO t VALUES (1, 12.34, 1234.5678, 123456.789012)") + .execute() + .await + .expect("INSERT failed"); + + #[derive(Debug, Row, Deserialize)] + struct DecRow { + id: u32, + d32: i32, + d128: i128, + d256: [u8; 32], // raw 32-byte LE + } + + let rows = client + .query("SELECT id, d32, d128, d256 FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].d32, 1234); // 12.34 × 100 + assert_eq!(rows[0].d128, 12345678i128); // 1234.5678 × 10^4 + // d256: 123456789012 (123456.789012 × 10^6) — check first bytes + let expected: i64 = 123_456_789_012; + let le_bytes = expected.to_le_bytes(); + assert_eq!(&rows[0].d256[..8], &le_bytes); + assert!(rows[0].d256[8..].iter().all(|&b| b == 0)); +} + +// --------------------------------------------------------------------------- +// Time and Time64 +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn native_time_types() { + let client = prepare_native_database("times").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, t Time, t64 Time64(3)) {}", + on_cluster(), + test_engine("tuple()") + )) + .execute() + .await + .expect("CREATE failed"); + + client + .query( + "INSERT INTO t VALUES \ + (1, '01:02:03', '01:02:03.456'), \ + (2, '00:00:00', '00:00:00.000')", + ) + .execute() + .await + .expect("INSERT failed"); + + // Time = i32 (seconds since midnight), Time64(3) = i64 (milliseconds since midnight) + #[derive(Debug, Row, Deserialize)] + struct TimeRow { + id: u32, + t: i32, + t64: i64, + } + + let rows = client + .query("SELECT id, t, t64 FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 2); + // 01:02:03 = 1*3600 + 2*60 + 3 = 3723 seconds + assert_eq!(rows[0].t, 3723); + // 01:02:03.456 = 3723 * 1000 + 456 = 3723456 ms + assert_eq!(rows[0].t64, 3_723_456); + assert_eq!(rows[1].t, 0); + assert_eq!(rows[1].t64, 0); +} + +// --------------------------------------------------------------------------- +// Geo types: Point +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn native_geo_types() { + let client = prepare_native_database("geo").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, pt Point) {}", + on_cluster(), + test_engine("tuple()") + )) + .execute() + .await + .expect("CREATE failed"); + + client + .query( + "INSERT INTO t VALUES \ + (1, (1.5, 2.5)), \ + (2, (0.0, -90.0))", + ) + .execute() + .await + .expect("INSERT failed"); + + // Point = 2 × Float64 LE (16 raw bytes). Read as [u8; 16] to avoid + // relying on serde tuple deserialization, then decode f64 values manually. + #[derive(Debug, Row, Deserialize)] + struct GeoRow { + id: u32, + pt: [u8; 16], + } + + let rows = client + .query("SELECT id, pt FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 2); + let x0 = f64::from_le_bytes(rows[0].pt[..8].try_into().unwrap()); + let y0 = f64::from_le_bytes(rows[0].pt[8..].try_into().unwrap()); + assert!((x0 - 1.5).abs() < f64::EPSILON); + assert!((y0 - 2.5).abs() < f64::EPSILON); + let x1 = f64::from_le_bytes(rows[1].pt[..8].try_into().unwrap()); + let y1 = f64::from_le_bytes(rows[1].pt[8..].try_into().unwrap()); + assert!((x1 - 0.0).abs() < f64::EPSILON); + assert!((y1 - (-90.0)).abs() < f64::EPSILON); +} + #[tokio::test] async fn native_insert_scalars() { let client = prepare_native_database("insert_scalars").await; From 5dfe86c38c00cfe4498f40beea9968571f9294b2 Mon Sep 17 00:00:00 2001 From: Derek Date: Tue, 10 Mar 2026 13:13:39 +1100 Subject: [PATCH 3/6] feat(native): connection pooling for native TCP transport MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each NativeClient now owns a lazily-initialised bounded connection pool (default size: 10) instead of opening a fresh TCP connection per operation. Pool model: - tokio Semaphore caps total (idle + in-use) connections to pool_size - Idle connections stored in Mutex (FIFO — reuses recently returned connections first) - Pool created on first acquire() call with a snapshot of the current client config; builder methods that change connection params (with_addr, with_database, with_user, with_password, with_setting, with_lz4) reset the pool so fresh connections use updated config - clone() shares the same pool across client copies PooledConnection guard: - DerefMut for transparent use at all call sites - Returns connection to pool on drop - discard() marks a connection as broken — it is closed and the semaphore slot is freed instead of queuing the conn as idle - Cursor aborts (exception / incomplete read) call discard() to avoid returning misaligned connections to the pool - NativeInsert::abort() discards the connection (incomplete INSERT) New public API: NativeClient::with_pool_size(n) Added native_pool_reuse integration test: 20 sequential pings on a pool capped to 1 connection verify that the same connection is reused. --- src/native/client.rs | 397 +++++++++++++++++++++++++++++++++++++++++++ src/native/cursor.rs | 119 +++++++++++++ src/native/insert.rs | 179 +++++++++++++++++++ src/native/mod.rs | 34 ++++ src/native/pool.rs | 169 ++++++++++++++++++ src/native/query.rs | 108 ++++++++++++ tests/it/native.rs | 13 ++ 7 files changed, 1019 insertions(+) create mode 100644 src/native/client.rs create mode 100644 src/native/cursor.rs create mode 100644 src/native/insert.rs create mode 100644 src/native/mod.rs create mode 100644 src/native/pool.rs create mode 100644 src/native/query.rs diff --git a/src/native/client.rs b/src/native/client.rs new file mode 100644 index 00000000..89c5f9ef --- /dev/null +++ b/src/native/client.rs @@ -0,0 +1,397 @@ +//! Public `NativeClient` — a ClickHouse client using the native TCP protocol. +//! +//! Mirrors the basic API of [`crate::Client`] so integration tests can switch +//! between transports with minimal changes. +//! +//! # Connection model +//! +//! Each `NativeClient` owns a lazily-initialised connection pool (default +//! size: 10). Connections are returned to the pool after each query/insert +//! and reused by subsequent operations. Use [`with_pool_size`] to tune the +//! cap. The pool is per-client-instance; clones share the same pool. +//! +//! Builder methods that affect connection parameters (`with_addr`, +//! `with_database`, `with_user`, `with_password`, `with_setting`, `with_lz4`) +//! reset the pool so the next `acquire` opens fresh connections with the +//! updated config. +//! +//! [`with_pool_size`]: NativeClient::with_pool_size + +use std::net::{SocketAddr, ToSocketAddrs}; +use std::sync::{Arc, OnceLock}; + +use crate::error::Result; +use crate::native::insert::NativeInsert; +use crate::native::inserter::NativeInserter; +use crate::native::pool::{NativePool, PoolConfig, PooledConnection}; +use crate::native::protocol::NativeCompressionMethod; +use crate::native::query::NativeQuery; +use crate::native::schema::NativeSchemaCache; +use crate::row::Row; + +/// A ClickHouse client using the native binary TCP protocol (port 9000). +/// +/// # Example +/// +/// ```no_run +/// # async fn example() -> clickhouse::error::Result<()> { +/// use clickhouse::native::NativeClient; +/// +/// let client = NativeClient::default() +/// .with_addr("localhost:9000") +/// .with_database("default") +/// .with_user("default") +/// .with_password(""); +/// +/// client.query("CREATE TABLE t (n UInt32) ENGINE = Memory").execute().await?; +/// # Ok(()) } +/// ``` +/// Default connection pool size. +const DEFAULT_POOL_SIZE: usize = 10; + +#[derive(Clone)] +pub struct NativeClient { + addr: SocketAddr, + database: String, + username: String, + password: String, + compression: NativeCompressionMethod, + /// Shared schema cache (TTL 300 s by default). + schema_cache: Arc, + /// Per-query settings sent with every query on this client. + settings: Arc>, + /// Maximum connections (idle + in-use) in the pool. + pool_size: usize, + /// Lazily-initialised pool; reset whenever connection params change. + /// + /// Wrapped in `Arc` so `Clone` shares the same pool across copies of a + /// fully-configured client. + pool: Arc>>, +} + +impl Default for NativeClient { + fn default() -> Self { + Self { + addr: "127.0.0.1:9000".parse().expect("valid default addr"), + database: "default".to_string(), + username: "default".to_string(), + password: String::new(), + compression: NativeCompressionMethod::None, + schema_cache: NativeSchemaCache::new(300), + settings: Arc::new(Vec::new()), + pool_size: DEFAULT_POOL_SIZE, + pool: Arc::new(OnceLock::new()), + } + } +} + +impl NativeClient { + /// Set the server address (host:port). + /// + /// # Panics + /// + /// If `addr` cannot be resolved to a socket address. + #[must_use] + pub fn with_addr(mut self, addr: impl ToSocketAddrs) -> Self { + self.addr = addr + .to_socket_addrs() + .expect("invalid address") + .next() + .expect("no address resolved"); + self.pool = Arc::new(OnceLock::new()); + self + } + + /// Set the database name. + #[must_use] + pub fn with_database(mut self, database: impl Into) -> Self { + self.database = database.into(); + self.pool = Arc::new(OnceLock::new()); + self + } + + /// Set the username. + #[must_use] + pub fn with_user(mut self, user: impl Into) -> Self { + self.username = user.into(); + self.pool = Arc::new(OnceLock::new()); + self + } + + /// Set the password. + #[must_use] + pub fn with_password(mut self, password: impl Into) -> Self { + self.password = password.into(); + self.pool = Arc::new(OnceLock::new()); + self + } + + /// Enable LZ4 compression for query data. + #[must_use] + pub fn with_lz4(mut self) -> Self { + self.compression = NativeCompressionMethod::Lz4; + self.pool = Arc::new(OnceLock::new()); + self + } + + /// Set the maximum number of connections (idle + in-use) in the pool. + /// + /// Defaults to 10. Must be called before the first query/insert — + /// changing it after the pool has been initialised has no effect. + #[must_use] + pub fn with_pool_size(mut self, size: usize) -> Self { + self.pool_size = size; + self.pool = Arc::new(OnceLock::new()); + self + } + + /// Add a session-level setting sent with every query on this client. + /// + /// Settings are sent in the query packet and apply to all query types + /// (SELECT, INSERT, DDL). Common examples: + /// + /// ```no_run + /// # use clickhouse::native::NativeClient; + /// let client = NativeClient::default() + /// // Read-after-write consistency on replicated tables: + /// .with_setting("select_sequential_consistency", "1") + /// // Require N replicas to acknowledge an INSERT before returning: + /// .with_setting("insert_quorum", "2"); + /// ``` + #[must_use] + pub fn with_setting( + mut self, + name: impl Into, + value: impl Into, + ) -> Self { + Arc::make_mut(&mut self.settings).push((name.into(), value.into())); + self.pool = Arc::new(OnceLock::new()); + self + } + + /// Return all session-level settings configured on this client. + pub(crate) fn settings(&self) -> &[(String, String)] { + &self.settings + } + + /// Start a query. + pub fn query(&self, sql: &str) -> NativeQuery { + NativeQuery::new(self.clone(), sql) + } + + /// Begin a single INSERT statement for rows of type `T`. + /// + /// The connection is opened lazily on the first call to + /// [`NativeInsert::write`]. Call [`NativeInsert::end`] to commit. + /// + /// ```no_run + /// # async fn example() -> clickhouse::error::Result<()> { + /// use clickhouse::{Row, native::NativeClient}; + /// use serde::Serialize; + /// + /// #[derive(Row, Serialize)] + /// struct Event { id: u64, name: String } + /// + /// let client = NativeClient::default(); + /// let mut insert = client.insert::("events"); + /// insert.write(&Event { id: 1, name: "foo".into() }).await?; + /// insert.end().await?; + /// # Ok(()) } + /// ``` + pub fn insert(&self, table: &str) -> NativeInsert { + NativeInsert::new(self.clone(), table) + } + + /// Create a multi-batch inserter for rows of type `T`. + /// + /// Mirrors [`crate::inserter::Inserter`] for the native transport. + /// + /// ```no_run + /// # async fn example() -> clickhouse::error::Result<()> { + /// use clickhouse::{Row, native::NativeClient}; + /// use serde::Serialize; + /// use std::time::Duration; + /// + /// #[derive(Row, Serialize)] + /// struct Event { id: u64, name: String } + /// + /// let client = NativeClient::default(); + /// let mut ins = client.inserter::("events") + /// .with_max_rows(100_000) + /// .with_period(Some(Duration::from_secs(5))); + /// + /// ins.write(&Event { id: 1, name: "foo".into() }).await?; + /// ins.commit().await?; + /// ins.end().await?; + /// # Ok(()) } + /// ``` + pub fn inserter(&self, table: &str) -> NativeInserter { + NativeInserter::new(self, table) + } + + /// Return the cached schema for `table` if it has been populated. + /// + /// The cache is populated automatically during INSERT operations when the + /// server sends column headers. To fetch the schema proactively, use + /// [`NativeClient::fetch_schema`]. + pub fn cached_schema(&self, table: &str) -> Option> { + self.schema_cache.get(table) + } + + /// Fetch column schema for `table` from `system.columns`, bypassing the cache. + /// + /// Parses the result at the RowBinary level so no serde derive is required. + /// The result is stored in the TTL cache for future calls to [`cached_schema`]. + /// + /// [`cached_schema`]: NativeClient::cached_schema + pub async fn fetch_schema( + &self, + table: &str, + ) -> Result> { + if let Some(cached) = self.schema_cache.get(table) { + return Ok(cached); + } + let db = &self.database; + let sql = format!( + "SELECT name, type \ + FROM system.columns \ + WHERE database = '{db}' AND table = '{table}' \ + ORDER BY position" + ); + let columns = fetch_string_pairs(self, &sql).await?; + self.schema_cache.insert(table.to_string(), columns.clone()); + Ok(columns) + } + + /// Remove `table`'s schema from the cache, forcing a refresh on next access. + pub fn clear_cached_schema(&self, table: &str) { + self.schema_cache.invalidate(table); + } + + /// Remove all cached schemas. + pub fn clear_all_cached_schemas(&self) { + self.schema_cache.invalidate_all(); + } + + /// Populate the schema cache entry for `table` from the given column headers. + /// + /// Called internally after a successful `begin_insert` to cache the schema + /// the server reported. + pub(crate) fn cache_schema( + &self, + table: &str, + columns: &[(String, String)], + ) { + self.schema_cache + .insert(table.to_string(), columns.to_vec()); + } + + /// Acquire a connection from the pool, opening a new one if needed. + /// + /// The pool is created lazily on first call with a snapshot of the + /// current connection parameters. + pub(crate) async fn acquire(&self) -> Result { + let pool = self.pool.get_or_init(|| { + NativePool::new( + PoolConfig { + addr: self.addr, + database: self.database.clone(), + username: self.username.clone(), + password: self.password.clone(), + compression: self.compression, + settings: self.settings.as_ref().clone(), + }, + self.pool_size, + ) + }); + pool.acquire().await + } + + /// Ping the server. + pub async fn ping(&self) -> Result<()> { + let mut conn = self.acquire().await?; + conn.ping().await + } +} + +/// Execute a query expected to return two `String` columns and collect all rows +/// as `Vec<(String, String)>`, parsing RowBinary directly without serde. +async fn fetch_string_pairs( + client: &NativeClient, + sql: &str, +) -> Result> { + use crate::native::reader::ServerPacket; + + let mut conn = client.acquire().await?; + let revision = conn.server_revision(); + let compression = conn.compression(); + + crate::native::writer::send_query( + conn.writer_mut(), + "", + sql, + client.settings(), + revision, + compression, + ) + .await?; + crate::native::writer::send_empty_block(conn.writer_mut(), compression).await?; + + let mut result = Vec::new(); + + loop { + let packet = crate::native::reader::read_packet( + conn.reader_mut(), + revision, + compression, + ) + .await?; + match packet { + ServerPacket::Data(block) if block.num_rows > 0 => { + // Each element in row_data is one complete RowBinary row. + // Two String columns: parse varuint(len)+bytes twice per row. + for row in &block.row_data { + let (a, rest) = rb_read_string(row)?; + let (b, _) = rb_read_string(rest)?; + result.push((a, b)); + } + } + ServerPacket::EndOfStream => break, + ServerPacket::Exception(err) => { + return Err(crate::error::Error::BadResponse(err.to_string())); + } + _ => {} + } + } + + Ok(result) +} + +/// Parse one RowBinary-encoded `String` from the start of `bytes`. +/// Returns `(value, remaining_bytes)`. +fn rb_read_string(bytes: &[u8]) -> crate::error::Result<(String, &[u8])> { + if bytes.is_empty() { + return Err(crate::error::Error::NotEnoughData); + } + let mut len = 0u64; + let mut shift = 0u32; + let mut i = 0usize; + loop { + if i >= bytes.len() { + return Err(crate::error::Error::NotEnoughData); + } + let b = bytes[i]; + i += 1; + len |= u64::from(b & 0x7F) << shift; + shift += 7; + if b & 0x80 == 0 { + break; + } + } + let len = len as usize; + if i + len > bytes.len() { + return Err(crate::error::Error::NotEnoughData); + } + let s = String::from_utf8_lossy(&bytes[i..i + len]).into_owned(); + Ok((s, &bytes[i + len..])) +} diff --git a/src/native/cursor.rs b/src/native/cursor.rs new file mode 100644 index 00000000..e3f157bb --- /dev/null +++ b/src/native/cursor.rs @@ -0,0 +1,119 @@ +//! Row cursor for native protocol query results. +//! +//! Reads native data blocks, transposes columnar data to RowBinary format, +//! and deserializes rows using the existing `rowbinary::deserialize_row` machinery. + +use std::collections::VecDeque; +use std::marker::PhantomData; + +use crate::error::{Error, Result}; +use crate::native::client::NativeClient; +use crate::native::pool::PooledConnection; +use crate::native::reader::ServerPacket; +use crate::row::{RowOwned, RowRead}; +use crate::rowbinary; + +/// A cursor that emits owned deserialized rows from a native TCP query. +/// +/// `T` must be [`RowOwned`] — i.e., the deserialized value must not borrow from +/// the network buffer. This covers the vast majority of use cases. +pub struct NativeRowCursor { + client: NativeClient, + sql: String, + /// Buffered row bytes from already-received blocks. + row_buf: VecDeque>, + state: CursorState, + _marker: PhantomData T>, +} + +enum CursorState { + /// Initial state — connection not yet acquired from pool. + NotStarted, + /// Connection open, reading packets. + Reading(Box), + /// EndOfStream received — no more data. + Done, +} + +impl NativeRowCursor { + pub(crate) fn new(client: NativeClient, sql: String) -> Self { + Self { + client, + sql, + row_buf: VecDeque::new(), + state: CursorState::NotStarted, + _marker: PhantomData, + } + } + + /// Return the next deserialized row, or `None` at end of stream. + /// + /// `T` must be [`RowOwned`], meaning the result does not borrow from + /// the network buffer. This is required for correctness with async streaming. + pub async fn next(&mut self) -> Result> { + loop { + // Return a buffered row if available. + if let Some(row_bytes) = self.row_buf.pop_front() { + let mut slice: &[u8] = &row_bytes; + let value = rowbinary::deserialize_row::(&mut slice, None)?; + return Ok(Some(value)); + } + + match &mut self.state { + CursorState::Done => return Ok(None), + + CursorState::NotStarted => { + let mut conn = self.client.acquire().await?; + let revision = conn.server_revision(); + let compression = conn.compression(); + crate::native::writer::send_query( + conn.writer_mut(), + "", + &self.sql, + self.client.settings(), + revision, + compression, + ) + .await?; + crate::native::writer::send_empty_block(conn.writer_mut(), compression).await?; + self.state = CursorState::Reading(Box::new(conn)); + } + + CursorState::Reading(conn) => { + let revision = conn.server_revision(); + let compression = conn.compression(); + let packet = crate::native::reader::read_packet( + conn.reader_mut(), + revision, + compression, + ) + .await?; + + match packet { + ServerPacket::EndOfStream => { + self.state = CursorState::Done; + // Connection is returned to pool automatically when + // the old CursorState::Reading is dropped here. + } + ServerPacket::Data(block) => { + if block.num_rows > 0 { + self.row_buf.extend(block.row_data); + } + } + ServerPacket::Exception(err) => { + // Discard the connection — the query didn't complete + // cleanly; subsequent reads on this conn would be misaligned. + if let CursorState::Reading(mut conn) = + std::mem::replace(&mut self.state, CursorState::Done) + { + conn.discard(); + } + return Err(Error::BadResponse(err.to_string())); + } + _ => {} + } + } + } + } + } +} diff --git a/src/native/insert.rs b/src/native/insert.rs new file mode 100644 index 00000000..d3243976 --- /dev/null +++ b/src/native/insert.rs @@ -0,0 +1,179 @@ +//! `NativeInsert` — a single INSERT statement over the native TCP protocol. +//! +//! Mirrors the public API of [`crate::insert::Insert`] so code using the HTTP +//! client can switch to the native transport with minimal changes. +//! +//! # Usage +//! +//! ```no_run +//! # async fn example() -> clickhouse::error::Result<()> { +//! use clickhouse::{Row, native::NativeClient}; +//! use serde::Serialize; +//! +//! #[derive(Row, Serialize)] +//! struct Event { id: u64, name: String } +//! +//! let client = NativeClient::default(); +//! let mut insert = client.insert::("events"); +//! insert.write(&Event { id: 1, name: "foo".into() }).await?; +//! insert.end().await?; +//! # Ok(()) } +//! ``` +//! +//! # Behaviour +//! +//! - Connection is opened lazily on the first call to [`write`](NativeInsert::write). +//! - Rows are serialised to RowBinary and buffered in memory. +//! - The buffer is flushed (transposed to native columnar format and sent) when +//! it exceeds ~256 KiB, and also when [`end`](NativeInsert::end) is called. +//! - [`end`](NativeInsert::end) must be called to commit the INSERT. Dropping +//! without calling `end` silently aborts (connection dropped). + +use std::marker::PhantomData; + +use bytes::BytesMut; + +use crate::error::{Error, Result}; +use crate::native::client::NativeClient; +use crate::native::encode::{ColumnSchema, encode_columns}; +use crate::native::pool::PooledConnection; +use crate::row::{self, Row, RowWrite}; +use crate::rowbinary::serialize_row_binary; + +/// Desired flush threshold (~256 KiB uncompressed). +const BUFFER_SIZE: usize = 256 * 1024; +/// Soft flush limit — slightly below `BUFFER_SIZE` to avoid one extra allocation. +const MIN_CHUNK_SIZE: usize = BUFFER_SIZE - 2048; + +/// A single in-flight native INSERT statement. +/// +/// Call [`write`](NativeInsert::write) for each row, then +/// [`end`](NativeInsert::end) to commit. Dropping without `end` aborts. +#[must_use] +pub struct NativeInsert { + client: NativeClient, + /// `INSERT INTO table(col1, col2, …) FORMAT Native` + sql: String, + /// Table name, used to populate the schema cache after handshake. + table: String, + /// Pooled connection; `None` until the first `write`. + conn: Option, + /// Column schema received from the server after `begin_insert`. + columns: Vec, + /// Buffered rows as RowBinary, one `Vec` per row. + row_buf: Vec>, + /// Total bytes across all buffered rows (used for flush threshold). + row_bytes: usize, + _marker: PhantomData T>, +} + +impl NativeInsert { + /// Create a new `NativeInsert`. Connection is deferred until first write. + pub(crate) fn new(client: NativeClient, table: &str) -> Self { + let fields = row::join_column_names::() + .expect("the row type must be a struct or a wrapper around it"); + let sql = format!("INSERT INTO {table}({fields}) FORMAT Native"); + Self { + client, + sql, + table: table.to_string(), + conn: None, + columns: Vec::new(), + row_buf: Vec::new(), + row_bytes: 0, + _marker: PhantomData, + } + } + + /// Serialise `row` into the internal buffer and flush if above threshold. + /// + /// The future does not borrow `row` after it returns. + pub async fn write(&mut self, row: &T::Value<'_>) -> Result<()> + where + T: RowWrite, + { + // Ensure connection is open and we have the column schema. + self.ensure_connected().await?; + + // Serialise to RowBinary. + let mut rb = BytesMut::new(); + if let Err(e) = serialize_row_binary(&mut rb, row) { + self.abort(); + return Err(e); + } + let rb = rb.freeze().to_vec(); + self.row_bytes += rb.len(); + self.row_buf.push(rb); + + // Flush when the buffer is large enough. + if self.row_bytes >= MIN_CHUNK_SIZE { + if let Err(e) = self.flush().await { + self.abort(); + return Err(e); + } + } + Ok(()) + } + + /// Flush remaining buffered rows and signal end of INSERT to the server. + /// + /// Must be called to commit the INSERT. On error the connection is dropped. + pub async fn end(mut self) -> Result<()> { + if self.conn.is_none() { + // Nothing was written — open a connection and immediately close it cleanly. + if let Err(e) = self.ensure_connected().await { + return Err(e); + } + } + if !self.row_buf.is_empty() { + if let Err(e) = self.flush().await { + return Err(e); + } + } + self.conn + .as_mut() + .expect("conn must be open") + .finish_insert() + .await + } + + async fn ensure_connected(&mut self) -> Result<()> { + if self.conn.is_some() { + return Ok(()); + } + let mut conn = self.client.acquire().await?; + let headers = conn.begin_insert(&self.sql).await?; + self.client.cache_schema(&self.table, &headers); + self.columns = ColumnSchema::from_headers(&headers).map_err(|e| { + Error::BadResponse(format!("native INSERT: bad schema from server: {e}")) + })?; + self.conn = Some(conn); + Ok(()) + } + + async fn flush(&mut self) -> Result<()> { + let rows = std::mem::take(&mut self.row_buf); + let n = rows.len(); + self.row_bytes = 0; + if n == 0 { + return Ok(()); + } + let conn = self.conn.as_mut().expect("conn must be open during flush"); + let revision = conn.server_revision(); + let column_bytes = encode_columns(&rows, &self.columns, revision)?; + conn.send_insert_block(&column_bytes, self.columns.len(), n).await + } + + /// Abort the INSERT: discard the connection and clear the buffer. + /// + /// The server-side INSERT is incomplete — we must not return this + /// connection to the pool as subsequent protocol exchanges would be + /// misaligned. + fn abort(&mut self) { + if let Some(mut conn) = self.conn.take() { + conn.discard(); + } + self.row_buf.clear(); + self.row_bytes = 0; + } +} diff --git a/src/native/mod.rs b/src/native/mod.rs new file mode 100644 index 00000000..c521f89a --- /dev/null +++ b/src/native/mod.rs @@ -0,0 +1,34 @@ +//! ClickHouse native TCP protocol (port 9000). +//! +//! Alternative transport to the default HTTP/RowBinary path. Ported and +//! extended by HYPERI PTY LIMITED from the HyperI `clickhouse-arrow` fork. +//! API names follow the ClickHouse Go client convention. + +// HyperI CTO moonlighting — ClickHouse Rust client needed love, so here we are. + +pub(crate) mod block_info; +pub(crate) mod client_info; +pub(crate) mod client; +pub(crate) mod columns; +pub(crate) mod compression; +pub(crate) mod connection; +pub(crate) mod cursor; +pub(crate) mod encode; +pub(crate) mod error_codes; +pub(crate) mod insert; +pub(crate) mod inserter; +pub(crate) mod pool; +pub(crate) mod io; +pub(crate) mod protocol; +pub(crate) mod query; +pub(crate) mod reader; +pub(crate) mod schema; +pub(crate) mod sparse; +pub(crate) mod tcp; +pub(crate) mod writer; + +pub use self::client::NativeClient; +pub use self::cursor::NativeRowCursor; +pub use self::insert::NativeInsert; +pub use self::inserter::NativeInserter; +pub use self::query::NativeQuery; diff --git a/src/native/pool.rs b/src/native/pool.rs new file mode 100644 index 00000000..b3f479a4 --- /dev/null +++ b/src/native/pool.rs @@ -0,0 +1,169 @@ +//! Connection pool for the native TCP transport. +//! +//! Holds a bounded set of idle [`NativeConnection`]s so successive queries +//! and inserts can reuse TCP connections instead of paying handshake overhead +//! on every operation. +//! +//! # Model +//! +//! - A [`Semaphore`] caps the total number of connections (idle + in-use) to +//! `max_size`. Callers block on [`NativePool::acquire`] when the pool is +//! full until a permit becomes available. +//! - Idle connections are stored in a `Mutex` (FIFO, so recently +//! used connections are preferred). +//! - On [`PooledConnection`] drop the connection is returned to the idle +//! queue unless [`PooledConnection::discard`] was called, in which case it +//! is closed and the semaphore slot is released. + +use std::collections::VecDeque; +use std::net::SocketAddr; +use std::ops::{Deref, DerefMut}; +use std::sync::{Arc, Mutex}; + +use tokio::sync::Semaphore; + +use crate::error::{Error, Result}; +use crate::native::connection::NativeConnection; +use crate::native::protocol::NativeCompressionMethod; + +/// Parameters needed to open a new connection. +pub(crate) struct PoolConfig { + pub(crate) addr: SocketAddr, + pub(crate) database: String, + pub(crate) username: String, + pub(crate) password: String, + pub(crate) compression: NativeCompressionMethod, + pub(crate) settings: Vec<(String, String)>, +} + +/// A bounded idle-connection pool. +pub(crate) struct NativePool { + idle: Mutex>, + /// Total in-use + idle connections must not exceed `max_size`. + semaphore: Semaphore, + config: PoolConfig, + max_size: usize, +} + +impl NativePool { + pub(crate) fn new(config: PoolConfig, max_size: usize) -> Arc { + Arc::new(Self { + idle: Mutex::new(VecDeque::new()), + semaphore: Semaphore::new(max_size), + config, + max_size, + }) + } + + /// Acquire a connection. Waits if the pool is at capacity. + /// + /// Tries an idle connection first; opens a new one if none are available. + pub(crate) async fn acquire(self: &Arc) -> Result { + let permit = self + .semaphore + .acquire() + .await + .map_err(|_| Error::Custom("connection pool closed".into()))?; + // We manage permits manually — forget the RAII guard. + permit.forget(); + + let conn = { + let mut idle = self.idle.lock().expect("pool mutex"); + idle.pop_front() + }; + + let conn = match conn { + Some(c) => c, + None => { + NativeConnection::open( + &self.config.addr, + &self.config.database, + &self.config.username, + &self.config.password, + self.config.compression, + self.config.settings.clone(), + ) + .await + .inspect_err(|_| { + // Opening failed — release the permit so the pool slot isn't lost. + self.semaphore.add_permits(1); + })? + } + }; + + Ok(PooledConnection { + conn: Some(conn), + pool: Arc::clone(self), + return_to_pool: true, + }) + } + + /// Return a connection to the idle queue and release its semaphore slot. + fn return_conn(&self, conn: NativeConnection) { + { + let mut idle = self.idle.lock().expect("pool mutex"); + // Guard against exceeding max_size in the idle queue. + if idle.len() < self.max_size { + idle.push_back(conn); + } + // If somehow over capacity, just drop the connection. + } + self.semaphore.add_permits(1); + } + + /// Release a semaphore permit without returning the connection (broken path). + fn release_permit(&self) { + self.semaphore.add_permits(1); + } +} + +/// A connection borrowed from a [`NativePool`]. +/// +/// Dereferences to [`NativeConnection`] for transparent method calls. +/// When dropped, the connection is returned to the pool — unless +/// [`discard`](PooledConnection::discard) was called, in which case the +/// connection is closed and the pool slot is freed. +pub(crate) struct PooledConnection { + conn: Option, + pool: Arc, + return_to_pool: bool, +} + +impl PooledConnection { + /// Mark this connection as broken — it will be closed on drop instead of + /// returned to the pool. Call this when an I/O error or incomplete + /// protocol exchange leaves the connection in an unrecoverable state. + pub(crate) fn discard(&mut self) { + self.return_to_pool = false; + } +} + +impl Deref for PooledConnection { + type Target = NativeConnection; + fn deref(&self) -> &NativeConnection { + self.conn + .as_ref() + .expect("PooledConnection invariant: conn is always Some while borrowed") + } +} + +impl DerefMut for PooledConnection { + fn deref_mut(&mut self) -> &mut NativeConnection { + self.conn + .as_mut() + .expect("PooledConnection invariant: conn is always Some while borrowed") + } +} + +impl Drop for PooledConnection { + fn drop(&mut self) { + if let Some(conn) = self.conn.take() { + if self.return_to_pool { + self.pool.return_conn(conn); + } else { + drop(conn); + self.pool.release_permit(); + } + } + } +} diff --git a/src/native/query.rs b/src/native/query.rs new file mode 100644 index 00000000..7c1f8bc2 --- /dev/null +++ b/src/native/query.rs @@ -0,0 +1,108 @@ +//! Native query builder — mirrors `crate::query::Query` for the native transport. + +use crate::error::{Error, Result}; +use crate::native::client::NativeClient; +use crate::native::cursor::NativeRowCursor; +use crate::row::{RowOwned, RowRead}; + +/// A query being built for native transport execution. +/// +/// Follows the same builder pattern as [`crate::query::Query`]. +#[must_use] +pub struct NativeQuery { + client: NativeClient, + sql: String, +} + +impl NativeQuery { + pub(crate) fn new(client: NativeClient, sql: &str) -> Self { + Self { + client, + sql: sql.to_string(), + } + } + + /// Bind a parameter using simple string substitution. + /// + /// Replaces the next `?` placeholder in the SQL string. + /// + /// For production use, prefer parameterized queries with ClickHouse's + /// `{name: Type}` syntax via the HTTP client. + pub fn bind(mut self, value: impl std::fmt::Display) -> Self { + if let Some(pos) = self.sql.find('?') { + self.sql = format!( + "{}{}{}", + &self.sql[..pos], + value, + &self.sql[pos + 1..] + ); + } + self + } + + /// Execute a DDL or non-SELECT query (CREATE, DROP, INSERT, etc.). + pub async fn execute(self) -> Result<()> { + let mut conn = self.client.acquire().await?; + conn.execute_query(&self.sql).await + } + + /// Execute a SELECT query, returning a cursor over deserialized rows. + /// + /// # Type support + /// + /// The native transport supports: Int/UInt 8/16/32/64/128/256, Float32/64, + /// String, FixedString(N), UUID, Date, Date32, DateTime, DateTime64, + /// Nullable(T), and LowCardinality(T). + /// + /// Complex types (Array, Map, Tuple) are not yet supported. + /// Execute a SELECT query, returning a cursor over deserialized rows. + /// + /// `T` must be [`RowOwned`] — the deserialized value must not borrow from + /// the network buffer. + /// + /// # Type support + /// + /// Supported: Int/UInt 8/16/32/64/128/256, Float32/64, String, FixedString(N), + /// UUID, Date, Date32, DateTime, DateTime64, Nullable(T), LowCardinality(T). + /// + /// Not yet supported: Array, Map, Tuple. + pub fn fetch(self) -> Result> + where + T: RowOwned + RowRead, + { + Ok(NativeRowCursor::new(self.client, self.sql)) + } + + /// Fetch a single row. + pub async fn fetch_one(self) -> Result + where + T: RowOwned + RowRead, + { + match self.fetch::()?.next().await { + Ok(Some(row)) => Ok(row), + Ok(None) => Err(Error::RowNotFound), + Err(err) => Err(err), + } + } + + /// Fetch all rows into a Vec. + pub async fn fetch_all(self) -> Result> + where + T: RowOwned + RowRead, + { + let mut result = Vec::new(); + let mut cursor = self.fetch::()?; + while let Some(row) = cursor.next().await? { + result.push(row); + } + Ok(result) + } + + /// Fetch at most one row. + pub async fn fetch_optional(self) -> Result> + where + T: RowOwned + RowRead, + { + self.fetch::()?.next().await + } +} diff --git a/tests/it/native.rs b/tests/it/native.rs index 57c1ca8b..544cb3a0 100644 --- a/tests/it/native.rs +++ b/tests/it/native.rs @@ -103,6 +103,19 @@ async fn native_ping() { client.ping().await.expect("ping failed"); } +/// Verify that the connection pool reuses connections across queries. +/// +/// Run 20 sequential pings on a pool capped to 1 connection. If pooling +/// works, all 20 succeed because the same connection is returned each time. +/// Without pooling each ping would open a new connection. +#[tokio::test] +async fn native_pool_reuse() { + let client = get_native_client().with_pool_size(1); + for _ in 0..20 { + client.ping().await.expect("ping failed"); + } +} + #[tokio::test] async fn native_ddl() { let client = prepare_native_database("ddl").await; From 40d8df60f2246688c9e91cf09bd9373bed8e5f61 Mon Sep 17 00:00:00 2001 From: Derek Date: Tue, 10 Mar 2026 14:03:05 +1100 Subject: [PATCH 4/6] =?UTF-8?q?feat(native):=20add=20native=20transport=20?= =?UTF-8?q?infrastructure=20=E2=80=94=20INSERT,=20encoder,=20schema=20cach?= =?UTF-8?q?e?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New files: - protocol.rs, io.rs, tcp.rs, connection.rs: low-level TCP + handshake - block_info.rs, client_info.rs: protocol framing helpers - writer.rs: packet serialisation (query, data block, ping, addendum) - compression.rs: LZ4 block decompression - error_codes.rs: server error parsing - sparse.rs: sparse column deserialisation (Bool / custom_ser=1) - encode.rs: RowBinary → columnar transpose for INSERT - schema.rs: TTL-backed schema cache (NativeSchemaCache) - inserter.rs: multi-batch NativeInserter (mirrors HTTP Inserter) - batcher.rs, tests/it/batcher.rs: HTTP batcher (separate feature) Core changes: - Cargo.toml: deadpool, zstd, socket2 under native-transport feature - src/lib.rs: expose native module + public re-exports - tests/it/main.rs: wire in native and batcher test modules --- .gitignore | 4 + Cargo.toml | 7 + src/batcher.rs | 187 +++++++++++++++ src/lib.rs | 5 + src/native/block_info.rs | 98 ++++++++ src/native/client_info.rs | 142 ++++++++++++ src/native/compression.rs | 314 +++++++++++++++++++++++++ src/native/connection.rs | 264 +++++++++++++++++++++ src/native/encode.rs | 474 ++++++++++++++++++++++++++++++++++++++ src/native/error_codes.rs | 118 ++++++++++ src/native/inserter.rs | 257 +++++++++++++++++++++ src/native/io.rs | 306 ++++++++++++++++++++++++ src/native/protocol.rs | 386 +++++++++++++++++++++++++++++++ src/native/schema.rs | 70 ++++++ src/native/sparse.rs | 327 ++++++++++++++++++++++++++ src/native/tcp.rs | 82 +++++++ src/native/writer.rs | 209 +++++++++++++++++ tests/it/batcher.rs | 169 ++++++++++++++ tests/it/main.rs | 4 + 19 files changed, 3423 insertions(+) create mode 100644 src/batcher.rs create mode 100644 src/native/block_info.rs create mode 100644 src/native/client_info.rs create mode 100644 src/native/compression.rs create mode 100644 src/native/connection.rs create mode 100644 src/native/encode.rs create mode 100644 src/native/error_codes.rs create mode 100644 src/native/inserter.rs create mode 100644 src/native/io.rs create mode 100644 src/native/protocol.rs create mode 100644 src/native/schema.rs create mode 100644 src/native/sparse.rs create mode 100644 src/native/tcp.rs create mode 100644 src/native/writer.rs create mode 100644 tests/it/batcher.rs diff --git a/.gitignore b/.gitignore index fbc9a58c..61222563 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,7 @@ .idea target Cargo.lock + +# HyperI-local: private session files not for public commits +CLAUDE.md +STATE.md diff --git a/Cargo.toml b/Cargo.toml index 080a9866..82bf9771 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -92,12 +92,16 @@ default = ["lz4"] test-util = ["hyper/server"] inserter = ["dep:quanta"] +batcher = ["inserter", "tokio/time"] uuid = ["dep:uuid"] time = ["dep:time"] lz4 = ["dep:lz4_flex", "dep:cityhash-rs"] chrono = ["dep:chrono"] futures03 = [] +## Native TCP protocol transport +native-transport = ["dep:cityhash-rs", "dep:lz4_flex", "dep:zstd", "dep:socket2", "dep:deadpool", "tokio/net", "tokio/io-util"] + ## TLS native-tls = ["dep:hyper-tls"] # ext: native-tls-alpn @@ -145,6 +149,8 @@ lz4_flex = { version = "0.11.3", default-features = false, features = [ "std", ], optional = true } cityhash-rs = { version = "=1.0.1", optional = true } # exact version for safety, this package has been stable for years +zstd = { version = "0.13", optional = true } +socket2 = { version = "0.6", features = ["all"], optional = true } uuid = { version = "1", optional = true } time = { version = "0.3", optional = true } chrono = { version = "0.4", optional = true, features = ["serde"] } @@ -153,6 +159,7 @@ quanta = { version = "0.12", optional = true } polonius-the-crab = "0.5.0" bnum = "0.13.0" +deadpool = { version = "0.12", features = ["managed", "rt_tokio_1"], optional = true } [dev-dependencies] clickhouse-macros = { version = "0.3.0", path = "macros" } diff --git a/src/batcher.rs b/src/batcher.rs new file mode 100644 index 00000000..edf42c08 --- /dev/null +++ b/src/batcher.rs @@ -0,0 +1,187 @@ +//! Per-table batch inserter with automatic flushing. +//! +//! [`TableBatcher`] wraps [`Inserter`][crate::inserter::Inserter] with +//! a shared buffer and a background task that handles period-based flushes, +//! so callers don't need to poll `time_left()`. +//! +//! API names follow the ClickHouse Go client's +//! [`Batch`](https://pkg.go.dev/github.com/ClickHouse/clickhouse-go/v2/lib/driver#Batch) +//! interface: [`append`][TableBatcher::append] / [`flush`][TableBatcher::flush] / +//! [`send`][TableBatcher::send]. +//! +//! A flush fires when **any** of these thresholds are crossed: +//! - serialised bytes reach [`BatchConfig::max_bytes`] +//! - row count reaches [`BatchConfig::max_rows`] +//! - [`BatchConfig::max_period`] elapses (background task) + +use std::sync::Arc; + +use tokio::sync::Mutex; +use tokio::time::Duration; + +use crate::{ + Client, + error::Result, + inserter::{Inserter, Quantities}, + row::{Row, RowWrite}, +}; + +/// Flush thresholds for [`TableBatcher`]. +/// +/// Defaults align with ClickHouse's async-insert defaults: +/// `async_insert_max_data_size` = 10 MiB, `max_rows` = 100 000 (upper end of +/// the recommended per-insert batch size to avoid MergeTree part fragmentation). +#[derive(Debug, Clone)] +pub struct BatchConfig { + /// Flush when this many rows have been buffered. Default: `100_000`. + pub max_rows: u64, + /// Flush when serialised bytes reach this size. Default: `10 MiB`. + pub max_bytes: u64, + /// Flush after this period regardless of row/byte counts. Default: `5 s`. + /// + /// `None` disables period-based flushing — no background task is spawned. + pub max_period: Option, +} + +impl Default for BatchConfig { + fn default() -> Self { + Self { + max_rows: 100_000, + max_bytes: 10 * 1024 * 1024, + max_period: Some(Duration::from_secs(5)), + } + } +} + +impl BatchConfig { + /// Override the row-count flush threshold. + pub fn with_max_rows(mut self, n: u64) -> Self { + self.max_rows = n; + self + } + + /// Override the byte-size flush threshold. + pub fn with_max_bytes(mut self, n: u64) -> Self { + self.max_bytes = n; + self + } + + /// Override the period-based flush interval. + pub fn with_max_period(mut self, d: Duration) -> Self { + self.max_period = Some(d); + self + } + + /// Disable period-based flushing (no background task is spawned). + pub fn without_period(mut self) -> Self { + self.max_period = None; + self + } +} + +// HyperI CTO moonlighting — dfe-loader needed this and no one else was going to write it. + +struct BatcherInner { + inserter: Inserter, +} + +/// Thread-safe, auto-flushing batch inserter for a single ClickHouse table. +/// +/// Wraps [`Inserter`][crate::inserter::Inserter] behind an `Arc` for +/// concurrent writes and spawns a background task to handle period-based flushes. +/// +/// Unlike `Inserter`, this type accepts `&self` on [`append`][Self::append] +/// and [`flush`][Self::flush], so it can be shared across tasks via [`Arc`]. +/// +/// Concurrent appends are serialised through the internal mutex. On the hot path +/// (limits not yet reached) the lock covers only RowBinary serialisation. A flush +/// holds the lock across the network round-trip, but that happens at most once per batch. +/// +/// For multi-table writes create one `TableBatcher` per table and share via `Arc`. +pub struct TableBatcher { + inner: Arc>>, + flush_task: tokio::task::JoinHandle<()>, +} + +impl TableBatcher +where + T: Row + RowWrite + Send + 'static, + for<'a> ::Value<'a>: Send, +{ + /// Create a new `TableBatcher` for `table` using `config` thresholds. + /// + /// If `config.max_period` is `Some`, a background tokio task is spawned + /// to handle periodic flushes. + pub fn new(client: &Client, table: &str, config: BatchConfig) -> Self { + let inserter = client + .inserter::(table) + .with_max_rows(config.max_rows) + .with_max_bytes(config.max_bytes) + .with_period(config.max_period); + + let inner = Arc::new(Mutex::new(BatcherInner { inserter })); + let inner_bg = Arc::clone(&inner); + let period = config.max_period; + + let flush_task = tokio::spawn(async move { + let Some(p) = period else { return }; + + let mut interval = tokio::time::interval(p); + // Skip ticks that arrive while a flush is already in progress. + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + // First tick fires immediately at t=0; skip it to avoid flushing an empty buffer. + interval.tick().await; + + loop { + interval.tick().await; + let mut guard = inner_bg.lock().await; + if let Err(_err) = guard.inserter.commit().await { + // Errors surface on the next explicit append/flush/send call. + // Background tasks can't propagate errors to callers. + } + } + }); + + Self { inner, flush_task } + } + + /// Add `row` to the buffer. Flushes automatically if a threshold is crossed. + pub async fn append(&self, row: &::Value<'_>) -> Result<()> { + let mut guard = self.inner.lock().await; + guard.inserter.write(row).await?; + guard.inserter.commit().await?; + Ok(()) + } + + /// Force-flush all pending rows to ClickHouse immediately. + /// + /// Returns the [`Quantities`] sent. Useful when shutting down a subsystem + /// while other clones of this batcher are still alive. + pub async fn flush(&self) -> Result { + let mut guard = self.inner.lock().await; + guard.inserter.force_commit().await + } + + /// Flush remaining rows and shut down the batcher. + /// + /// Aborts the background flush task, waits for it to exit (so its `Arc` + /// clone is dropped), then finalises the INSERT via + /// [`Inserter::end`][crate::inserter::Inserter::end]. + /// + /// All other `Arc` holders over the same inner buffer must be dropped before + /// calling `send`. If any remain after the abort, a force-flush is done + /// instead of a clean `end`. + pub async fn send(self) -> Result { + self.flush_task.abort(); + let _ = self.flush_task.await; + + match Arc::try_unwrap(self.inner) { + Ok(mutex) => mutex.into_inner().inserter.end().await, + Err(arc) => { + // Another Arc clone still exists — force-commit what we can. + let mut guard = arc.lock().await; + guard.inserter.force_commit().await + } + } + } +} diff --git a/src/lib.rs b/src/lib.rs index 23892f15..60b09d58 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,6 +22,8 @@ pub mod insert; pub mod insert_formatted; #[cfg(feature = "inserter")] pub mod inserter; +#[cfg(feature = "batcher")] +pub mod batcher; pub mod query; pub mod serde; pub mod sql; @@ -43,6 +45,9 @@ mod rowbinary; #[cfg(feature = "inserter")] mod ticks; +#[cfg(feature = "native-transport")] +pub mod native; + /// A client containing HTTP pool. /// /// ### Cloning behavior diff --git a/src/native/block_info.rs b/src/native/block_info.rs new file mode 100644 index 00000000..c348d8c7 --- /dev/null +++ b/src/native/block_info.rs @@ -0,0 +1,98 @@ +//! Block metadata for ClickHouse native protocol. +//! +//! Each data block carries overflow/bucket info used by the server for +//! aggregation and distributed query routing. + +use tokio::io::{AsyncReadExt, AsyncWriteExt}; + +use crate::error::{Error, Result}; +use crate::native::io::{ClickHouseRead, ClickHouseWrite}; + +/// Metadata about a native protocol data block. +#[derive(Debug, Clone, Copy)] +pub(crate) struct BlockInfo { + pub(crate) is_overflows: bool, + pub(crate) bucket_num: i32, +} + +impl Default for BlockInfo { + fn default() -> Self { + BlockInfo { + is_overflows: false, + bucket_num: -1, + } + } +} + +impl BlockInfo { + pub(crate) async fn read_async(reader: &mut R) -> Result { + let mut info = Self::default(); + loop { + let field_num = reader.read_var_uint().await?; + match field_num { + 0 => break, + 1 => { + info.is_overflows = reader.read_u8().await? != 0; + } + 2 => { + info.bucket_num = reader.read_i32_le().await?; + } + n => { + return Err(Error::BadResponse(format!( + "native protocol: unknown block info field: {n}" + ))); + } + } + } + Ok(info) + } + + pub(crate) async fn write_async(&self, writer: &mut W) -> Result<()> { + writer.write_var_uint(1).await?; + writer + .write_u8(if self.is_overflows { 1 } else { 0 }) + .await?; + writer.write_var_uint(2).await?; + writer.write_i32_le(self.bucket_num).await?; + writer.write_var_uint(0).await?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::io::Cursor; + + use super::*; + + #[tokio::test] + async fn test_block_info_roundtrip() { + let info = BlockInfo { + is_overflows: true, + bucket_num: 42, + }; + + let mut buf = Vec::new(); + info.write_async(&mut buf).await.unwrap(); + + let mut reader = Cursor::new(buf); + let decoded = BlockInfo::read_async(&mut reader).await.unwrap(); + + assert!(decoded.is_overflows); + assert_eq!(decoded.bucket_num, 42); + } + + #[tokio::test] + async fn test_block_info_default_roundtrip() { + let info = BlockInfo::default(); + + let mut buf = Vec::new(); + info.write_async(&mut buf).await.unwrap(); + + let mut reader = Cursor::new(buf); + let decoded = BlockInfo::read_async(&mut reader).await.unwrap(); + + assert!(!decoded.is_overflows); + assert_eq!(decoded.bucket_num, -1); + } +} diff --git a/src/native/client_info.rs b/src/native/client_info.rs new file mode 100644 index 00000000..974907b2 --- /dev/null +++ b/src/native/client_info.rs @@ -0,0 +1,142 @@ +//! Client information sent during query execution. +//! +//! Version-gated fields are written conditionally based on the negotiated +//! protocol revision with the server. + +use tokio::io::AsyncWriteExt; + +use crate::error::Result; +use crate::native::io::ClickHouseWrite; +use crate::native::protocol::{ + DBMS_MIN_PROTOCOL_VERSION_WITH_DISTRIBUTED_DEPTH, + DBMS_MIN_PROTOCOL_VERSION_WITH_PARALLEL_REPLICAS, + DBMS_MIN_PROTOCOL_VERSION_WITH_QUERY_START_TIME, DBMS_MIN_REVISION_WITH_JWT_IN_INTERSERVER, + DBMS_MIN_REVISION_WITH_OPENTELEMETRY, DBMS_MIN_REVISION_WITH_QUERY_AND_LINE_NUMBERS, + DBMS_MIN_REVISION_WITH_QUOTA_KEY_IN_CLIENT_INFO, DBMS_MIN_REVISION_WITH_VERSION_PATCH, + DBMS_TCP_PROTOCOL_VERSION, +}; + +// Client version derived from this crate's Cargo.toml +const CLIENT_VERSION_MAJOR: u64 = 0; +const CLIENT_VERSION_MINOR: u64 = 14; +const CLIENT_VERSION_PATCH: u64 = 2; + +#[repr(u8)] +#[derive(PartialEq, Clone, Copy, Debug)] +#[allow(unused)] +pub(crate) enum QueryKind { + NoQuery, + InitialQuery, + SecondaryQuery, +} + +#[derive(Debug)] +pub(crate) struct ClientInfo<'a> { + pub(crate) kind: QueryKind, + pub(crate) initial_user: &'a str, + pub(crate) initial_query_id: &'a str, + pub(crate) initial_address: &'a str, + pub(crate) os_user: &'a str, + pub(crate) client_hostname: &'a str, + pub(crate) client_name: &'a str, + pub(crate) client_version_major: u64, + pub(crate) client_version_minor: u64, + pub(crate) client_version_patch: u64, + pub(crate) client_tcp_protocol_version: u64, + pub(crate) query_start_time: u64, + pub(crate) quota_key: &'a str, + pub(crate) distributed_depth: u64, +} + +impl Default for ClientInfo<'_> { + fn default() -> Self { + #[allow(clippy::cast_possible_truncation)] + let query_start_time = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or(std::time::Duration::from_secs(0)) + .as_micros() as u64; + + ClientInfo { + kind: QueryKind::InitialQuery, + initial_user: "", + initial_query_id: "", + initial_address: "0.0.0.0:0", + os_user: "", + client_hostname: "localhost", + client_name: "clickhouse-rs", + client_version_major: CLIENT_VERSION_MAJOR, + client_version_minor: CLIENT_VERSION_MINOR, + client_version_patch: CLIENT_VERSION_PATCH, + client_tcp_protocol_version: DBMS_TCP_PROTOCOL_VERSION, + query_start_time, + quota_key: "", + distributed_depth: 1, + } + } +} + +impl ClientInfo<'_> { + pub(crate) async fn write( + &self, + writer: &mut W, + revision: u64, + ) -> Result<()> { + writer.write_u8(self.kind as u8).await?; + if self.kind == QueryKind::NoQuery { + return Ok(()); + } + + writer.write_string(self.initial_user).await?; + writer.write_string(self.initial_query_id).await?; + writer.write_string(self.initial_address).await?; + + if revision >= DBMS_MIN_PROTOCOL_VERSION_WITH_QUERY_START_TIME { + writer.write_u64_le(self.query_start_time).await?; + } + + // interface = TCP = 1 + writer.write_u8(1).await?; + + writer.write_string(self.os_user).await?; + writer.write_string(self.client_hostname).await?; + writer.write_string(self.client_name).await?; + + writer + .write_var_uint(self.client_version_major) + .await?; + writer + .write_var_uint(self.client_version_minor) + .await?; + writer + .write_var_uint(self.client_tcp_protocol_version) + .await?; + + if revision >= DBMS_MIN_REVISION_WITH_QUOTA_KEY_IN_CLIENT_INFO { + writer.write_string(self.quota_key).await?; + } + if revision >= DBMS_MIN_PROTOCOL_VERSION_WITH_DISTRIBUTED_DEPTH { + writer.write_var_uint(self.distributed_depth).await?; + } + if revision >= DBMS_MIN_REVISION_WITH_VERSION_PATCH { + writer.write_var_uint(self.client_version_patch).await?; + } + if revision >= DBMS_MIN_REVISION_WITH_OPENTELEMETRY { + // No OpenTelemetry support in MVP + writer.write_u8(0).await?; + } + if revision >= DBMS_MIN_PROTOCOL_VERSION_WITH_PARALLEL_REPLICAS { + writer.write_var_uint(0).await?; // collaborate_with_initiator + writer.write_var_uint(0).await?; // count_participating_replicas + writer.write_var_uint(0).await?; // number_of_current_replica + } + if revision >= DBMS_MIN_REVISION_WITH_QUERY_AND_LINE_NUMBERS { + writer.write_var_uint(0).await?; // script_query_number + writer.write_var_uint(0).await?; // script_line_number + } + if revision >= DBMS_MIN_REVISION_WITH_JWT_IN_INTERSERVER { + writer.write_u8(0).await?; + } + + Ok(()) + } +} diff --git a/src/native/compression.rs b/src/native/compression.rs new file mode 100644 index 00000000..63349581 --- /dev/null +++ b/src/native/compression.rs @@ -0,0 +1,314 @@ +//! Compression/decompression for ClickHouse native protocol. +//! +//! LZ4 and ZSTD support with ClickHouse's custom frame format: +//! - 16 bytes: CityHash128 checksum +//! - 1 byte: compression method (0x82=LZ4, 0x90=ZSTD) +//! - 4 bytes: compressed size (incl. 9-byte header) +//! - 4 bytes: decompressed size +//! - N bytes: payload +//! +//! Checksum covers method+sizes+payload. + +use std::future::Future; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use futures_util::FutureExt; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt, ReadBuf}; + +use crate::error::{Error, Result}; +use crate::native::io::{ClickHouseRead, ClickHouseWrite}; +use crate::native::protocol::NativeCompressionMethod; + +/// Compress and write data in ClickHouse native chunk format. +#[allow(clippy::cast_possible_truncation)] +pub(crate) async fn compress_data( + writer: &mut W, + raw: &[u8], + compression: NativeCompressionMethod, +) -> Result<()> { + let decompressed_size = raw.len(); + let compressed_payload = match compression { + NativeCompressionMethod::Zstd => zstd::bulk::compress(raw, 1) + .map_err(|e| Error::Compression(Box::new(e)))?, + NativeCompressionMethod::Lz4 => lz4_flex::compress(raw), + NativeCompressionMethod::None => return Ok(()), + }; + + // Build header: method(1) + compressed_size(4) + decompressed_size(4) + payload + let mut frame = Vec::with_capacity(compressed_payload.len() + 9); + frame.push(compression.byte()); + frame.extend_from_slice(&(compressed_payload.len() as u32 + 9).to_le_bytes()); + frame.extend_from_slice(&(decompressed_size as u32).to_le_bytes()); + frame.extend_from_slice(&compressed_payload); + + let hash = cityhash_rs::cityhash_102_128(&frame); + writer.write_u64_le((hash >> 64) as u64).await?; + writer.write_u64_le(hash as u64).await?; + writer.write_all(&frame).await?; + + Ok(()) +} + +/// Read and decompress a single chunk. Validates CityHash128 checksum. +pub(crate) async fn decompress_data( + reader: &mut R, + compression: NativeCompressionMethod, +) -> Result> { + // Read checksum (16 bytes) + let checksum_high = reader + .read_u64_le() + .await + .map_err(|e| Error::Decompression(Box::new(e)))?; + let checksum_low = reader + .read_u64_le() + .await + .map_err(|e| Error::Decompression(Box::new(e)))?; + let checksum = (u128::from(checksum_high) << 64) | u128::from(checksum_low); + + // Read compression header (9 bytes) + let type_byte = reader + .read_u8() + .await + .map_err(|e| Error::Decompression(Box::new(e)))?; + if type_byte != compression.byte() { + return Err(Error::Decompression( + format!( + "unexpected compression algorithm for {compression}: 0x{type_byte:02x}" + ) + .into(), + )); + } + + let compressed_size = reader + .read_u32_le() + .await + .map_err(|e| Error::Decompression(Box::new(e)))?; + let decompressed_size = reader + .read_u32_le() + .await + .map_err(|e| Error::Decompression(Box::new(e)))?; + + // Sanity checks + if compressed_size > 100_000_000 || decompressed_size > 1_000_000_000 { + return Err(Error::Decompression("chunk size too large".into())); + } + + // Build the complete compressed block for checksum validation + let mut compressed = vec![0u8; compressed_size as usize]; + reader + .read_exact(&mut compressed[9..]) + .await + .map_err(|e| Error::Decompression(Box::new(e)))?; + compressed[0] = type_byte; + compressed[1..5].copy_from_slice(&compressed_size.to_le_bytes()); + compressed[5..9].copy_from_slice(&decompressed_size.to_le_bytes()); + + // Validate checksum + let calc_checksum = cityhash_rs::cityhash_102_128(&compressed); + if calc_checksum != checksum { + return Err(Error::Decompression( + format!("checksum mismatch: expected {checksum:032x}, got {calc_checksum:032x}") + .into(), + )); + } + + // Decompress + match compression { + NativeCompressionMethod::Lz4 => lz4_flex::decompress(&compressed[9..], decompressed_size as usize) + .map_err(|e| Error::Decompression(Box::new(e))), + NativeCompressionMethod::Zstd => zstd::bulk::decompress(&compressed[9..], decompressed_size as usize) + .map_err(|e| Error::Decompression(Box::new(e))), + NativeCompressionMethod::None => { + Err(Error::Decompression("attempted to decompress uncompressed data".into())) + } + } +} + +type BlockReadingFuture<'a, R> = + Pin, &'a mut R)>> + Send + Sync + 'a>>; + +/// Async reader that decompresses ClickHouse native protocol blocks on-the-fly. +pub(crate) struct DecompressionReader<'a, R: ClickHouseRead + 'static> { + mode: NativeCompressionMethod, + inner: Option<&'a mut R>, + decompressed: Vec, + position: usize, + block_reading_future: Option>, +} + +impl<'a, R: ClickHouseRead> DecompressionReader<'a, R> { + /// Create decompressor. Reads first chunk immediately. + pub(crate) async fn new(mode: NativeCompressionMethod, inner: &'a mut R) -> Result { + let decompressed = decompress_data(inner, mode).await?; + Ok(Self { + mode, + inner: Some(inner), + decompressed, + position: 0, + block_reading_future: None, + }) + } +} + +impl AsyncRead for DecompressionReader<'_, R> { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + if buf.remaining() == 0 { + return Poll::Ready(Ok(())); + } + + // Check if we have a pending decompression future + if let Some(block_reading_future) = self.block_reading_future.as_mut() { + match block_reading_future.poll_unpin(cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Ok((value, inner))) => { + drop(self.block_reading_future.take()); + self.decompressed = value; + self.position = 0; + self.inner = Some(inner); + } + Poll::Ready(Err(e)) => { + drop(self.block_reading_future.take()); + return Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + e, + ))); + } + } + } + + // Serve available data + let available = self.decompressed.len() - self.position; + if available > 0 { + let to_serve = available.min(buf.remaining()); + buf.put_slice(&self.decompressed[self.position..self.position + to_serve]); + self.position += to_serve; + return Poll::Ready(Ok(())); + } + + // Need more data — start reading next chunk + if let Some(inner) = self.inner.take() { + let mode = self.mode; + self.block_reading_future = Some(Box::pin(async move { + let value = decompress_data(inner, mode).await?; + Ok((value, inner)) + })); + return self.poll_read(cx, buf); + } + + // EOF + Poll::Ready(Ok(())) + } +} + +#[cfg(test)] +mod tests { + use std::io::Cursor; + + use tokio::io::AsyncReadExt; + + use super::*; + + #[tokio::test] + async fn test_roundtrip_lz4() { + let data = b"test data for LZ4 compression".to_vec(); + let mut buffer = Vec::new(); + compress_data(&mut buffer, &data, NativeCompressionMethod::Lz4) + .await + .unwrap(); + assert!(!buffer.is_empty()); + + let mut reader = Cursor::new(buffer); + let decompressed = decompress_data(&mut reader, NativeCompressionMethod::Lz4) + .await + .unwrap(); + assert_eq!(decompressed, data); + } + + #[tokio::test] + async fn test_roundtrip_zstd() { + let data = b"test data for ZSTD compression".to_vec(); + let mut buffer = Vec::new(); + compress_data(&mut buffer, &data, NativeCompressionMethod::Zstd) + .await + .unwrap(); + assert!(!buffer.is_empty()); + + let mut reader = Cursor::new(buffer); + let decompressed = decompress_data(&mut reader, NativeCompressionMethod::Zstd) + .await + .unwrap(); + assert_eq!(decompressed, data); + } + + #[tokio::test] + async fn test_compress_none_is_noop() { + let data = b"test data no compression".to_vec(); + let mut buffer = Vec::new(); + compress_data(&mut buffer, &data, NativeCompressionMethod::None) + .await + .unwrap(); + assert!(buffer.is_empty()); + } + + #[tokio::test] + async fn test_checksum_validation() { + let data = b"test data for checksum validation".to_vec(); + let mut buffer = Vec::new(); + compress_data(&mut buffer, &data, NativeCompressionMethod::Lz4) + .await + .unwrap(); + + // Corrupt the checksum + buffer[0] ^= 0xFF; + + let mut reader = Cursor::new(buffer); + let result = decompress_data(&mut reader, NativeCompressionMethod::Lz4).await; + assert!(result.is_err()); + let err_msg = result.unwrap_err().to_string(); + assert!(err_msg.contains("checksum mismatch"), "got: {err_msg}"); + } + + #[tokio::test] + async fn test_decompression_reader_single_chunk() { + let data = b"test data for single chunk reading".to_vec(); + let expected_len = data.len(); + + let mut buffer = Vec::new(); + compress_data(&mut buffer, &data, NativeCompressionMethod::Lz4) + .await + .unwrap(); + + let mut reader = Cursor::new(buffer); + let mut decomp_reader = + DecompressionReader::new(NativeCompressionMethod::Lz4, &mut reader) + .await + .unwrap(); + + let mut result = vec![0u8; expected_len]; + decomp_reader.read_exact(&mut result).await.unwrap(); + assert_eq!(result, data); + } + + #[tokio::test] + async fn test_roundtrip_both_algorithms() { + let original = b"This is a longer piece of test data that should compress well \ + with both LZ4 and ZSTD algorithms" + .to_vec(); + + for compression in [NativeCompressionMethod::Lz4, NativeCompressionMethod::Zstd] { + let mut compressed_buffer = Vec::new(); + compress_data(&mut compressed_buffer, &original, compression) + .await + .unwrap(); + + let mut reader = Cursor::new(compressed_buffer); + let decompressed = decompress_data(&mut reader, compression).await.unwrap(); + assert_eq!(decompressed, original, "round trip failed for {compression}"); + } + } +} diff --git a/src/native/connection.rs b/src/native/connection.rs new file mode 100644 index 00000000..a34b9d20 --- /dev/null +++ b/src/native/connection.rs @@ -0,0 +1,264 @@ +//! Connection management for ClickHouse native TCP protocol. +//! +//! Single-connection MVP — handles handshake, query execution, and packet +//! reading over a buffered TCP stream. + +use std::net::SocketAddr; +use std::pin::Pin; +use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; + +use tokio::io::{AsyncRead, BufReader, BufWriter, ReadBuf}; +use tokio::net::TcpStream; + +use crate::error::{Error, Result}; +use crate::native::protocol::{ + ChunkedProtocolMode, NativeCompressionMethod, ServerHello, DBMS_TCP_PROTOCOL_VERSION, +}; +use crate::native::reader::{self, ServerPacket}; +use crate::native::tcp::{self, CONN_READ_BUFFER, CONN_WRITE_BUFFER}; +use crate::native::writer; + +/// A single native TCP connection to ClickHouse. +pub(crate) struct NativeConnection { + reader: BufReader>, + writer: BufWriter>, + server_hello: ServerHello, + compression: NativeCompressionMethod, + settings: Vec<(String, String)>, + /// Set to `true` by [`crate::native::pool::PooledConnection::discard`] to + /// prevent this connection being returned to the idle pool on drop. + pub(crate) poisoned: bool, +} + +impl NativeConnection { + /// Connect and perform the handshake. + pub(crate) async fn open( + addr: &SocketAddr, + database: &str, + username: &str, + password: &str, + compression: NativeCompressionMethod, + settings: Vec<(String, String)>, + ) -> Result { + let stream = tcp::connect(addr).await?; + let (read_half, write_half) = tokio::io::split(stream); + let mut reader = BufReader::with_capacity(CONN_READ_BUFFER, read_half); + let mut writer = BufWriter::with_capacity(CONN_WRITE_BUFFER, write_half); + + // Send hello + writer::send_hello(&mut writer, database, username, password).await?; + + // Read hello response + let chunked_modes = ( + ChunkedProtocolMode::default(), + ChunkedProtocolMode::default(), + ); + let server_hello = + reader::read_hello(&mut reader, DBMS_TCP_PROTOCOL_VERSION, chunked_modes).await?; + + // Send addendum + writer::send_addendum(&mut writer, &server_hello).await?; + + Ok(Self { + reader, + writer, + server_hello, + compression, + settings, + poisoned: false, + }) + } + + /// Returns `true` if this connection has been marked as broken and should + /// not be returned to the idle pool. + pub(crate) fn is_poisoned(&self) -> bool { + self.poisoned + } + + /// Non-blocking liveness check for pool recycling. + /// + /// Returns `false` (connection should be discarded) if: + /// - the connection is poisoned + /// - the `BufReader` has unread bytes (leftover data from a previous query) + /// - the TCP socket reports EOF (server closed the connection) + /// - the TCP socket has unexpected data ready (protocol misalignment) + /// + /// Returns `true` only when the socket is clean and idle (no pending bytes). + pub(crate) fn check_alive(&mut self) -> bool { + if self.poisoned { + return false; + } + // Leftover bytes in the read buffer mean a previous query didn't drain + // completely — the connection is in an unknown state. + if !self.reader.buffer().is_empty() { + return false; + } + // Non-blocking poll: detect EOF or unexpected data without blocking. + // A Pending result means the socket is idle → connection is alive. + let mut buf = [0u8; 1]; + let mut read_buf = ReadBuf::new(&mut buf); + let waker = noop_waker(); + let mut cx = Context::from_waker(&waker); + match Pin::new(&mut self.reader).poll_read(&mut cx, &mut read_buf) { + Poll::Pending => true, // idle — connection is healthy + Poll::Ready(_) => false, // EOF or unexpected data — discard + } + } + + /// Get the server hello info. + #[allow(unused)] + pub(crate) fn server_hello(&self) -> &ServerHello { + &self.server_hello + } + + /// Negotiated server revision. + pub(crate) fn server_revision(&self) -> u64 { + self.server_hello.revision_version + } + + /// Compression method in use. + pub(crate) fn compression(&self) -> NativeCompressionMethod { + self.compression + } + + /// Mutable access to the write half for sending packets. + pub(crate) fn writer_mut(&mut self) -> &mut BufWriter> { + &mut self.writer + } + + /// Mutable access to the read half for receiving packets. + pub(crate) fn reader_mut(&mut self) -> &mut BufReader> { + &mut self.reader + } + + /// Execute a query and read all response packets until EndOfStream. + pub(crate) async fn execute_query(&mut self, query: &str) -> Result<()> { + let revision = self.server_hello.revision_version; + let compression = self.compression; + + writer::send_query(&mut self.writer, "", query, &self.settings, revision, compression).await?; + writer::send_empty_block(&mut self.writer, compression).await?; + + loop { + let packet = + reader::read_packet(&mut self.reader, revision, compression).await?; + match packet { + ServerPacket::EndOfStream => break, + ServerPacket::Exception(err) => { + return Err(Error::BadResponse(err.to_string())); + } + _ => {} + } + } + + Ok(()) + } + + /// Begin an INSERT operation. + /// + /// Sends `INSERT INTO table(cols) FORMAT Native` + an empty data block, + /// then reads server packets until the schema Data block (0 rows) arrives. + /// Returns the column headers `(name, type_name)` declared by the server. + pub(crate) async fn begin_insert( + &mut self, + query: &str, + ) -> Result> { + let revision = self.server_hello.revision_version; + let compression = self.compression; + + writer::send_query(&mut self.writer, "", query, &self.settings, revision, compression).await?; + writer::send_empty_block(&mut self.writer, compression).await?; + + loop { + let packet = + reader::read_packet(&mut self.reader, revision, compression).await?; + match packet { + reader::ServerPacket::Data(block) => { + return Ok(block + .column_headers + .into_iter() + .map(|h| (h.name, h.type_name)) + .collect()); + } + reader::ServerPacket::Exception(err) => { + return Err(Error::BadResponse(err.to_string())); + } + _ => {} // skip Progress, ProfileInfo, etc. + } + } + } + + /// Send one data block during an INSERT. + /// + /// `column_bytes` must be produced by [`crate::native::encode::encode_columns`]. + pub(crate) async fn send_insert_block( + &mut self, + column_bytes: &[u8], + num_columns: usize, + num_rows: usize, + ) -> Result<()> { + writer::send_data_block( + &mut self.writer, + num_columns, + num_rows, + column_bytes, + self.compression, + ) + .await + } + + /// Finish an INSERT: send the empty terminator block and consume until + /// `EndOfStream` (or surface any server exception). + pub(crate) async fn finish_insert(&mut self) -> Result<()> { + let revision = self.server_hello.revision_version; + let compression = self.compression; + + writer::send_empty_block(&mut self.writer, compression).await?; + + loop { + let packet = + reader::read_packet(&mut self.reader, revision, compression).await?; + match packet { + reader::ServerPacket::EndOfStream => return Ok(()), + reader::ServerPacket::Exception(err) => { + return Err(Error::BadResponse(err.to_string())); + } + _ => {} // skip Progress, ProfileInfo, etc. + } + } + } + + /// Send ping and wait for pong. + pub(crate) async fn ping(&mut self) -> Result<()> { + let revision = self.server_hello.revision_version; + let compression = self.compression; + + writer::send_ping(&mut self.writer).await?; + loop { + let packet = + reader::read_packet(&mut self.reader, revision, compression).await?; + match packet { + ServerPacket::Pong => return Ok(()), + ServerPacket::Exception(err) => { + return Err(Error::BadResponse(err.to_string())); + } + _ => continue, + } + } + } +} + +/// A no-op [`Waker`] used for non-blocking `poll_read` calls in `check_alive`. +/// +/// The waker never schedules anything — it is used purely to drive a single +/// synchronous poll without registering for wake-up notifications. +fn noop_waker() -> Waker { + const VTABLE: RawWakerVTable = RawWakerVTable::new( + |p| RawWaker::new(p, &VTABLE), // clone + |_| {}, // wake + |_| {}, // wake_by_ref + |_| {}, // drop + ); + // SAFETY: the vtable is a no-op; the data pointer is never dereferenced. + unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &VTABLE)) } +} diff --git a/src/native/encode.rs b/src/native/encode.rs new file mode 100644 index 00000000..66c0e33c --- /dev/null +++ b/src/native/encode.rs @@ -0,0 +1,474 @@ +//! Columnar block encoder for native INSERT. +//! +//! Transposes row-oriented RowBinary data (one `Vec` per row) into the +//! native columnar wire format used by ClickHouse data blocks. +//! +//! # Supported types for INSERT +//! +//! All scalar fixed-size types, String, FixedString(N), Nullable(T), +//! Array(T), Map(K, V), Tuple(T1..Tn), and nested combinations thereof. +//! LowCardinality is stripped to its inner type (ClickHouse accepts plain values). +//! Variant, Dynamic, and JSON are not yet supported. + +use crate::error::{Error, Result}; +use crate::native::columns::ColumnType; +use crate::native::io::ClickHouseBytesWrite; +use crate::native::protocol::DBMS_MIN_PROTOCOL_VERSION_WITH_CUSTOM_SERIALIZATION; + +/// Column schema entry for a native INSERT block. +#[derive(Debug, Clone)] +pub(crate) struct ColumnSchema { + /// Column name as declared to the server. + pub(crate) name: String, + /// Type name string sent on the wire (LowCardinality stripped). + pub(crate) type_name: String, + /// Parsed column type used for encoding decisions. + pub(crate) col_type: ColumnType, +} + +impl ColumnSchema { + /// Build a `ColumnSchema` list from server-provided `(name, type_name)` pairs. + /// + /// LowCardinality wrappers are stripped — the inner type is sent on wire, + /// which ClickHouse accepts transparently. + pub(crate) fn from_headers(headers: &[(String, String)]) -> Result> { + headers + .iter() + .map(|(name, type_name)| { + let col_type = + ColumnType::parse(type_name).ok_or_else(|| { + Error::BadResponse(format!( + "native INSERT: unsupported column type '{type_name}' \ + for column '{name}'" + )) + })?; + // Strip LowCardinality: send inner type bytes, CH handles encoding + let (effective_type, effective_name) = + strip_low_cardinality(col_type, type_name); + Ok(ColumnSchema { + name: name.clone(), + type_name: effective_name, + col_type: effective_type, + }) + }) + .collect() + } +} + +/// Recursively strip `LowCardinality(...)` returning the inner `(ColumnType, type_name)`. +fn strip_low_cardinality(col_type: ColumnType, type_name: &str) -> (ColumnType, String) { + match col_type { + ColumnType::LowCardinality(inner) => { + let inner_name = extract_inner(type_name, "LowCardinality"); + strip_low_cardinality(*inner, inner_name) + } + other => (other, type_name.to_string()), + } +} + +fn extract_inner<'a>(s: &'a str, wrapper: &str) -> &'a str { + let prefix = format!("{wrapper}("); + if let Some(rest) = s.strip_prefix(prefix.as_str()) { + if let Some(inner) = rest.strip_suffix(')') { + return inner; + } + } + s +} + +/// Encode buffered RowBinary rows into native columnar block column bytes. +/// +/// Returns a flat byte buffer containing, for each column in order: +/// - `string(column_name)` +/// - `string(column_type_name)` +/// - optional custom-serialization flag byte (0x00) for newer servers +/// - column data (native columnar encoding, recursively for Array/Map/Tuple) +/// +/// This output is written directly after the block header +/// (`num_columns` + `num_rows`) in a Data packet. +/// +/// # Errors +/// +/// Returns `Error::BadResponse` if any row's RowBinary data is truncated or +/// contains an unsupported type for INSERT. +pub(crate) fn encode_columns( + rows: &[Vec], + columns: &[ColumnSchema], + revision: u64, +) -> Result> { + let has_custom_ser = revision >= DBMS_MIN_PROTOCOL_VERSION_WITH_CUSTOM_SERIALIZATION; + if columns.is_empty() { + return Ok(Vec::new()); + } + + // Pass 1 — extract per-column raw RowBinary value bytes (one per row). + let n = rows.len(); + let mut per_col: Vec>> = vec![Vec::with_capacity(n); columns.len()]; + for row in rows { + let mut pos = 0; + for (ci, col) in columns.iter().enumerate() { + let start = pos; + rb_advance(row, &mut pos, &col.col_type)?; + per_col[ci].push(row[start..pos].to_vec()); + } + } + + // Pass 2 — emit header + native-encoded data for each column. + let mut out = Vec::new(); + for (ci, col) in columns.iter().enumerate() { + out.put_string(col.name.as_bytes()); + out.put_string(col.type_name.as_bytes()); + // Newer servers expect a custom-serialization flag byte (0 = normal) per column. + if has_custom_ser { + out.push(0u8); + } + write_col_values(&per_col[ci], &col.col_type, &mut out)?; + } + + Ok(out) +} + +/// Recursively write native columnar data for `values` (one `Vec` per row, +/// containing raw RowBinary bytes for a single value). +fn write_col_values(values: &[Vec], col_type: &ColumnType, out: &mut Vec) -> Result<()> { + // Fixed-size scalars, String, and FixedString: RowBinary bytes == native bytes. + if col_type.fixed_size().is_some() + || matches!( + col_type, + ColumnType::String | ColumnType::FixedString(_) | ColumnType::Json + ) + { + for v in values { + out.extend_from_slice(v); + } + return Ok(()); + } + + match col_type { + ColumnType::Nullable(inner) => { + // Native: u8[n] null flags, then inner_type[n] values (zero for nulls). + let mut inner_vals: Vec> = Vec::with_capacity(values.len()); + for v in values { + if v.is_empty() { + return Err(rb_truncated()); + } + let flag = v[0]; // RowBinary: 0 = has value, 1 = null + out.push(flag); + if flag == 0 { + inner_vals.push(v[1..].to_vec()); // value bytes follow flag + } else { + let mut def = Vec::new(); + rb_write_default(&mut def, inner); // zero bytes for null slot + inner_vals.push(def); + } + } + write_col_values(&inner_vals, inner, out)?; + } + + ColumnType::LowCardinality(inner) => { + // Stripped at schema level — just encode as inner type. + write_col_values(values, inner, out)?; + } + + ColumnType::Array(inner) => { + // Native: u64[n] cumulative offsets, then all elements as a sub-column. + let mut cum: u64 = 0; + let mut offsets: Vec = Vec::with_capacity(values.len()); + let mut all_elems: Vec> = Vec::new(); + + for v in values { + let mut pos = 0; + let (count, hdr) = rb_read_varuint(v, pos)?; + pos += hdr; + for _ in 0..count { + let start = pos; + rb_advance(v, &mut pos, inner)?; + all_elems.push(v[start..pos].to_vec()); + } + cum += count; + offsets.push(cum); + } + + for off in &offsets { + out.extend_from_slice(&off.to_le_bytes()); + } + write_col_values(&all_elems, inner, out)?; + } + + ColumnType::Map(key_type, val_type) => { + // Native: u64[n] cumulative offsets, then key sub-column, then value sub-column. + let mut cum: u64 = 0; + let mut offsets: Vec = Vec::with_capacity(values.len()); + let mut all_keys: Vec> = Vec::new(); + let mut all_vals: Vec> = Vec::new(); + + for v in values { + let mut pos = 0; + let (count, hdr) = rb_read_varuint(v, pos)?; + pos += hdr; + for _ in 0..count { + let ks = pos; + rb_advance(v, &mut pos, key_type)?; + all_keys.push(v[ks..pos].to_vec()); + let vs = pos; + rb_advance(v, &mut pos, val_type)?; + all_vals.push(v[vs..pos].to_vec()); + } + cum += count; + offsets.push(cum); + } + + for off in &offsets { + out.extend_from_slice(&off.to_le_bytes()); + } + write_col_values(&all_keys, key_type, out)?; + write_col_values(&all_vals, val_type, out)?; + } + + ColumnType::Tuple(fields) => { + // Native: each field is a separate sub-column in definition order. + let mut field_vals: Vec>> = + vec![Vec::with_capacity(values.len()); fields.len()]; + for v in values { + let mut pos = 0; + for (fi, field_type) in fields.iter().enumerate() { + let start = pos; + rb_advance(v, &mut pos, field_type)?; + field_vals[fi].push(v[start..pos].to_vec()); + } + } + for (fi, field_type) in fields.iter().enumerate() { + write_col_values(&field_vals[fi], field_type, out)?; + } + } + + unsupported => { + return Err(Error::BadResponse(format!( + "native INSERT: column type {unsupported:?} is not supported for INSERT" + ))); + } + } + + Ok(()) +} + +/// Advance `pos` past one RowBinary-encoded value of `col_type`. +/// +/// RowBinary and native wire formats are identical for all scalar types. +/// Only `Nullable` differs: RowBinary has a per-row flag followed by the +/// value (or nothing for null), while native packs flags and values separately. +fn rb_advance(data: &[u8], pos: &mut usize, col_type: &ColumnType) -> Result<()> { + // Fixed-size types: same byte count in RowBinary and native. + if let Some(size) = col_type.fixed_size() { + if *pos + size > data.len() { + return Err(rb_truncated()); + } + *pos += size; + return Ok(()); + } + + match col_type { + ColumnType::String | ColumnType::Json => { + let (len, hdr) = rb_read_varuint(data, *pos)?; + let end = *pos + hdr + len as usize; + if end > data.len() { + return Err(rb_truncated()); + } + *pos = end; + } + ColumnType::FixedString(n) => { + if *pos + n > data.len() { + return Err(rb_truncated()); + } + *pos += n; + } + ColumnType::Nullable(inner) => { + if *pos >= data.len() { + return Err(rb_truncated()); + } + let flag = data[*pos]; + *pos += 1; + if flag == 0 { + rb_advance(data, pos, inner)?; + } + } + ColumnType::LowCardinality(inner) => { + // RowBinary serialises LowCardinality transparently as the inner type. + rb_advance(data, pos, inner)?; + } + ColumnType::Array(inner) => { + let (count, hdr) = rb_read_varuint(data, *pos)?; + *pos += hdr; + for _ in 0..count { + rb_advance(data, pos, inner)?; + } + } + ColumnType::Tuple(fields) => { + for field in fields { + rb_advance(data, pos, field)?; + } + } + ColumnType::Map(key_type, val_type) => { + let (count, hdr) = rb_read_varuint(data, *pos)?; + *pos += hdr; + for _ in 0..count { + rb_advance(data, pos, key_type)?; + rb_advance(data, pos, val_type)?; + } + } + unsupported => { + return Err(Error::BadResponse(format!( + "native INSERT: column type {unsupported:?} is not supported for INSERT" + ))); + } + } + Ok(()) +} + +/// Write the default (zero) native encoding for `col_type`. +/// +/// Used to fill the value slot for NULL rows in a Nullable column — +/// the native protocol requires value bytes even when the null flag is set. +fn rb_write_default(out: &mut Vec, col_type: &ColumnType) { + if let Some(size) = col_type.fixed_size() { + out.extend(std::iter::repeat_n(0u8, size)); + return; + } + match col_type { + ColumnType::String | ColumnType::Json => { + out.put_var_uint(0); // empty string: single 0x00 varuint + } + ColumnType::FixedString(n) => { + out.extend(std::iter::repeat_n(0u8, *n)); + } + ColumnType::LowCardinality(inner) => { + rb_write_default(out, inner); + } + _ => { + // Best-effort: empty string for unknown variable-length types + out.put_var_uint(0); + } + } +} + +/// Read a varuint from `data` starting at `pos`, returning `(value, bytes_consumed)`. +fn rb_read_varuint(data: &[u8], pos: usize) -> Result<(u64, usize)> { + let mut out = 0u64; + let mut shift = 0u32; + let mut i = pos; + loop { + if i >= data.len() { + return Err(rb_truncated()); + } + let b = data[i]; + i += 1; + out |= u64::from(b & 0x7F) << shift; + shift += 7; + if b & 0x80 == 0 { + break; + } + if shift >= 64 { + return Err(Error::BadResponse( + "native INSERT: varuint overflow in RowBinary".to_string(), + )); + } + } + Ok((out, i - pos)) +} + +fn rb_truncated() -> Error { + Error::BadResponse( + "native INSERT: RowBinary row data is truncated; \ + does the row struct match the table schema?" + .to_string(), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn u8_col() -> ColumnSchema { + ColumnSchema { + name: "n".to_string(), + type_name: "UInt8".to_string(), + col_type: ColumnType::UInt8, + } + } + + fn str_col() -> ColumnSchema { + ColumnSchema { + name: "s".to_string(), + type_name: "String".to_string(), + col_type: ColumnType::String, + } + } + + fn nullable_u8_col() -> ColumnSchema { + ColumnSchema { + name: "n".to_string(), + type_name: "Nullable(UInt8)".to_string(), + col_type: ColumnType::Nullable(Box::new(ColumnType::UInt8)), + } + } + + #[test] + fn test_encode_single_u8_column() { + // Two rows: UInt8 values 1 and 2 + let rows = vec![vec![1u8], vec![2u8]]; + let cols = vec![u8_col()]; + let out = encode_columns(&rows, &cols, 0).unwrap(); + + // string("n") = varuint(1) + "n" + // string("UInt8") = varuint(5) + "UInt8" + // data = [1, 2] + let expected_name = b"\x01n"; + let expected_type = b"\x05UInt8"; + let expected_data = b"\x01\x02"; + assert!(out.starts_with(expected_name)); + let after_name = &out[expected_name.len()..]; + assert!(after_name.starts_with(expected_type)); + let after_type = &after_name[expected_type.len()..]; + assert_eq!(after_type, expected_data); + } + + #[test] + fn test_encode_string_column() { + // One row: String "hi" + let mut row = Vec::new(); + row.push(0x02u8); // varuint(2) + row.extend_from_slice(b"hi"); + let rows = vec![row.clone()]; + let cols = vec![str_col()]; + let out = encode_columns(&rows, &cols, 0).unwrap(); + // After header: the string bytes from RowBinary are passed through unchanged + let after_hdr = out[b"\x01s\x06String".len()..].to_vec(); + assert_eq!(after_hdr, row); + } + + #[test] + fn test_encode_nullable_u8_not_null() { + // One row: Nullable(UInt8) = Some(42) + // RowBinary: [0x00 (not null), 42] + let rows = vec![vec![0x00u8, 42u8]]; + let cols = vec![nullable_u8_col()]; + let out = encode_columns(&rows, &cols, 0).unwrap(); + // After header: [0x00 (null flag)] then [42 (value)] + let hdr_len = b"\x01n\x10Nullable(UInt8)".len(); + let data = &out[hdr_len..]; + assert_eq!(data, &[0x00u8, 42u8]); // flag then value + } + + #[test] + fn test_encode_nullable_u8_null() { + // One row: Nullable(UInt8) = None + // RowBinary: [0x01 (null)] + let rows = vec![vec![0x01u8]]; + let cols = vec![nullable_u8_col()]; + let out = encode_columns(&rows, &cols, 0).unwrap(); + // After header: [0x01 (null flag)] then [0x00 (zero default value)] + let hdr_len = b"\x01n\x10Nullable(UInt8)".len(); + let data = &out[hdr_len..]; + assert_eq!(data, &[0x01u8, 0x00u8]); + } +} diff --git a/src/native/error_codes.rs b/src/native/error_codes.rs new file mode 100644 index 00000000..2ad85023 --- /dev/null +++ b/src/native/error_codes.rs @@ -0,0 +1,118 @@ +//! Server exception mapping for ClickHouse native protocol. +//! +//! Maps ClickHouse error codes to severity levels to distinguish +//! fatal server errors from recoverable client/query errors. + +use std::fmt; + +use crate::native::protocol::ServerException; + +/// Severity classification for server exceptions. +#[derive(Debug, Clone)] +pub(crate) enum Severity { + /// Fatal server-side error — connection should be dropped. + Server(ServerErrorKind), + /// Non-fatal query/client error — connection can be reused. + Client(ClientErrorKind), +} + +impl fmt::Display for Severity { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Severity::Server(kind) => write!(f, "Server({kind:?})"), + Severity::Client(kind) => write!(f, "Client({kind:?})"), + } + } +} + +#[derive(Debug, Clone)] +#[allow(unused)] +pub(crate) enum ServerErrorKind { + Internal, + Timeout, + ResourceExhausted, + Other, +} + +#[derive(Debug, Clone)] +#[allow(unused)] +pub(crate) enum ClientErrorKind { + Syntax, + Type, + NotFound, + Auth, + Other, +} + +/// A mapped server error with severity classification. +#[derive(Debug, Clone)] +pub(crate) struct ServerError { + pub(crate) severity: Severity, + pub(crate) code: i32, + pub(crate) name: String, + pub(crate) message: String, + pub(crate) stack_trace: String, +} + +impl ServerError { + pub(crate) fn is_fatal(&self) -> bool { + matches!(self.severity, Severity::Server(_)) + } +} + +impl fmt::Display for ServerError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "ClickHouse exception: {severity} code={code} {name}: {message}", + severity = self.severity, + code = self.code, + name = self.name, + message = self.message, + )?; + if !self.stack_trace.is_empty() { + write!(f, "\nStack trace:\n")?; + for line in self.stack_trace.lines() { + writeln!(f, " {line}")?; + } + } + Ok(()) + } +} + +/// Map a raw server exception to a classified `ServerError`. +pub(crate) fn map_exception_to_error(exception: ServerException) -> ServerError { + let severity = map_error_code(exception.code); + ServerError { + severity, + code: exception.code, + name: exception.name, + message: exception.message, + stack_trace: exception.stack_trace, + } +} + +/// Classify an error code into severity. +/// +/// Based on ClickHouse error codes from ErrorCodes.h. +/// Only the most common codes are mapped; everything else defaults to Client(Other). +fn map_error_code(code: i32) -> Severity { + match code { + // Server-side fatal errors + 1 => Severity::Server(ServerErrorKind::Internal), // UNSUPPORTED_METHOD + 48 => Severity::Server(ServerErrorKind::Internal), // NOT_IMPLEMENTED + 76 => Severity::Server(ServerErrorKind::Internal), // LOGICAL_ERROR + 159 => Severity::Server(ServerErrorKind::Timeout), // TIMEOUT_EXCEEDED + 241 => Severity::Server(ServerErrorKind::ResourceExhausted), // MEMORY_LIMIT_EXCEEDED + 252 => Severity::Server(ServerErrorKind::ResourceExhausted), // TOO_MANY_SIMULTANEOUS_QUERIES + // Client / query errors + 27 => Severity::Client(ClientErrorKind::NotFound), // UNKNOWN_DATABASE + 36 => Severity::Client(ClientErrorKind::Type), // TYPE_MISMATCH + 47 => Severity::Client(ClientErrorKind::Syntax), // UNKNOWN_IDENTIFIER + 60 => Severity::Client(ClientErrorKind::NotFound), // UNKNOWN_TABLE + 62 => Severity::Client(ClientErrorKind::Syntax), // SYNTAX_ERROR + 192 => Severity::Client(ClientErrorKind::Auth), // AUTHENTICATION_FAILED + 516 => Severity::Client(ClientErrorKind::Auth), // AUTHENTICATION_FAILED (v2) + _ => Severity::Client(ClientErrorKind::Other), + } +} diff --git a/src/native/inserter.rs b/src/native/inserter.rs new file mode 100644 index 00000000..4784d213 --- /dev/null +++ b/src/native/inserter.rs @@ -0,0 +1,257 @@ +//! `NativeInserter` — multi-batch INSERT wrapper for the native transport. +//! +//! Mirrors the public API of [`crate::inserter::Inserter`] (HTTP transport) +//! without requiring the `inserter` crate feature. + +use std::mem; +use std::time::{Duration, Instant}; + +use crate::error::Result; +use crate::native::client::NativeClient; +use crate::native::insert::NativeInsert; +use crate::row::{Row, RowWrite}; + +/// Statistics about pending or inserted data. +/// +/// Mirrors [`crate::inserter::Quantities`] for the native transport. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Quantities { + /// Approximate number of uncompressed bytes (RowBinary representation). + pub bytes: u64, + /// Number of rows written since the last commit. + pub rows: u64, + /// Number of non-empty transactions (INSERT statements) committed. + pub transactions: u64, +} + +impl Quantities { + /// All-zero quantities. + pub const ZERO: Quantities = Quantities { + bytes: 0, + rows: 0, + transactions: 0, + }; +} + +/// Simple wall-clock period tracker used by [`NativeInserter`]. +struct NativeTicks { + period: Option, + next_at: Option, +} + +impl Default for NativeTicks { + fn default() -> Self { + Self { + period: None, + next_at: None, + } + } +} + +impl NativeTicks { + fn set_period(&mut self, period: Option) { + self.period = period; + } + + fn set_period_bias(&mut self, _bias: f64) { + // Bias (jitter) is accepted for API compatibility; not applied in native MVP. + } + + fn reschedule(&mut self) { + self.next_at = self.period.map(|p| Instant::now() + p); + } + + fn time_left(&mut self) -> Option { + let next = self.next_at?; + let now = Instant::now(); + Some(if next > now { next - now } else { Duration::ZERO }) + } + + fn reached(&self) -> bool { + self.next_at + .map(|next| Instant::now() >= next) + .unwrap_or(false) + } +} + +/// Multi-batch native INSERT manager. +/// +/// Wraps [`NativeInsert`] to produce multiple consecutive INSERT statements +/// bounded by configurable thresholds. See [`crate::inserter::Inserter`] for +/// the equivalent HTTP version and full documentation. +#[must_use] +pub struct NativeInserter { + client: NativeClient, + table: String, + max_bytes: u64, + max_rows: u64, + insert: Option>, + ticks: NativeTicks, + pending: Quantities, + in_transaction: bool, + #[allow(clippy::type_complexity)] + on_commit: Option>, +} + +impl NativeInserter { + pub(crate) fn new(client: &NativeClient, table: &str) -> Self { + Self { + client: client.clone(), + table: table.into(), + max_bytes: u64::MAX, + max_rows: u64::MAX, + insert: None, + ticks: NativeTicks::default(), + pending: Quantities::ZERO, + in_transaction: false, + on_commit: None, + } + } + + /// Maximum uncompressed bytes per INSERT statement (soft limit). + pub fn with_max_bytes(mut self, threshold: u64) -> Self { + self.max_bytes = threshold; + self + } + + /// Maximum rows per INSERT statement (soft limit). + pub fn with_max_rows(mut self, threshold: u64) -> Self { + self.max_rows = threshold; + self + } + + /// Maximum elapsed time between INSERT commits. + pub fn with_period(mut self, period: Option) -> Self { + self.ticks.set_period(period); + self.ticks.reschedule(); + self + } + + /// Add a bias to the period for jitter (API-compatible; bias is ignored in MVP). + pub fn with_period_bias(mut self, bias: f64) -> Self { + self.ticks.set_period_bias(bias); + self + } + + /// Register a callback invoked after each successful non-empty commit. + pub fn with_commit_callback( + mut self, + callback: impl FnMut(&Quantities) + Send + 'static, + ) -> Self { + self.on_commit = Some(Box::new(callback)); + self + } + + /// See [`with_max_bytes`](Self::with_max_bytes). + pub fn set_max_bytes(&mut self, threshold: u64) { + self.max_bytes = threshold; + } + + /// See [`with_max_rows`](Self::with_max_rows). + pub fn set_max_rows(&mut self, threshold: u64) { + self.max_rows = threshold; + } + + /// See [`with_period`](Self::with_period). + pub fn set_period(&mut self, period: Option) { + self.ticks.set_period(period); + self.ticks.reschedule(); + } + + /// See [`with_period_bias`](Self::with_period_bias). + pub fn set_period_bias(&mut self, bias: f64) { + self.ticks.set_period_bias(bias); + } + + /// How much time remains until the next tick. `None` if no period is set. + pub fn time_left(&mut self) -> Option { + self.ticks.time_left() + } + + /// Statistics about rows/bytes not yet committed. + pub fn pending(&self) -> &Quantities { + &self.pending + } + + /// Serialise `row` into the internal buffer. + /// + /// Flushes to the network when the active `NativeInsert`'s buffer is full. + /// Call [`commit`](Self::commit) or [`force_commit`](Self::force_commit) + /// to check limits and end the current INSERT. + pub async fn write(&mut self, row: &T::Value<'_>) -> Result<()> + where + T: RowWrite, + { + if self.insert.is_none() { + self.init_insert(); + } + + match self.insert.as_mut().unwrap().write(row).await { + Ok(()) => { + self.pending.rows += 1; + if !self.in_transaction { + self.pending.transactions += 1; + self.in_transaction = true; + } + Ok(()) + } + Err(e) => { + self.pending = Quantities::ZERO; + self.insert = None; + Err(e) + } + } + } + + /// Check limits; if reached, end the active INSERT. + /// + /// Returns [`Quantities::ZERO`] when limits have not been reached. + pub async fn commit(&mut self) -> Result { + if !self.limits_reached() { + self.in_transaction = false; + return Ok(Quantities::ZERO); + } + self.force_commit().await + } + + /// End the active INSERT unconditionally, regardless of limits. + pub async fn force_commit(&mut self) -> Result { + let q = self.do_commit().await?; + self.ticks.reschedule(); + Ok(q) + } + + /// End the active INSERT and consume the `NativeInserter`. + /// + /// Must be called to flush the final batch. + pub async fn end(mut self) -> Result { + self.do_commit().await + } + + fn limits_reached(&self) -> bool { + self.pending.rows >= self.max_rows || self.ticks.reached() + } + + async fn do_commit(&mut self) -> Result { + self.in_transaction = false; + let quantities = mem::replace(&mut self.pending, Quantities::ZERO); + + if let Some(insert) = self.insert.take() { + insert.end().await?; + } + + if let Some(cb) = &mut self.on_commit { + if quantities.transactions > 0 { + (cb)(&quantities); + } + } + + Ok(quantities) + } + + #[inline(never)] + fn init_insert(&mut self) { + debug_assert!(self.insert.is_none()); + self.insert = Some(NativeInsert::new(self.client.clone(), &self.table)); + } +} diff --git a/src/native/io.rs b/src/native/io.rs new file mode 100644 index 00000000..db7b4dc8 --- /dev/null +++ b/src/native/io.rs @@ -0,0 +1,306 @@ +//! Extension traits for reading/writing ClickHouse native wire protocol primitives. +//! +//! Provides VarUInt and length-prefixed string encoding used by the native TCP protocol. + +use std::io::IoSlice; + +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; + +use crate::error::{Error, Result}; +use crate::native::protocol::MAX_STRING_SIZE; + +/// Extension trait on AsyncRead for ClickHouse wire protocol. +pub(crate) trait ClickHouseRead: AsyncRead + Unpin + Send + Sync { + fn read_var_uint(&mut self) -> impl Future> + Send + '_; + + fn read_string(&mut self) -> impl Future>> + Send + '_; + + fn read_utf8_string(&mut self) -> impl Future> + Send + '_ { + async { + let bytes = self.read_string().await?; + String::from_utf8(bytes) + .map_err(|e| Error::BadResponse(format!("native protocol: invalid utf8: {e}"))) + } + } +} + +impl ClickHouseRead for T { + async fn read_var_uint(&mut self) -> Result { + let mut out = 0u64; + for i in 0..9u64 { + let mut octet = [0u8]; + self.read_exact(&mut octet[..]).await?; + out |= u64::from(octet[0] & 0x7F) << (7 * i); + if (octet[0] & 0x80) == 0 { + break; + } + } + Ok(out) + } + + async fn read_string(&mut self) -> Result> { + #[allow(clippy::cast_possible_truncation)] + let len = self.read_var_uint().await? as usize; + if len > MAX_STRING_SIZE { + return Err(Error::BadResponse(format!( + "native protocol: string too large: {len} > {MAX_STRING_SIZE}" + ))); + } + if len == 0 { + return Ok(vec![]); + } + let mut buf = vec![0u8; len]; + self.read_exact(&mut buf).await?; + Ok(buf) + } +} + +/// Extension trait on AsyncWrite for ClickHouse wire protocol. +pub(crate) trait ClickHouseWrite: AsyncWrite + Unpin + Send + Sync { + fn write_var_uint(&mut self, value: u64) -> impl Future> + Send + '_; + + fn write_string + Send>( + &mut self, + value: V, + ) -> impl Future> + Send + use<'_, Self, V>; + + /// Write multiple buffers in one syscall (vectored I/O). + fn write_vectored_all<'a>( + &'a mut self, + bufs: &'a mut [IoSlice<'a>], + ) -> impl Future> + Send + 'a; +} + +impl ClickHouseWrite for T { + async fn write_var_uint(&mut self, mut value: u64) -> Result<()> { + let mut buf = [0u8; 9]; // Max 9 bytes for u64 + let mut pos = 0; + + #[allow(clippy::cast_possible_truncation)] + while pos < 9 { + let mut byte = value & 0x7F; + value >>= 7; + if value > 0 { + byte |= 0x80; + } + buf[pos] = byte as u8; + pos += 1; + if value == 0 { + break; + } + } + self.write_all(&buf[..pos]).await?; + Ok(()) + } + + async fn write_string + Send>(&mut self, value: V) -> Result<()> { + let value = value.as_ref(); + self.write_var_uint(value.len() as u64).await?; + self.write_all(value).await?; + Ok(()) + } + + async fn write_vectored_all<'a>(&'a mut self, bufs: &'a mut [IoSlice<'a>]) -> Result<()> { + let total: usize = bufs.iter().map(|b| b.len()).sum(); + if total == 0 { + return Ok(()); + } + + let mut written = 0usize; + while written < total { + let mut remaining_bufs: Vec> = + bufs.iter().skip_while(|b| b.is_empty()).map(|b| IoSlice::new(b)).collect(); + + if remaining_bufs.is_empty() { + break; + } + + let mut to_skip = written; + for buf in &mut remaining_bufs { + if to_skip == 0 { + break; + } + let buf_len = buf.len(); + if to_skip >= buf_len { + to_skip -= buf_len; + *buf = IoSlice::new(&[]); + } else { + break; + } + } + + let active_bufs: Vec> = + remaining_bufs.into_iter().filter(|b| !b.is_empty()).collect(); + + if active_bufs.is_empty() { + break; + } + + match self.write_vectored(&active_bufs).await { + Ok(0) => { + return Err(Error::Network(Box::new(std::io::Error::new( + std::io::ErrorKind::WriteZero, + "write_vectored returned 0", + )))); + } + Ok(n) => written += n, + Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {} + Err(e) => return Err(e.into()), + } + } + + Ok(()) + } +} + +/// Sync extension trait on `bytes::Buf` for ClickHouse wire protocol. +pub(crate) trait ClickHouseBytesRead: bytes::Buf { + fn try_get_var_uint(&mut self) -> Result; + fn try_get_string(&mut self) -> Result; +} + +impl ClickHouseBytesRead for T { + #[inline] + fn try_get_var_uint(&mut self) -> Result { + if !self.has_remaining() { + return Err(Error::NotEnoughData); + } + let b = self.get_u8(); + let mut out = u64::from(b & 0x7F); + if (b & 0x80) == 0 { + return Ok(out); + } + + for i in 1..9 { + if !self.has_remaining() { + return Err(Error::NotEnoughData); + } + let b = self.get_u8(); + out |= u64::from(b & 0x7F) << (7 * i); + if (b & 0x80) == 0 { + return Ok(out); + } + } + + Ok(out) + } + + #[inline] + fn try_get_string(&mut self) -> Result { + #[allow(clippy::cast_possible_truncation)] + let len = self.try_get_var_uint()? as usize; + + if len > MAX_STRING_SIZE { + return Err(Error::BadResponse(format!( + "native protocol: string too large: {len}" + ))); + } + + if len == 0 { + return Ok(bytes::Bytes::new()); + } + + if self.remaining() < len { + return Err(Error::NotEnoughData); + } + + Ok(self.copy_to_bytes(len)) + } +} + +/// Sync extension trait on `bytes::BufMut` for ClickHouse wire protocol. +pub(crate) trait ClickHouseBytesWrite: bytes::BufMut { + fn put_var_uint(&mut self, value: u64); + fn put_string>(&mut self, value: V); +} + +impl ClickHouseBytesWrite for T { + fn put_var_uint(&mut self, mut value: u64) { + let mut buf = [0u8; 9]; + let mut pos = 0; + + #[allow(clippy::cast_possible_truncation)] + while pos < 9 { + let mut byte = value & 0x7F; + value >>= 7; + if value > 0 { + byte |= 0x80; + } + buf[pos] = byte as u8; + pos += 1; + if value == 0 { + break; + } + } + + self.put_slice(&buf[..pos]); + } + + fn put_string>(&mut self, value: V) { + let value = value.as_ref(); + self.put_var_uint(value.len() as u64); + self.put_slice(value); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use bytes::{Bytes, BytesMut}; + + #[test] + fn test_var_uint_roundtrip_sync() { + // Note: ClickHouse varint uses 7 bits × 9 bytes = 63 bits max + let test_values: &[u64] = &[0, 1, 127, 128, 255, 256, 16383, 16384, (1 << 63) - 1]; + for &val in test_values { + let mut buf = BytesMut::new(); + buf.put_var_uint(val); + let mut reader = buf.freeze(); + let decoded = reader.try_get_var_uint().unwrap(); + assert_eq!(val, decoded, "roundtrip failed for {val}"); + } + } + + #[test] + fn test_string_roundtrip_sync() { + let test_strings: &[&[u8]] = &[b"", b"hello", b"hello world", &[0u8; 1000]]; + for &val in test_strings { + let mut buf = BytesMut::new(); + buf.put_string(val); + let mut reader = buf.freeze(); + let decoded = reader.try_get_string().unwrap(); + assert_eq!(val, &decoded[..], "roundtrip failed"); + } + } + + #[tokio::test] + async fn test_var_uint_roundtrip_async() { + let test_values: &[u64] = &[0, 1, 127, 128, 255, 256, 16383, 16384, (1 << 63) - 1]; + for &val in test_values { + let mut buf = Vec::new(); + buf.write_var_uint(val).await.unwrap(); + let mut reader = std::io::Cursor::new(buf); + let decoded = reader.read_var_uint().await.unwrap(); + assert_eq!(val, decoded, "async roundtrip failed for {val}"); + } + } + + #[tokio::test] + async fn test_string_roundtrip_async() { + let test_strings: &[&[u8]] = &[b"", b"hello", b"hello world"]; + for &val in test_strings { + let mut buf = Vec::new(); + buf.write_string(val).await.unwrap(); + let mut reader = std::io::Cursor::new(buf); + let decoded = reader.read_string().await.unwrap(); + assert_eq!(val, &decoded[..], "async roundtrip failed"); + } + } + + #[test] + fn test_eof_returns_error() { + let mut buf = Bytes::new(); + let result = buf.try_get_var_uint(); + assert!(result.is_err()); + } +} diff --git a/src/native/protocol.rs b/src/native/protocol.rs new file mode 100644 index 00000000..9b38047a --- /dev/null +++ b/src/native/protocol.rs @@ -0,0 +1,386 @@ +//! ClickHouse native TCP protocol definitions. +//! +//! Packet IDs, handshake structures, version constants, and compression methods +//! for the native binary protocol (port 9000). + +use std::str::FromStr; + +use crate::error::{Error, Result}; + +// === Protocol version constants === + +pub(crate) const DBMS_MIN_REVISION_WITH_CLIENT_INFO: u64 = 54032; +pub(crate) const DBMS_MIN_REVISION_WITH_SERVER_TIMEZONE: u64 = 54058; +pub(crate) const DBMS_MIN_REVISION_WITH_QUOTA_KEY_IN_CLIENT_INFO: u64 = 54060; +pub(crate) const DBMS_MIN_REVISION_WITH_SERVER_DISPLAY_NAME: u64 = 54372; +pub(crate) const DBMS_MIN_REVISION_WITH_VERSION_PATCH: u64 = 54401; +pub(crate) const DBMS_MIN_REVISION_WITH_SERVER_LOGS: u64 = 54406; +pub(crate) const DBMS_MIN_REVISION_WITH_CLIENT_WRITE_INFO: u64 = 54420; +pub(crate) const DBMS_MIN_REVISION_WITH_SETTINGS_SERIALIZED_AS_STRINGS: u64 = 54429; +pub(crate) const DBMS_MIN_REVISION_WITH_OPENTELEMETRY: u64 = 54442; +pub(crate) const DBMS_MIN_REVISION_WITH_INTERSERVER_SECRET: u64 = 54441; +pub(crate) const DBMS_MIN_PROTOCOL_VERSION_WITH_DISTRIBUTED_DEPTH: u64 = 54448; +pub(crate) const DBMS_MIN_PROTOCOL_VERSION_WITH_QUERY_START_TIME: u64 = 54449; +pub(crate) const DBMS_MIN_PROTOCOL_VERSION_WITH_PARALLEL_REPLICAS: u64 = 54453; +pub(crate) const DBMS_MIN_PROTOCOL_VERSION_WITH_CUSTOM_SERIALIZATION: u64 = 54454; +pub(crate) const DBMS_MIN_PROTOCOL_VERSION_WITH_PROFILE_EVENTS_IN_INSERT: u64 = 54456; +pub(crate) const DBMS_MIN_PROTOCOL_VERSION_WITH_ADDENDUM: u64 = 54458; +pub(crate) const DBMS_MIN_PROTOCOL_VERSION_WITH_QUOTA_KEY: u64 = 54458; +pub(crate) const DBMS_MIN_PROTOCOL_VERSION_WITH_PARAMETERS: u64 = 54459; +pub(crate) const DBMS_MIN_PROTOCOL_VERSION_WITH_SERVER_QUERY_TIME_IN_PROGRESS: u64 = 54460; +pub(crate) const DBMS_MIN_PROTOCOL_VERSION_WITH_PASSWORD_COMPLEXITY_RULES: u64 = 54461; +pub(crate) const DBMS_MIN_REVISION_WITH_INTERSERVER_SECRET_V2: u64 = 54462; +pub(crate) const DBMS_MIN_PROTOCOL_VERSION_WITH_TOTAL_BYTES_IN_PROGRESS: u64 = 54463; +pub(crate) const DBMS_MIN_REVISION_WITH_ROWS_BEFORE_AGGREGATION: u64 = 54469; +pub(crate) const DBMS_MIN_PROTOCOL_VERSION_WITH_CHUNKED_PACKETS: u64 = 54470; +pub(crate) const DBMS_MIN_REVISION_WITH_VERSIONED_PARALLEL_REPLICAS_PROTOCOL: u64 = 54471; +pub(crate) const DBMS_MIN_PROTOCOL_VERSION_WITH_INTERSERVER_EXTERNALLY_GRANTED_ROLES: u64 = 54472; +pub(crate) const DBMS_MIN_REVISION_WITH_SERVER_SETTINGS: u64 = 54474; +pub(crate) const DBMS_MIN_REVISION_WITH_QUERY_AND_LINE_NUMBERS: u64 = 54475; +pub(crate) const DBMS_MIN_REVISION_WITH_JWT_IN_INTERSERVER: u64 = 54476; +pub(crate) const DBMS_MIN_REVISION_WITH_QUERY_PLAN_SERIALIZATION: u64 = 54477; +pub(crate) const DBMS_MIN_REVISION_WITH_VERSIONED_CLUSTER_FUNCTION_PROTOCOL: u64 = 54479; + +/// Active protocol version this client advertises. +pub(crate) const DBMS_TCP_PROTOCOL_VERSION: u64 = + DBMS_MIN_REVISION_WITH_VERSIONED_CLUSTER_FUNCTION_PROTOCOL; + +pub(crate) const DBMS_PARALLEL_REPLICAS_PROTOCOL_VERSION: u64 = 4; + +/// Maximum string size over the native protocol (1 GiB). +pub(crate) const MAX_STRING_SIZE: usize = 1 << 30; + +// === Query processing stage === + +#[repr(u64)] +#[derive(Clone, Copy, Debug)] +#[allow(unused)] +pub(crate) enum QueryProcessingStage { + FetchColumns, + WithMergeableState, + Complete, + WithMergableStateAfterAggregation, +} + +// === Client packets === + +#[allow(unused)] +#[repr(u64)] +#[derive(Clone, Copy, Debug)] +pub(crate) enum ClientPacketId { + Hello = 0, + Query = 1, + Data = 2, + Cancel = 3, + Ping = 4, + TablesStatusRequest = 5, + KeepAlive = 6, + Scalar = 7, + IgnoredPartUUIDs = 8, + ReadTaskResponse = 9, + MergeTreeReadTaskResponse = 10, + SSHChallengeRequest = 11, + SSHChallengeResponse = 12, + QueryPlan = 13, +} + +pub(crate) struct ClientHello { + pub(crate) default_database: String, + pub(crate) username: String, + pub(crate) password: String, +} + +// === Server packets === + +#[repr(u64)] +#[derive(Clone, Copy, Debug)] +pub(crate) enum ServerPacketId { + Hello = 0, + Data = 1, + Exception = 2, + Progress = 3, + Pong = 4, + EndOfStream = 5, + ProfileInfo = 6, + Totals = 7, + Extremes = 8, + TablesStatusResponse = 9, + Log = 10, + TableColumns = 11, + PartUUIDs = 12, + ReadTaskRequest = 13, + ProfileEvents = 14, + MergeTreeAllRangesAnnouncement = 15, + MergeTreeReadTaskRequest = 16, + TimezoneUpdate = 17, + SSHChallenge = 18, +} + +impl ServerPacketId { + pub(crate) fn from_u64(i: u64) -> Result { + Ok(match i { + 0 => ServerPacketId::Hello, + 1 => ServerPacketId::Data, + 2 => ServerPacketId::Exception, + 3 => ServerPacketId::Progress, + 4 => ServerPacketId::Pong, + 5 => ServerPacketId::EndOfStream, + 6 => ServerPacketId::ProfileInfo, + 7 => ServerPacketId::Totals, + 8 => ServerPacketId::Extremes, + 9 => ServerPacketId::TablesStatusResponse, + 10 => ServerPacketId::Log, + 11 => ServerPacketId::TableColumns, + 12 => ServerPacketId::PartUUIDs, + 13 => ServerPacketId::ReadTaskRequest, + 14 => ServerPacketId::ProfileEvents, + 15 => ServerPacketId::MergeTreeAllRangesAnnouncement, + 16 => ServerPacketId::MergeTreeReadTaskRequest, + 17 => ServerPacketId::TimezoneUpdate, + 18 => ServerPacketId::SSHChallenge, + x => { + return Err(Error::BadResponse(format!( + "native protocol: unknown server packet id {x}" + ))); + } + }) + } +} + +// === Server response structures === + +#[derive(Debug, Clone, Default)] +pub(crate) struct ServerHello { + pub(crate) server_name: String, + pub(crate) version: (u64, u64, u64), + pub(crate) revision_version: u64, + pub(crate) timezone: Option, + pub(crate) display_name: Option, + pub(crate) chunked_send: ChunkedProtocolMode, + pub(crate) chunked_recv: ChunkedProtocolMode, +} + +impl ServerHello { + #[allow(unused)] + pub(crate) fn supports_chunked_send(&self) -> bool { + matches!( + self.chunked_send, + ChunkedProtocolMode::Chunked | ChunkedProtocolMode::ChunkedOptional + ) + } + + #[allow(unused)] + pub(crate) fn supports_chunked_recv(&self) -> bool { + matches!( + self.chunked_recv, + ChunkedProtocolMode::Chunked | ChunkedProtocolMode::ChunkedOptional + ) + } +} + +#[derive(Debug, Clone)] +pub(crate) struct ServerException { + pub(crate) code: i32, + pub(crate) name: String, + pub(crate) message: String, + pub(crate) stack_trace: String, + pub(crate) has_nested: bool, +} + +#[allow(unused)] +#[derive(Debug, Clone)] +pub(crate) struct ProfileInfo { + pub(crate) rows: u64, + pub(crate) blocks: u64, + pub(crate) bytes: u64, + pub(crate) applied_limit: bool, + pub(crate) rows_before_limit: u64, + pub(crate) calculated_rows_before_limit: bool, + pub(crate) applied_aggregation: bool, + pub(crate) rows_before_aggregation: u64, +} + +#[allow(unused)] +#[derive(Debug, Clone)] +pub(crate) struct TableColumns { + pub(crate) name: String, + pub(crate) description: String, +} + +// === Progress === + +#[derive(Debug, Clone, Default)] +pub(crate) struct Progress { + pub(crate) read_rows: u64, + pub(crate) read_bytes: u64, + pub(crate) total_rows_to_read: u64, + pub(crate) total_bytes_to_read: Option, + pub(crate) written_rows: Option, + pub(crate) written_bytes: Option, + pub(crate) elapsed_ns: Option, +} + +impl std::ops::Add for Progress { + type Output = Self; + + fn add(self, rhs: Self) -> Self { + Self { + read_rows: self.read_rows + rhs.read_rows, + read_bytes: self.read_bytes + rhs.read_bytes, + total_rows_to_read: self.total_rows_to_read + rhs.total_rows_to_read, + total_bytes_to_read: match (self.total_bytes_to_read, rhs.total_bytes_to_read) { + (Some(a), Some(b)) => Some(a + b), + (a, b) => a.or(b), + }, + written_rows: match (self.written_rows, rhs.written_rows) { + (Some(a), Some(b)) => Some(a + b), + (a, b) => a.or(b), + }, + written_bytes: match (self.written_bytes, rhs.written_bytes) { + (Some(a), Some(b)) => Some(a + b), + (a, b) => a.or(b), + }, + elapsed_ns: match (self.elapsed_ns, rhs.elapsed_ns) { + (Some(a), Some(b)) => Some(a + b), + (a, b) => a.or(b), + }, + } + } +} + +impl std::ops::AddAssign for Progress { + fn add_assign(&mut self, rhs: Self) { + *self = std::mem::take(self) + rhs; + } +} + +// === Chunked protocol negotiation === + +#[derive(Clone, Default, Copy, Debug, PartialEq, Eq, Hash)] +pub(crate) enum ChunkedProtocolMode { + #[default] + ChunkedOptional, + Chunked, + NotChunkedOptional, + NotChunked, +} + +impl ChunkedProtocolMode { + /// Negotiates chunked protocol between client and server (based on C++ `is_chunked` function). + pub(crate) fn negotiate( + server_mode: ChunkedProtocolMode, + client_mode: ChunkedProtocolMode, + direction: &str, + ) -> Result { + let server_chunked = matches!( + server_mode, + ChunkedProtocolMode::Chunked | ChunkedProtocolMode::ChunkedOptional + ); + let server_optional = matches!( + server_mode, + ChunkedProtocolMode::ChunkedOptional | ChunkedProtocolMode::NotChunkedOptional + ); + let client_chunked = matches!( + client_mode, + ChunkedProtocolMode::Chunked | ChunkedProtocolMode::ChunkedOptional + ); + let client_optional = matches!( + client_mode, + ChunkedProtocolMode::ChunkedOptional | ChunkedProtocolMode::NotChunkedOptional + ); + let result_chunked = if server_optional { + client_chunked + } else if client_optional { + server_chunked + } else if client_chunked != server_chunked { + return Err(Error::BadResponse(format!( + "native protocol: incompatible chunked mode for {direction}: \ + client={}, server={}", + if client_chunked { "chunked" } else { "notchunked" }, + if server_chunked { "chunked" } else { "notchunked" }, + ))); + } else { + server_chunked + }; + + Ok(if result_chunked { + ChunkedProtocolMode::Chunked + } else { + ChunkedProtocolMode::NotChunked + }) + } +} + +impl FromStr for ChunkedProtocolMode { + type Err = Error; + + fn from_str(s: &str) -> Result { + Ok(match s { + "chunked" => Self::Chunked, + "chunked_optional" => Self::ChunkedOptional, + "notchunked" => Self::NotChunked, + "notchunked_optional" => Self::NotChunkedOptional, + _ => { + return Err(Error::BadResponse(format!( + "native protocol: unexpected chunked mode: {s}" + ))); + } + }) + } +} + +impl std::fmt::Display for ChunkedProtocolMode { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::ChunkedOptional => write!(f, "chunked_optional"), + Self::Chunked => write!(f, "chunked"), + Self::NotChunkedOptional => write!(f, "notchunked_optional"), + Self::NotChunked => write!(f, "notchunked"), + } + } +} + +// === Compression method for native protocol === + +#[derive(Clone, Default, Copy, Debug, PartialEq, Eq, Hash)] +pub(crate) enum NativeCompressionMethod { + None, + #[default] + Lz4, + Zstd, +} + +impl NativeCompressionMethod { + pub(crate) fn byte(self) -> u8 { + match self { + NativeCompressionMethod::None => 0x02, + NativeCompressionMethod::Lz4 => 0x82, + NativeCompressionMethod::Zstd => 0x90, + } + } +} + +impl std::fmt::Display for NativeCompressionMethod { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + NativeCompressionMethod::None => write!(f, "None"), + NativeCompressionMethod::Lz4 => write!(f, "LZ4"), + NativeCompressionMethod::Zstd => write!(f, "ZSTD"), + } + } +} + +impl FromStr for NativeCompressionMethod { + type Err = Error; + + fn from_str(s: &str) -> Result { + match s { + "lz4" | "LZ4" => Ok(NativeCompressionMethod::Lz4), + "zstd" | "ZSTD" => Ok(NativeCompressionMethod::Zstd), + "none" | "None" => Ok(NativeCompressionMethod::None), + _ => Err(Error::BadResponse(format!( + "native protocol: unknown compression method: {s}" + ))), + } + } +} diff --git a/src/native/schema.rs b/src/native/schema.rs new file mode 100644 index 00000000..ff6fcc14 --- /dev/null +++ b/src/native/schema.rs @@ -0,0 +1,70 @@ +//! Schema cache for the native transport. +//! +//! Caches `(column_name, type_name)` pairs fetched from `system.columns` so +//! that consumers (e.g. dfe-loader) can inspect table schemas without a +//! round-trip on every insert. +//! +//! The cache is shared across clones of [`crate::native::NativeClient`] via `Arc`. + +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; +use std::time::{Duration, Instant}; + +/// A cached schema entry. +struct Entry { + /// Ordered `(name, type_name)` pairs. + columns: Vec<(String, String)>, + fetched_at: Instant, +} + +/// TTL-based schema cache shared across [`crate::native::NativeClient`] clones. +pub(crate) struct NativeSchemaCache { + inner: RwLock>, + ttl: Duration, +} + +impl NativeSchemaCache { + /// Create a new cache wrapped in `Arc`. + /// + /// A TTL of 300 s (5 minutes) is a sensible default. + pub(crate) fn new(ttl_secs: u64) -> Arc { + Arc::new(Self { + inner: RwLock::new(HashMap::new()), + ttl: Duration::from_secs(ttl_secs), + }) + } + + /// Return cached columns if the entry exists and has not expired. + pub(crate) fn get(&self, table: &str) -> Option> { + let guard = self.inner.read().unwrap(); + guard.get(table).and_then(|e| { + if e.fetched_at.elapsed() < self.ttl { + Some(e.columns.clone()) + } else { + None + } + }) + } + + /// Insert or refresh a schema entry. + pub(crate) fn insert(&self, table: String, columns: Vec<(String, String)>) { + let mut guard = self.inner.write().unwrap(); + guard.insert( + table, + Entry { + columns, + fetched_at: Instant::now(), + }, + ); + } + + /// Remove a schema from the cache, forcing a refresh on next access. + pub(crate) fn invalidate(&self, table: &str) { + self.inner.write().unwrap().remove(table); + } + + /// Remove all cached schemas. + pub(crate) fn invalidate_all(&self) { + self.inner.write().unwrap().clear(); + } +} diff --git a/src/native/sparse.rs b/src/native/sparse.rs new file mode 100644 index 00000000..4eca1f38 --- /dev/null +++ b/src/native/sparse.rs @@ -0,0 +1,327 @@ +//! Sparse serialization for ClickHouse native protocol. +//! +//! Optimization for columns with many default values — only non-default values +//! are stored along with their positions. Wire format: +//! +//! 1. Offsets: VarUInt group sizes (count of defaults before each non-default) +//! - Final group has `END_OF_GRANULE_FLAG` (2^62) ORed in +//! 2. Values: Only the non-default values +//! +//! Example: `[0, 0, 5, 0, 3, 0, 0, 0]` → offsets [2, 1, 3|END], values [5, 3] + +use crate::error::Result; +use crate::native::io::{ClickHouseBytesRead, ClickHouseRead}; + +/// End-of-granule marker (bit 62). When set, this is the final VarUInt in the offsets stream. +pub(crate) const END_OF_GRANULE_FLAG: u64 = 1 << 62; + +/// State for sparse deserialization across multiple reads. +#[derive(Debug, Default, Clone)] +pub(crate) struct SparseDeserializeState { + /// Trailing defaults from previous read that haven't been consumed yet. + pub(crate) num_trailing_defaults: u64, + /// Non-default value pending after the trailing defaults. + pub(crate) has_value_after_defaults: bool, +} + +/// Read sparse offsets from an async stream. Returns positions of non-default values. +/// +/// Must loop until `END_OF_GRANULE_FLAG` — can't stop early even if we have enough +/// rows, or the stream will be misaligned for the next column. +#[allow(clippy::cast_possible_truncation)] +pub(crate) async fn read_sparse_offsets( + reader: &mut R, + num_rows: usize, + state: &mut SparseDeserializeState, +) -> Result> { + let mut offsets = Vec::new(); + let mut current_position: u64 = 0; + + // Handle state carried over from previous read + if state.num_trailing_defaults > 0 { + current_position += state.num_trailing_defaults; + state.num_trailing_defaults = 0; + } + if state.has_value_after_defaults { + if (current_position as usize) < num_rows { + offsets.push(current_position as usize); + } + current_position += 1; + state.has_value_after_defaults = false; + } + + loop { + let group_size = reader.read_var_uint().await?; + + let is_end_of_granule = (group_size & END_OF_GRANULE_FLAG) != 0; + let actual_group_size = group_size & !END_OF_GRANULE_FLAG; + + current_position += actual_group_size; + + if is_end_of_granule { + if current_position > num_rows as u64 { + state.num_trailing_defaults = current_position - num_rows as u64; + } + break; + } + + if (current_position as usize) < num_rows { + offsets.push(current_position as usize); + current_position += 1; + } else { + state.has_value_after_defaults = true; + } + } + + Ok(offsets) +} + +/// Sync version of `read_sparse_offsets` for `bytes::Buf` readers. +#[allow(clippy::cast_possible_truncation)] +pub(crate) fn read_sparse_offsets_sync( + reader: &mut R, + num_rows: usize, + state: &mut SparseDeserializeState, +) -> Result> { + let mut offsets = Vec::new(); + let mut current_position: u64 = 0; + + if state.num_trailing_defaults > 0 { + current_position += state.num_trailing_defaults; + state.num_trailing_defaults = 0; + } + if state.has_value_after_defaults { + if (current_position as usize) < num_rows { + offsets.push(current_position as usize); + } + current_position += 1; + state.has_value_after_defaults = false; + } + + loop { + let group_size = reader.try_get_var_uint()?; + + let is_end_of_granule = (group_size & END_OF_GRANULE_FLAG) != 0; + let actual_group_size = group_size & !END_OF_GRANULE_FLAG; + + current_position += actual_group_size; + + if is_end_of_granule { + if current_position > num_rows as u64 { + state.num_trailing_defaults = current_position - num_rows as u64; + } + break; + } + + if (current_position as usize) < num_rows { + offsets.push(current_position as usize); + current_position += 1; + } else { + state.has_value_after_defaults = true; + } + } + + Ok(offsets) +} + +#[cfg(test)] +mod tests { + use bytes::Bytes; + + use super::*; + + fn encode_var_uint(value: u64) -> Vec { + let mut result = Vec::new(); + let mut v = value; + loop { + let byte = (v & 0x7f) as u8; + v >>= 7; + if v == 0 { + result.push(byte); + break; + } + result.push(byte | 0x80); + } + result + } + + #[test] + fn test_read_sparse_offsets_simple() { + // Column: [default, default, value, default, value, default, default, default] + // Positions of non-defaults: [2, 4] + let mut data = Vec::new(); + data.extend(encode_var_uint(2)); // 2 defaults before first value + data.extend(encode_var_uint(1)); // 1 default before second value + data.extend(encode_var_uint(3 | END_OF_GRANULE_FLAG)); // 3 trailing defaults + + let mut bytes = Bytes::from(data); + let mut state = SparseDeserializeState::default(); + let offsets = read_sparse_offsets_sync(&mut bytes, 8, &mut state).unwrap(); + + assert_eq!(offsets, vec![2, 4]); + } + + #[test] + fn test_read_sparse_offsets_all_defaults() { + let mut data = Vec::new(); + data.extend(encode_var_uint(4 | END_OF_GRANULE_FLAG)); + + let mut bytes = Bytes::from(data); + let mut state = SparseDeserializeState::default(); + let offsets = read_sparse_offsets_sync(&mut bytes, 4, &mut state).unwrap(); + + assert!(offsets.is_empty()); + } + + #[test] + fn test_read_sparse_offsets_no_defaults() { + // All non-default values + let mut data = Vec::new(); + data.extend(encode_var_uint(0)); // value at 0 + data.extend(encode_var_uint(0)); // value at 1 + data.extend(encode_var_uint(0)); // value at 2 + data.extend(encode_var_uint(END_OF_GRANULE_FLAG)); // 0 trailing defaults + + let mut bytes = Bytes::from(data); + let mut state = SparseDeserializeState::default(); + let offsets = read_sparse_offsets_sync(&mut bytes, 3, &mut state).unwrap(); + + assert_eq!(offsets, vec![0, 1, 2]); + } + + #[test] + fn test_read_sparse_offsets_first_is_value() { + // [value, default, default, value] + let mut data = Vec::new(); + data.extend(encode_var_uint(0)); // 0 defaults before first value + data.extend(encode_var_uint(2)); // 2 defaults before second value + data.extend(encode_var_uint(END_OF_GRANULE_FLAG)); // 0 trailing defaults + + let mut bytes = Bytes::from(data); + let mut state = SparseDeserializeState::default(); + let offsets = read_sparse_offsets_sync(&mut bytes, 4, &mut state).unwrap(); + + assert_eq!(offsets, vec![0, 3]); + } + + #[tokio::test] + async fn test_read_sparse_offsets_async() { + let mut data = Vec::new(); + data.extend(encode_var_uint(2)); + data.extend(encode_var_uint(1)); + data.extend(encode_var_uint(3 | END_OF_GRANULE_FLAG)); + + let mut reader = std::io::Cursor::new(data); + let mut state = SparseDeserializeState::default(); + let offsets = read_sparse_offsets(&mut reader, 8, &mut state).await.unwrap(); + + assert_eq!(offsets, vec![2, 4]); + } + + #[test] + fn test_single_row_default() { + // One row, it's the default value. + let mut data = Vec::new(); + data.extend(encode_var_uint(1 | END_OF_GRANULE_FLAG)); + + let mut bytes = Bytes::from(data); + let mut state = SparseDeserializeState::default(); + let offsets = read_sparse_offsets_sync(&mut bytes, 1, &mut state).unwrap(); + + assert!(offsets.is_empty()); + } + + #[test] + fn test_single_row_non_default() { + // One row, it's a non-default value. + let mut data = Vec::new(); + data.extend(encode_var_uint(0)); // 0 defaults before the value + data.extend(encode_var_uint(END_OF_GRANULE_FLAG)); // 0 trailing + + let mut bytes = Bytes::from(data); + let mut state = SparseDeserializeState::default(); + let offsets = read_sparse_offsets_sync(&mut bytes, 1, &mut state).unwrap(); + + assert_eq!(offsets, vec![0]); + } + + #[test] + fn test_large_gap_value_at_end() { + // 1000 rows, only the very last is non-default. + // Sparse stream: offset group = 999 defaults, then the value, then END. + let mut data = Vec::new(); + data.extend(encode_var_uint(999)); // 999 defaults before position 999 + data.extend(encode_var_uint(END_OF_GRANULE_FLAG)); // 0 trailing defaults + + let mut bytes = Bytes::from(data); + let mut state = SparseDeserializeState::default(); + let offsets = read_sparse_offsets_sync(&mut bytes, 1000, &mut state).unwrap(); + + assert_eq!(offsets, vec![999]); + } + + #[test] + fn test_value_at_position_zero_only() { + // 100 rows, only position 0 is non-default. + let mut data = Vec::new(); + data.extend(encode_var_uint(0)); // 0 defaults before position 0 + data.extend(encode_var_uint(99 | END_OF_GRANULE_FLAG)); // 99 trailing defaults + + let mut bytes = Bytes::from(data); + let mut state = SparseDeserializeState::default(); + let offsets = read_sparse_offsets_sync(&mut bytes, 100, &mut state).unwrap(); + + assert_eq!(offsets, vec![0]); + } + + #[test] + fn test_consecutive_non_defaults() { + // [T, T, T, F, F, T] — positions 0, 1, 2, 5 are non-default. + let mut data = Vec::new(); + data.extend(encode_var_uint(0)); // position 0 + data.extend(encode_var_uint(0)); // position 1 + data.extend(encode_var_uint(0)); // position 2 + data.extend(encode_var_uint(2)); // 2 defaults → position 5 + data.extend(encode_var_uint(END_OF_GRANULE_FLAG)); // 0 trailing + + let mut bytes = Bytes::from(data); + let mut state = SparseDeserializeState::default(); + let offsets = read_sparse_offsets_sync(&mut bytes, 6, &mut state).unwrap(); + + assert_eq!(offsets, vec![0, 1, 2, 5]); + } + + #[test] + fn test_state_carry_trailing_defaults() { + // Simulate state from a previous partial read that left trailing defaults. + // State: 2 unconsumed defaults from the previous read. + // + // New read covers only 3 rows of a block that the sparse stream + // encodes as covering 4 positions (2 carried + 1 value + 1 trailing). + // The trailing default beyond num_rows must be saved in state for the + // next caller. + // + // In practice our code always creates fresh state per column per block, + // but the carry-over paths must still be correct. + let mut state = SparseDeserializeState { + num_trailing_defaults: 2, // 2 unconsumed defaults carried in + has_value_after_defaults: false, + }; + // Stream: 0 more defaults before next value (→ pos 2), then END with 1 trailing. + let mut data = Vec::new(); + data.extend(encode_var_uint(0)); // 0 more defaults → value at position 2 + data.extend(encode_var_uint(1 | END_OF_GRANULE_FLAG)); // 1 trailing default + + let mut bytes = Bytes::from(data); + // Only 3 rows in this block — the trailing default goes past the end. + let offsets = read_sparse_offsets_sync(&mut bytes, 3, &mut state).unwrap(); + + // Carried 2 defaults → position 2 is the value. + // But num_rows=3, so position 2 is inside the block. + assert_eq!(offsets, vec![2]); + // The 1 trailing default puts current_position at 4, which is > num_rows(3). + // state.num_trailing_defaults = 4 - 3 = 1. + assert_eq!(state.num_trailing_defaults, 1); + assert!(!state.has_value_after_defaults); + } +} diff --git a/src/native/tcp.rs b/src/native/tcp.rs new file mode 100644 index 00000000..0361c90f --- /dev/null +++ b/src/native/tcp.rs @@ -0,0 +1,82 @@ +//! TCP connection setup for ClickHouse native protocol. +//! +//! Configures socket options (keepalive, buffer sizes, nodelay) via `socket2` +//! for high-throughput data transfer on port 9000. + +use std::net::SocketAddr; +use std::time::Duration; + +use tokio::net::TcpStream; + +use crate::error::{Error, Result}; + +// Socket configuration constants +const TCP_READ_BUFFER_SIZE: usize = 128 * 1024; +const TCP_WRITE_BUFFER_SIZE: usize = 8 * 1024 * 1024; +const TCP_CONNECT_TIMEOUT_SECS: u64 = 30; +const TCP_KEEP_ALIVE_SECS: u64 = 60; +const TCP_KEEP_ALIVE_INTERVAL: u64 = 10; +const TCP_KEEP_ALIVE_RETRIES: u32 = 6; + +// Buffered I/O sizes for the connection +pub(crate) const CONN_READ_BUFFER: usize = 1024 * 1024; +pub(crate) const CONN_WRITE_BUFFER: usize = 10 * 1024 * 1024; + +/// Connect to ClickHouse via TCP with configured socket options. +pub(crate) async fn connect(addr: &SocketAddr) -> Result { + let domain = if addr.is_ipv4() { + socket2::Domain::IPV4 + } else { + socket2::Domain::IPV6 + }; + + let socket = + socket2::Socket::new(domain, socket2::Type::STREAM, Some(socket2::Protocol::TCP)) + .map_err(|e| Error::Network(Box::new(e)))?; + + socket + .set_nonblocking(true) + .map_err(|e| Error::Network(Box::new(e)))?; + socket + .set_recv_buffer_size(TCP_READ_BUFFER_SIZE) + .map_err(|e| Error::Network(Box::new(e)))?; + socket + .set_send_buffer_size(TCP_WRITE_BUFFER_SIZE) + .map_err(|e| Error::Network(Box::new(e)))?; + + let keepalive = socket2::TcpKeepalive::new() + .with_time(Duration::from_secs(TCP_KEEP_ALIVE_SECS)) + .with_interval(Duration::from_secs(TCP_KEEP_ALIVE_INTERVAL)) + .with_retries(TCP_KEEP_ALIVE_RETRIES); + socket + .set_tcp_keepalive(&keepalive) + .map_err(|e| Error::Network(Box::new(e)))?; + + let sock_addr = socket2::SockAddr::from(*addr); + socket + .connect_timeout(&sock_addr, Duration::from_secs(TCP_CONNECT_TIMEOUT_SECS)) + .map_err(|e| Error::Network(Box::new(e)))?; + + let stream = std::net::TcpStream::from(socket); + stream + .set_nodelay(true) + .map_err(|e| Error::Network(Box::new(e)))?; + stream + .set_nonblocking(true) + .map_err(|e| Error::Network(Box::new(e)))?; + + TcpStream::from_std(stream).map_err(|e| Error::Network(Box::new(e))) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_connect_refused() { + // Connecting to a port with nothing listening should fail + let addr: SocketAddr = "127.0.0.1:19999".parse().unwrap(); + let result = connect(&addr).await; + assert!(result.is_err()); + } +} diff --git a/src/native/writer.rs b/src/native/writer.rs new file mode 100644 index 00000000..0ee91ec1 --- /dev/null +++ b/src/native/writer.rs @@ -0,0 +1,209 @@ +//! Packet writer for ClickHouse native protocol. +//! +//! Sends client packets: hello, query, data, addendum, ping. + +use tokio::io::AsyncWriteExt; + +use crate::error::Result; +use crate::native::block_info::BlockInfo; +use crate::native::client_info::ClientInfo; +use crate::native::compression::compress_data; +use crate::native::io::ClickHouseWrite; +use crate::native::protocol::{ + ClientPacketId, NativeCompressionMethod, QueryProcessingStage, ServerHello, + DBMS_MIN_PROTOCOL_VERSION_WITH_CHUNKED_PACKETS, + DBMS_MIN_PROTOCOL_VERSION_WITH_INTERSERVER_EXTERNALLY_GRANTED_ROLES, + DBMS_MIN_PROTOCOL_VERSION_WITH_PARAMETERS, DBMS_MIN_PROTOCOL_VERSION_WITH_QUOTA_KEY, + DBMS_MIN_REVISION_WITH_CLIENT_INFO, DBMS_MIN_REVISION_WITH_INTERSERVER_SECRET, + DBMS_MIN_REVISION_WITH_VERSIONED_PARALLEL_REPLICAS_PROTOCOL, + DBMS_PARALLEL_REPLICAS_PROTOCOL_VERSION, DBMS_TCP_PROTOCOL_VERSION, +}; + +/// Send client hello packet. +pub(crate) async fn send_hello( + writer: &mut W, + database: &str, + username: &str, + password: &str, +) -> Result<()> { + writer + .write_var_uint(ClientPacketId::Hello as u64) + .await?; + writer + .write_string(format!( + "clickhouse-rs native {}", + env!("CARGO_PKG_VERSION") + )) + .await?; + // Client version (major, minor, revision) + writer.write_var_uint(0).await?; // major + writer.write_var_uint(14).await?; // minor + writer + .write_var_uint(DBMS_TCP_PROTOCOL_VERSION) + .await?; + writer.write_string(database).await?; + writer.write_string(username).await?; + writer.write_string(password).await?; + writer.flush().await?; + Ok(()) +} + +/// Send a query for execution. +pub(crate) async fn send_query( + writer: &mut W, + query_id: &str, + query: &str, + settings: &[(String, String)], + revision: u64, + compression: NativeCompressionMethod, +) -> Result<()> { + writer + .write_var_uint(ClientPacketId::Query as u64) + .await?; + writer.write_string(query_id).await?; + + if revision >= DBMS_MIN_REVISION_WITH_CLIENT_INFO { + let info = ClientInfo::default(); + info.write(writer, revision).await?; + } + + // Settings: (name, is_important u8, value) per entry, terminated by empty name. + for (name, value) in settings { + writer.write_string(name).await?; + writer.write_u8(0).await?; // not important + writer.write_string(value).await?; + } + writer.write_string("").await?; // end marker + + if revision >= DBMS_MIN_PROTOCOL_VERSION_WITH_INTERSERVER_EXTERNALLY_GRANTED_ROLES { + writer.write_string("").await?; + } + + if revision >= DBMS_MIN_REVISION_WITH_INTERSERVER_SECRET { + writer.write_string("").await?; + } + + writer + .write_var_uint(QueryProcessingStage::Complete as u64) + .await?; + + // Compression flag + let use_compression = !matches!(compression, NativeCompressionMethod::None); + writer.write_u8(u8::from(use_compression)).await?; + + writer.write_string(query).await?; + + if revision >= DBMS_MIN_PROTOCOL_VERSION_WITH_PARAMETERS { + writer.write_string("").await?; // end of params + } + + writer.flush().await?; + Ok(()) +} + +/// Send an empty data block (signals end of client data). +/// +/// When `compression` is not `None`, the block body (info + counts) is +/// wrapped in a single ClickHouse compressed chunk. +pub(crate) async fn send_empty_block( + writer: &mut W, + compression: NativeCompressionMethod, +) -> Result<()> { + writer + .write_var_uint(ClientPacketId::Data as u64) + .await?; + writer.write_string("").await?; // table name (always uncompressed) + + if matches!(compression, NativeCompressionMethod::None) { + let info = BlockInfo::default(); + info.write_async(writer).await?; + writer.write_var_uint(0).await?; // 0 columns + writer.write_var_uint(0).await?; // 0 rows + } else { + let mut body: Vec = Vec::new(); + BlockInfo::default().write_async(&mut body).await?; + body.write_var_uint(0).await?; // 0 columns + body.write_var_uint(0).await?; // 0 rows + compress_data(writer, &body, compression).await?; + } + + writer.flush().await?; + Ok(()) +} + +/// Send addendum after hello exchange. +pub(crate) async fn send_addendum( + writer: &mut W, + server_hello: &ServerHello, +) -> Result<()> { + let revision = server_hello.revision_version; + + if revision >= DBMS_MIN_PROTOCOL_VERSION_WITH_QUOTA_KEY { + writer.write_string("").await?; + } + + if revision >= DBMS_MIN_PROTOCOL_VERSION_WITH_CHUNKED_PACKETS { + writer + .write_string(server_hello.chunked_send.to_string()) + .await?; + writer + .write_string(server_hello.chunked_recv.to_string()) + .await?; + } + + if revision >= DBMS_MIN_REVISION_WITH_VERSIONED_PARALLEL_REPLICAS_PROTOCOL { + writer + .write_var_uint(DBMS_PARALLEL_REPLICAS_PROTOCOL_VERSION) + .await?; + } + + writer.flush().await?; + Ok(()) +} + +/// Send a data block containing encoded column bytes. +/// +/// `column_bytes` must be pre-encoded by [`crate::native::encode::encode_columns`]: +/// for each column in order, it contains `string(name) + string(type) + column_data`. +/// +/// When `compression` is not `None`, the block body (info + counts + column_bytes) +/// is wrapped in a single ClickHouse compressed chunk. +pub(crate) async fn send_data_block( + writer: &mut W, + num_columns: usize, + num_rows: usize, + column_bytes: &[u8], + compression: NativeCompressionMethod, +) -> Result<()> { + writer + .write_var_uint(ClientPacketId::Data as u64) + .await?; + writer.write_string("").await?; // temp table name (always uncompressed) + + if matches!(compression, NativeCompressionMethod::None) { + let info = BlockInfo::default(); + info.write_async(writer).await?; + writer.write_var_uint(num_columns as u64).await?; + writer.write_var_uint(num_rows as u64).await?; + writer.write_all(column_bytes).await?; + } else { + let mut body: Vec = Vec::with_capacity(32 + column_bytes.len()); + BlockInfo::default().write_async(&mut body).await?; + body.write_var_uint(num_columns as u64).await?; + body.write_var_uint(num_rows as u64).await?; + body.extend_from_slice(column_bytes); + compress_data(writer, &body, compression).await?; + } + + writer.flush().await?; + Ok(()) +} + +/// Send ping. +pub(crate) async fn send_ping(writer: &mut W) -> Result<()> { + writer + .write_var_uint(ClientPacketId::Ping as u64) + .await?; + writer.flush().await?; + Ok(()) +} diff --git a/tests/it/batcher.rs b/tests/it/batcher.rs new file mode 100644 index 00000000..f6fdc128 --- /dev/null +++ b/tests/it/batcher.rs @@ -0,0 +1,169 @@ +use serde::{Deserialize, Serialize}; + +use clickhouse::batcher::{BatchConfig, TableBatcher}; +use clickhouse::{Client, Row}; + +#[derive(Debug, PartialEq, Eq, Row, Serialize, Deserialize)] +struct MyRow { + id: u32, + data: String, +} + +async fn create_table(client: &Client) { + client + .query( + "CREATE TABLE test(id UInt32, data String) \ + ENGINE = MergeTree ORDER BY id", + ) + .execute() + .await + .unwrap(); +} + +async fn count_rows(client: &Client) -> u64 { + client + .query("SELECT count() FROM test") + .fetch_one::() + .await + .unwrap() +} + +// ── Basic append + send ─────────────────────────────────────────────────────── + +#[tokio::test] +async fn batcher_basic() { + let client = prepare_database!(); + create_table(&client).await; + + let batcher = TableBatcher::::new( + &client, + "test", + BatchConfig::default().without_period(), + ); + + for i in 0..100u32 { + batcher + .append(&MyRow { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + batcher.send().await.unwrap(); + + assert_eq!(count_rows(&client).await, 100); +} + +// ── flush() mid-stream ──────────────────────────────────────────────────────── + +#[tokio::test] +async fn batcher_explicit_flush() { + let client = prepare_database!(); + create_table(&client).await; + + let batcher = TableBatcher::::new( + &client, + "test", + BatchConfig::default().without_period(), + ); + + for i in 0..50u32 { + batcher + .append(&MyRow { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + let q = batcher.flush().await.unwrap(); + assert_eq!(q.rows, 50); + assert_eq!(count_rows(&client).await, 50); + + for i in 50..100u32 { + batcher + .append(&MyRow { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + batcher.send().await.unwrap(); + assert_eq!(count_rows(&client).await, 100); +} + +// ── max_rows threshold ──────────────────────────────────────────────────────── + +#[tokio::test] +async fn batcher_max_rows_flush() { + let client = prepare_database!(); + create_table(&client).await; + + // Flush every 10 rows. + let batcher = TableBatcher::::new( + &client, + "test", + BatchConfig::default().with_max_rows(10).without_period(), + ); + + for i in 0..35u32 { + batcher + .append(&MyRow { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + // 3 automatic flushes of 10 rows = 30 committed; 5 still buffered. + // send() flushes the remaining 5. + batcher.send().await.unwrap(); + + assert_eq!(count_rows(&client).await, 35); +} + +// ── period-based background flush ───────────────────────────────────────────── + +#[tokio::test] +async fn batcher_period_flush() { + let client = prepare_database!(); + create_table(&client).await; + + let batcher = TableBatcher::::new( + &client, + "test", + BatchConfig::default() + .with_max_rows(u64::MAX) + .with_max_bytes(u64::MAX) + .with_max_period(tokio::time::Duration::from_millis(200)), + ); + + for i in 0..20u32 { + batcher + .append(&MyRow { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + // Wait long enough for two background flush ticks. + tokio::time::sleep(tokio::time::Duration::from_millis(600)).await; + + // Data should already be in ClickHouse from the background task. + assert_eq!(count_rows(&client).await, 20); + + batcher.send().await.unwrap(); + assert_eq!(count_rows(&client).await, 20); +} + +// ── empty send() is safe ────────────────────────────────────────────────────── + +#[tokio::test] +async fn batcher_empty_send() { + let client = prepare_database!(); + create_table(&client).await; + + let batcher = TableBatcher::::new( + &client, + "test", + BatchConfig::default().without_period(), + ); + + // No appends — send() must not panic or error. + batcher.send().await.unwrap(); + + assert_eq!(count_rows(&client).await, 0); +} diff --git a/tests/it/main.rs b/tests/it/main.rs index a137381d..a5f57a03 100644 --- a/tests/it/main.rs +++ b/tests/it/main.rs @@ -249,6 +249,8 @@ pub(crate) mod decimals { mod chrono; mod cloud_jwt; mod compression; +#[cfg(feature = "native-transport")] +mod native; mod cursor_error; mod cursor_stats; mod fetch_bytes; @@ -257,6 +259,8 @@ mod insert; mod insert_formatted; #[cfg(feature = "inserter")] mod inserter; +#[cfg(feature = "batcher")] +mod batcher; mod int128; mod int256; mod ip; From 971e98c8949bb4ecd995f320b221148cf2e81af8 Mon Sep 17 00:00:00 2001 From: Derek Date: Tue, 10 Mar 2026 14:05:15 +1100 Subject: [PATCH 5/6] refactor(native): deadpool pool, cursor drain, connection health check; edge-case tests Pool: - Replace hand-rolled Semaphore/Mutex pool with deadpool::managed - NativeConnectionManager implements Manager::create/recycle - Poisoned connections rejected in recycle(); dropped instead of recycled - NativeClient stores NativePool directly (Arc-backed); builder methods call rebuild_pool() on config changes Connection: - Add poisoned field + is_poisoned(); check_alive() for recycle health check - check_alive(): rejects if poisoned, BufReader has leftover bytes, or non-blocking poll_read returns EOF/unexpected data (noop waker) Cursor / Query: - NativeRowCursor: add drain() to consume packets until EndOfStream - NativeRowCursor: Drop impl discards connection if stream not fully consumed - fetch_one / fetch_optional: call drain() after retrieving row so the connection is returned to the pool in a clean state (fixes RowNotFound race when pool connections were reused mid-stream) Insert: - NativeInsert: add Drop impl that calls abort() to discard connection if end() was never called - end(): take conn out on success so Drop is a no-op after clean commit Tests (59 integration tests, 11 sparse unit tests): - Pool: concurrent (10 tasks / pool_size=2), error recovery, insert+ping - Bool/sparse: stream alignment, multi-column, large gap (1000 rows), single true at start, single true at end - Sparse unit: single-row, large gap, consecutive non-defaults, state carry - Cursor: fetch_one on empty, fetch_optional, bind() edge cases - INSERT: large batch, abort, sequential, bool insert - Schema cache: miss, clear_all --- src/native/client.rs | 101 ++++-- src/native/cursor.rs | 54 +++ src/native/insert.rs | 22 +- src/native/pool.rs | 180 ++++----- src/native/query.rs | 21 +- tests/it/native.rs | 844 +++++++++++++++++++++++++++++++++++++++++++ 6 files changed, 1058 insertions(+), 164 deletions(-) diff --git a/src/native/client.rs b/src/native/client.rs index 89c5f9ef..d21b409b 100644 --- a/src/native/client.rs +++ b/src/native/client.rs @@ -18,12 +18,12 @@ //! [`with_pool_size`]: NativeClient::with_pool_size use std::net::{SocketAddr, ToSocketAddrs}; -use std::sync::{Arc, OnceLock}; +use std::sync::Arc; -use crate::error::Result; +use crate::error::{Error, Result}; use crate::native::insert::NativeInsert; use crate::native::inserter::NativeInserter; -use crate::native::pool::{NativePool, PoolConfig, PooledConnection}; +use crate::native::pool::{NativePool, PoolConfig, PooledConnection, build_pool}; use crate::native::protocol::NativeCompressionMethod; use crate::native::query::NativeQuery; use crate::native::schema::NativeSchemaCache; @@ -62,30 +62,61 @@ pub struct NativeClient { settings: Arc>, /// Maximum connections (idle + in-use) in the pool. pool_size: usize, - /// Lazily-initialised pool; reset whenever connection params change. - /// - /// Wrapped in `Arc` so `Clone` shares the same pool across copies of a - /// fully-configured client. - pool: Arc>>, + /// Deadpool-backed connection pool. Already Arc-backed internally, so + /// cloning this client shares the same pool across all copies. + pool: NativePool, } impl Default for NativeClient { fn default() -> Self { + let addr: SocketAddr = "127.0.0.1:9000".parse().expect("valid default addr"); + let database = "default".to_string(); + let username = "default".to_string(); + let password = String::new(); + let compression = NativeCompressionMethod::None; + let settings: Vec<(String, String)> = Vec::new(); + let pool = build_pool( + PoolConfig { + addr, + database: database.clone(), + username: username.clone(), + password: password.clone(), + compression, + settings: settings.clone(), + }, + DEFAULT_POOL_SIZE, + ); Self { - addr: "127.0.0.1:9000".parse().expect("valid default addr"), - database: "default".to_string(), - username: "default".to_string(), - password: String::new(), - compression: NativeCompressionMethod::None, + addr, + database, + username, + password, + compression, schema_cache: NativeSchemaCache::new(300), - settings: Arc::new(Vec::new()), + settings: Arc::new(settings), pool_size: DEFAULT_POOL_SIZE, - pool: Arc::new(OnceLock::new()), + pool, } } } impl NativeClient { + /// Rebuild the connection pool from the current client configuration. + /// Called internally whenever a connection parameter changes. + fn rebuild_pool(&mut self) { + self.pool = build_pool( + PoolConfig { + addr: self.addr, + database: self.database.clone(), + username: self.username.clone(), + password: self.password.clone(), + compression: self.compression, + settings: self.settings.as_ref().clone(), + }, + self.pool_size, + ); + } + /// Set the server address (host:port). /// /// # Panics @@ -98,7 +129,7 @@ impl NativeClient { .expect("invalid address") .next() .expect("no address resolved"); - self.pool = Arc::new(OnceLock::new()); + self.rebuild_pool(); self } @@ -106,7 +137,7 @@ impl NativeClient { #[must_use] pub fn with_database(mut self, database: impl Into) -> Self { self.database = database.into(); - self.pool = Arc::new(OnceLock::new()); + self.rebuild_pool(); self } @@ -114,7 +145,7 @@ impl NativeClient { #[must_use] pub fn with_user(mut self, user: impl Into) -> Self { self.username = user.into(); - self.pool = Arc::new(OnceLock::new()); + self.rebuild_pool(); self } @@ -122,7 +153,7 @@ impl NativeClient { #[must_use] pub fn with_password(mut self, password: impl Into) -> Self { self.password = password.into(); - self.pool = Arc::new(OnceLock::new()); + self.rebuild_pool(); self } @@ -130,7 +161,7 @@ impl NativeClient { #[must_use] pub fn with_lz4(mut self) -> Self { self.compression = NativeCompressionMethod::Lz4; - self.pool = Arc::new(OnceLock::new()); + self.rebuild_pool(); self } @@ -141,7 +172,7 @@ impl NativeClient { #[must_use] pub fn with_pool_size(mut self, size: usize) -> Self { self.pool_size = size; - self.pool = Arc::new(OnceLock::new()); + self.rebuild_pool(); self } @@ -165,7 +196,7 @@ impl NativeClient { value: impl Into, ) -> Self { Arc::make_mut(&mut self.settings).push((name.into(), value.into())); - self.pool = Arc::new(OnceLock::new()); + self.rebuild_pool(); self } @@ -287,24 +318,16 @@ impl NativeClient { } /// Acquire a connection from the pool, opening a new one if needed. - /// - /// The pool is created lazily on first call with a snapshot of the - /// current connection parameters. pub(crate) async fn acquire(&self) -> Result { - let pool = self.pool.get_or_init(|| { - NativePool::new( - PoolConfig { - addr: self.addr, - database: self.database.clone(), - username: self.username.clone(), - password: self.password.clone(), - compression: self.compression, - settings: self.settings.as_ref().clone(), - }, - self.pool_size, - ) - }); - pool.acquire().await + use deadpool::managed::PoolError; + self.pool + .get() + .await + .map(PooledConnection::new) + .map_err(|e| match e { + PoolError::Backend(e) => e, + e => Error::Custom(format!("pool: {e}")), + }) } /// Ping the server. diff --git a/src/native/cursor.rs b/src/native/cursor.rs index e3f157bb..cbccfa5b 100644 --- a/src/native/cursor.rs +++ b/src/native/cursor.rs @@ -35,6 +35,23 @@ enum CursorState { Done, } +impl Drop for NativeRowCursor { + /// Discard the connection if the stream was never fully consumed. + /// + /// Dropping a cursor mid-stream (e.g. after `fetch_one`) without calling + /// `drain()` first would return a connection with unread bytes to the pool. + /// Marking it poisoned here ensures deadpool drops it instead of recycling. + fn drop(&mut self) { + if let CursorState::Reading(conn) = + std::mem::replace(&mut self.state, CursorState::Done) + { + // We can't async-drain here, so discard the connection. + let mut conn = conn; + conn.discard(); + } + } +} + impl NativeRowCursor { pub(crate) fn new(client: NativeClient, sql: String) -> Self { Self { @@ -46,6 +63,43 @@ impl NativeRowCursor { } } + /// Consume all remaining packets until `EndOfStream`, allowing the + /// underlying connection to be returned to the pool in a clean state. + /// + /// Must be called after a partial read (e.g. after `fetch_one` got its row) + /// to prevent the half-read connection from being recycled with unread data. + pub(crate) async fn drain(&mut self) -> Result<()> { + loop { + match &self.state { + CursorState::Done | CursorState::NotStarted => return Ok(()), + CursorState::Reading(_) => {} + } + let CursorState::Reading(conn) = &mut self.state else { + unreachable!() + }; + let revision = conn.server_revision(); + let compression = conn.compression(); + let packet = + crate::native::reader::read_packet(conn.reader_mut(), revision, compression) + .await; + match packet { + Ok(ServerPacket::EndOfStream) => { + self.state = CursorState::Done; + return Ok(()); + } + Ok(_) => {} + Err(e) => { + if let CursorState::Reading(mut conn) = + std::mem::replace(&mut self.state, CursorState::Done) + { + conn.discard(); + } + return Err(e); + } + } + } + } + /// Return the next deserialized row, or `None` at end of stream. /// /// `T` must be [`RowOwned`], meaning the result does not borrow from diff --git a/src/native/insert.rs b/src/native/insert.rs index d3243976..def87039 100644 --- a/src/native/insert.rs +++ b/src/native/insert.rs @@ -130,11 +130,18 @@ impl NativeInsert { return Err(e); } } - self.conn + let result = self + .conn .as_mut() .expect("conn must be open") .finish_insert() - .await + .await; + if result.is_ok() { + // Take the connection out so our Drop impl does not discard it. + // Dropping the PooledConnection here returns it to the idle pool. + let _ = self.conn.take(); + } + result } async fn ensure_connected(&mut self) -> Result<()> { @@ -164,6 +171,9 @@ impl NativeInsert { conn.send_insert_block(&column_bytes, self.columns.len(), n).await } +} + +impl NativeInsert { /// Abort the INSERT: discard the connection and clear the buffer. /// /// The server-side INSERT is incomplete — we must not return this @@ -177,3 +187,11 @@ impl NativeInsert { self.row_bytes = 0; } } + +impl Drop for NativeInsert { + /// If [`end`](NativeInsert::end) was not called, discard the connection so + /// it is never returned to the pool mid-INSERT. + fn drop(&mut self) { + self.abort(); + } +} diff --git a/src/native/pool.rs b/src/native/pool.rs index b3f479a4..3dc6c1fa 100644 --- a/src/native/pool.rs +++ b/src/native/pool.rs @@ -1,26 +1,22 @@ //! Connection pool for the native TCP transport. //! -//! Holds a bounded set of idle [`NativeConnection`]s so successive queries -//! and inserts can reuse TCP connections instead of paying handshake overhead -//! on every operation. +//! Thin wrapper around [`deadpool::managed`]. A [`NativeConnectionManager`] +//! teaches deadpool how to open and recycle [`NativeConnection`]s; all pool +//! mechanics (semaphore, idle queue, timeouts, metrics) are handled by +//! deadpool. //! -//! # Model +//! # Discard pattern //! -//! - A [`Semaphore`] caps the total number of connections (idle + in-use) to -//! `max_size`. Callers block on [`NativePool::acquire`] when the pool is -//! full until a permit becomes available. -//! - Idle connections are stored in a `Mutex` (FIFO, so recently -//! used connections are preferred). -//! - On [`PooledConnection`] drop the connection is returned to the idle -//! queue unless [`PooledConnection::discard`] was called, in which case it -//! is closed and the semaphore slot is released. +//! When an I/O error or incomplete protocol exchange leaves a connection in an +//! unrecoverable state, call [`PooledConnection::discard`]. This sets the +//! `poisoned` flag on the underlying [`NativeConnection`]; deadpool's +//! `recycle()` hook sees the flag and drops the connection instead of +//! returning it to the idle queue. -use std::collections::VecDeque; use std::net::SocketAddr; use std::ops::{Deref, DerefMut}; -use std::sync::{Arc, Mutex}; -use tokio::sync::Semaphore; +use deadpool::managed::{self, RecycleError, RecycleResult}; use crate::error::{Error, Result}; use crate::native::connection::NativeConnection; @@ -36,134 +32,84 @@ pub(crate) struct PoolConfig { pub(crate) settings: Vec<(String, String)>, } -/// A bounded idle-connection pool. -pub(crate) struct NativePool { - idle: Mutex>, - /// Total in-use + idle connections must not exceed `max_size`. - semaphore: Semaphore, +/// deadpool [`Manager`](managed::Manager) for [`NativeConnection`]. +pub(crate) struct NativeConnectionManager { config: PoolConfig, - max_size: usize, } -impl NativePool { - pub(crate) fn new(config: PoolConfig, max_size: usize) -> Arc { - Arc::new(Self { - idle: Mutex::new(VecDeque::new()), - semaphore: Semaphore::new(max_size), - config, - max_size, - }) +impl managed::Manager for NativeConnectionManager { + type Type = NativeConnection; + type Error = Error; + + async fn create(&self) -> Result { + NativeConnection::open( + &self.config.addr, + &self.config.database, + &self.config.username, + &self.config.password, + self.config.compression, + self.config.settings.clone(), + ) + .await } - /// Acquire a connection. Waits if the pool is at capacity. - /// - /// Tries an idle connection first; opens a new one if none are available. - pub(crate) async fn acquire(self: &Arc) -> Result { - let permit = self - .semaphore - .acquire() - .await - .map_err(|_| Error::Custom("connection pool closed".into()))?; - // We manage permits manually — forget the RAII guard. - permit.forget(); - - let conn = { - let mut idle = self.idle.lock().expect("pool mutex"); - idle.pop_front() - }; - - let conn = match conn { - Some(c) => c, - None => { - NativeConnection::open( - &self.config.addr, - &self.config.database, - &self.config.username, - &self.config.password, - self.config.compression, - self.config.settings.clone(), - ) - .await - .inspect_err(|_| { - // Opening failed — release the permit so the pool slot isn't lost. - self.semaphore.add_permits(1); - })? - } - }; - - Ok(PooledConnection { - conn: Some(conn), - pool: Arc::clone(self), - return_to_pool: true, - }) - } - - /// Return a connection to the idle queue and release its semaphore slot. - fn return_conn(&self, conn: NativeConnection) { - { - let mut idle = self.idle.lock().expect("pool mutex"); - // Guard against exceeding max_size in the idle queue. - if idle.len() < self.max_size { - idle.push_back(conn); - } - // If somehow over capacity, just drop the connection. + async fn recycle( + &self, + conn: &mut NativeConnection, + _: &managed::Metrics, + ) -> RecycleResult { + if !conn.check_alive() { + return Err(RecycleError::message("connection dead or dirty")); } - self.semaphore.add_permits(1); + Ok(()) } +} - /// Release a semaphore permit without returning the connection (broken path). - fn release_permit(&self) { - self.semaphore.add_permits(1); - } +/// A bounded connection pool backed by deadpool. +pub(crate) type NativePool = managed::Pool; + +/// Build a new pool with the given config and connection cap. +pub(crate) fn build_pool(config: PoolConfig, max_size: usize) -> NativePool { + let mgr = NativeConnectionManager { config }; + managed::Pool::builder(mgr) + .max_size(max_size) + .build() + .expect("pool config is always valid") } -/// A connection borrowed from a [`NativePool`]. +/// A connection borrowed from the pool. /// -/// Dereferences to [`NativeConnection`] for transparent method calls. -/// When dropped, the connection is returned to the pool — unless -/// [`discard`](PooledConnection::discard) was called, in which case the -/// connection is closed and the pool slot is freed. +/// Dereferences to [`NativeConnection`] for transparent method access. +/// Returns the connection to the idle queue on drop, unless +/// [`discard`](PooledConnection::discard) was called first. pub(crate) struct PooledConnection { - conn: Option, - pool: Arc, - return_to_pool: bool, + inner: managed::Object, } impl PooledConnection { - /// Mark this connection as broken — it will be closed on drop instead of - /// returned to the pool. Call this when an I/O error or incomplete - /// protocol exchange leaves the connection in an unrecoverable state. + pub(crate) fn new(inner: managed::Object) -> Self { + Self { inner } + } + + /// Mark this connection as broken. + /// + /// The connection will be closed on drop rather than returned to the pool. + /// Call this after any I/O error or incomplete protocol exchange that + /// leaves the connection in an unrecoverable state. pub(crate) fn discard(&mut self) { - self.return_to_pool = false; + self.inner.poisoned = true; } } impl Deref for PooledConnection { type Target = NativeConnection; fn deref(&self) -> &NativeConnection { - self.conn - .as_ref() - .expect("PooledConnection invariant: conn is always Some while borrowed") + &self.inner } } impl DerefMut for PooledConnection { fn deref_mut(&mut self) -> &mut NativeConnection { - self.conn - .as_mut() - .expect("PooledConnection invariant: conn is always Some while borrowed") - } -} - -impl Drop for PooledConnection { - fn drop(&mut self) { - if let Some(conn) = self.conn.take() { - if self.return_to_pool { - self.pool.return_conn(conn); - } else { - drop(conn); - self.pool.release_permit(); - } - } + &mut self.inner } } diff --git a/src/native/query.rs b/src/native/query.rs index 7c1f8bc2..bdc46c73 100644 --- a/src/native/query.rs +++ b/src/native/query.rs @@ -78,11 +78,16 @@ impl NativeQuery { where T: RowOwned + RowRead, { - match self.fetch::()?.next().await { - Ok(Some(row)) => Ok(row), - Ok(None) => Err(Error::RowNotFound), - Err(err) => Err(err), - } + let mut cursor = self.fetch::()?; + let row = match cursor.next().await { + Ok(Some(row)) => row, + Ok(None) => return Err(Error::RowNotFound), + Err(err) => return Err(err), + }; + // Drain remaining packets so the connection is returned to the pool + // in a clean state rather than mid-stream. + cursor.drain().await?; + Ok(row) } /// Fetch all rows into a Vec. @@ -103,6 +108,10 @@ impl NativeQuery { where T: RowOwned + RowRead, { - self.fetch::()?.next().await + let mut cursor = self.fetch::()?; + let row = cursor.next().await?; + // Drain if we got a row without reaching EndOfStream. + cursor.drain().await?; + Ok(row) } } diff --git a/tests/it/native.rs b/tests/it/native.rs index 544cb3a0..c8e118ec 100644 --- a/tests/it/native.rs +++ b/tests/it/native.rs @@ -1592,3 +1592,847 @@ async fn native_insert_lz4() { assert_eq!(rows[1].id, 2); assert_eq!(rows[1].name, "bob"); } + +// --------------------------------------------------------------------------- +// Pool edge cases +// --------------------------------------------------------------------------- + +/// Pool size 2, 10 concurrent tasks — all must succeed. +/// Verifies that tasks waiting for a connection are eventually served. +#[tokio::test] +async fn native_pool_concurrent() { + let client = get_native_client().with_pool_size(2); + let tasks: Vec<_> = (0..10u8) + .map(|i| { + let c = client.clone(); + tokio::spawn(async move { + let n: u8 = c + .query(&format!("SELECT {i}")) + .fetch_one::() + .await + .expect("concurrent query failed"); + assert_eq!(n, i); + }) + }) + .collect(); + for task in tasks { + task.await.expect("task panicked"); + } +} + +/// A server exception must not permanently break the pool. +/// The next query after an error must succeed on a fresh/recycled connection. +#[tokio::test] +async fn native_pool_error_recovery() { + let client = get_native_client(); + + // Trigger a server exception (table does not exist). + let result = client + .query("SELECT * FROM _this_table_does_not_exist_clickhouse_rs_test") + .fetch_all::() + .await; + assert!(result.is_err(), "expected error from bad query"); + + // Pool must still be usable after the error. + let n: u8 = client + .query("SELECT 99") + .fetch_one::() + .await + .expect("query after error must succeed"); + assert_eq!(n, 99); +} + +/// Pool size 1 + many concurrent inserts — verifies no deadlock when the +/// INSERT holds the sole connection and another task waits for it. +#[tokio::test] +async fn native_pool_insert_wait() { + let client = prepare_native_database("pool_insert_wait").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + let small_client = client.clone().with_pool_size(1); + + #[derive(Debug, Row, Serialize)] + struct R { + id: u32, + } + + // Task A holds the connection in an INSERT. + // Task B tries to ping at the same time — it must wait, not deadlock. + let client_a = small_client.clone(); + let client_b = small_client.clone(); + + let insert_task = tokio::spawn(async move { + let mut ins = client_a.insert::("t"); + for i in 0..100u32 { + ins.write(&R { id: i }).await.expect("write failed"); + } + ins.end().await.expect("end failed"); + }); + let ping_task = tokio::spawn(async move { + client_b.ping().await.expect("ping failed while insert held connection"); + }); + + insert_task.await.expect("insert task panicked"); + ping_task.await.expect("ping task panicked"); +} + +// --------------------------------------------------------------------------- +// Bool / sparse-serialization edge cases +// --------------------------------------------------------------------------- + +/// All rows false — sparse format sends 0 non-default values. +#[tokio::test] +async fn native_bool_all_false() { + let client = prepare_native_database("bool_all_false").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, flag Bool) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + for i in 1..=8u32 { + client + .query(&format!("INSERT INTO t VALUES ({i}, false)")) + .execute() + .await + .expect("INSERT failed"); + } + + #[derive(Debug, Row, Deserialize, PartialEq)] + struct BoolRow { + id: u32, + flag: bool, + } + + let rows = client + .query("SELECT id, flag FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 8); + for row in &rows { + assert!(!row.flag, "expected false for id={}", row.id); + } +} + +/// All rows true — sparse format stores every row as a non-default value. +#[tokio::test] +async fn native_bool_all_true() { + let client = prepare_native_database("bool_all_true").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, flag Bool) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + for i in 1..=8u32 { + client + .query(&format!("INSERT INTO t VALUES ({i}, true)")) + .execute() + .await + .expect("INSERT failed"); + } + + #[derive(Debug, Row, Deserialize, PartialEq)] + struct BoolRow { + id: u32, + flag: bool, + } + + let rows = client + .query("SELECT id, flag FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 8); + for row in &rows { + assert!(row.flag, "expected true for id={}", row.id); + } +} + +/// Mixed true/false across many rows — exercises sparse offset groups. +#[tokio::test] +async fn native_bool_many_rows() { + let client = prepare_native_database("bool_many_rows").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, flag Bool) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + // Insert 200 rows in one batch: alternating true/false, then a run of + // 50 trues, then 50 falses — exercises multiple sparse offset groups. + let vals: String = (0..200u32) + .map(|i| { + let b = if i < 100 { i % 2 == 0 } else { i < 150 }; + format!("({i}, {})", b) + }) + .collect::>() + .join(", "); + client + .query(&format!("INSERT INTO t VALUES {vals}")) + .execute() + .await + .expect("INSERT failed"); + + #[derive(Debug, Row, Deserialize)] + struct BoolRow { + id: u32, + flag: bool, + } + + let rows = client + .query("SELECT id, flag FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 200); + for row in &rows { + let expected = if row.id < 100 { + row.id % 2 == 0 + } else { + row.id < 150 + }; + assert_eq!(row.flag, expected, "mismatch at id={}", row.id); + } +} + +/// Nullable(Bool): Some(true), Some(false), NULL. +#[tokio::test] +async fn native_bool_nullable() { + let client = prepare_native_database("bool_nullable").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, flag Nullable(Bool)) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + client + .query("INSERT INTO t VALUES (1, true), (2, false), (3, NULL)") + .execute() + .await + .expect("INSERT failed"); + + #[derive(Debug, Row, Deserialize, PartialEq)] + struct BoolRow { + id: u32, + flag: Option, + } + + let rows = client + .query("SELECT id, flag FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 3); + assert_eq!(rows[0], BoolRow { id: 1, flag: Some(true) }); + assert_eq!(rows[1], BoolRow { id: 2, flag: Some(false) }); + assert_eq!(rows[2], BoolRow { id: 3, flag: None }); +} + +/// INSERT Bool via `NativeInsert` (not SQL VALUES) — tests the encoder path. +#[tokio::test] +async fn native_insert_bool() { + let client = prepare_native_database("insert_bool").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, flag Bool) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Serialize, Deserialize, PartialEq)] + struct BoolRow { + id: u32, + flag: bool, + } + + let mut insert = client.insert::("t"); + insert.write(&BoolRow { id: 1, flag: true }).await.expect("write 1 failed"); + insert.write(&BoolRow { id: 2, flag: false }).await.expect("write 2 failed"); + insert.write(&BoolRow { id: 3, flag: true }).await.expect("write 3 failed"); + insert.end().await.expect("end failed"); + + let rows = client + .query("SELECT id, flag FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 3); + assert_eq!(rows[0], BoolRow { id: 1, flag: true }); + assert_eq!(rows[1], BoolRow { id: 2, flag: false }); + assert_eq!(rows[2], BoolRow { id: 3, flag: true }); +} + +/// Bool column alongside non-sparse UInt32 and String columns. +/// +/// Verifies that the sparse decoder does not misalign the stream — after reading +/// the Bool column's sparse offsets + values, the reader must be positioned +/// exactly at the next column's data. +#[tokio::test] +async fn native_bool_sparse_stream_alignment() { + let client = prepare_native_database("bool_sparse_align").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, flag Bool, name String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + client + .query("INSERT INTO t VALUES (1, true, 'alice'), (2, false, 'bob'), (3, true, 'carol'), (4, false, 'dave')") + .execute() + .await + .expect("INSERT failed"); + + #[derive(Debug, Row, Deserialize, PartialEq)] + struct R { + id: u32, + flag: bool, + name: String, + } + + let rows = client + .query("SELECT id, flag, name FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 4); + assert_eq!(rows[0], R { id: 1, flag: true, name: "alice".into() }); + assert_eq!(rows[1], R { id: 2, flag: false, name: "bob".into() }); + assert_eq!(rows[2], R { id: 3, flag: true, name: "carol".into() }); + assert_eq!(rows[3], R { id: 4, flag: false, name: "dave".into() }); +} + +/// Two consecutive Bool columns — each must decode its own sparse stream +/// independently without cross-contamination. +#[tokio::test] +async fn native_bool_multi_sparse_columns() { + let client = prepare_native_database("bool_multi_sparse").await; + + client + .query(&format!( + "CREATE TABLE t{} (a Bool, b Bool) {}", + on_cluster(), + test_engine("a"), + )) + .execute() + .await + .expect("CREATE failed"); + + // a: T F T F T, b: F F T T F → different sparse patterns. + client + .query("INSERT INTO t VALUES (true,false),(false,false),(true,true),(false,true),(true,false)") + .execute() + .await + .expect("INSERT failed"); + + #[derive(Debug, Row, Deserialize, PartialEq)] + struct R { a: bool, b: bool } + + let rows = client + .query("SELECT a, b FROM t ORDER BY (a,b)") + .fetch_all::() + .await + .expect("fetch failed"); + + // Sort-order independent check: collect (a,b) pairs. + let mut got: Vec<(bool, bool)> = rows.iter().map(|r| (r.a, r.b)).collect(); + got.sort(); + // (T,F),(F,F),(T,T),(F,T),(T,F) sorted: (F,F),(F,T),(T,F),(T,F),(T,T) + assert_eq!( + got, + vec![(false,false),(false,true),(true,false),(true,false),(true,true)] + ); +} + +/// 1 000 rows, only the last one is `true`. +/// +/// Exercises large VarUInt offsets in the sparse stream (offset group = 999). +#[tokio::test] +async fn native_bool_sparse_large_gap() { + let client = prepare_native_database("bool_sparse_large_gap").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, flag Bool) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Serialize)] + struct W { id: u32, flag: bool } + #[derive(Debug, Row, Deserialize)] + struct R { id: u32, flag: bool } + + const N: u32 = 1000; + let mut insert = client.insert::("t"); + for i in 0..N { + insert.write(&W { id: i, flag: i == N - 1 }).await.expect("write failed"); + } + insert.end().await.expect("end failed"); + + let rows = client + .query("SELECT id, flag FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), N as usize); + for (i, row) in rows.iter().enumerate() { + assert_eq!(row.id, i as u32); + assert_eq!(row.flag, i == (N - 1) as usize, + "row {i}: expected flag={}", i == (N - 1) as usize); + } +} + +/// 1 000 rows, only position 0 is `true` — zero-offset sparse group. +#[tokio::test] +async fn native_bool_sparse_single_at_start() { + let client = prepare_native_database("bool_sparse_single_start").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, flag Bool) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Serialize)] + struct W { id: u32, flag: bool } + #[derive(Debug, Row, Deserialize)] + struct R { id: u32, flag: bool } + + const N: u32 = 1000; + let mut insert = client.insert::("t"); + for i in 0..N { + insert.write(&W { id: i, flag: i == 0 }).await.expect("write failed"); + } + insert.end().await.expect("end failed"); + + let rows = client + .query("SELECT id, flag FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), N as usize); + assert!(rows[0].flag, "row 0 should be true"); + for row in &rows[1..] { + assert!(!row.flag, "row {} should be false", row.id); + } +} + +// --------------------------------------------------------------------------- +// INSERT edge cases +// --------------------------------------------------------------------------- + +/// Write 50 000 rows — enough to trigger multiple intermediate flushes at the +/// 256 KiB threshold. Verifies all rows arrive after end(). +#[tokio::test] +async fn native_insert_large_batch() { + let client = prepare_native_database("insert_large_batch").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt64) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Serialize)] + struct NumRow { + id: u64, + } + + const N: u64 = 50_000; + let mut insert = client.insert::("t"); + for i in 0..N { + insert.write(&NumRow { id: i }).await.expect("write failed"); + } + insert.end().await.expect("end failed"); + + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one::() + .await + .expect("count failed"); + + assert_eq!(count, N, "row count mismatch after large batch"); +} + +/// Drop `NativeInsert` without calling `end()` — must not commit any data, +/// and must not leave the pool connection in a broken state. +#[tokio::test] +async fn native_insert_abort() { + let client = prepare_native_database("insert_abort").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Serialize)] + struct R { + id: u32, + } + + // Write two rows then drop without end() — aborts the INSERT. + { + let mut insert = client.insert::("t"); + insert.write(&R { id: 1 }).await.expect("write 1 failed"); + insert.write(&R { id: 2 }).await.expect("write 2 failed"); + // dropped here — connection must be discarded, not returned to pool + } + + // The pool must still work after the aborted insert. + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one::() + .await + .expect("count after abort failed"); + + assert_eq!(count, 0, "aborted insert must not commit data"); +} + +/// Two sequential inserts into the same table to verify pool reuse between +/// INSERT operations does not misalign the protocol. +#[tokio::test] +async fn native_insert_sequential() { + let client = prepare_native_database("insert_sequential").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Serialize)] + struct R { + id: u32, + } + + // First INSERT + let mut ins = client.insert::("t"); + ins.write(&R { id: 1 }).await.expect("write 1a failed"); + ins.write(&R { id: 2 }).await.expect("write 1b failed"); + ins.end().await.expect("end 1 failed"); + + // Second INSERT reuses the same connection from the pool. + let mut ins2 = client.insert::("t"); + ins2.write(&R { id: 3 }).await.expect("write 2a failed"); + ins2.end().await.expect("end 2 failed"); + + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one::() + .await + .expect("count failed"); + + assert_eq!(count, 3); +} + +// --------------------------------------------------------------------------- +// Query API edge cases +// --------------------------------------------------------------------------- + +/// `fetch_one` on an empty result set must return `RowNotFound`. +#[tokio::test] +async fn native_query_fetch_one_empty() { + let client = prepare_native_database("fetch_one_empty").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + let result = client + .query("SELECT id FROM t") + .fetch_one::() + .await; + + assert!( + matches!(result, Err(clickhouse::error::Error::RowNotFound)), + "expected RowNotFound, got {result:?}" + ); +} + +/// `fetch_optional` returns `None` when no rows match, `Some` when one does. +#[tokio::test] +async fn native_query_fetch_optional() { + let client = prepare_native_database("fetch_optional").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + // Empty table → None. + let none: Option = client + .query("SELECT id FROM t") + .fetch_optional::() + .await + .expect("fetch_optional failed"); + assert!(none.is_none()); + + // Insert one row. + client + .query("INSERT INTO t VALUES (42)") + .execute() + .await + .expect("INSERT failed"); + + // One row → Some. + let some: Option = client + .query("SELECT id FROM t LIMIT 1") + .fetch_optional::() + .await + .expect("fetch_optional failed"); + assert_eq!(some, Some(42u32)); +} + +/// `bind()` with multiple `?` placeholders — each replaces the next occurrence. +#[tokio::test] +async fn native_query_bind_multiple() { + let client = get_native_client(); + + // ClickHouse infers UInt8 for small literals; match the inferred type. + let result: u8 = client + .query("SELECT ? + ?") + .bind(10u8) + .bind(32u8) + .fetch_one::() + .await + .expect("fetch failed"); + + assert_eq!(result, 42u8); +} + +/// `bind()` when the SQL has no `?` — should be a no-op (query unchanged). +#[tokio::test] +async fn native_query_bind_no_placeholder() { + let client = get_native_client(); + + let result: u8 = client + .query("SELECT 1") + .bind(999u32) // no placeholder — ignored + .fetch_one::() + .await + .expect("fetch failed"); + + assert_eq!(result, 1u8); +} + +// --------------------------------------------------------------------------- +// Nullable edge cases +// --------------------------------------------------------------------------- + +/// Column that is NULL for every row. +#[tokio::test] +async fn native_nullable_all_null() { + let client = prepare_native_database("nullable_all_null").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, val Nullable(Int64)) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + client + .query("INSERT INTO t VALUES (1, NULL), (2, NULL), (3, NULL)") + .execute() + .await + .expect("INSERT failed"); + + #[derive(Debug, Row, Deserialize, PartialEq)] + struct R { + id: u32, + val: Option, + } + + let rows = client + .query("SELECT id, val FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 3); + for row in &rows { + assert!(row.val.is_none(), "expected NULL for id={}", row.id); + } +} + +/// Array(Nullable(String)) — nulls inside an array. +#[tokio::test] +async fn native_array_of_nullable() { + let client = prepare_native_database("array_nullable").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, tags Array(Nullable(String))) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + client + .query("INSERT INTO t VALUES (1, ['a', NULL, 'b']), (2, [])") + .execute() + .await + .expect("INSERT failed"); + + #[derive(Debug, Row, Deserialize)] + struct R { + id: u32, + tags: Vec>, + } + + let rows = client + .query("SELECT id, tags FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 2); + assert_eq!( + rows[0].tags, + vec![Some("a".into()), None, Some("b".into())] + ); + assert_eq!(rows[1].tags, Vec::>::new()); +} + +// --------------------------------------------------------------------------- +// Schema cache edge cases +// --------------------------------------------------------------------------- + +/// `fetch_schema` on a non-existent table must return an error (empty result). +#[tokio::test] +async fn native_schema_cache_miss() { + let client = get_native_client(); + + // Non-existent table returns empty schema (not an error from ClickHouse). + let schema = client + .fetch_schema("_this_table_does_not_exist_xyz_clickhouse_rs") + .await + .expect("fetch_schema should not error on missing table"); + + assert!( + schema.is_empty(), + "expected empty schema for non-existent table, got {schema:?}" + ); +} + +/// `clear_all_cached_schemas` removes all entries; subsequent access re-fetches. +#[tokio::test] +async fn native_schema_cache_clear_all() { + let client = prepare_native_database("schema_clear_all").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, name String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Serialize)] + struct TestRow { + id: u32, + name: String, + } + + // Populate the cache via an INSERT. + let mut ins = client.insert::("t"); + ins.write(&TestRow { id: 1, name: "x".into() }).await.expect("write failed"); + ins.end().await.expect("end failed"); + + assert!(client.cached_schema("t").is_some(), "cache should be populated after INSERT"); + + // Clear all — cache must be empty. + client.clear_all_cached_schemas(); + assert!(client.cached_schema("t").is_none(), "cache should be empty after clear_all"); + + // fetch_schema re-populates. + let schema = client.fetch_schema("t").await.expect("fetch_schema failed"); + assert!(!schema.is_empty()); + assert!(client.cached_schema("t").is_some(), "cache should be re-populated after fetch_schema"); +} From 1688e1100eb8e55231d14ec66f8e52034991da41 Mon Sep 17 00:00:00 2001 From: Derek Date: Tue, 10 Mar 2026 14:42:02 +1100 Subject: [PATCH 6/6] feat(native): LowCardinality INSERT + fix LC(Nullable(T)) reader MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Encoder (encode.rs): - Remove strip_low_cardinality() — send original type name to server - Implement proper LowCardinality wire encoding in write_col_values: * version=1, HAS_ADDITIONAL_KEYS flag, dict + indices * For LC(Nullable(T)): dict type is T (not Nullable(T)); index 0 is the null sentinel (default T value); null inputs → index 0 * For LC(T): plain T dict; indices map rows to unique values * Chooses smallest index type (U8/U16/U32/U64) based on dict size Reader (columns.rs): - Fix read_low_cardinality_column for LC(Nullable(T)): * Dict is now read using T (not Nullable(T)) — matches ClickHouse wire * Index 0 → RowBinary null [0x01]; other indices → [0x00, T bytes] * Non-nullable LC unchanged Tests: - native_insert_low_cardinality: LowCardinality(String) INSERT + readback - native_insert_low_cardinality_nullable: LC(Nullable(String)) with None values --- src/native/columns.rs | 43 +++++++++++---- src/native/encode.rs | 123 ++++++++++++++++++++++++++++++------------ tests/it/native.rs | 81 ++++++++++++++++++++++++++++ 3 files changed, 205 insertions(+), 42 deletions(-) diff --git a/src/native/columns.rs b/src/native/columns.rs index 482408b5..2e9238f8 100644 --- a/src/native/columns.rs +++ b/src/native/columns.rs @@ -531,6 +531,15 @@ async fn read_low_cardinality_column( // Bit 9: HAS_ADDITIONAL_KEYS — server sends per-block additional keys let has_additional_keys = (state & 0x200) != 0; + // For LowCardinality(Nullable(T)), the dictionary on the wire is of type T + // (not Nullable(T)). Index 0 is a special null-sentinel entry (the default + // T value, e.g. "" for String). All other indices reference non-null T values. + let (dict_type, is_nullable_inner) = if let ColumnType::Nullable(t) = inner { + (t.as_ref(), true) + } else { + (inner, false) + }; + // Indices 0.. reference additional_keys first, then global_dict. // Build combined dict in that order. let mut additional: ColumnData = Vec::new(); @@ -538,12 +547,12 @@ async fn read_low_cardinality_column( if has_global_dict { let sz = reader.read_u64_le().await?; - global = read_column(reader, inner, sz).await?; + global = read_column(reader, dict_type, sz).await?; } if has_additional_keys { let sz = reader.read_u64_le().await?; - additional = read_column(reader, inner, sz).await?; + additional = read_column(reader, dict_type, sz).await?; } // Combined dict: additional_keys first (indices 0..additional.len()), @@ -556,7 +565,7 @@ async fn read_low_cardinality_column( // (older / simpler LowCardinality without shared dictionaries). if !has_global_dict && !has_additional_keys { let sz = reader.read_u64_le().await?; - dict.extend(read_column(reader, inner, sz).await?); + dict.extend(read_column(reader, dict_type, sz).await?); } let num_indices = reader.read_u64_le().await?; @@ -582,12 +591,28 @@ async fn read_low_cardinality_column( let mut result = Vec::with_capacity(n); for _ in 0..n { let idx = read_index(reader, index_bytes).await? as usize; - let value = dict.get(idx).ok_or_else(|| { - Error::BadResponse(format!( - "native protocol: LowCardinality index {idx} out of range (dict size {dict_size})" - )) - })?; - result.push(value.clone()); + if is_nullable_inner { + // Index 0 = null sentinel → RowBinary null; other indices = Some(T). + if idx == 0 { + result.push(vec![0x01u8]); // RowBinary Nullable null flag + } else { + let value = dict.get(idx).ok_or_else(|| { + Error::BadResponse(format!( + "native protocol: LowCardinality index {idx} out of range (dict size {dict_size})" + )) + })?; + let mut rb = vec![0x00u8]; // RowBinary not-null flag + rb.extend_from_slice(value); + result.push(rb); + } + } else { + let value = dict.get(idx).ok_or_else(|| { + Error::BadResponse(format!( + "native protocol: LowCardinality index {idx} out of range (dict size {dict_size})" + )) + })?; + result.push(value.clone()); + } } Ok(result) } diff --git a/src/native/encode.rs b/src/native/encode.rs index 66c0e33c..5f42a2c9 100644 --- a/src/native/encode.rs +++ b/src/native/encode.rs @@ -6,8 +6,8 @@ //! # Supported types for INSERT //! //! All scalar fixed-size types, String, FixedString(N), Nullable(T), -//! Array(T), Map(K, V), Tuple(T1..Tn), and nested combinations thereof. -//! LowCardinality is stripped to its inner type (ClickHouse accepts plain values). +//! LowCardinality(T), Array(T), Map(K, V), Tuple(T1..Tn), and nested combinations. +//! LowCardinality is fully encoded with a per-block dictionary + indices. //! Variant, Dynamic, and JSON are not yet supported. use crate::error::{Error, Result}; @@ -28,9 +28,6 @@ pub(crate) struct ColumnSchema { impl ColumnSchema { /// Build a `ColumnSchema` list from server-provided `(name, type_name)` pairs. - /// - /// LowCardinality wrappers are stripped — the inner type is sent on wire, - /// which ClickHouse accepts transparently. pub(crate) fn from_headers(headers: &[(String, String)]) -> Result> { headers .iter() @@ -42,40 +39,16 @@ impl ColumnSchema { for column '{name}'" )) })?; - // Strip LowCardinality: send inner type bytes, CH handles encoding - let (effective_type, effective_name) = - strip_low_cardinality(col_type, type_name); Ok(ColumnSchema { name: name.clone(), - type_name: effective_name, - col_type: effective_type, + type_name: type_name.clone(), + col_type, }) }) .collect() } } -/// Recursively strip `LowCardinality(...)` returning the inner `(ColumnType, type_name)`. -fn strip_low_cardinality(col_type: ColumnType, type_name: &str) -> (ColumnType, String) { - match col_type { - ColumnType::LowCardinality(inner) => { - let inner_name = extract_inner(type_name, "LowCardinality"); - strip_low_cardinality(*inner, inner_name) - } - other => (other, type_name.to_string()), - } -} - -fn extract_inner<'a>(s: &'a str, wrapper: &str) -> &'a str { - let prefix = format!("{wrapper}("); - if let Some(rest) = s.strip_prefix(prefix.as_str()) { - if let Some(inner) = rest.strip_suffix(')') { - return inner; - } - } - s -} - /// Encode buffered RowBinary rows into native columnar block column bytes. /// /// Returns a flat byte buffer containing, for each column in order: @@ -166,8 +139,92 @@ fn write_col_values(values: &[Vec], col_type: &ColumnType, out: &mut Vec } ColumnType::LowCardinality(inner) => { - // Stripped at schema level — just encode as inner type. - write_col_values(values, inner, out)?; + // LowCardinality wire format (ClickHouse native INSERT): + // u64 version = 1 + // u64 flags = HAS_ADDITIONAL_KEYS (bit 9) | index_type (bits 0-1) + // u64 dict_size + dict_size values (of dict_type) + // u64 num_indices + indices (1/2/4/8 bytes each) + // + // For LowCardinality(Nullable(T)), the DICTIONARY type is T (not Nullable(T)). + // ClickHouse stores nullable LC as a T-typed dict with index 0 always + // pointing to the default T value (representing NULL). + // + // For LowCardinality(T) (non-nullable), dict type is T directly. + + // Determine dict type and extract RowBinary key bytes from each value. + let (dict_type, is_nullable_inner) = + if let ColumnType::Nullable(t_inner) = inner.as_ref() { + (t_inner.as_ref(), true) + } else { + (inner.as_ref(), false) + }; + + let mut dict: Vec> = Vec::new(); + let mut seen: std::collections::HashMap, u32> = + std::collections::HashMap::new(); + + if is_nullable_inner { + // Index 0 = default T value, represents NULL. + let mut default_val = Vec::new(); + rb_write_default(&mut default_val, dict_type); + seen.insert(default_val.clone(), 0); + dict.push(default_val); + } + + let mut indices: Vec = Vec::with_capacity(values.len()); + for v in values { + // For Nullable inner: strip the Nullable RowBinary wrapper. + // [0x01] = NULL → index 0; [0x00, bytes...] = Some(v) → extract bytes. + let key: Option> = if is_nullable_inner { + if v.is_empty() || v[0] == 0x01 { + None // NULL + } else { + Some(v[1..].to_vec()) // extract T bytes + } + } else { + Some(v.clone()) + }; + + let idx = match key { + None => 0, // NULL → index 0 + Some(bytes) => { + if let Some(&i) = seen.get(&bytes) { + i + } else { + let i = dict.len() as u32; + seen.insert(bytes.clone(), i); + dict.push(bytes); + i + } + } + }; + indices.push(idx); + } + + // Choose smallest index type that fits all dict indices. + let index_type: u64 = if dict.len() <= 0x100 { + 0 // U8 + } else if dict.len() <= 0x1_0000 { + 1 // U16 + } else if (dict.len() as u64) <= 0x1_0000_0000 { + 2 // U32 + } else { + 3 // U64 + }; + + // ClickHouse requires HAS_ADDITIONAL_KEYS (bit 9 = 0x200) for client INSERT blocks. + const HAS_ADDITIONAL_KEYS: u64 = 1 << 9; + let flags = HAS_ADDITIONAL_KEYS | index_type; + + out.extend_from_slice(&1u64.to_le_bytes()); // version + out.extend_from_slice(&flags.to_le_bytes()); // flags + out.extend_from_slice(&(dict.len() as u64).to_le_bytes()); // dict_size + write_col_values(&dict, dict_type, out)?; // dict values (type = T, not Nullable(T)) + out.extend_from_slice(&(indices.len() as u64).to_le_bytes()); // num_indices + let ibytes = [1usize, 2, 4, 8][index_type as usize]; + for idx in &indices { + out.extend_from_slice(&idx.to_le_bytes()[..ibytes]); + } } ColumnType::Array(inner) => { diff --git a/tests/it/native.rs b/tests/it/native.rs index c8e118ec..edcfa4d8 100644 --- a/tests/it/native.rs +++ b/tests/it/native.rs @@ -1196,6 +1196,87 @@ async fn native_insert_strings() { assert_eq!(rows[1], StringRow { id: 2, name: "Bob".into() }); } +/// INSERT into a table with a LowCardinality(String) column. +/// +/// The encoder strips LowCardinality to its inner type and sends plain String +/// bytes; ClickHouse accepts this via implicit type conversion. +#[tokio::test] +async fn native_insert_low_cardinality() { + let client = prepare_native_database("insert_lc").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, tag LowCardinality(String)) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Serialize, Deserialize, PartialEq)] + struct Row { + id: u32, + tag: String, + } + + let mut insert = client.insert::("t"); + insert.write(&Row { id: 1, tag: "foo".into() }).await.expect("write 1 failed"); + insert.write(&Row { id: 2, tag: "bar".into() }).await.expect("write 2 failed"); + insert.write(&Row { id: 3, tag: "foo".into() }).await.expect("write 3 failed"); + insert.end().await.expect("end failed"); + + let rows = client + .query("SELECT id, tag FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 3); + assert_eq!(rows[0], Row { id: 1, tag: "foo".into() }); + assert_eq!(rows[1], Row { id: 2, tag: "bar".into() }); + assert_eq!(rows[2], Row { id: 3, tag: "foo".into() }); +} + +/// INSERT into a table with LowCardinality(Nullable(String)). +#[tokio::test] +async fn native_insert_low_cardinality_nullable() { + let client = prepare_native_database("insert_lc_nullable").await; + + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, tag LowCardinality(Nullable(String))) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .expect("CREATE failed"); + + #[derive(Debug, Row, Serialize, Deserialize, PartialEq)] + struct Row { + id: u32, + tag: Option, + } + + let mut insert = client.insert::("t"); + insert.write(&Row { id: 1, tag: Some("alpha".into()) }).await.expect("write 1"); + insert.write(&Row { id: 2, tag: None }).await.expect("write 2"); + insert.write(&Row { id: 3, tag: Some("beta".into()) }).await.expect("write 3"); + insert.end().await.expect("end failed"); + + let rows = client + .query("SELECT id, tag FROM t ORDER BY id ASC") + .fetch_all::() + .await + .expect("fetch failed"); + + assert_eq!(rows.len(), 3); + assert_eq!(rows[0], Row { id: 1, tag: Some("alpha".into()) }); + assert_eq!(rows[1], Row { id: 2, tag: None }); + assert_eq!(rows[2], Row { id: 3, tag: Some("beta".into()) }); +} + #[tokio::test] async fn native_insert_nullable() { let client = prepare_native_database("insert_nullable").await;