diff --git a/Cargo.toml b/Cargo.toml index 080a9866..110aefe8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -92,12 +92,17 @@ default = ["lz4"] test-util = ["hyper/server"] inserter = ["dep:quanta"] +async-inserter = ["inserter", "tokio/time", "tokio/sync"] +batcher = ["async-inserter"] 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 @@ -141,18 +146,20 @@ hyper-rustls = { version = "0.27.3", default-features = false, features = [ url = "2.1.1" futures-util = { version = "0.3.5", default-features = false, features = ["sink", "io"] } futures-channel = { version = "0.3.30", features = ["sink"] } -lz4_flex = { version = "0.11.3", default-features = false, features = [ +lz4_flex = { version = "0.11.6", 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"] } bstr = { version = "1.11.0", default-features = false } 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 } +serde_json = "1" [dev-dependencies] clickhouse-macros = { version = "0.3.0", path = "macros" } @@ -161,8 +168,7 @@ serde = { version = "1.0.106", features = ["derive"] } tokio = { version = "1.0.1", features = ["full", "test-util", "io-util"] } hyper = { version = "1.1", features = ["server"] } indexmap = { version = "2.10.0", features = ["serde"] } -linked-hash-map = { version = "0.5.6", features = ["serde_impl"] } -fxhash = { version = "0.2.1" } +rustc-hash = "2" serde_bytes = "0.11.4" serde_json = "1" serde_repr = "0.1.7" diff --git a/README.md b/README.md index d9ce586a..8999cb9d 100644 --- a/README.md +++ b/README.md @@ -602,6 +602,42 @@ The functionality can be enabled with the `test-util` feature. Use it **only** i See [the example](https://github.com/ClickHouse/clickhouse-rs/tree/main/examples/mock.rs). +## Native TCP Transport (HyperI Fork) + +This fork adds a native TCP protocol client (`feature = "native-transport"`) +that connects on port 9000 — the same binary protocol used by `clickhouse-client` +and the Go client. + +```rust +use clickhouse::native::NativeClient; + +let client = NativeClient::default() + .with_addr("localhost:9000") + .with_database("default") + .with_lz4(); +``` + +The same `#[derive(Row)]` structs work with both HTTP and native transports. + +### Additional feature flags (HyperI) + +| Feature | Description | +|---|---| +| `native-transport` | Native TCP client with connection pooling, SELECT + INSERT | +| `async-inserter` | `AsyncInserter` — concurrent MPSC-based inserter (HTTP + native) | +| `batcher` | `TableBatcher` — Go-style `append`/`flush`/`send` wrapper | + +### Extended documentation + +| Guide | Description | +|---|---| +| [Native Transport](docs/native-transport.md) | Connect, query, and insert over TCP | +| [Connection Pooling](docs/connection-pooling.md) | Deadpool pool, health checks, recycling | +| [Batching](docs/batching.md) | AsyncInserter, AsyncNativeInserter, TableBatcher | +| [Types](docs/types.md) | Full type coverage matrix (HTTP + native) | +| [Wire Format](docs/wire-format.md) | LowCardinality, Dynamic, Variant encoding internals | +| [Migration](docs/migration.md) | HTTP vs native trade-offs, switching guide | + ## Support Policies ### Minimum Supported Rust Version (MSRV) diff --git a/docs/batching.md b/docs/batching.md new file mode 100644 index 00000000..38147375 --- /dev/null +++ b/docs/batching.md @@ -0,0 +1,166 @@ +# Batching and Concurrent Inserters + +This crate provides a hierarchy of inserter types for different concurrency and +transport needs. All share the same three-threshold flush policy: **row count**, +**byte size**, and **time period**. + +## Inserter hierarchy + +```mermaid +graph TD + subgraph HTTP Transport + TB["TableBatcher<T>
feature = batcher
Go-style append/flush/send"] + AI["AsyncInserter<T>
feature = async-inserter
MPSC channel + background task"] + I["Inserter<T>
feature = inserter
Single-owner &mut self"] + TB -- delegates to --> AI + AI -- wraps --> I + end + + subgraph Native TCP Transport + ANI["AsyncNativeInserter<T>
feature = native-transport
MPSC channel + background task"] + NI["NativeInserter<T>
feature = native-transport
Single-owner &mut self"] + ANI -- wraps --> NI + end +``` + +## Choosing an inserter + +| Type | Transport | Concurrency | Use case | +|---|---|---|---| +| `Insert` / `NativeInsert` | HTTP / Native | Single owner | One-shot batch, manual control | +| `Inserter` / `NativeInserter` | HTTP / Native | Single owner (`&mut self`) | Long-running pipeline, single task | +| `AsyncInserter` | HTTP | Multi-task (`&self` + handles) | Fan-in from many producers | +| `AsyncNativeInserter` | Native | Multi-task (`&self` + handles) | Fan-in from many producers | +| `TableBatcher` | HTTP | Multi-task (Go-style API) | Drop-in replacement for Go `Batch` | + +## AsyncInserter / AsyncNativeInserter + +Both share identical architecture — an MPSC channel feeding a background tokio +task that owns the underlying `Inserter` or `NativeInserter`: + +```mermaid +graph TD + A["Task A
tx.send()"] --> CH{{"bounded mpsc channel
(default: 8192 slots)"}} + B["Task B
tx.send()"] --> CH + C["Task C
tx.send()"] --> CH + CH --> BG["Background Task

select! {
  cmd = rx.recv()   ← biased
  _ = interval.tick() ← periodic flush
}

serialize → buffer
check limits → flush"] + BG --> CK[("ClickHouse server")] +``` + +### Key properties + +- **`&self` on `write()` and `flush()`** — safe to call from multiple tasks + without external synchronization. +- **Backpressure** — the bounded channel blocks producers when the background + task can't keep up. Tune with `with_channel_capacity()`. +- **Error propagation** — each `write()` returns a `Result<()>` via a oneshot + channel. Serialization or network errors are surfaced to the caller. +- **`RowOwned` requirement** — rows cross a channel boundary, so `T` must own + its data (no borrowed `&str` fields). This is automatic for `#[derive(Row)]` + structs with owned fields. + +### Configuration + +```rust +use clickhouse::async_inserter::{AsyncInserter, AsyncInserterConfig}; + +let config = AsyncInserterConfig::default() + .with_max_rows(100_000) // flush at 100K rows + .with_max_bytes(10_485_760) // flush at 10 MiB + .with_max_period(Duration::from_secs(5)) // flush every 5s + .with_channel_capacity(4096); // backpressure at 4K pending + +let inserter = AsyncInserter::::new(&client, "my_table", config); +``` + +### Handles + +Handles are cheap clones of the channel sender. Use them to fan out writes +across tasks: + +```rust +let inserter = AsyncInserter::::new(&client, "my_table", config); + +for i in 0..10 { + let h = inserter.handle(); + tokio::spawn(async move { + h.write(MyRow { id: i, data: format!("task-{i}") }).await.unwrap(); + }); +} + +let stats = inserter.end().await?; // graceful shutdown +``` + +### Lifecycle + +1. **`new()`** — spawns background task immediately. +2. **`write(row)`** — sends row over channel; blocks if full. +3. **`flush()`** — forces immediate flush of buffered rows. +4. **`end()`** — sends shutdown command, waits for final flush, joins task. + +Dropping without `end()` causes the background task to flush and exit when all +senders are dropped (including handles). + +## TableBatcher + +`TableBatcher` is a thin wrapper over `AsyncInserter` with Go client +naming conventions: + +| TableBatcher | AsyncInserter | Go `Batch` | +|---|---|---| +| `append(row)` | `write(row)` | `Append(args...)` | +| `flush()` | `flush()` | `Flush()` | +| `send()` | `end()` | `Send()` | + +```rust +use clickhouse::batcher::{TableBatcher, BatchConfig}; + +let config = BatchConfig { + max_rows: 100_000, + max_bytes: 10 * 1024 * 1024, + max_period: Some(Duration::from_secs(5)), +}; + +let batcher = TableBatcher::::new(&client, "my_table", config); +batcher.append(MyRow { id: 1, data: "foo".into() }).await?; +batcher.append(MyRow { id: 2, data: "bar".into() }).await?; + +let stats = batcher.send().await?; // final flush + shutdown +``` + +## Default thresholds + +All inserter configs share these defaults, aligned with ClickHouse server +settings and the Go client: + +| Setting | Default | Rationale | +|---|---|---| +| `max_rows` | 100,000 | Upper end of recommended per-INSERT batch size | +| `max_bytes` | 10 MiB | Matches `async_insert_max_data_size` server default | +| `max_period` | 5 seconds | Balances latency vs throughput | +| `channel_capacity` | 8,192 | Backpressure threshold for concurrent inserters | + +### Why client-side batching? + +ClickHouse's server-side `async_insert` is convenient for distributed agents but +has drawbacks for high-throughput pipelines: + +- **OOM risk** — server buffers data in memory until thresholds are met +- **Delayed visibility** — data is not queryable until the server flushes +- **Error opacity** — insert errors may be lost or delayed +- **No client control** — can't tune batch size per table or per source + +Client-side batching (as implemented here) gives you immediate query visibility, +per-table tuning, and explicit error handling at the cost of managing batch +state in the client. + +### MergeTree part fragmentation + +Each INSERT creates one part per partition in MergeTree. Too many small INSERTs +cause part fragmentation: + +- `parts_to_delay_insert` (default: 150) — slows down inserts +- `parts_to_throw_insert` (default: 300) — hard "Too many parts" error + +The defaults here (100K rows, 10 MiB) are designed to produce reasonably-sized +parts. For tables with multiple partitions, adjust downward. diff --git a/docs/connection-pooling.md b/docs/connection-pooling.md new file mode 100644 index 00000000..c3e5bedf --- /dev/null +++ b/docs/connection-pooling.md @@ -0,0 +1,95 @@ +# Connection Pooling + +The native transport uses [deadpool](https://docs.rs/deadpool) for connection +pooling. Each `NativeClient` owns a pool; clones share it (the pool is +`Arc`-backed internally). + +## Configuration + +```rust +use clickhouse::native::NativeClient; + +let client = NativeClient::default() + .with_pool_size(20); // max 20 connections (default: 10) +``` + +The pool is bounded — when all connections are in use, `acquire()` waits until +one is returned. There is no idle timeout; connections persist until the client +is dropped or they fail a health check. + +## Health checks (recycle) + +When a connection is returned to the pool, deadpool calls `recycle()` which +runs `check_alive()`: + +1. **Poisoned flag** — if `discard()` was called, the connection is dropped + unconditionally. +2. **Buffered data** — if the `BufReader` has leftover bytes from an incomplete + read, the connection is dropped (stale protocol state). +3. **Non-blocking poll** — a non-blocking `poll_read` detects EOF or unexpected + server data. If either is found, the connection is dropped. + +Connections that pass all three checks are returned to the idle queue for reuse. + +## The discard pattern + +When an I/O error or incomplete protocol exchange leaves a connection in an +unrecoverable state, call `discard()` on the `PooledConnection`. This sets a +`poisoned` flag that causes `recycle()` to drop the connection rather than +returning it to the pool. + +This is used internally by: +- `NativeRowCursor::Drop` — if a cursor is dropped mid-stream (e.g. after + `fetch_one` without draining), the connection is discarded. +- Error paths in `NativeInsert` — if an INSERT fails mid-stream, the + connection is discarded rather than risk protocol desync. + +## Cursor drain + +`NativeRowCursor` provides a `drain()` method that reads and discards all +remaining server packets until `EndOfStream`. This is called automatically by +`fetch_one` and `fetch_optional` to clean up the connection before returning it +to the pool: + +```mermaid +graph LR + F["fetch_one()"] --> N["next().await?
get first row"] + N --> D["drain().await?
consume remaining packets"] + D --> R["return row
connection returned to pool"] +``` + +If `drain()` is not called (e.g. cursor dropped early), the `Drop` impl +discards the connection as a safety net. + +## Pool rebuild on config change + +Builder methods that affect connection parameters trigger `rebuild_pool()`, +which creates a new pool instance. Existing connections from the old pool are +not immediately closed — they drain naturally as they're returned and not +recycled into the new pool. + +Affected methods: +- `with_addr()` +- `with_database()` +- `with_user()` / `with_password()` +- `with_lz4()` +- `with_pool_size()` +- `with_setting()` + +## Architecture + +```mermaid +graph TD + NC["NativeClient"] --> POOL["pool: NativePool
(deadpool::managed::Pool)"] + NC --> SC["schema_cache: Arc<NativeSchemaCache>
HashMap<table, (columns, expires_at)>"] + NC --> SET["settings: Arc<Vec<(key, value)>>"] + + POOL --> MGR["NativeConnectionManager"] + MGR -->|"create()"| OPEN["NativeConnection::open()"] + MGR -->|"recycle()"| CHECK["check_alive()"] + + POOL --> IDLE["Idle queue
(bounded semaphore)"] + IDLE --> C1["conn 1"] + IDLE --> C2["conn 2"] + IDLE --> C3["..."] +``` diff --git a/docs/migration.md b/docs/migration.md new file mode 100644 index 00000000..43ea3eaa --- /dev/null +++ b/docs/migration.md @@ -0,0 +1,148 @@ +# Migration Guide + +## HTTP vs Native: choosing a transport + +Both transports use the same `Row` derive macro and serde machinery. The +primary differences are in connection model and feature coverage. + +### Feature comparison + +| Feature | HTTP (`Client`) | Native (`NativeClient`) | +|---|---|---| +| Transport | HTTP/1.1 (port 8123) | TCP (port 9000) | +| Compression | LZ4 stream-level | LZ4 block-level | +| Connection pooling | HTTP keep-alive (hyper) | Deadpool bounded pool | +| Load balancer support | Yes (stateless) | No (stateful TCP) | +| TLS | `native-tls` / `rustls-tls` features | Not yet implemented | +| SELECT | `fetch`, `fetch_one`, `fetch_all` | `fetch`, `fetch_one`, `fetch_optional`, `fetch_all` | +| Single INSERT | `Insert` | `NativeInsert` | +| Multi-batch INSERT | `Inserter` | `NativeInserter` | +| Concurrent INSERT | `AsyncInserter` | `AsyncNativeInserter` | +| Batch wrapper | `TableBatcher` | (use `AsyncNativeInserter` directly) | +| Query bind parameters | `?` placeholders, `?fields` | Not yet implemented | +| Mocking (`test-util`) | Yes | Not yet implemented | +| Validation | `RowBinaryWithNamesAndTypes` | Schema cache (TTL-based) | +| `serde::uuid`, `serde::ipv4`, etc. | Yes | Yes (via RowBinary bridge) | + +### When to use HTTP + +- Behind a load balancer or HTTP proxy +- Need TLS (not yet available on native) +- Need query bind parameters (`?` placeholders) +- Need mock testing (`test-util` feature) +- General-purpose use + +### When to use Native + +- Direct connection to ClickHouse (co-located, same network) +- High-throughput INSERT pipelines (less overhead per block) +- Need types only available on native (BFloat16, Time, Time64) +- Need connection pooling with health checks +- Want the same protocol as `clickhouse-client` and the Go client + +## Switching from HTTP to Native + +### Client creation + +```rust +// HTTP +use clickhouse::Client; +let client = Client::default() + .with_url("http://localhost:8123") + .with_database("default"); + +// Native +use clickhouse::native::NativeClient; +let client = NativeClient::default() + .with_addr("localhost:9000") + .with_database("default"); +``` + +### Row types — no changes needed + +```rust +use clickhouse::Row; +use serde::{Serialize, Deserialize}; + +#[derive(Row, Serialize, Deserialize)] +struct MyRow { + id: u64, + name: String, +} +``` + +The same `#[derive(Row)]` struct works with both transports. + +### SELECT + +```rust +// HTTP +let mut cursor = client.query("SELECT ?fields FROM t") + .fetch::()?; + +// Native — no ?fields support yet, list columns explicitly +let mut cursor = client.query("SELECT id, name FROM t") + .fetch::()?; + +// Both use the same cursor API +while let Some(row) = cursor.next().await? { + // ... +} +``` + +### INSERT + +```rust +// HTTP +let mut insert = client.insert::("t").await?; +insert.write(&row).await?; +insert.end().await?; + +// Native — note: no .await on insert creation +let mut insert = client.insert::("t"); +insert.write(&row).await?; +insert.end().await?; +``` + +The key difference: `client.insert()` is `async` on HTTP (opens connection +immediately) but synchronous on native (connection is lazy, opened on first +`write`). + +### DDL + +```rust +// HTTP +client.query("CREATE TABLE ...").execute().await?; + +// Native — identical +client.query("CREATE TABLE ...").execute().await?; +``` + +## Differences from upstream clickhouse-rs + +This fork (HyperI) adds the following on top of upstream v0.14.2: + +| Feature | Upstream | HyperI Fork | +|---|---|---| +| Native TCP transport | Planned | Implemented (`feature = "native-transport"`) | +| Connection pooling | N/A (HTTP) | Deadpool-based for native | +| LowCardinality INSERT | N/A | Full dictionary encoding | +| AsyncInserter | N/A | MPSC-based concurrent inserter | +| TableBatcher | N/A | Go-style batch wrapper | +| BFloat16, Time, Time64 | No | Yes (native only) | +| Variant/Dynamic/JSON (native) | No | Yes (SELECT, emitted as JSON strings) | +| Schema cache | N/A | TTL-based, per-client | +| LZ4 for INSERT blocks | N/A | Yes (native) | + +### Branch structure + +```mermaid +graph LR + M["main
(upstream v0.14.2)"] --> NT["hyperi/native-transport
native SELECT + INSERT"] + NT --> CP["hyperi/connection-pooling
deadpool, cursor drain, health checks"] + CP --> LC["hyperi/lc-insert
LowCardinality INSERT + LC(Nullable) fix"] + LC --> AI["hyperi/async-inserter
AsyncInserter, TableBatcher"] +``` + +Branches are designed to be merged in order. `hyperi/native-transport` is the +base PR; each subsequent branch stacks cleanly on top. diff --git a/docs/native-transport.md b/docs/native-transport.md new file mode 100644 index 00000000..1d862e39 --- /dev/null +++ b/docs/native-transport.md @@ -0,0 +1,217 @@ +# Native Transport + +The native TCP transport (`feature = "native-transport"`) connects to ClickHouse +on port 9000 using the same binary protocol as `clickhouse-client` and the Go +client (`clickhouse-go`). + +## When to use native vs HTTP + +| | Native (port 9000) | HTTP (port 8123) | +|---|---|---| +| Protocol | Binary, columnar blocks | Text/RowBinary over HTTP | +| Compression | LZ4 block-level (per data block) | LZ4 stream-level | +| Connection model | Persistent TCP, pooled | HTTP/1.1 keep-alive | +| Load balancer friendly | No (stateful TCP) | Yes (stateless HTTP) | +| Best for | High-throughput pipelines, co-located apps | General use, through proxies/LBs | + +## Creating a client + +```rust +use clickhouse::native::NativeClient; + +let client = NativeClient::default() // 127.0.0.1:9000 + .with_addr("clickhouse.internal:9000") + .with_database("analytics") + .with_user("writer") + .with_password("secret") + .with_lz4() // enable LZ4 compression + .with_pool_size(20); // max 20 connections (default: 10) +``` + +`NativeClient` is `Clone` — clones share the same connection pool. Builder +methods that change connection parameters (`with_addr`, `with_database`, etc.) +rebuild the pool so the next `acquire()` opens fresh connections. + +### Per-query settings + +```rust +let client = NativeClient::default() + .with_setting("select_sequential_consistency", "1") + .with_setting("insert_quorum", "2"); +``` + +Settings are sent in every query packet and apply to SELECT, INSERT, and DDL. + +## Querying (SELECT) + +```rust +use clickhouse::Row; +use serde::Deserialize; + +#[derive(Row, Deserialize)] +struct Event { + id: u64, + name: String, +} + +// Cursor — streaming, row-by-row +let mut cursor = client + .query("SELECT id, name FROM events WHERE id > 100") + .fetch::()?; + +while let Some(row) = cursor.next().await? { + println!("{}: {}", row.id, row.name); +} +``` + +### Convenience methods + +```rust +// Single row (or error if none) +let row = client + .query("SELECT id, name FROM events WHERE id = 42") + .fetch_one::() + .await?; + +// Optional single row +let maybe = client + .query("SELECT id, name FROM events WHERE id = 42") + .fetch_optional::() + .await?; + +// All rows into a Vec +let all = client + .query("SELECT id, name FROM events ORDER BY id LIMIT 1000") + .fetch_all::() + .await?; +``` + +### DDL and other statements + +```rust +client.query("CREATE TABLE t (n UInt32) ENGINE = Memory") + .execute() + .await?; +``` + +## Inserting + +### Single INSERT + +```rust +use clickhouse::Row; +use serde::Serialize; + +#[derive(Row, Serialize)] +struct Event { id: u64, name: String } + +let mut insert = client.insert::("events"); +insert.write(&Event { id: 1, name: "foo".into() }).await?; +insert.write(&Event { id: 2, name: "bar".into() }).await?; +insert.end().await?; // commit — dropping without end() aborts +``` + +Rows are serialised to RowBinary internally, buffered, and flushed as native +columnar blocks when the buffer exceeds ~256 KiB or when `end()` is called. + +### Multi-batch inserter (NativeInserter) + +For long-running pipelines, `NativeInserter` automatically commits when +row/byte/period thresholds are reached — producing multiple INSERT statements: + +```rust +use std::time::Duration; + +let mut ins = client.inserter::("events") + .with_max_rows(100_000) + .with_max_bytes(10 * 1024 * 1024) // 10 MiB + .with_period(Some(Duration::from_secs(5))); + +for event in events { + ins.write(&event).await?; + ins.commit().await?; // ends INSERT only if limits are reached +} +ins.end().await?; // final flush +``` + +### Concurrent inserter (AsyncNativeInserter) + +For multi-task writers, `AsyncNativeInserter` moves serialisation and I/O to +a background tokio task with an MPSC channel for backpressure: + +```rust +use clickhouse::native::async_inserter::{AsyncNativeInserter, AsyncNativeInserterConfig}; + +let config = AsyncNativeInserterConfig::default() + .with_max_rows(100_000) + .with_channel_capacity(4096); + +let inserter = AsyncNativeInserter::::new(&client, "events", config); + +// Multiple tasks can write concurrently via handles: +let handle = inserter.handle(); +tokio::spawn(async move { + handle.write(Event { id: 3, name: "baz".into() }).await.unwrap(); +}); + +// Graceful shutdown +let stats = inserter.end().await?; +``` + +See [Batching](batching.md) for the full concurrent inserter architecture. + +## LZ4 Compression + +Enable with `.with_lz4()` on the client. This compresses: +- INSERT data blocks (both payload blocks and empty terminator blocks) +- Query result data blocks from the server + +**Important**: ClickHouse sends `Log` and `ProfileEvents` blocks **uncompressed** +even when compression is negotiated. The reader handles this automatically. + +## Schema Cache + +The client maintains a TTL-based schema cache (default: 300 seconds) that is +populated automatically during INSERT operations when the server returns column +headers. You can also manage it explicitly: + +```rust +// Pre-fetch schema from system.columns +let schema = client.fetch_schema("events").await?; +// Returns Vec<(column_name, column_type)> + +// Check cache +if let Some(cached) = client.cached_schema("events") { + println!("cached {} columns", cached.len()); +} + +// Invalidate +client.clear_cached_schema("events"); +client.clear_all_cached_schemas(); +``` + +## Connection lifecycle + +```mermaid +graph TD + NC["NativeClient::default()"] --> Pool["Pool
(deadpool, max_size=10)"] + + Pool --> ACQ["acquire()"] + ACQ --> CREATE["Create or reuse connection"] + CREATE --> TCP["TCP connect"] + TCP --> HELLO["Hello handshake"] + HELLO --> PC["Return PooledConnection"] + + Pool --> REC["recycle()"] + REC --> CA["check_alive()"] + CA -->|poisoned?| DROP1["Drop connection"] + CA -->|buffered data?| DROP2["Drop connection"] + CA -->|EOF?| DROP3["Drop connection"] + CA -->|ok| IDLE["Return to idle queue"] + + Pool --> DISC["discard()"] + DISC --> POISON["Set poisoned flag"] + POISON -.-> DROP1 +``` + +See [Connection Pooling](connection-pooling.md) for details. diff --git a/docs/types.md b/docs/types.md new file mode 100644 index 00000000..80c9ff26 --- /dev/null +++ b/docs/types.md @@ -0,0 +1,132 @@ +# Type Coverage + +This document lists all ClickHouse column types and their support status across +both the HTTP and native TCP transports. + +## Scalars + +| ClickHouse Type | Rust Type | Wire Size | HTTP | Native SELECT | Native INSERT | +|---|---|---|---|---|---| +| UInt8 | `u8` | 1 | Yes | Yes | Yes | +| UInt16 | `u16` | 2 | Yes | Yes | Yes | +| UInt32 | `u32` | 4 | Yes | Yes | Yes | +| UInt64 | `u64` | 8 | Yes | Yes | Yes | +| UInt128 | `u128` | 16 | Yes | Yes | Yes | +| UInt256 | `clickhouse::types::UInt256` (`[u8; 32]`) | 32 | Yes | Yes | Yes | +| Int8 | `i8` | 1 | Yes | Yes | Yes | +| Int16 | `i16` | 2 | Yes | Yes | Yes | +| Int32 | `i32` | 4 | Yes | Yes | Yes | +| Int64 | `i64` | 8 | Yes | Yes | Yes | +| Int128 | `i128` | 16 | Yes | Yes | Yes | +| Int256 | `clickhouse::types::Int256` (`[u8; 32]`) | 32 | Yes | Yes | Yes | +| Float32 | `f32` | 4 | Yes | Yes | Yes | +| Float64 | `f64` | 8 | Yes | Yes | Yes | +| BFloat16 | `u16` (raw bits) | 2 | No | Yes | Yes | +| Boolean | `bool` | 1 | Yes | Yes | Yes | +| Decimal32(S) | `i32` | 4 | Yes | Yes | Yes | +| Decimal64(S) | `i64` | 8 | Yes | Yes | Yes | +| Decimal128(S) | `i128` | 16 | Yes | Yes | Yes | +| Decimal256(S) | `[u8; 32]` | 32 | Yes | Yes | Yes | +| Date | `u16` (days since epoch) | 2 | Yes | Yes | Yes | +| Date32 | `i32` (days since epoch) | 4 | Yes | Yes | Yes | +| DateTime | `u32` (seconds since epoch) | 4 | Yes | Yes | Yes | +| DateTime64(P) | `i64` (scaled since epoch) | 8 | Yes | Yes | Yes | +| Time | `i32` (seconds since midnight) | 4 | No | Yes | Yes | +| Time64(P) | `i64` (scaled since midnight) | 8 | No | Yes | Yes | +| UUID | `uuid::Uuid` (with `serde::uuid`) | 16 | Yes | Yes | Yes | +| IPv4 | `Ipv4Addr` (with `serde::ipv4`) | 4 | Yes | Yes | Yes | +| IPv6 | `Ipv6Addr` | 16 | Yes | Yes | Yes | +| Enum8 | `#[repr(i8)]` enum | 1 | Yes | Yes | Yes | +| Enum16 | `#[repr(i16)]` enum | 2 | Yes | Yes | Yes | +| Point | `(f64, f64)` | 16 | Yes | Yes | Yes | + +### Notes on scalars + +- **Decimal**: scale is part of the type string but ignored on the wire. Map to + the corresponding integer type or use [fixnum](https://docs.rs/fixnum). +- **DateTime/DateTime64**: timezone and precision are in the type string but do + not affect wire encoding. Use `serde::time` or `serde::chrono` helpers for + ergonomic date/time types. +- **BFloat16, Time, Time64**: only available on the native transport. These are + newer ClickHouse types not yet supported by the HTTP RowBinary path. + +## String types + +| ClickHouse Type | Rust Type | HTTP | Native SELECT | Native INSERT | +|---|---|---|---|---| +| String | `String`, `&str`, `Vec`, `&[u8]` | Yes | Yes | Yes | +| FixedString(N) | `[u8; N]` | Yes | Yes | Yes | + +### FixedString encoding + +- **HTTP (RowBinary)**: `varuint(N)` + N bytes +- **Native wire**: N raw bytes (no length prefix) +- **Native INSERT**: strips the varuint prefix from RowBinary before sending + +## Composite types + +| ClickHouse Type | Rust Type | HTTP | Native SELECT | Native INSERT | +|---|---|---|---|---| +| Nullable(T) | `Option` | Yes | Yes | Yes | +| LowCardinality(T) | same as T | Yes | Yes | Yes | +| Array(T) | `Vec`, `&[T]` | Yes | Yes | Yes | +| Tuple(T1, ..., Tn) | `(T1, ..., Tn)` | Yes | Yes | Yes | +| Map(K, V) | `HashMap`, `Vec<(K, V)>` | Yes | Yes | Yes | +| Nested(col1 T1, ...) | multiple `Vec` with `#[serde(rename)]` | Yes | Yes | No | +| SimpleAggregateFunction(f, T) | same as T | Yes | Yes | No | + +### LowCardinality + +LowCardinality wraps String, FixedString, or Nullable variants of these. On +the wire it uses a dictionary encoding: + +- **SELECT**: dictionary + index array decoded transparently +- **INSERT**: values are dictionary-encoded automatically, including + `LowCardinality(Nullable(T))` where index 0 is the null sentinel + +See [Wire Format](wire-format.md) for encoding details. + +## Modern types (ClickHouse 24.x+) + +| ClickHouse Type | Rust Type | HTTP | Native SELECT | Native INSERT | +|---|---|---|---|---| +| Variant(T1, ..., Tn) | `enum` or `String` (JSON) | Yes | Yes | No | +| Dynamic | `String` (JSON) | No | Yes | No | +| JSON (new, 24.10+) | `String` (JSON) | Yes | Yes | No | +| Object('json') (legacy) | `String` | Yes | Yes | No | + +### Variant output + +On the native transport, Variant cells are read from their per-discriminator +sub-columns and emitted as JSON-encoded strings. On HTTP, Variant maps to a +Rust enum with variants in alphabetical order matching the ClickHouse type +definition. + +### Dynamic / JSON output + +The new JSON type (ClickHouse 24.10+) uses the Dynamic wire format internally. +On the native transport, Dynamic/JSON cells are decoded from discriminator + +per-type sub-columns and emitted as JSON strings. Users map these to `String` +or deserialize with `serde_json::Value`. + +## Geo types + +| ClickHouse Type | Rust Type | HTTP | Native SELECT | Native INSERT | +|---|---|---|---|---| +| Point | `(f64, f64)` | Yes | Yes | Yes | +| Ring | `Vec<(f64, f64)>` | Yes | No | No | +| Polygon | `Vec>` | Yes | No | No | +| MultiPolygon | `Vec>>` | Yes | No | No | +| LineString | `Vec<(f64, f64)>` | Yes | No | No | +| MultiLineString | `Vec>` | Yes | No | No | + +Ring, Polygon, MultiPolygon, LineString, and MultiLineString are composed of +Arrays of Points. They work on HTTP via the standard Array machinery but are +not yet implemented as named types on the native transport. + +## Not yet supported + +| ClickHouse Type | Notes | +|---|---| +| AggregateFunction(...) | Opaque binary blob; complex intermediate state | +| Sparse serialization | Per-column `custom_ser = 1` flag; returns error on native transport | diff --git a/docs/wire-format.md b/docs/wire-format.md new file mode 100644 index 00000000..61a60d30 --- /dev/null +++ b/docs/wire-format.md @@ -0,0 +1,211 @@ +# Wire Format Reference + +Internal documentation for the native TCP protocol wire encoding of complex +column types. This is intended for contributors and anyone debugging protocol +issues. + +For the authoritative source, see the ClickHouse C++ code in +`src/DataTypes/Serializations/`. + +## Native protocol overview + +Data is exchanged in **blocks** — each block contains N rows across all columns. +Within a block, data is columnar: all N values for column 1, then all N values +for column 2, etc. + +```text +Block header: + varuint block_info.field1 (0) + u8 block_info.is_overflows (0) + varuint block_info.field2 (0) + i32 block_info.bucket_num (-1) + varuint 0 (end of block info) + varuint num_columns + varuint num_rows + +Per column (repeated num_columns times): + String column_name + String column_type + [u8] custom_serialization flag (if revision >= 54454) + [bytes] column data (num_rows values) +``` + +### Custom serialization flag + +ClickHouse servers with revision >= 54454 (`DBMS_MIN_PROTOCOL_VERSION_WITH_CUSTOM_SERIALIZATION`) +send a `u8` flag after each column's type string: +- `0x00` = normal serialization +- `0x01` = sparse serialization (offsets + values) + +The INSERT encoder must also write this flag. Omitting it causes the server to +misinterpret the first data byte as the flag, leading to hangs or corrupt data. + +## LowCardinality + +Wire format (verified working for both SELECT and INSERT): + +```text +u64 version = 1 ← serialization version prefix + +Per-block: + u64 flags + bits 0-1: index type (0=U8, 1=U16, 2=U32, 3=U64) + bit 8: NEED_GLOBAL_DICTIONARY (0x100) + bit 9: HAS_ADDITIONAL_KEYS (0x200) + + [if NEED_GLOBAL_DICTIONARY] + u64 global_dict_size + global_dict_size × T values + + [if HAS_ADDITIONAL_KEYS] + u64 additional_keys_size + additional_keys_size × T values + + [if neither flag] + u64 dict_size + dict_size × T values + + u64 num_indices (= num_rows) + num_indices × index_type bytes +``` + +### Index space + +Indices reference the combined dictionary: +- Indices `0..additional_keys_size` → additional keys +- Indices `additional_keys_size..` → global dictionary + +In practice, ClickHouse almost always uses `HAS_ADDITIONAL_KEYS` without a +global dictionary, so the additional keys *are* the entire dictionary. + +### LowCardinality(Nullable(T)) + +When the inner type is `Nullable(T)`: +- The **dictionary type is `T`** (not `Nullable(T)`) — no null flags in the dict +- Index 0 is a **null sentinel** — the value at dict position 0 is the default + value of T (e.g. empty string), but any row with index 0 should be treated as + NULL +- For SELECT: index 0 → emit RowBinary null (`0x01`); other indices → emit + `0x00` (not-null) + T value bytes +- For INSERT: null inputs → index 0; `Some(v)` → extract T bytes (strip null + flag) and dictionary-encode normally + +### INSERT encoding + +The INSERT encoder builds the dictionary by collecting unique values: + +1. Collect all unique T-values (for Nullable: strip the `0x00`/`0x01` null flag) +2. Assign index 0 as the null sentinel (if Nullable) +3. Choose index type based on dictionary size (U8 if ≤256, U16 if ≤65536, etc.) +4. Write: version=1, flags with `HAS_ADDITIONAL_KEYS`, dict values, indices + +## Array(T) + +```text +u64[num_rows] cumulative offsets (last offset = total elements) +T[total] element values as a sub-column +``` + +The offsets are cumulative — the i-th array contains elements from +`offsets[i-1]` (or 0 for i=0) to `offsets[i]`. + +Arrays of arrays (e.g. `Array(Array(String))`) nest recursively: the outer +offsets point into the inner offset array. + +## Map(K, V) + +Maps are encoded as arrays of key-value pairs: + +```text +u64[num_rows] cumulative offsets (same as Array) +K[total] key sub-column +V[total] value sub-column +``` + +## Tuple(T1, ..., Tn) + +Each element is a separate sub-column in definition order: + +```text +T1[num_rows] first element values +T2[num_rows] second element values +... +Tn[num_rows] nth element values +``` + +## Nullable(T) + +```text +u8[num_rows] null flags (1 = null, 0 = not null) +T[num_rows] values (null slots contain default/zero T values) +``` + +All N values are always present on the wire — null rows have zero-initialized +values that are ignored by the reader. + +## Variant(T1, ..., Tn) + +```text +u64 version = 0 ← different from LowCardinality! +u8[num_rows] discriminators (255 = NULL) +T1[count1] values for discriminator 0 +T2[count2] values for discriminator 1 +... +Tn[countn] values for discriminator n-1 +``` + +Types are always in the order specified in the Variant definition (which +ClickHouse sorts alphabetically). The count for each type is derived by +counting its discriminator value in the discriminator array. + +## Dynamic + +```text +u64 version = 1 +varuint num_prefix_types (usually 0) +String[] prefix type names (if num_prefix_types > 0) + +Per-block: + varuint num_types + String[] type_names + u8[num_rows] discriminators (255 = NULL) + for each type in order: + T[count] values for that discriminator +``` + +The new JSON type (ClickHouse 24.10+) uses this same wire format. Legacy +`Object('json')` is a plain String on the wire. + +## FixedString(N) + +- **Native wire**: N raw bytes, no length prefix +- **RowBinary**: `varuint(N)` + N bytes (length-prefixed) + +The native transport reader emits FixedString as RowBinary (prepends varuint +length) for compatibility with the serde deserializer. The INSERT encoder +strips the varuint prefix before sending. + +## String + +```text +varuint(len) length prefix +u8[len] UTF-8 bytes +``` + +Identical encoding in both native wire format and RowBinary. + +## Compression (LZ4) + +When compression is enabled, data blocks are wrapped: + +```text +u8 checksum[16] (CityHash128 of the rest) +u8 method (0x82 = LZ4) +u32 compressed_size (including this 9-byte header) +u32 uncompressed_size +u8[] LZ4-compressed payload +``` + +**Important**: `Log` and `ProfileEvents` packets from the server are always +sent uncompressed, even when compression is negotiated. The reader detects +these packet types and reads them without decompression. diff --git a/rustfmt.toml b/rustfmt.toml index ef4162c2..33a75456 100644 --- a/rustfmt.toml +++ b/rustfmt.toml @@ -1,2 +1,2 @@ -edition = "2021" +edition = "2024" merge_derives = false diff --git a/src/async_inserter.rs b/src/async_inserter.rs new file mode 100644 index 00000000..051fbb6c --- /dev/null +++ b/src/async_inserter.rs @@ -0,0 +1,330 @@ +//! Concurrent, auto-flushing inserter with background task (HTTP transport). +//! +//! [`AsyncInserter`] moves serialisation, limit-checking, and periodic +//! flushing into a dedicated tokio task that communicates with callers via an +//! MPSC channel. Multiple tasks can call [`write`][AsyncInserter::write] +//! concurrently — the bounded channel provides natural backpressure. +//! +//! Ported from the HyperI DFE Loader project (`dfe-loader/src/buffer/`) +//! where a similar architecture (per-table buffer + background flush task + +//! orchestrator select! loop) was used to feed ClickHouse from Kafka at +//! sustained throughput. The key improvement here is that the batching +//! policy is embedded in the library rather than requiring each consumer to +//! re-implement the orchestrator pattern. +//! +//! # Architecture +//! +//! ```text +//! ┌─ Task A ──┐ ┌─ Task B ──┐ ┌─ Task C ──┐ +//! │ tx.send() │ │ tx.send() │ │ tx.send() │ +//! └─────┬─────┘ └─────┬─────┘ └─────┬─────┘ +//! └───────────────┴───────────────┘ +//! │ +//! bounded mpsc channel +//! │ +//! ┌───────────▼────────────┐ +//! │ Background Task │ +//! │ │ +//! │ select! { │ +//! │ cmd = rx.recv() │ +//! │ _ = interval.tick() │ +//! │ } │ +//! │ │ +//! │ serialize → buffer │ +//! │ check limits → flush │ +//! └──────────┬─────────────┘ +//! │ HTTP +//! ▼ +//! ClickHouse :8123 +//! ``` +//! +//! The Go ClickHouse client (`clickhouse-go`) keeps batch inserts purely +//! caller-driven (no background goroutines). This design goes further — +//! providing the concurrent, auto-flushing inserter that Go users typically +//! build themselves with goroutines and channels. + +use tokio::sync::{mpsc, oneshot}; +use tokio::time::Duration; + +use crate::{ + Client, + error::Result, + inserter::{Inserter, Quantities}, + row::{RowOwned, RowWrite}, +}; + +const DEFAULT_CHANNEL_CAPACITY: usize = 8192; + +// --------------------------------------------------------------------------- +// Commands sent over the MPSC channel +// --------------------------------------------------------------------------- + +enum Command { + Write(T, oneshot::Sender>), + Flush(oneshot::Sender>), + End(oneshot::Sender>), +} + +// --------------------------------------------------------------------------- +// Configuration +// --------------------------------------------------------------------------- + +/// Configuration for [`AsyncInserter`]. +/// +/// Defaults align with ClickHouse's recommended batch sizes and the +/// `async_insert_max_data_size` server setting. +#[derive(Debug, Clone)] +pub struct AsyncInserterConfig { + /// 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. + pub max_period: Option, + /// Bounded channel capacity. Default: `8192`. + /// + /// Controls backpressure: producers block when the channel is full. + pub channel_capacity: usize, +} + +impl Default for AsyncInserterConfig { + fn default() -> Self { + Self { + max_rows: 100_000, + max_bytes: 10 * 1024 * 1024, + max_period: Some(Duration::from_secs(5)), + channel_capacity: DEFAULT_CHANNEL_CAPACITY, + } + } +} + +impl AsyncInserterConfig { + /// 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. + pub fn without_period(mut self) -> Self { + self.max_period = None; + self + } + + /// Override the bounded channel capacity. + pub fn with_channel_capacity(mut self, cap: usize) -> Self { + self.channel_capacity = cap; + self + } +} + +// --------------------------------------------------------------------------- +// AsyncInserter — HTTP transport +// --------------------------------------------------------------------------- + +/// Concurrent, auto-flushing inserter for a single ClickHouse table (HTTP). +/// +/// Unlike [`Inserter`][crate::inserter::Inserter], this type: +/// +/// - Accepts `&self` on [`write`][Self::write] and [`flush`][Self::flush], +/// so it can be shared across tasks via `Arc` (or via cheap +/// [`handle()`][Self::handle] clones). +/// - Moves serialisation and network I/O to a background tokio task. +/// - Flushes automatically when row/byte/period limits are reached. +/// - Provides backpressure via a bounded MPSC channel. +/// +/// # Note: `RowOwned` requirement +/// +/// Because rows are sent over an MPSC channel, `T` must be [`RowOwned`] +/// (i.e. `T::Value<'a> = T` for all lifetimes). This is automatically +/// satisfied by any `#[derive(Row)]` struct that owns its fields. +/// +/// Ported from HyperI DFE Loader's per-table buffer + orchestrator pattern. +pub struct AsyncInserter { + tx: mpsc::Sender>, + handle: tokio::task::JoinHandle<()>, +} + +/// A cheap, clonable handle for writing rows to an [`AsyncInserter`]. +/// +/// Obtained via [`AsyncInserter::handle`]. Multiple handles can write +/// concurrently. The background task exits when all handles and the +/// original `AsyncInserter` are dropped. +#[derive(Clone)] +pub struct AsyncInserterHandle { + tx: mpsc::Sender>, +} + +fn channel_closed_err() -> crate::error::Error { + crate::error::Error::Custom("AsyncInserter background task gone".into()) +} + +impl AsyncInserter +where + T: RowOwned + RowWrite + Send + Sync + 'static, +{ + /// Create a new `AsyncInserter` for `table` using `config` thresholds. + /// + /// Spawns a background tokio task immediately. + pub fn new(client: &Client, table: &str, config: AsyncInserterConfig) -> Self { + let (tx, rx) = mpsc::channel(config.channel_capacity); + + let inserter = client + .inserter::(table) + .with_max_rows(config.max_rows) + .with_max_bytes(config.max_bytes) + .with_period(config.max_period); + + let period = config.max_period; + let handle = tokio::spawn(background_task(inserter, rx, period)); + + Self { tx, handle } + } + + /// Obtain a cheap, clonable write handle. + pub fn handle(&self) -> AsyncInserterHandle { + AsyncInserterHandle { + tx: self.tx.clone(), + } + } + + /// Serialize and buffer a row. + /// + /// Blocks (asynchronously) if the channel is full (backpressure). + /// Returns once the row has been serialised into the internal buffer. + pub async fn write(&self, row: T) -> Result<()> { + let (resp_tx, resp_rx) = oneshot::channel(); + self.tx + .send(Command::Write(row, resp_tx)) + .await + .map_err(|_| channel_closed_err())?; + resp_rx.await.map_err(|_| channel_closed_err())? + } + + /// Force-flush all buffered rows to ClickHouse. + pub async fn flush(&self) -> Result { + let (resp_tx, resp_rx) = oneshot::channel(); + self.tx + .send(Command::Flush(resp_tx)) + .await + .map_err(|_| channel_closed_err())?; + resp_rx.await.map_err(|_| channel_closed_err())? + } + + /// Graceful shutdown: flush remaining rows, end the current INSERT, + /// and stop the background task. + /// + /// Consumes `self`. All cloned handles become inert after this call. + pub async fn end(self) -> Result { + let (resp_tx, resp_rx) = oneshot::channel(); + if self.tx.send(Command::End(resp_tx)).await.is_err() { + return Ok(Quantities::ZERO); + } + drop(self.tx); + let result = resp_rx.await.map_err(|_| channel_closed_err())?; + let _ = self.handle.await; + result + } +} + +impl AsyncInserterHandle +where + T: RowOwned + RowWrite + Send + Sync + 'static, +{ + /// Serialize and buffer a row (same as [`AsyncInserter::write`]). + pub async fn write(&self, row: T) -> Result<()> { + let (resp_tx, resp_rx) = oneshot::channel(); + self.tx + .send(Command::Write(row, resp_tx)) + .await + .map_err(|_| channel_closed_err())?; + resp_rx.await.map_err(|_| channel_closed_err())? + } + + /// Force-flush all buffered rows to ClickHouse. + pub async fn flush(&self) -> Result { + let (resp_tx, resp_rx) = oneshot::channel(); + self.tx + .send(Command::Flush(resp_tx)) + .await + .map_err(|_| channel_closed_err())?; + resp_rx.await.map_err(|_| channel_closed_err())? + } +} + +// --------------------------------------------------------------------------- +// Background task +// --------------------------------------------------------------------------- + +async fn background_task( + mut inserter: Inserter, + mut rx: mpsc::Receiver>, + period: Option, +) where + T: RowOwned + RowWrite + Send + Sync + 'static, +{ + let mut interval = period.map(|p| { + let mut iv = tokio::time::interval(p); + iv.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + iv + }); + + // Skip the immediate first tick. + if let Some(ref mut iv) = interval { + iv.tick().await; + } + + loop { + let tick_fut = async { + match interval { + Some(ref mut iv) => iv.tick().await, + None => std::future::pending().await, + } + }; + + tokio::select! { + biased; + + cmd = rx.recv() => { + match cmd { + Some(Command::Write(row, resp)) => { + let result = inserter.write(&row).await; + if result.is_ok() { + let _ = inserter.commit().await; + } + let _ = resp.send(result); + } + Some(Command::Flush(resp)) => { + let _ = resp.send(inserter.force_commit().await); + } + Some(Command::End(resp)) => { + let _ = resp.send(inserter.end().await); + return; + } + None => { + let _ = inserter.end().await; + return; + } + } + } + + _ = tick_fut => { + let _ = inserter.commit().await; + } + } + } +} diff --git a/src/batcher.rs b/src/batcher.rs new file mode 100644 index 00000000..745933ab --- /dev/null +++ b/src/batcher.rs @@ -0,0 +1,149 @@ +//! Per-table batch inserter with automatic flushing. +//! +//! [`TableBatcher`] is a thin convenience wrapper over +//! [`AsyncInserter`][crate::async_inserter::AsyncInserter] that provides +//! ClickHouse Go client–style naming ([`append`][TableBatcher::append] / +//! [`flush`][TableBatcher::flush] / [`send`][TableBatcher::send]) and +//! sensible defaults. +//! +//! # Architecture +//! +//! ```text +//! TableBatcher (thin wrapper over AsyncInserter) +//! ┌───────────────────────────────────────────┐ +//! │ append(row) ──→ AsyncInserter.write(row) │ +//! │ flush() ──→ AsyncInserter.flush() │ +//! │ send() ──→ AsyncInserter.end() │ +//! └──────────────────────┬────────────────────┘ +//! │ mpsc channel +//! ▼ +//! Background Task (select!) +//! │ +//! Inserter +//! │ HTTP +//! ▼ +//! ClickHouse :8123 +//! ``` +//! +//! 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) +//! +//! Ported from the HyperI DFE Loader project (`dfe-loader/src/buffer/`). + +use tokio::time::Duration; + +use crate::{ + Client, + async_inserter::{AsyncInserter, AsyncInserterConfig}, + error::Result, + inserter::Quantities, + row::{RowOwned, 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. + +/// Thread-safe, auto-flushing batch inserter for a single ClickHouse table. +/// +/// Thin wrapper over [`AsyncInserter`][crate::async_inserter::AsyncInserter] +/// with Go client–style naming. +/// +/// Unlike `Inserter`, this type accepts `&self` on [`append`][Self::append] +/// and [`flush`][Self::flush], so it can be shared across tasks via [`std::sync::Arc`]. +/// +/// For multi-table writes create one `TableBatcher` per table. +pub struct TableBatcher { + inner: AsyncInserter, +} + +impl TableBatcher +where + T: RowOwned + RowWrite + Send + Sync + 'static, +{ + /// Create a new `TableBatcher` for `table` using `config` thresholds. + pub fn new(client: &Client, table: &str, config: BatchConfig) -> Self { + let ai_config = AsyncInserterConfig { + max_rows: config.max_rows, + max_bytes: config.max_bytes, + max_period: config.max_period, + ..AsyncInserterConfig::default() + }; + + Self { + inner: AsyncInserter::new(client, table, ai_config), + } + } + + /// Add `row` to the buffer. Flushes automatically if a threshold is crossed. + pub async fn append(&self, row: T) -> Result<()> { + self.inner.write(row).await + } + + /// Force-flush all pending rows to ClickHouse immediately. + /// + /// Returns the [`Quantities`] sent. + pub async fn flush(&self) -> Result { + self.inner.flush().await + } + + /// Flush remaining rows and shut down the batcher. + /// + /// Consumes `self`. + pub async fn send(self) -> Result { + self.inner.end().await + } +} diff --git a/src/cursors/row.rs b/src/cursors/row.rs index ea622b5c..df5ed35a 100644 --- a/src/cursors/row.rs +++ b/src/cursors/row.rs @@ -12,7 +12,6 @@ use crate::{ use bytes::Buf; use clickhouse_types::error::TypesError; use clickhouse_types::parse_rbwnat_columns_header; -use polonius_the_crab::prelude::*; use std::marker::PhantomData; use std::pin::Pin; use std::task::{Context, Poll, ready}; @@ -100,6 +99,38 @@ impl RowCursor { Next::new(self).await } + // ----------------------------------------------------------------------- + // Why the unsafe reborrow? + // + // We hate unsafe. Genuinely. But NLL (the current borrow checker) can't + // see that `bytes` is dead in the NotEnoughData branch of this loop. + // The returned value borrows from `bytes`, so NLL extends that borrow + // to the function's return lifetime — blocking the `bytes.extend()` + // that only runs when no value exists. Classic Polonius limitation: + // https://github.com/rust-lang/rust/issues/51132 + // + // This used to be the `polonius-the-crab` crate, which wraps the exact + // same raw-pointer reborrow behind a macro. We dropped it because + // polonius-the-crab has so many abandonment issues it needs therapy: + // - `paste` transitive dep: RUSTSEC-2024-0436 (unmaintained) + // - `polonius-the-crab` itself: no meaningful commits in 12+ months + // - `higher-kinded-types`, `macro_rules_attribute`: same story + // Four stagnant crates, two RustSec advisories, all for a macro that + // expands to one line of unsafe. Two lines of unsafe instead of four + // crates is a good return. + // + // We properly tried to avoid this: + // - TryRow enum (borrow still escapes via return type — same error) + // - async-only next() + poll_next_owned for Stream (same NLL issue) + // - interior mutability in BytesExt via UnsafeCell (3x the diff, + // same amount of actual unsafe, just hidden — not actually better) + // - double deserialisation / probe-then-extract (~2x deser cost on + // the happy path — non-starter for a perf-sensitive cursor) + // None compiled without unsafe somewhere, or had unacceptable costs. + // + // When Polonius lands in stable rustc, rip this out. We'll buy it a beer. + // ----------------------------------------------------------------------- + #[inline] fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll>>> where @@ -110,27 +141,34 @@ impl RowCursor { debug_assert!(self.row_metadata.is_some()); } - let mut bytes = &mut self.bytes; + let bytes = &mut self.bytes; loop { - polonius!(|bytes| -> Poll>>> { - if bytes.remaining() > 0 { - let mut slice = bytes.slice(); - let result = rowbinary::deserialize_row::>( - &mut slice, - self.row_metadata.as_ref(), - ); - - match result { - Ok(value) => { - bytes.set_remaining(slice.len()); - polonius_return!(Poll::Ready(Ok(Some(value)))) - } - Err(Error::NotEnoughData) => {} - Err(err) => polonius_return!(Poll::Ready(Err(err))), + // SAFETY: we create a second &mut to `bytes` via raw pointer so the + // borrow checker releases the original. This is sound because: + // - On Ok: we return immediately — only one &mut is live. + // - On NotEnoughData: the deserialized value doesn't exist, the + // reborrow is dead, and we fall through to extend(). + // - On Err: we return immediately. + // Polonius would prove this automatically. NLL can't (yet). + let reborrowed = unsafe { &mut *(bytes as *mut BytesExt) }; + + if reborrowed.remaining() > 0 { + let mut slice = reborrowed.slice(); + let result = rowbinary::deserialize_row::>( + &mut slice, + self.row_metadata.as_ref(), + ); + + match result { + Ok(value) => { + reborrowed.set_remaining(slice.len()); + return Poll::Ready(Ok(Some(value))); } + Err(Error::NotEnoughData) => {} + Err(err) => return Poll::Ready(Err(err)), } - }); + } match ready!(self.raw.poll_next(cx))? { Some(chunk) => bytes.extend(chunk), @@ -196,18 +234,22 @@ where #[inline] fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - // Temporarily take the cursor out in order for `cursor.poll_next` to return a value with - // the correct lifetime `'a` rather than the unnamed lifetime of `&mut self`. - let mut cursor = self.cursor.take().expect("Future polled after completion"); - - polonius!(|cursor| -> Poll>>> { - match cursor.poll_next(cx) { - Poll::Ready(value) => polonius_return!(Poll::Ready(value)), - Poll::Pending => {} + // Take cursor out so poll_next's return value gets lifetime 'a + // (not the anonymous reborrow lifetime of &mut self). + let cursor = self.cursor.take().expect("Future polled after completion"); + + // SAFETY: same pattern as poll_next above — we create a second &mut + // via raw pointer. On Ready the reborrow escapes via the return value + // and cursor is consumed. On Pending the reborrow is dead and we put + // cursor back. Sound for the same reasons; Polonius would accept this. + let reborrowed = unsafe { &mut *(cursor as *mut RowCursor) }; + + match reborrowed.poll_next(cx) { + Poll::Ready(value) => Poll::Ready(value), + Poll::Pending => { + self.cursor = Some(cursor); + Poll::Pending } - }); - - self.cursor = Some(cursor); - Poll::Pending + } } } diff --git a/src/dynamic/batcher.rs b/src/dynamic/batcher.rs new file mode 100644 index 00000000..0277db5c --- /dev/null +++ b/src/dynamic/batcher.rs @@ -0,0 +1,357 @@ +//! Async auto-flushing dynamic inserter with background task. +//! +//! `DynamicBatcher` is the async, multi-producer variant of `DynamicInsert`. +//! It moves schema fetch, RowBinary encoding, and periodic flushing into a +//! dedicated tokio task that communicates with callers via a bounded MPSC +//! channel. Multiple tasks can call `write_map()` concurrently — the bounded +//! channel provides natural backpressure. +//! +//! # Schema Recovery +//! +//! On schema mismatch errors from ClickHouse, the background task: +//! 1. Invalidates the cached schema +//! 2. Re-fetches from `system.columns` +//! 3. Retries the current batch with the new schema +//! 4. Resumes normal operation +//! +//! One retry attempt per mismatch — prevents infinite loops on genuine +//! data errors. +//! +//! # Architecture +//! +//! ```text +//! ┌─ Task A ──┐ ┌─ Task B ──┐ ┌─ Task C ──┐ +//! │ write_map()│ │ write_map()│ │ write_map()│ +//! └─────┬─────┘ └─────┬─────┘ └─────┬─────┘ +//! └───────────────┴───────────────┘ +//! │ +//! bounded mpsc channel +//! │ +//! ┌───────────▼────────────┐ +//! │ Background Task │ +//! │ select! { │ +//! │ cmd = rx.recv() │ +//! │ _ = interval.tick() │ +//! │ } │ +//! │ encode → RowBinary │ +//! │ buffer → flush │ +//! └──────────┬─────────────┘ +//! │ HTTP RowBinary +//! ▼ +//! ClickHouse :8123 +//! ``` + +use std::sync::Arc; + +use serde_json::{Map, Value}; +use tokio::sync::{mpsc, oneshot}; +use tokio::time::Duration; + +use crate::Client; + +use super::error::DynamicError; +use super::schema::DynamicSchemaCache; + +const DEFAULT_CHANNEL_CAPACITY: usize = 8192; + +/// Configuration for [`DynamicBatcher`]. +#[derive(Debug, Clone)] +pub struct DynamicBatchConfig { + /// Flush when this many rows have been buffered. Default: `10_000`. + pub max_rows: u64, + /// Flush after this period regardless of row count. Default: `5s`. + pub max_period: Duration, + /// Bounded channel capacity. Default: `8192`. + pub channel_capacity: usize, +} + +impl Default for DynamicBatchConfig { + fn default() -> Self { + Self { + max_rows: 10_000, + max_period: Duration::from_secs(5), + channel_capacity: DEFAULT_CHANNEL_CAPACITY, + } + } +} + +// --------------------------------------------------------------------------- +// Commands over the MPSC channel +// --------------------------------------------------------------------------- + +enum Cmd { + Write(Map, oneshot::Sender>), + Flush(oneshot::Sender>), + End(oneshot::Sender>), +} + +// --------------------------------------------------------------------------- +// DynamicBatcher +// --------------------------------------------------------------------------- + +/// Async auto-flushing dynamic inserter for a single ClickHouse table. +/// +/// Push `Map`, the batcher encodes to RowBinary in a background +/// task and flushes to ClickHouse when row count or time thresholds are reached. +/// +/// Schema is fetched lazily from `system.columns` and cached. On schema +/// mismatch, the batcher automatically invalidates and re-fetches. +pub struct DynamicBatcher { + tx: mpsc::Sender, + handle: tokio::task::JoinHandle<()>, +} + +/// Cheap clonable handle for writing rows to a [`DynamicBatcher`]. +#[derive(Clone)] +pub struct DynamicBatcherHandle { + tx: mpsc::Sender, +} + +fn channel_closed() -> DynamicError { + DynamicError::EncodingError { + column: String::new(), + message: "DynamicBatcher background task gone".to_string(), + } +} + +impl DynamicBatcher { + /// Create a new `DynamicBatcher`. Spawns a background tokio task immediately. + pub fn new( + client: &Client, + database: &str, + table: &str, + config: DynamicBatchConfig, + ) -> Self { + let (tx, rx) = mpsc::channel(config.channel_capacity); + let client = client.clone(); + let database = database.to_string(); + let table = table.to_string(); + let schema_cache = client.dynamic_schema_cache.clone(); + + let handle = tokio::spawn(background_task( + client, + database, + table, + schema_cache, + config, + rx, + )); + + Self { tx, handle } + } + + /// Get a cheap, clonable write handle. + pub fn handle(&self) -> DynamicBatcherHandle { + DynamicBatcherHandle { + tx: self.tx.clone(), + } + } + + /// Buffer a row. Blocks (async) if the channel is full (backpressure). + pub async fn write_map(&self, row: Map) -> Result<(), DynamicError> { + let (resp_tx, resp_rx) = oneshot::channel(); + self.tx + .send(Cmd::Write(row, resp_tx)) + .await + .map_err(|_| channel_closed())?; + resp_rx.await.map_err(|_| channel_closed())? + } + + /// Force-flush all buffered rows to ClickHouse. + pub async fn flush(&self) -> Result { + let (resp_tx, resp_rx) = oneshot::channel(); + self.tx + .send(Cmd::Flush(resp_tx)) + .await + .map_err(|_| channel_closed())?; + resp_rx.await.map_err(|_| channel_closed())? + } + + /// Flush remaining rows and shut down the batcher. Consumes self. + pub async fn end(self) -> Result { + let (resp_tx, resp_rx) = oneshot::channel(); + if self.tx.send(Cmd::End(resp_tx)).await.is_err() { + return Ok(0); + } + drop(self.tx); + let result = resp_rx.await.map_err(|_| channel_closed())?; + let _ = self.handle.await; + result + } +} + +impl DynamicBatcherHandle { + /// Buffer a row (same as [`DynamicBatcher::write_map`]). + pub async fn write_map(&self, row: Map) -> Result<(), DynamicError> { + let (resp_tx, resp_rx) = oneshot::channel(); + self.tx + .send(Cmd::Write(row, resp_tx)) + .await + .map_err(|_| channel_closed())?; + resp_rx.await.map_err(|_| channel_closed())? + } + + /// Force-flush all buffered rows. + pub async fn flush(&self) -> Result { + let (resp_tx, resp_rx) = oneshot::channel(); + self.tx + .send(Cmd::Flush(resp_tx)) + .await + .map_err(|_| channel_closed())?; + resp_rx.await.map_err(|_| channel_closed())? + } +} + +// --------------------------------------------------------------------------- +// Background task +// --------------------------------------------------------------------------- + +async fn background_task( + client: Client, + database: String, + table: String, + schema_cache: Arc, + config: DynamicBatchConfig, + mut rx: mpsc::Receiver, +) { + let mut buffer: Vec> = Vec::with_capacity(config.max_rows as usize); + let mut total_rows: u64 = 0; + + let mut interval = tokio::time::interval(config.max_period); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + // Skip the immediate first tick + interval.tick().await; + + loop { + tokio::select! { + biased; + + cmd = rx.recv() => { + match cmd { + Some(Cmd::Write(row, resp)) => { + buffer.push(row); + if buffer.len() as u64 >= config.max_rows { + let flushed = flush_buffer( + &client, &database, &table, &schema_cache, + &mut buffer, + ).await; + match flushed { + Ok(n) => { + total_rows += n; + let _ = resp.send(Ok(())); + } + Err(e) => { + let _ = resp.send(Err(e)); + } + } + } else { + let _ = resp.send(Ok(())); + } + } + Some(Cmd::Flush(resp)) => { + let flushed = flush_buffer( + &client, &database, &table, &schema_cache, + &mut buffer, + ).await; + match flushed { + Ok(n) => { + total_rows += n; + let _ = resp.send(Ok(total_rows)); + } + Err(e) => { + let _ = resp.send(Err(e)); + } + } + } + Some(Cmd::End(resp)) => { + let flushed = flush_buffer( + &client, &database, &table, &schema_cache, + &mut buffer, + ).await; + match flushed { + Ok(n) => { + total_rows += n; + let _ = resp.send(Ok(total_rows)); + } + Err(e) => { + let _ = resp.send(Err(e)); + } + } + return; + } + None => { + // All senders dropped — flush and exit + let _ = flush_buffer( + &client, &database, &table, &schema_cache, + &mut buffer, + ).await; + return; + } + } + } + + _ = interval.tick() => { + if !buffer.is_empty() { + match flush_buffer( + &client, &database, &table, &schema_cache, + &mut buffer, + ).await { + Ok(n) => total_rows += n, + Err(_e) => { + // Timer-triggered flush errors are logged but not fatal + // The next write_map will surface errors to callers + } + } + } + } + } + } +} + +/// Flush buffered rows via DynamicInsert. +/// +/// On schema mismatch, invalidates cache and retries once with fresh schema. +async fn flush_buffer( + client: &Client, + database: &str, + table: &str, + schema_cache: &Arc, + buffer: &mut Vec>, +) -> Result { + if buffer.is_empty() { + return Ok(0); + } + + let rows = std::mem::take(buffer); + let count = rows.len() as u64; + + match try_insert(client, database, table, &rows).await { + Ok(()) => Ok(count), + Err(DynamicError::SchemaMismatch { .. }) => { + // Schema changed — invalidate and retry once + let full_table = format!("{database}.{table}"); + schema_cache.invalidate(&full_table); + + // Retry with fresh schema + try_insert(client, database, table, &rows) + .await + .map(|()| count) + } + Err(e) => Err(e), + } +} + +/// Attempt to insert rows via DynamicInsert. +async fn try_insert( + client: &Client, + database: &str, + table: &str, + rows: &[Map], +) -> Result<(), DynamicError> { + let mut insert = client.dynamic_insert(database, table); + for row in rows { + insert.write_map(row).await?; + } + insert.end().await?; + Ok(()) +} diff --git a/src/dynamic/encode.rs b/src/dynamic/encode.rs new file mode 100644 index 00000000..12a5c9aa --- /dev/null +++ b/src/dynamic/encode.rs @@ -0,0 +1,542 @@ +//! Runtime RowBinary encoder for `serde_json::Value`. +//! +//! Converts a JSON map to RowBinary bytes using a [`DynamicSchema`]. +//! This is the bridge between dynamic schemas (`Map`) +//! and the efficient binary wire format that ClickHouse expects. +//! +//! **Performance:** avoids the JSON text overhead of JSONEachRow. +//! ClickHouse receives pre-columnarised binary — zero server-side parsing. +//! +//! # Encoding Rules +//! +//! - Columns are written in schema order +//! - Missing columns with server-side defaults are skipped +//! - Missing columns without defaults get a type-appropriate zero value +//! - Nullable columns: `0x01` for NULL, `0x00` + value for non-NULL +//! - Strings: varint length prefix + UTF-8 bytes +//! - Integers: little-endian fixed-width +//! - UUID: two little-endian u64 (high, low) + +use serde_json::{Map, Value}; + +use super::error::DynamicError; +use super::schema::{ColumnDef, DynamicSchema}; + +/// Encode a JSON row map to RowBinary bytes according to the schema. +/// +/// Columns are written in schema order. Missing columns with server-side +/// defaults are omitted (the INSERT column list excludes them). Missing +/// columns WITHOUT defaults get a type-appropriate zero value. +pub fn encode_dynamic_row( + row: &Map, + _schema: &DynamicSchema, + columns_to_send: &[&ColumnDef], +) -> Result, DynamicError> { + let mut buf = Vec::with_capacity(256); + + for col in columns_to_send { + let value = row.get(&col.name).unwrap_or(&Value::Null); + encode_value(value, col, &mut buf)?; + } + + Ok(buf) +} + +/// Determine which columns to include in the INSERT column list. +/// +/// Includes columns that are present in the row OR that have no default +/// (must send something). Columns with defaults that aren't in the row +/// are omitted — ClickHouse fills them server-side. +pub fn columns_to_send<'a>( + row: &Map, + schema: &'a DynamicSchema, +) -> Vec<&'a ColumnDef> { + schema + .columns + .iter() + .filter(|col| row.contains_key(&col.name) || !col.has_default) + .collect() +} + +// --------------------------------------------------------------------------- +// Core encoding +// --------------------------------------------------------------------------- + +fn encode_value(value: &Value, col: &ColumnDef, buf: &mut Vec) -> Result<(), DynamicError> { + let pt = &col.parsed_type; + + // Handle Nullable wrapper + if pt.nullable { + if value.is_null() { + buf.push(1); // is_null = true + return Ok(()); + } + buf.push(0); // is_null = false + } else if value.is_null() { + // Non-nullable column with null value — write type default + write_default(pt, buf); + return Ok(()); + } + + encode_typed(value, pt, &col.name, buf) +} + +fn encode_typed( + value: &Value, + pt: &super::parsed_type::ParsedType, + col_name: &str, + buf: &mut Vec, +) -> Result<(), DynamicError> { + match pt.base.as_str() { + "String" => { + let s = value_to_string(value); + write_string(s.as_bytes(), buf); + } + "FixedString" => { + let s = value_to_string(value); + let n = pt.fixed_size.unwrap_or(1); + let bytes = s.as_bytes(); + if bytes.len() <= n { + buf.extend_from_slice(bytes); + buf.resize(buf.len() + (n - bytes.len()), 0); + } else { + buf.extend_from_slice(&bytes[..n]); + } + } + "UInt8" | "Bool" => { + buf.push(as_u64(value, col_name)? as u8); + } + "UInt16" => { + buf.extend_from_slice(&(as_u64(value, col_name)? as u16).to_le_bytes()); + } + "UInt32" | "DateTime" => { + buf.extend_from_slice(&(as_u64(value, col_name)? as u32).to_le_bytes()); + } + "UInt64" => { + buf.extend_from_slice(&as_u64(value, col_name)?.to_le_bytes()); + } + "Int8" | "Enum8" => { + buf.extend_from_slice(&(as_i64(value, col_name)? as i8).to_le_bytes()); + } + "Int16" | "Enum16" | "Date" => { + buf.extend_from_slice(&(as_i64(value, col_name)? as i16).to_le_bytes()); + } + "Int32" | "Date32" | "Decimal32" => { + buf.extend_from_slice(&(as_i64(value, col_name)? as i32).to_le_bytes()); + } + "Int64" | "DateTime64" | "Decimal64" => { + buf.extend_from_slice(&as_i64(value, col_name)?.to_le_bytes()); + } + "Float32" => { + buf.extend_from_slice(&(as_f64(value, col_name)? as f32).to_le_bytes()); + } + "Float64" => { + buf.extend_from_slice(&as_f64(value, col_name)?.to_le_bytes()); + } + "UUID" => encode_uuid(value, col_name, buf)?, + "IPv4" => encode_ipv4(value, col_name, buf)?, + "IPv6" => encode_ipv6(value, col_name, buf)?, + "Array" => { + let elem = pt + .array_element + .as_ref() + .ok_or_else(|| enc_err(col_name, "Array without element type"))?; + encode_array(value, elem, col_name, buf)?; + } + "Map" => { + let (kt, vt) = pt + .map_types + .as_ref() + .ok_or_else(|| enc_err(col_name, "Map without key/value types"))?; + encode_map(value, kt, vt, col_name, buf)?; + } + "JSON" => { + // JSON type — send as length-prefixed JSON string + let json_str = value.to_string(); + write_string(json_str.as_bytes(), buf); + } + other => { + // Unknown type — try as string (forward-compatible) + let s = value_to_string(value); + write_string(s.as_bytes(), buf); + // Log but don't fail — ClickHouse may accept it + #[cfg(feature = "tracing")] + tracing::debug!( + column = col_name, + r#type = other, + "encoding unknown type as String" + ); + let _ = other; + } + } + Ok(()) +} + +// --------------------------------------------------------------------------- +// Wire format helpers +// --------------------------------------------------------------------------- + +fn write_string(bytes: &[u8], buf: &mut Vec) { + write_varint(bytes.len() as u64, buf); + buf.extend_from_slice(bytes); +} + +fn write_varint(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); + } +} + +fn write_default(pt: &super::parsed_type::ParsedType, buf: &mut Vec) { + if let Some(size) = pt.fixed_byte_size() { + buf.extend(std::iter::repeat_n(0u8, size)); + } else { + // Variable-length: empty string / empty array / empty map + write_varint(0, buf); + } +} + +// --------------------------------------------------------------------------- +// Value coercion helpers +// --------------------------------------------------------------------------- + +fn value_to_string(value: &Value) -> String { + match value { + Value::String(s) => s.clone(), + Value::Number(n) => n.to_string(), + Value::Bool(b) => b.to_string(), + Value::Null => String::new(), + other => other.to_string(), + } +} + +fn as_u64(value: &Value, col: &str) -> Result { + match value { + Value::Number(n) => n + .as_u64() + .or_else(|| n.as_i64().map(|v| v as u64)) + .or_else(|| n.as_f64().map(|v| v as u64)) + .ok_or_else(|| enc_err(col, "not a valid unsigned integer")), + Value::Bool(b) => Ok(u64::from(*b)), + Value::String(s) => s + .parse::() + .map_err(|_| enc_err(col, "string not parseable as u64")), + _ => Err(enc_err(col, "expected number")), + } +} + +fn as_i64(value: &Value, col: &str) -> Result { + match value { + Value::Number(n) => n + .as_i64() + .or_else(|| n.as_u64().map(|v| v as i64)) + .or_else(|| n.as_f64().map(|v| v as i64)) + .ok_or_else(|| enc_err(col, "not a valid integer")), + Value::Bool(b) => Ok(i64::from(*b)), + Value::String(s) => s + .parse::() + .map_err(|_| enc_err(col, "string not parseable as i64")), + _ => Err(enc_err(col, "expected number")), + } +} + +fn as_f64(value: &Value, col: &str) -> Result { + match value { + Value::Number(n) => n.as_f64().ok_or_else(|| enc_err(col, "not a valid float")), + Value::String(s) => s + .parse::() + .map_err(|_| enc_err(col, "string not parseable as f64")), + _ => Err(enc_err(col, "expected number")), + } +} + +// --------------------------------------------------------------------------- +// Complex type encoders +// --------------------------------------------------------------------------- + +fn encode_uuid(value: &Value, col: &str, buf: &mut Vec) -> Result<(), DynamicError> { + let s = value_to_string(value); + let hex: String = s.chars().filter(|c| c.is_ascii_hexdigit()).collect(); + if hex.len() != 32 { + return Err(enc_err(col, "invalid UUID length")); + } + // ClickHouse RowBinary UUID: two LE u64 (high word first, then low) + let high = u64::from_str_radix(&hex[..16], 16).map_err(|_| enc_err(col, "invalid UUID hex"))?; + let low = u64::from_str_radix(&hex[16..], 16).map_err(|_| enc_err(col, "invalid UUID hex"))?; + buf.extend_from_slice(&high.to_le_bytes()); + buf.extend_from_slice(&low.to_le_bytes()); + Ok(()) +} + +fn encode_ipv4(value: &Value, col: &str, buf: &mut Vec) -> Result<(), DynamicError> { + let s = value_to_string(value); + let addr: std::net::Ipv4Addr = s.parse().map_err(|_| enc_err(col, "invalid IPv4"))?; + // ClickHouse stores IPv4 as UInt32 little-endian + buf.extend_from_slice(&u32::from(addr).to_le_bytes()); + Ok(()) +} + +fn encode_ipv6(value: &Value, col: &str, buf: &mut Vec) -> Result<(), DynamicError> { + let s = value_to_string(value); + let addr: std::net::Ipv6Addr = s.parse().map_err(|_| enc_err(col, "invalid IPv6"))?; + buf.extend_from_slice(&addr.octets()); + Ok(()) +} + +fn encode_array( + value: &Value, + elem_type: &super::parsed_type::ParsedType, + col_name: &str, + buf: &mut Vec, +) -> Result<(), DynamicError> { + let arr = match value { + Value::Array(a) => a, + _ => return Err(enc_err(col_name, "expected array")), + }; + write_varint(arr.len() as u64, buf); + let dummy_col = ColumnDef { + name: col_name.to_string(), + raw_type: String::new(), + parsed_type: elem_type.clone(), + default_kind: String::new(), + has_default: false, + }; + for item in arr { + encode_value(item, &dummy_col, buf)?; + } + Ok(()) +} + +fn encode_map( + value: &Value, + key_type: &super::parsed_type::ParsedType, + val_type: &super::parsed_type::ParsedType, + col_name: &str, + buf: &mut Vec, +) -> Result<(), DynamicError> { + let obj = match value { + Value::Object(m) => m, + _ => return Err(enc_err(col_name, "expected object for Map")), + }; + write_varint(obj.len() as u64, buf); + let key_col = ColumnDef { + name: format!("{col_name}.key"), + raw_type: String::new(), + parsed_type: key_type.clone(), + default_kind: String::new(), + has_default: false, + }; + let val_col = ColumnDef { + name: format!("{col_name}.value"), + raw_type: String::new(), + parsed_type: val_type.clone(), + default_kind: String::new(), + has_default: false, + }; + for (k, v) in obj { + encode_value(&Value::String(k.clone()), &key_col, buf)?; + encode_value(v, &val_col, buf)?; + } + Ok(()) +} + +fn enc_err(col: &str, msg: &str) -> DynamicError { + DynamicError::EncodingError { + column: col.to_string(), + message: msg.to_string(), + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::super::parsed_type::ParsedType; + use super::super::schema::DynamicSchema; + use super::*; + + fn col(name: &str, type_str: &str) -> ColumnDef { + ColumnDef { + name: name.to_string(), + raw_type: type_str.to_string(), + parsed_type: ParsedType::parse(type_str), + default_kind: String::new(), + has_default: false, + } + } + + fn col_default(name: &str, type_str: &str) -> ColumnDef { + ColumnDef { + name: name.to_string(), + raw_type: type_str.to_string(), + parsed_type: ParsedType::parse(type_str), + default_kind: "DEFAULT".to_string(), + has_default: true, + } + } + + #[test] + fn test_encode_string() { + let schema = DynamicSchema::from_columns("t", vec![col("name", "String")]); + let cols = columns_to_send( + &serde_json::json!({"name": "hello"}).as_object().unwrap(), + &schema, + ); + let row = serde_json::json!({"name": "hello"}); + let bytes = encode_dynamic_row(row.as_object().unwrap(), &schema, &cols).unwrap(); + // varint(5) + "hello" + assert_eq!(bytes, vec![5, b'h', b'e', b'l', b'l', b'o']); + } + + #[test] + fn test_encode_uint32() { + let schema = DynamicSchema::from_columns("t", vec![col("id", "UInt32")]); + let row = serde_json::json!({"id": 42}); + let cols = columns_to_send(row.as_object().unwrap(), &schema); + let bytes = encode_dynamic_row(row.as_object().unwrap(), &schema, &cols).unwrap(); + assert_eq!(bytes, 42u32.to_le_bytes().to_vec()); + } + + #[test] + fn test_encode_int64() { + let schema = DynamicSchema::from_columns("t", vec![col("val", "Int64")]); + let row = serde_json::json!({"val": -100}); + let cols = columns_to_send(row.as_object().unwrap(), &schema); + let bytes = encode_dynamic_row(row.as_object().unwrap(), &schema, &cols).unwrap(); + assert_eq!(bytes, (-100i64).to_le_bytes().to_vec()); + } + + #[test] + fn test_encode_float64() { + let schema = DynamicSchema::from_columns("t", vec![col("f", "Float64")]); + let row = serde_json::json!({"f": 3.14}); + let cols = columns_to_send(row.as_object().unwrap(), &schema); + let bytes = encode_dynamic_row(row.as_object().unwrap(), &schema, &cols).unwrap(); + assert_eq!(bytes, 3.14f64.to_le_bytes().to_vec()); + } + + #[test] + fn test_encode_bool() { + let schema = DynamicSchema::from_columns("t", vec![col("b", "Bool")]); + let row = serde_json::json!({"b": true}); + let cols = columns_to_send(row.as_object().unwrap(), &schema); + let bytes = encode_dynamic_row(row.as_object().unwrap(), &schema, &cols).unwrap(); + assert_eq!(bytes, vec![1]); + } + + #[test] + fn test_encode_nullable_null() { + let schema = DynamicSchema::from_columns("t", vec![col("n", "Nullable(String)")]); + let row = serde_json::json!({"n": null}); + let cols = columns_to_send(row.as_object().unwrap(), &schema); + let bytes = encode_dynamic_row(row.as_object().unwrap(), &schema, &cols).unwrap(); + // 0x01 = is_null + assert_eq!(bytes, vec![1]); + } + + #[test] + fn test_encode_nullable_non_null() { + let schema = DynamicSchema::from_columns("t", vec![col("n", "Nullable(String)")]); + let row = serde_json::json!({"n": "hi"}); + let cols = columns_to_send(row.as_object().unwrap(), &schema); + let bytes = encode_dynamic_row(row.as_object().unwrap(), &schema, &cols).unwrap(); + // 0x00 = not_null, varint(2), "hi" + assert_eq!(bytes, vec![0, 2, b'h', b'i']); + } + + #[test] + fn test_encode_missing_column_with_default_skipped() { + let schema = DynamicSchema::from_columns( + "t", + vec![col("id", "UInt32"), col_default("ts", "DateTime64(3)")], + ); + let row = serde_json::json!({"id": 1}); + let cols = columns_to_send(row.as_object().unwrap(), &schema); + // Only id should be in the column list (ts has default and is absent) + assert_eq!(cols.len(), 1); + assert_eq!(cols[0].name, "id"); + let bytes = encode_dynamic_row(row.as_object().unwrap(), &schema, &cols).unwrap(); + assert_eq!(bytes, 1u32.to_le_bytes().to_vec()); + } + + #[test] + fn test_encode_missing_non_nullable_gets_zero() { + let schema = DynamicSchema::from_columns("t", vec![col("x", "UInt32")]); + let row = serde_json::json!({}); + let cols = columns_to_send(row.as_object().unwrap(), &schema); + let bytes = encode_dynamic_row(row.as_object().unwrap(), &schema, &cols).unwrap(); + assert_eq!(bytes, 0u32.to_le_bytes().to_vec()); + } + + #[test] + fn test_encode_array() { + let schema = DynamicSchema::from_columns("t", vec![col("a", "Array(UInt32)")]); + let row = serde_json::json!({"a": [1, 2, 3]}); + let cols = columns_to_send(row.as_object().unwrap(), &schema); + let bytes = encode_dynamic_row(row.as_object().unwrap(), &schema, &cols).unwrap(); + let mut expected = vec![3u8]; // varint(3) + expected.extend_from_slice(&1u32.to_le_bytes()); + expected.extend_from_slice(&2u32.to_le_bytes()); + expected.extend_from_slice(&3u32.to_le_bytes()); + assert_eq!(bytes, expected); + } + + #[test] + fn test_encode_multi_column() { + let schema = + DynamicSchema::from_columns("t", vec![col("id", "UInt32"), col("name", "String")]); + let row = serde_json::json!({"id": 42, "name": "test"}); + let cols = columns_to_send(row.as_object().unwrap(), &schema); + let bytes = encode_dynamic_row(row.as_object().unwrap(), &schema, &cols).unwrap(); + let mut expected = Vec::new(); + expected.extend_from_slice(&42u32.to_le_bytes()); + expected.extend_from_slice(&[4, b't', b'e', b's', b't']); // varint(4) + "test" + assert_eq!(bytes, expected); + } + + #[test] + fn test_encode_fixed_string() { + let schema = DynamicSchema::from_columns("t", vec![col("f", "FixedString(4)")]); + let row = serde_json::json!({"f": "ab"}); + let cols = columns_to_send(row.as_object().unwrap(), &schema); + let bytes = encode_dynamic_row(row.as_object().unwrap(), &schema, &cols).unwrap(); + assert_eq!(bytes, vec![b'a', b'b', 0, 0]); // padded with zeros + } + + #[test] + fn test_encode_string_from_number() { + // Numbers should coerce to string + let schema = DynamicSchema::from_columns("t", vec![col("s", "String")]); + let row = serde_json::json!({"s": 42}); + let cols = columns_to_send(row.as_object().unwrap(), &schema); + let bytes = encode_dynamic_row(row.as_object().unwrap(), &schema, &cols).unwrap(); + assert_eq!(bytes, vec![2, b'4', b'2']); + } + + #[test] + fn test_varint_encoding() { + let mut buf = Vec::new(); + write_varint(0, &mut buf); + assert_eq!(buf, vec![0]); + + buf.clear(); + write_varint(127, &mut buf); + assert_eq!(buf, vec![127]); + + buf.clear(); + write_varint(128, &mut buf); + assert_eq!(buf, vec![0x80, 0x01]); + + buf.clear(); + write_varint(300, &mut buf); + assert_eq!(buf, vec![0xAC, 0x02]); + } +} diff --git a/src/dynamic/error.rs b/src/dynamic/error.rs new file mode 100644 index 00000000..8831daf0 --- /dev/null +++ b/src/dynamic/error.rs @@ -0,0 +1,51 @@ +//! Error types for dynamic (schema-driven) inserts. + +use std::fmt; + +/// Errors specific to dynamic schema-driven inserts. +#[derive(Debug)] +pub enum DynamicError { + /// Column type string could not be parsed. + UnsupportedType { column: String, type_str: String }, + /// Value could not be encoded for the target column type. + EncodingError { column: String, message: String }, + /// Schema mismatch detected — server rejected the insert. + SchemaMismatch { table: String, message: String }, + /// Schema fetch from system.columns failed. + SchemaFetch { + table: String, + source: crate::error::Error, + }, + /// Table has no columns (or does not exist). + EmptySchema { table: String }, +} + +impl fmt::Display for DynamicError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::UnsupportedType { column, type_str } => { + write!(f, "unsupported type '{type_str}' for column '{column}'") + } + Self::EncodingError { column, message } => { + write!(f, "encoding error for column '{column}': {message}") + } + Self::SchemaMismatch { table, message } => { + write!(f, "schema mismatch for table '{table}': {message}") + } + Self::SchemaFetch { table, source } => { + write!(f, "failed to fetch schema for '{table}': {source}") + } + Self::EmptySchema { table } => { + write!(f, "table '{table}' has no columns or does not exist") + } + } + } +} + +impl std::error::Error for DynamicError {} + +impl From for crate::error::Error { + fn from(e: DynamicError) -> Self { + crate::error::Error::Custom(e.to_string()) + } +} diff --git a/src/dynamic/insert.rs b/src/dynamic/insert.rs new file mode 100644 index 00000000..0d7231a1 --- /dev/null +++ b/src/dynamic/insert.rs @@ -0,0 +1,190 @@ +//! Single-table dynamic insert with automatic schema fetch and recovery. +//! +//! `DynamicInsert` fetches the table schema from `system.columns` on first use, +//! encodes `Map` to RowBinary, and sends via HTTP `InsertFormatted`. +//! +//! On schema mismatch errors (e.g. `ALTER TABLE ADD COLUMN`), it automatically +//! invalidates the cached schema so the next insert re-fetches. The caller's +//! retry/salvage logic handles re-sending failed rows. +//! +//! # End-to-End Efficiency +//! +//! This replaces the "push JSON, let ClickHouse parse it" approach with +//! schema-reflected binary. Your app does roughly the same work (binary encoding +//! instead of JSON serialisation), but the ClickHouse cluster does zero parsing +//! on ingest. Think total CPU across client + cluster, not just your app. + +use std::sync::Arc; + +use serde_json::{Map, Value}; + +use crate::Client; + +use super::encode::{columns_to_send, encode_dynamic_row}; +use super::error::DynamicError; +use super::schema::{fetch_dynamic_schema, ColumnDef, DynamicSchema, DynamicSchemaCache}; + +/// Dynamic insert for a single table. +/// +/// Encodes `Map` to RowBinary using a schema fetched from +/// `system.columns`. As simple to use as JSONEachRow, but binary wire format. +/// +/// # Schema Recovery +/// +/// If ClickHouse rejects an insert due to schema mismatch (column added/removed, +/// type changed), call [`invalidate_schema()`][Self::invalidate_schema] and create +/// a new `DynamicInsert`. The next insert will re-fetch the schema automatically. +/// +/// For automatic recovery in a pipeline context, use `DynamicBatcher` which +/// handles this transparently. +pub struct DynamicInsert { + client: Client, + database: String, + table: String, + schema_cache: Arc, + schema: Option, + /// Column list for the current INSERT (determined from first row). + insert_columns: Option>, + /// The active HTTP insert (created lazily on first write_map). + insert: Option, + rows_written: u64, +} + +impl DynamicInsert { + /// Create a new `DynamicInsert`. Schema is fetched lazily on first `write_map()`. + pub(crate) fn new( + client: Client, + database: String, + table: String, + schema_cache: Arc, + ) -> Self { + Self { + client, + database, + table, + schema_cache, + schema: None, + insert_columns: None, + insert: None, + rows_written: 0, + } + } + + /// Ensure schema is loaded (from cache or system.columns). + async fn ensure_schema(&mut self) -> Result<&DynamicSchema, DynamicError> { + if self.schema.is_none() { + let full_table = format!("{}.{}", self.database, self.table); + let schema = if let Some(cached) = self.schema_cache.get(&full_table) { + cached + } else { + let fetched = + fetch_dynamic_schema(&self.client, &self.database, &self.table).await?; + self.schema_cache.insert(&full_table, fetched.clone()); + fetched + }; + self.schema = Some(schema); + } + Ok(self.schema.as_ref().unwrap()) + } + + /// Encode and buffer a row for insert. + /// + /// The row is encoded to RowBinary and written to the HTTP insert buffer. + /// On first call, fetches the schema and creates the INSERT statement. + pub async fn write_map(&mut self, row: &Map) -> Result<(), DynamicError> { + // Ensure schema is loaded + if self.schema.is_none() { + self.ensure_schema().await?; + } + let schema = self.schema.as_ref().unwrap(); + + // On first row, determine the column list and create the INSERT + if self.insert.is_none() { + let cols = columns_to_send(row, schema); + let col_names: Vec = cols.iter().map(|c| c.name.clone()).collect(); + let col_list = col_names.join(", "); + let sql = format!( + "INSERT INTO {}.{} ({col_list}) FORMAT RowBinary", + self.database, self.table + ); + self.insert = Some(self.client.insert_formatted_with(sql).buffered()); + self.insert_columns = Some(col_names); + } + + // Build the column def refs for encoding based on stored column names + let col_defs: Vec<&ColumnDef> = self + .insert_columns + .as_ref() + .unwrap() + .iter() + .filter_map(|name| schema.column(name)) + .collect(); + + // Encode row to RowBinary + let rb_bytes = encode_dynamic_row(row, schema, &col_defs)?; + + // Write to the HTTP insert buffer + let insert = self.insert.as_mut().unwrap(); + insert + .write(&rb_bytes) + .await + .map_err(|e| classify_error(&self.database, &self.table, e))?; + + self.rows_written += 1; + Ok(()) + } + + /// Flush the buffer and finalise the INSERT. + /// + /// Returns the number of rows written. + pub async fn end(mut self) -> Result { + if let Some(mut insert) = self.insert.take() { + insert + .end() + .await + .map_err(|e| classify_error(&self.database, &self.table, e))?; + } + Ok(self.rows_written) + } + + /// Invalidate the cached schema, forcing a re-fetch on next insert. + /// + /// Call this after a schema mismatch error before creating a new + /// `DynamicInsert` for the same table. + pub fn invalidate_schema(&mut self) { + let full_table = format!("{}.{}", self.database, self.table); + self.schema_cache.invalidate(&full_table); + self.schema = None; + } + + /// Number of rows written so far. + pub fn rows_written(&self) -> u64 { + self.rows_written + } + + /// Get the current schema (if loaded). + pub fn schema(&self) -> Option<&DynamicSchema> { + self.schema.as_ref() + } +} + +/// Classify a ClickHouse error as schema mismatch or generic encoding error. +fn classify_error(database: &str, table: &str, e: crate::error::Error) -> DynamicError { + let msg = e.to_string(); + if msg.contains("UNKNOWN_IDENTIFIER") + || msg.contains("NO_SUCH_COLUMN") + || msg.contains("THERE_IS_NO_COLUMN") + || msg.contains("TYPE_MISMATCH") + || msg.contains("ILLEGAL_COLUMN") + { + DynamicError::SchemaMismatch { + table: format!("{database}.{table}"), + message: msg, + } + } else { + DynamicError::EncodingError { + column: String::new(), + message: msg, + } + } +} diff --git a/src/dynamic/mod.rs b/src/dynamic/mod.rs new file mode 100644 index 00000000..dc90721c --- /dev/null +++ b/src/dynamic/mod.rs @@ -0,0 +1,47 @@ +//! Runtime schema-driven inserts for dynamic schemas. +//! +//! Use this when table schemas are not known at compile time. +//! `DynamicInsert` fetches the schema from `system.columns` and encodes +//! `Map` directly to RowBinary — same ease as JSONEachRow +//! but without the server-side JSON parsing overhead. +//! +//! # Why This Exists — End-to-End CPU Savings +//! +//! JSONEachRow is easy: push JSON text, ClickHouse parses it. But at scale, +//! the ClickHouse cluster itself pays the CPU cost of parsing every JSON row +//! on ingest. That's not "someone else's problem" — it's your total solution +//! budget. If your ClickHouse cluster is CPU-loaded because every INSERT runs +//! through a JSON parser, that's capacity you can't use for queries. +//! +//! Schema-reflected RowBinary shifts the work to the client: fetch the schema +//! once, encode binary directly, ClickHouse receives pre-columnarised data +//! with zero parsing. The client does roughly the same work (binary encoding +//! instead of JSON serialisation), but the server does dramatically less. +//! The big picture: total CPU across client + cluster drops significantly. +//! +//! "Hey, my app works — if the CH cluster is loaded, that's the infra team's +//! problem" is exactly the mindset this module replaces. Think end-to-end. +//! +//! # Three Insert Tiers +//! +//! | Tier | API | Use When | +//! |------|-----|----------| +//! | 1 | `Insert` / `Inserter` | Compile-time schema, `#[derive(Row)]` | +//! | 2 | `DynamicInsert` | Runtime schema, `Map` (this module) | +//! | 3 | `InsertFormatted` | Raw bytes, any format (JSONEachRow, CSV) | +//! +//! Tier 2 gives you the ergonomics of Tier 3 (push any JSON map) with the +//! performance profile of Tier 1 (ClickHouse skips JSON parsing entirely). + +pub mod batcher; +pub mod encode; +pub mod error; +pub mod insert; +pub mod parsed_type; +pub mod schema; + +pub use batcher::{DynamicBatchConfig, DynamicBatcher, DynamicBatcherHandle}; +pub use error::DynamicError; +pub use insert::DynamicInsert; +pub use parsed_type::ParsedType; +pub use schema::{ColumnDef, DynamicSchema, DynamicSchemaCache, fetch_dynamic_schema}; diff --git a/src/dynamic/parsed_type.rs b/src/dynamic/parsed_type.rs new file mode 100644 index 00000000..835bc226 --- /dev/null +++ b/src/dynamic/parsed_type.rs @@ -0,0 +1,478 @@ +//! Rich ClickHouse type parser for runtime schema-driven inserts. +//! +//! Parses ClickHouse type strings from `system.columns` into a structured +//! `ParsedType` AST. Handles Nullable, LowCardinality, Array, Map, +//! DateTime64(precision, timezone), Decimal(precision, scale), FixedString(n), +//! Enum8/Enum16, and all scalar types. +//! +//! Lifted from the HyperI DFE Loader project — generic enough for any +//! clickhouse-rs user with dynamic schemas. + +use std::fmt; + +/// Parsed ClickHouse type information. +/// +/// Runtime representation of a ClickHouse column type. Unknown types are +/// preserved as-is for forward compatibility. +/// +/// # Examples +/// +/// ``` +/// use clickhouse::dynamic::ParsedType; +/// +/// let t = ParsedType::parse("LowCardinality(Nullable(String))"); +/// assert_eq!(t.base, "String"); +/// assert!(t.nullable); +/// assert!(t.low_cardinality); +/// +/// let t = ParsedType::parse("DateTime64(3, 'UTC')"); +/// assert_eq!(t.base, "DateTime64"); +/// assert_eq!(t.precision, Some(3)); +/// assert_eq!(t.timezone.as_deref(), Some("UTC")); +/// ``` +#[derive(Debug, Clone, PartialEq)] +pub struct ParsedType { + /// Original type string from ClickHouse. + pub raw: String, + /// Base type name (e.g., "String", "Int64", "DateTime64"). + pub base: String, + /// Whether wrapped in Nullable(). + pub nullable: bool, + /// Whether wrapped in LowCardinality(). + pub low_cardinality: bool, + /// For Array types, the element type. + pub array_element: Option>, + /// For Map types, (key_type, value_type). + pub map_types: Option<(Box, Box)>, + /// Extended info: precision for DateTime64, Decimal, etc. + pub precision: Option, + /// Extended info: scale for Decimal types. + pub scale: Option, + /// Extended info: timezone for DateTime64. + pub timezone: Option, + /// Extended info: size for FixedString. + pub fixed_size: Option, +} + +impl ParsedType { + /// Parse a ClickHouse type string into a structured `ParsedType`. + #[must_use] + pub fn parse(type_str: &str) -> Self { + let type_str = type_str.trim(); + Self::parse_inner(type_str, type_str.to_string()) + } + + fn parse_inner(type_str: &str, raw: String) -> Self { + let mut result = Self { + raw, + base: String::new(), + nullable: false, + low_cardinality: false, + array_element: None, + map_types: None, + precision: None, + scale: None, + timezone: None, + fixed_size: None, + }; + + let mut type_str = type_str.trim().to_string(); + + // Unwrap wrappers in a loop (handles LowCardinality(Nullable(...)) etc.) + loop { + let (unwrapped, is_nullable) = Self::unwrap_wrapper(&type_str, "Nullable"); + if is_nullable { + result.nullable = true; + type_str = unwrapped; + continue; + } + + let (unwrapped, is_lc) = Self::unwrap_wrapper(&type_str, "LowCardinality"); + if is_lc { + result.low_cardinality = true; + type_str = unwrapped; + continue; + } + + break; + } + + // Check for Array + if let Some(inner) = Self::extract_wrapper(&type_str, "Array") { + result.base = "Array".to_string(); + result.array_element = Some(Box::new(Self::parse(&inner))); + return result; + } + + // Check for Map + if let Some(inner) = Self::extract_wrapper(&type_str, "Map") + && let Some((key, value)) = Self::split_type_args(&inner) + { + result.base = "Map".to_string(); + result.map_types = Some((Box::new(Self::parse(&key)), Box::new(Self::parse(&value)))); + return result; + } + + // Check for DateTime64(precision, 'timezone') + if type_str.starts_with("DateTime64") { + result.base = "DateTime64".to_string(); + if let Some(inner) = Self::extract_wrapper(&type_str, "DateTime64") { + let parts: Vec<&str> = inner.splitn(2, ',').collect(); + result.precision = parts.first().and_then(|p| p.trim().parse().ok()); + result.timezone = parts + .get(1) + .map(|tz| tz.trim().trim_matches('\'').trim_matches('"').to_string()); + } + return result; + } + + // Check for FixedString(N) + if let Some(inner) = Self::extract_wrapper(&type_str, "FixedString") { + result.base = "FixedString".to_string(); + result.fixed_size = inner.trim().parse().ok(); + return result; + } + + // Check for Decimal(P, S) or Decimal32/64/128/256(S) + if type_str.starts_with("Decimal") { + result.base = Self::parse_decimal_base(&type_str); + if let Some(inner) = Self::extract_parens(&type_str) { + let parts: Vec<&str> = inner.split(',').collect(); + if parts.len() == 2 { + result.precision = parts[0].trim().parse().ok(); + result.scale = parts[1].trim().parse().ok(); + } else if parts.len() == 1 { + result.scale = parts[0].trim().parse().ok(); + } + } + return result; + } + + // Check for Enum8/Enum16 + if type_str.starts_with("Enum8") || type_str.starts_with("Enum16") { + result.base = if type_str.starts_with("Enum8") { + "Enum8".to_string() + } else { + "Enum16".to_string() + }; + return result; + } + + // Simple type + result.base = type_str.to_string(); + result + } + + fn unwrap_wrapper(type_str: &str, wrapper: &str) -> (String, bool) { + let prefix = format!("{wrapper}("); + if let Some(rest) = type_str.strip_prefix(&prefix) + && let Some(inner) = rest.strip_suffix(')') + { + return (inner.to_string(), true); + } + (type_str.to_string(), false) + } + + fn extract_wrapper(type_str: &str, wrapper: &str) -> Option { + let prefix = format!("{wrapper}("); + type_str + .strip_prefix(&prefix) + .and_then(|rest| rest.strip_suffix(')')) + .map(std::string::ToString::to_string) + } + + fn extract_parens(type_str: &str) -> Option { + let start = type_str.find('(')?; + let end = type_str.rfind(')')?; + if start < end { + Some(type_str[start + 1..end].to_string()) + } else { + None + } + } + + fn parse_decimal_base(type_str: &str) -> String { + if type_str.starts_with("Decimal256") { + "Decimal256".to_string() + } else if type_str.starts_with("Decimal128") { + "Decimal128".to_string() + } else if type_str.starts_with("Decimal64") { + "Decimal64".to_string() + } else if type_str.starts_with("Decimal32") { + "Decimal32".to_string() + } else { + "Decimal".to_string() + } + } + + /// Split Map(K, V) or similar two-arg types, handling nested parens. + fn split_type_args(inner: &str) -> Option<(String, String)> { + let mut depth = 0; + for (i, c) in inner.char_indices() { + match c { + '(' => depth += 1, + ')' => depth -= 1, + ',' if depth == 0 => { + return Some(( + inner[..i].trim().to_string(), + inner[i + 1..].trim().to_string(), + )); + } + _ => {} + } + } + None + } + + /// Get the type category for coercion/encoding decisions. + /// + /// Maps ClickHouse types to categories. Unknown types map to "String". + #[must_use] + pub fn category(&self) -> &str { + match self.base.as_str() { + "String" | "FixedString" => "String", + "Int8" | "Int16" | "Int32" | "Int64" | "Int128" | "Int256" => "Int", + "UInt8" | "UInt16" | "UInt32" | "UInt64" | "UInt128" | "UInt256" => "UInt", + "Float32" | "Float64" => "Float", + "Decimal" | "Decimal32" | "Decimal64" | "Decimal128" | "Decimal256" => "Decimal", + "Bool" => "Bool", + "Date" | "Date32" => "Date", + "DateTime" => "DateTime", + "DateTime64" => "DateTime64", + "UUID" => "UUID", + "IPv4" => "IPv4", + "IPv6" => "IPv6", + "Array" => "Array", + "Map" => "Map", + "Tuple" => "Tuple", + "JSON" | "Object" => "JSON", + "Variant" => "Variant", + "Dynamic" => "Dynamic", + "Enum8" | "Enum16" => "Enum", + "Point" | "Ring" | "Polygon" | "MultiPolygon" | "LineString" + | "MultiLineString" => "Geo", + _ => "String", + } + } + + /// Check if this is a numeric type. + #[must_use] + pub fn is_numeric(&self) -> bool { + matches!(self.category(), "Int" | "UInt" | "Float" | "Decimal") + } + + /// Check if this is a string type. + #[must_use] + pub fn is_string(&self) -> bool { + self.category() == "String" + } + + /// Check if this is a date/time type. + #[must_use] + pub fn is_datetime(&self) -> bool { + matches!(self.category(), "Date" | "DateTime" | "DateTime64") + } + + /// Check if this is an IP address type. + #[must_use] + pub fn is_ip(&self) -> bool { + matches!(self.category(), "IPv4" | "IPv6") + } + + /// Byte size of this type's fixed-width representation, if applicable. + /// + /// Returns `None` for variable-length types (String, Array, Map, JSON). + #[must_use] + pub fn fixed_byte_size(&self) -> Option { + match self.base.as_str() { + "UInt8" | "Int8" | "Bool" | "Enum8" => Some(1), + "UInt16" | "Int16" | "Enum16" | "Date" => Some(2), + "UInt32" | "Int32" | "Date32" | "DateTime" | "Float32" | "Decimal32" | "IPv4" => { + Some(4) + } + "UInt64" | "Int64" | "Float64" | "Decimal64" | "DateTime64" => Some(8), + "Int128" | "UInt128" | "Decimal128" | "UUID" | "IPv6" => Some(16), + "Int256" | "UInt256" | "Decimal256" => Some(32), + "FixedString" => self.fixed_size, + _ => None, + } + } +} + +impl fmt::Display for ParsedType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.raw) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_simple_types() { + let t = ParsedType::parse("String"); + assert_eq!(t.base, "String"); + assert!(!t.nullable); + assert!(!t.low_cardinality); + + let t = ParsedType::parse("Int64"); + assert_eq!(t.base, "Int64"); + assert_eq!(t.category(), "Int"); + + let t = ParsedType::parse("Float64"); + assert_eq!(t.base, "Float64"); + assert_eq!(t.category(), "Float"); + } + + #[test] + fn test_parse_nullable() { + let t = ParsedType::parse("Nullable(String)"); + assert_eq!(t.base, "String"); + assert!(t.nullable); + assert!(!t.low_cardinality); + } + + #[test] + fn test_parse_low_cardinality() { + let t = ParsedType::parse("LowCardinality(String)"); + assert_eq!(t.base, "String"); + assert!(!t.nullable); + assert!(t.low_cardinality); + } + + #[test] + fn test_parse_nullable_low_cardinality() { + let t = ParsedType::parse("LowCardinality(Nullable(String))"); + assert_eq!(t.base, "String"); + assert!(t.nullable); + assert!(t.low_cardinality); + } + + #[test] + fn test_parse_array() { + let t = ParsedType::parse("Array(Int64)"); + assert_eq!(t.base, "Array"); + assert!(t.array_element.is_some()); + let elem = t.array_element.as_ref().unwrap(); + assert_eq!(elem.base, "Int64"); + } + + #[test] + fn test_parse_map() { + let t = ParsedType::parse("Map(String, Int64)"); + assert_eq!(t.base, "Map"); + assert!(t.map_types.is_some()); + let (key, value) = t.map_types.as_ref().unwrap(); + assert_eq!(key.base, "String"); + assert_eq!(value.base, "Int64"); + } + + #[test] + fn test_parse_datetime64() { + let t = ParsedType::parse("DateTime64(3)"); + assert_eq!(t.base, "DateTime64"); + assert_eq!(t.precision, Some(3)); + assert!(t.timezone.is_none()); + + let t = ParsedType::parse("DateTime64(6, 'UTC')"); + assert_eq!(t.base, "DateTime64"); + assert_eq!(t.precision, Some(6)); + assert_eq!(t.timezone, Some("UTC".to_string())); + } + + #[test] + fn test_parse_fixed_string() { + let t = ParsedType::parse("FixedString(32)"); + assert_eq!(t.base, "FixedString"); + assert_eq!(t.fixed_size, Some(32)); + } + + #[test] + fn test_parse_decimal() { + let t = ParsedType::parse("Decimal(18, 6)"); + assert_eq!(t.base, "Decimal"); + assert_eq!(t.precision, Some(18)); + assert_eq!(t.scale, Some(6)); + + let t = ParsedType::parse("Decimal64(4)"); + assert_eq!(t.base, "Decimal64"); + assert_eq!(t.scale, Some(4)); + } + + #[test] + fn test_parse_enum() { + let t = ParsedType::parse("Enum8('a' = 1, 'b' = 2)"); + assert_eq!(t.base, "Enum8"); + + let t = ParsedType::parse("Enum16('x' = 100)"); + assert_eq!(t.base, "Enum16"); + } + + #[test] + fn test_categories() { + assert_eq!(ParsedType::parse("String").category(), "String"); + assert_eq!(ParsedType::parse("Int64").category(), "Int"); + assert_eq!(ParsedType::parse("UInt32").category(), "UInt"); + assert_eq!(ParsedType::parse("Float64").category(), "Float"); + assert_eq!(ParsedType::parse("Bool").category(), "Bool"); + assert_eq!(ParsedType::parse("DateTime").category(), "DateTime"); + assert_eq!(ParsedType::parse("UUID").category(), "UUID"); + assert_eq!(ParsedType::parse("IPv4").category(), "IPv4"); + assert_eq!(ParsedType::parse("JSON").category(), "JSON"); + assert_eq!(ParsedType::parse("SomeNewType").category(), "String"); + } + + #[test] + fn test_is_helpers() { + assert!(ParsedType::parse("Int64").is_numeric()); + assert!(ParsedType::parse("Float64").is_numeric()); + assert!(!ParsedType::parse("String").is_numeric()); + + assert!(ParsedType::parse("String").is_string()); + assert!(ParsedType::parse("FixedString(10)").is_string()); + + assert!(ParsedType::parse("DateTime").is_datetime()); + assert!(ParsedType::parse("DateTime64(3)").is_datetime()); + assert!(ParsedType::parse("Date").is_datetime()); + + assert!(ParsedType::parse("IPv4").is_ip()); + assert!(ParsedType::parse("IPv6").is_ip()); + } + + #[test] + fn test_fixed_byte_size() { + assert_eq!(ParsedType::parse("UInt8").fixed_byte_size(), Some(1)); + assert_eq!(ParsedType::parse("Int32").fixed_byte_size(), Some(4)); + assert_eq!(ParsedType::parse("Float64").fixed_byte_size(), Some(8)); + assert_eq!(ParsedType::parse("UUID").fixed_byte_size(), Some(16)); + assert_eq!(ParsedType::parse("Int256").fixed_byte_size(), Some(32)); + assert_eq!( + ParsedType::parse("FixedString(32)").fixed_byte_size(), + Some(32) + ); + assert_eq!(ParsedType::parse("String").fixed_byte_size(), None); + assert_eq!(ParsedType::parse("Array(Int64)").fixed_byte_size(), None); + } + + #[test] + fn test_nested_array_map() { + let t = ParsedType::parse("Array(Nullable(String))"); + assert_eq!(t.base, "Array"); + let elem = t.array_element.as_ref().unwrap(); + assert_eq!(elem.base, "String"); + assert!(elem.nullable); + + let t = ParsedType::parse("Map(String, Array(UInt64))"); + assert_eq!(t.base, "Map"); + let (k, v) = t.map_types.as_ref().unwrap(); + assert_eq!(k.base, "String"); + assert_eq!(v.base, "Array"); + } + + #[test] + fn test_display() { + let t = ParsedType::parse("LowCardinality(Nullable(String))"); + assert_eq!(t.to_string(), "LowCardinality(Nullable(String))"); + } +} diff --git a/src/dynamic/schema.rs b/src/dynamic/schema.rs new file mode 100644 index 00000000..cabd82a5 --- /dev/null +++ b/src/dynamic/schema.rs @@ -0,0 +1,321 @@ +//! Schema reflection for dynamic inserts. +//! +//! Fetches column definitions from `system.columns` and caches them with TTL. +//! The schema drives runtime RowBinary encoding — each column's [`ParsedType`] +//! determines how `serde_json::Value` is converted to binary. +//! +//! # Usage +//! +//! ```rust,ignore +//! use clickhouse::dynamic::schema::{fetch_dynamic_schema, DynamicSchemaCache}; +//! +//! let cache = DynamicSchemaCache::new(Duration::from_secs(300)); +//! let schema = fetch_dynamic_schema(&client, "mydb", "mytable").await?; +//! cache.insert("mydb.mytable", schema); +//! ``` + +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; +use std::time::{Duration, Instant}; + +use super::error::DynamicError; +use super::parsed_type::ParsedType; + +/// Column definition from `system.columns`. +#[derive(Debug, Clone)] +pub struct ColumnDef { + /// Column name. + pub name: String, + /// Raw type string from ClickHouse (e.g. "LowCardinality(Nullable(String))"). + pub raw_type: String, + /// Parsed type with full structure. + pub parsed_type: ParsedType, + /// Default kind: "", "DEFAULT", "MATERIALIZED", "ALIAS", "EPHEMERAL". + pub default_kind: String, + /// Whether this column can be omitted from INSERT (has a server-side default). + pub has_default: bool, +} + +/// Schema for a single table — ordered list of column definitions. +#[derive(Debug, Clone)] +pub struct DynamicSchema { + /// Fully qualified table name (database.table). + pub table: String, + /// Columns in position order. + pub columns: Vec, + /// Lookup by column name for O(1) access during encoding. + column_index: HashMap, +} + +impl DynamicSchema { + /// Build from a list of column definitions. + pub fn from_columns(table: &str, columns: Vec) -> Self { + let column_index = columns + .iter() + .enumerate() + .map(|(i, c)| (c.name.clone(), i)) + .collect(); + Self { + table: table.to_string(), + columns, + column_index, + } + } + + /// Look up a column by name. + pub fn column(&self, name: &str) -> Option<&ColumnDef> { + self.column_index.get(name).map(|&i| &self.columns[i]) + } + + /// Columns that MUST appear in INSERT (no server-side default). + pub fn required_columns(&self) -> impl Iterator { + self.columns.iter().filter(|c| !c.has_default) + } + + /// Columns that CAN be omitted (have DEFAULT/MATERIALIZED/ALIAS). + pub fn optional_columns(&self) -> impl Iterator { + self.columns.iter().filter(|c| c.has_default) + } + + /// Number of columns. + pub fn len(&self) -> usize { + self.columns.len() + } + + /// Whether the schema has no columns. + pub fn is_empty(&self) -> bool { + self.columns.is_empty() + } +} + +// --------------------------------------------------------------------------- +// Schema Cache +// --------------------------------------------------------------------------- + +/// TTL-based schema cache with invalidation. +/// +/// Thread-safe via `RwLock`. Designed to be shared across insert instances +/// via `Arc`. +pub struct DynamicSchemaCache { + inner: RwLock>, + ttl: Duration, +} + +struct CacheEntry { + schema: DynamicSchema, + fetched_at: Instant, +} + +impl DynamicSchemaCache { + /// Create a new cache wrapped in `Arc`. + pub fn new(ttl: Duration) -> Arc { + Arc::new(Self { + inner: RwLock::new(HashMap::new()), + ttl, + }) + } + + /// Get cached schema if not expired. + pub fn get(&self, table: &str) -> Option { + let guard = self.inner.read().ok()?; + guard.get(table).and_then(|e| { + if e.fetched_at.elapsed() < self.ttl { + Some(e.schema.clone()) + } else { + None + } + }) + } + + /// Insert or refresh a schema entry. + pub fn insert(&self, table: &str, schema: DynamicSchema) { + if let Ok(mut guard) = self.inner.write() { + guard.insert( + table.to_string(), + CacheEntry { + schema, + fetched_at: Instant::now(), + }, + ); + } + } + + /// Invalidate a single table (forces re-fetch on next access). + pub fn invalidate(&self, table: &str) { + if let Ok(mut guard) = self.inner.write() { + guard.remove(table); + } + } + + /// Invalidate all cached schemas. + pub fn invalidate_all(&self) { + if let Ok(mut guard) = self.inner.write() { + guard.clear(); + } + } +} + +impl std::fmt::Debug for DynamicSchemaCache { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let count = self.inner.read().map(|g| g.len()).unwrap_or(0); + f.debug_struct("DynamicSchemaCache") + .field("ttl", &self.ttl) + .field("entries", &count) + .finish() + } +} + +// --------------------------------------------------------------------------- +// Schema fetch +// --------------------------------------------------------------------------- + +/// Fetch table schema from `system.columns` via the HTTP client. +/// +/// Parses each column's type string into a full [`ParsedType`]. +pub async fn fetch_dynamic_schema( + client: &crate::Client, + database: &str, + table: &str, +) -> Result { + let full_table = format!("{database}.{table}"); + + // Build query using the crate's SQL escaping + let mut sql = + String::from("SELECT name, type, default_kind FROM system.columns WHERE database = "); + crate::sql::escape::string(database, &mut sql).map_err(|e| DynamicError::SchemaFetch { + table: full_table.clone(), + source: crate::error::Error::Custom(e.to_string()), + })?; + sql.push_str(" AND table = "); + crate::sql::escape::string(table, &mut sql).map_err(|e| DynamicError::SchemaFetch { + table: full_table.clone(), + source: crate::error::Error::Custom(e.to_string()), + })?; + sql.push_str(" ORDER BY position"); + + // Fetch as positional tuples to avoid derive(Row) macro issues inside the crate + let mut cursor = client + .query(&sql) + .fetch::<(String, String, String)>() + .map_err(|e| DynamicError::SchemaFetch { + table: full_table.clone(), + source: e, + })?; + + let mut columns = Vec::new(); + while let Some((name, col_type, default_kind)) = + cursor.next().await.map_err(|e| DynamicError::SchemaFetch { + table: full_table.clone(), + source: e, + })? + { + let parsed_type = ParsedType::parse(&col_type); + let has_default = !default_kind.is_empty(); + columns.push(ColumnDef { + name, + raw_type: col_type, + parsed_type, + default_kind, + has_default, + }); + } + + if columns.is_empty() { + return Err(DynamicError::EmptySchema { table: full_table }); + } + + Ok(DynamicSchema::from_columns(&full_table, columns)) +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + + fn make_col(name: &str, type_str: &str, default_kind: &str) -> ColumnDef { + ColumnDef { + name: name.to_string(), + raw_type: type_str.to_string(), + parsed_type: ParsedType::parse(type_str), + default_kind: default_kind.to_string(), + has_default: !default_kind.is_empty(), + } + } + + #[test] + fn test_dynamic_schema_basic() { + let schema = DynamicSchema::from_columns( + "db.test", + vec![ + make_col("id", "UInt64", ""), + make_col("name", "String", ""), + make_col("created_at", "DateTime64(3)", "DEFAULT"), + ], + ); + + assert_eq!(schema.len(), 3); + assert!(!schema.is_empty()); + assert!(schema.column("id").is_some()); + assert!(schema.column("missing").is_none()); + assert_eq!(schema.required_columns().count(), 2); + assert_eq!(schema.optional_columns().count(), 1); + } + + #[test] + fn test_column_def_has_default() { + let col = make_col("ts", "DateTime64(3)", "DEFAULT"); + assert!(col.has_default); + + let col = make_col("id", "UInt64", ""); + assert!(!col.has_default); + + let col = make_col("mv", "String", "MATERIALIZED"); + assert!(col.has_default); + } + + #[test] + fn test_schema_cache_basic() { + let cache = DynamicSchemaCache::new(Duration::from_secs(300)); + let schema = DynamicSchema::from_columns("db.test", vec![make_col("id", "UInt64", "")]); + + assert!(cache.get("db.test").is_none()); + cache.insert("db.test", schema.clone()); + assert!(cache.get("db.test").is_some()); + + cache.invalidate("db.test"); + assert!(cache.get("db.test").is_none()); + } + + #[test] + fn test_schema_cache_ttl() { + let cache = DynamicSchemaCache::new(Duration::from_millis(1)); + let schema = DynamicSchema::from_columns("db.test", vec![make_col("id", "UInt64", "")]); + + cache.insert("db.test", schema); + // Immediately should still be cached + assert!(cache.get("db.test").is_some()); + + // After TTL expires + std::thread::sleep(Duration::from_millis(10)); + assert!(cache.get("db.test").is_none()); + } + + #[test] + fn test_schema_cache_invalidate_all() { + let cache = DynamicSchemaCache::new(Duration::from_secs(300)); + let schema1 = DynamicSchema::from_columns("db.t1", vec![make_col("id", "UInt64", "")]); + let schema2 = DynamicSchema::from_columns("db.t2", vec![make_col("id", "UInt64", "")]); + + cache.insert("db.t1", schema1); + cache.insert("db.t2", schema2); + assert!(cache.get("db.t1").is_some()); + assert!(cache.get("db.t2").is_some()); + + cache.invalidate_all(); + assert!(cache.get("db.t1").is_none()); + assert!(cache.get("db.t2").is_none()); + } +} diff --git a/src/lib.rs b/src/lib.rs index 23892f15..feb35545 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,6 +22,10 @@ pub mod insert; pub mod insert_formatted; #[cfg(feature = "inserter")] pub mod inserter; +#[cfg(feature = "async-inserter")] +pub mod async_inserter; +#[cfg(feature = "batcher")] +pub mod batcher; pub mod query; pub mod serde; pub mod sql; @@ -43,6 +47,11 @@ mod rowbinary; #[cfg(feature = "inserter")] mod ticks; +#[cfg(feature = "native-transport")] +pub mod native; + +pub mod dynamic; + /// A client containing HTTP pool. /// /// ### Cloning behavior @@ -64,6 +73,7 @@ pub struct Client { products_info: Vec, validation: bool, insert_metadata_cache: Arc, + pub(crate) dynamic_schema_cache: Arc, #[cfg(feature = "test-util")] mocked: bool, @@ -129,6 +139,9 @@ impl Client { products_info: Vec::default(), validation: true, insert_metadata_cache: Arc::new(InsertMetadataCache::default()), + dynamic_schema_cache: dynamic::DynamicSchemaCache::new( + std::time::Duration::from_secs(300), + ), #[cfg(feature = "test-util")] mocked: false, } @@ -455,6 +468,57 @@ impl Client { insert_formatted::InsertFormatted::new(self, sql.into()) } + /// Start a dynamic INSERT for a table with runtime schema. + /// + /// Fetches the schema from `system.columns` (cached with TTL) and encodes + /// `Map` to RowBinary. As simple as JSONEachRow to use, but + /// ClickHouse skips JSON parsing entirely — significant CPU savings on the + /// cluster at scale. + /// + /// # Example + /// + /// ```rust,ignore + /// let mut insert = client.dynamic_insert("mydb", "mytable"); + /// insert.write_map(&row).await?; + /// insert.write_map(&row2).await?; + /// let rows_written = insert.end().await?; + /// ``` + pub fn dynamic_insert( + &self, + database: &str, + table: &str, + ) -> dynamic::insert::DynamicInsert { + dynamic::insert::DynamicInsert::new( + self.clone(), + database.to_string(), + table.to_string(), + self.dynamic_schema_cache.clone(), + ) + } + + /// Start an async auto-flushing dynamic batcher for a table. + /// + /// Same as [`dynamic_insert`][Self::dynamic_insert] but with a background + /// task that auto-flushes on row count and time thresholds. Multiple tasks + /// can write concurrently via [`DynamicBatcherHandle`][dynamic::DynamicBatcherHandle]. + /// + /// # Example + /// + /// ```rust,ignore + /// let batcher = client.dynamic_batcher("mydb", "mytable", Default::default()); + /// let handle = batcher.handle(); + /// handle.write_map(row).await?; + /// batcher.end().await?; + /// ``` + pub fn dynamic_batcher( + &self, + database: &str, + table: &str, + config: dynamic::DynamicBatchConfig, + ) -> dynamic::DynamicBatcher { + dynamic::DynamicBatcher::new(self, database, table, config) + } + /// Starts a new SELECT/DDL query. pub fn query(&self, query: &str) -> query::Query { query::Query::new(self, query) diff --git a/src/native/async_inserter.rs b/src/native/async_inserter.rs new file mode 100644 index 00000000..5177ed4e --- /dev/null +++ b/src/native/async_inserter.rs @@ -0,0 +1,302 @@ +//! Concurrent, auto-flushing inserter with background task (native TCP transport). +//! +//! [`AsyncNativeInserter`] is the native TCP equivalent of +//! [`crate::async_inserter::AsyncInserter`] (HTTP). It wraps +//! [`NativeInserter`] in a background tokio task with an MPSC channel, +//! providing concurrent writes, backpressure, and automatic periodic flushing. +//! +//! Ported from the HyperI DFE Loader project (`dfe-loader/src/buffer/`). +//! +//! # Architecture +//! +//! ```text +//! ┌─ Task A ──┐ ┌─ Task B ──┐ ┌─ Task C ──┐ +//! │ tx.send() │ │ tx.send() │ │ tx.send() │ +//! └─────┬─────┘ └─────┬─────┘ └─────┬─────┘ +//! └───────────────┴───────────────┘ +//! │ +//! bounded mpsc channel +//! │ +//! ┌───────────▼────────────┐ +//! │ Background Task │ +//! │ │ +//! │ select! { │ +//! │ cmd = rx.recv() │ +//! │ _ = interval.tick() │ +//! │ } │ +//! │ │ +//! │ serialize → buffer │ +//! │ check limits → flush │ +//! └──────────┬─────────────┘ +//! │ native TCP +//! ▼ +//! ClickHouse :9000 +//! ``` + +use std::time::Duration; + +use tokio::sync::{mpsc, oneshot}; + +use crate::error::Result; +use crate::native::client::NativeClient; +use crate::native::inserter::{NativeInserter, Quantities}; +use crate::row::{RowOwned, RowWrite}; + +const DEFAULT_CHANNEL_CAPACITY: usize = 8192; + +// --------------------------------------------------------------------------- +// Commands +// --------------------------------------------------------------------------- + +enum Command { + Write(T, oneshot::Sender>), + Flush(oneshot::Sender>), + End(oneshot::Sender>), +} + +// --------------------------------------------------------------------------- +// Configuration +// --------------------------------------------------------------------------- + +/// Configuration for [`AsyncNativeInserter`]. +/// +/// Same defaults as [`crate::async_inserter::AsyncInserterConfig`]. +#[derive(Debug, Clone)] +pub struct AsyncNativeInserterConfig { + /// 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. + pub max_period: Option, + /// Bounded channel capacity. Default: `8192`. + pub channel_capacity: usize, +} + +impl Default for AsyncNativeInserterConfig { + fn default() -> Self { + Self { + max_rows: 100_000, + max_bytes: 10 * 1024 * 1024, + max_period: Some(Duration::from_secs(5)), + channel_capacity: DEFAULT_CHANNEL_CAPACITY, + } + } +} + +impl AsyncNativeInserterConfig { + /// 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. + pub fn without_period(mut self) -> Self { + self.max_period = None; + self + } + + /// Override the bounded channel capacity. + pub fn with_channel_capacity(mut self, cap: usize) -> Self { + self.channel_capacity = cap; + self + } +} + +// --------------------------------------------------------------------------- +// AsyncNativeInserter — native TCP transport +// --------------------------------------------------------------------------- + +/// Concurrent, auto-flushing inserter for a single ClickHouse table (native TCP). +/// +/// This is the native transport equivalent of +/// [`AsyncInserter`][crate::async_inserter::AsyncInserter]. +/// +/// - Accepts `&self` on [`write`][Self::write] and [`flush`][Self::flush]. +/// - Moves serialisation and network I/O to a background tokio task. +/// - Flushes automatically when row/byte/period limits are reached. +/// - Provides backpressure via a bounded MPSC channel. +/// +/// Ported from HyperI DFE Loader's per-table buffer + orchestrator pattern. +pub struct AsyncNativeInserter { + tx: mpsc::Sender>, + handle: tokio::task::JoinHandle<()>, +} + +/// A cheap, clonable handle for writing rows to an [`AsyncNativeInserter`]. +#[derive(Clone)] +pub struct AsyncNativeInserterHandle { + tx: mpsc::Sender>, +} + +fn channel_closed_err() -> crate::error::Error { + crate::error::Error::Custom("AsyncNativeInserter background task gone".into()) +} + +impl AsyncNativeInserter +where + T: RowOwned + RowWrite + Send + Sync + 'static, +{ + /// Create a new `AsyncNativeInserter` for `table` using `config` thresholds. + /// + /// Spawns a background tokio task immediately. + pub fn new(client: &NativeClient, table: &str, config: AsyncNativeInserterConfig) -> Self { + let (tx, rx) = mpsc::channel(config.channel_capacity); + + let inserter = client + .inserter::(table) + .with_max_rows(config.max_rows) + .with_max_bytes(config.max_bytes) + .with_period(config.max_period); + + let period = config.max_period; + let handle = tokio::spawn(background_task(inserter, rx, period)); + + Self { tx, handle } + } + + /// Obtain a cheap, clonable write handle. + pub fn handle(&self) -> AsyncNativeInserterHandle { + AsyncNativeInserterHandle { + tx: self.tx.clone(), + } + } + + /// Serialize and buffer a row. + /// + /// Blocks (asynchronously) if the channel is full (backpressure). + pub async fn write(&self, row: T) -> Result<()> { + let (resp_tx, resp_rx) = oneshot::channel(); + self.tx + .send(Command::Write(row, resp_tx)) + .await + .map_err(|_| channel_closed_err())?; + resp_rx.await.map_err(|_| channel_closed_err())? + } + + /// Force-flush all buffered rows to ClickHouse. + pub async fn flush(&self) -> Result { + let (resp_tx, resp_rx) = oneshot::channel(); + self.tx + .send(Command::Flush(resp_tx)) + .await + .map_err(|_| channel_closed_err())?; + resp_rx.await.map_err(|_| channel_closed_err())? + } + + /// Graceful shutdown: flush remaining rows, end the current INSERT, + /// and stop the background task. + pub async fn end(self) -> Result { + let (resp_tx, resp_rx) = oneshot::channel(); + if self.tx.send(Command::End(resp_tx)).await.is_err() { + return Ok(Quantities::ZERO); + } + drop(self.tx); + let result = resp_rx.await.map_err(|_| channel_closed_err())?; + let _ = self.handle.await; + result + } +} + +impl AsyncNativeInserterHandle +where + T: RowOwned + RowWrite + Send + Sync + 'static, +{ + /// Serialize and buffer a row (same as [`AsyncNativeInserter::write`]). + pub async fn write(&self, row: T) -> Result<()> { + let (resp_tx, resp_rx) = oneshot::channel(); + self.tx + .send(Command::Write(row, resp_tx)) + .await + .map_err(|_| channel_closed_err())?; + resp_rx.await.map_err(|_| channel_closed_err())? + } + + /// Force-flush all buffered rows to ClickHouse. + pub async fn flush(&self) -> Result { + let (resp_tx, resp_rx) = oneshot::channel(); + self.tx + .send(Command::Flush(resp_tx)) + .await + .map_err(|_| channel_closed_err())?; + resp_rx.await.map_err(|_| channel_closed_err())? + } +} + +// --------------------------------------------------------------------------- +// Background task +// --------------------------------------------------------------------------- + +async fn background_task( + mut inserter: NativeInserter, + mut rx: mpsc::Receiver>, + period: Option, +) where + T: RowOwned + RowWrite + Send + Sync + 'static, +{ + let mut interval = period.map(|p| { + let mut iv = tokio::time::interval(tokio::time::Duration::from(p)); + iv.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + iv + }); + + // Skip the immediate first tick. + if let Some(ref mut iv) = interval { + iv.tick().await; + } + + loop { + let tick_fut = async { + match interval { + Some(ref mut iv) => iv.tick().await, + None => std::future::pending().await, + } + }; + + tokio::select! { + biased; + + cmd = rx.recv() => { + match cmd { + Some(Command::Write(row, resp)) => { + let result = inserter.write(&row).await; + if result.is_ok() { + let _ = inserter.commit().await; + } + let _ = resp.send(result); + } + Some(Command::Flush(resp)) => { + let _ = resp.send(inserter.force_commit().await); + } + Some(Command::End(resp)) => { + let _ = resp.send(inserter.end().await); + return; + } + None => { + let _ = inserter.end().await; + return; + } + } + } + + _ = tick_fut => { + let _ = inserter.commit().await; + } + } + } +} 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.rs b/src/native/client.rs new file mode 100644 index 00000000..d21b409b --- /dev/null +++ b/src/native/client.rs @@ -0,0 +1,420 @@ +//! 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; + +use crate::error::{Error, Result}; +use crate::native::insert::NativeInsert; +use crate::native::inserter::NativeInserter; +use crate::native::pool::{NativePool, PoolConfig, PooledConnection, build_pool}; +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, + /// 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, + database, + username, + password, + compression, + schema_cache: NativeSchemaCache::new(300), + settings: Arc::new(settings), + pool_size: DEFAULT_POOL_SIZE, + 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 + /// + /// 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.rebuild_pool(); + self + } + + /// Set the database name. + #[must_use] + pub fn with_database(mut self, database: impl Into) -> Self { + self.database = database.into(); + self.rebuild_pool(); + self + } + + /// Set the username. + #[must_use] + pub fn with_user(mut self, user: impl Into) -> Self { + self.username = user.into(); + self.rebuild_pool(); + self + } + + /// Set the password. + #[must_use] + pub fn with_password(mut self, password: impl Into) -> Self { + self.password = password.into(); + self.rebuild_pool(); + self + } + + /// Enable LZ4 compression for query data. + #[must_use] + pub fn with_lz4(mut self) -> Self { + self.compression = NativeCompressionMethod::Lz4; + self.rebuild_pool(); + 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.rebuild_pool(); + 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.rebuild_pool(); + 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. + pub(crate) async fn acquire(&self) -> Result { + 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. + 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/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/columns.rs b/src/native/columns.rs new file mode 100644 index 00000000..2e9238f8 --- /dev/null +++ b/src/native/columns.rs @@ -0,0 +1,1494 @@ +//! 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), + 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 + // Point is Tuple(Float64, Float64) in columnar format — not a flat 16-byte blob. + | Self::Point => 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 => { + 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, + 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; + + // 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(); + let mut global: ColumnData = Vec::new(); + + if has_global_dict { + let sz = reader.read_u64_le().await?; + global = read_column(reader, dict_type, sz).await?; + } + + if has_additional_keys { + let sz = reader.read_u64_le().await?; + additional = read_column(reader, dict_type, 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, dict_type, 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; + 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) +} + +/// 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/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/cursor.rs b/src/native/cursor.rs new file mode 100644 index 00000000..cbccfa5b --- /dev/null +++ b/src/native/cursor.rs @@ -0,0 +1,173 @@ +//! 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 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 { + client, + sql, + row_buf: VecDeque::new(), + state: CursorState::NotStarted, + _marker: PhantomData, + } + } + + /// 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 + /// 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/encode.rs b/src/native/encode.rs new file mode 100644 index 00000000..5f42a2c9 --- /dev/null +++ b/src/native/encode.rs @@ -0,0 +1,531 @@ +//! 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), +//! 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}; +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. + 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}'" + )) + })?; + Ok(ColumnSchema { + name: name.clone(), + type_name: type_name.clone(), + col_type, + }) + }) + .collect() + } +} + +/// 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) => { + // 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) => { + // 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/insert.rs b/src/native/insert.rs new file mode 100644 index 00000000..def87039 --- /dev/null +++ b/src/native/insert.rs @@ -0,0 +1,197 @@ +//! `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); + } + } + let result = self + .conn + .as_mut() + .expect("conn must be open") + .finish_insert() + .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<()> { + 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 + } + +} + +impl NativeInsert { + /// 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; + } +} + +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/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/mod.rs b/src/native/mod.rs new file mode 100644 index 00000000..0760827f --- /dev/null +++ b/src/native/mod.rs @@ -0,0 +1,36 @@ +//! 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 async_inserter; +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::async_inserter::{AsyncNativeInserter, AsyncNativeInserterConfig, AsyncNativeInserterHandle}; +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..3dc6c1fa --- /dev/null +++ b/src/native/pool.rs @@ -0,0 +1,115 @@ +//! Connection pool for the native TCP transport. +//! +//! 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. +//! +//! # Discard pattern +//! +//! 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::net::SocketAddr; +use std::ops::{Deref, DerefMut}; + +use deadpool::managed::{self, RecycleError, RecycleResult}; + +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)>, +} + +/// deadpool [`Manager`](managed::Manager) for [`NativeConnection`]. +pub(crate) struct NativeConnectionManager { + config: PoolConfig, +} + +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 + } + + async fn recycle( + &self, + conn: &mut NativeConnection, + _: &managed::Metrics, + ) -> RecycleResult { + if !conn.check_alive() { + return Err(RecycleError::message("connection dead or dirty")); + } + Ok(()) + } +} + +/// 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 the pool. +/// +/// 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 { + inner: managed::Object, +} + +impl PooledConnection { + 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.inner.poisoned = true; + } +} + +impl Deref for PooledConnection { + type Target = NativeConnection; + fn deref(&self) -> &NativeConnection { + &self.inner + } +} + +impl DerefMut for PooledConnection { + fn deref_mut(&mut self) -> &mut NativeConnection { + &mut self.inner + } +} 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/query.rs b/src/native/query.rs new file mode 100644 index 00000000..bdc46c73 --- /dev/null +++ b/src/native/query.rs @@ -0,0 +1,117 @@ +//! 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, + { + 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. + 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, + { + 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/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/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/async_inserter.rs b/tests/it/async_inserter.rs new file mode 100644 index 00000000..be846a01 --- /dev/null +++ b/tests/it/async_inserter.rs @@ -0,0 +1,1209 @@ +use serde::{Deserialize, Serialize}; + +use clickhouse::async_inserter::{AsyncInserter, AsyncInserterConfig}; +use clickhouse::{Client, Row}; + +#[derive(Debug, Clone, 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() +} + +// ═══════════════════════════════════════════════════════════════════════════ +// Happy-path tests +// ═══════════════════════════════════════════════════════════════════════════ + +#[tokio::test] +async fn async_inserter_basic() { + let client = prepare_database!(); + create_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test", + AsyncInserterConfig::default().without_period(), + ); + + for i in 0..100u32 { + inserter + .write(MyRow { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + + assert_eq!(count_rows(&client).await, 100); +} + +#[tokio::test] +async fn async_inserter_flush() { + let client = prepare_database!(); + create_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test", + AsyncInserterConfig::default().without_period(), + ); + + for i in 0..50u32 { + inserter + .write(MyRow { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + let q = inserter.flush().await.unwrap(); + assert_eq!(q.rows, 50); + assert_eq!(count_rows(&client).await, 50); + + for i in 50..100u32 { + inserter + .write(MyRow { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + assert_eq!(count_rows(&client).await, 100); +} + +#[tokio::test] +async fn async_inserter_max_rows() { + let client = prepare_database!(); + create_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test", + AsyncInserterConfig::default() + .with_max_rows(10) + .without_period(), + ); + + for i in 0..35u32 { + inserter + .write(MyRow { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + assert_eq!(count_rows(&client).await, 35); +} + +#[tokio::test] +async fn async_inserter_period_flush() { + let client = prepare_database!(); + create_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test", + AsyncInserterConfig::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 { + inserter + .write(MyRow { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + tokio::time::sleep(tokio::time::Duration::from_millis(600)).await; + + assert_eq!(count_rows(&client).await, 20); + + inserter.end().await.unwrap(); + assert_eq!(count_rows(&client).await, 20); +} + +#[tokio::test] +async fn async_inserter_empty_end() { + let client = prepare_database!(); + create_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test", + AsyncInserterConfig::default().without_period(), + ); + + inserter.end().await.unwrap(); + assert_eq!(count_rows(&client).await, 0); +} + +#[tokio::test] +async fn async_inserter_concurrent_handles() { + let client = prepare_database!(); + create_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test", + AsyncInserterConfig::default().without_period(), + ); + + let mut tasks = Vec::new(); + for chunk_start in (0..100u32).step_by(10) { + let handle = inserter.handle(); + tasks.push(tokio::spawn(async move { + for i in chunk_start..chunk_start + 10 { + handle + .write(MyRow { id: i, data: i.to_string() }) + .await + .unwrap(); + } + })); + } + + for task in tasks { + task.await.unwrap(); + } + + inserter.end().await.unwrap(); + assert_eq!(count_rows(&client).await, 100); +} + +// ═══════════════════════════════════════════════════════════════════════════ +// Edge cases +// ═══════════════════════════════════════════════════════════════════════════ + +/// Writing a single row should work. +#[tokio::test] +async fn async_inserter_single_row() { + let client = prepare_database!(); + create_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test", + AsyncInserterConfig::default().without_period(), + ); + + inserter + .write(MyRow { + id: 42, + data: "hello".into(), + }) + .await + .unwrap(); + + inserter.end().await.unwrap(); + assert_eq!(count_rows(&client).await, 1); +} + +/// Flush on an empty buffer should return zero quantities (not error). +#[tokio::test] +async fn async_inserter_flush_empty() { + let client = prepare_database!(); + create_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test", + AsyncInserterConfig::default().without_period(), + ); + + let q = inserter.flush().await.unwrap(); + assert_eq!(q.rows, 0); + assert_eq!(q.bytes, 0); + + inserter.end().await.unwrap(); + assert_eq!(count_rows(&client).await, 0); +} + +/// Multiple flushes in a row without writes between them. +#[tokio::test] +async fn async_inserter_double_flush() { + let client = prepare_database!(); + create_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test", + AsyncInserterConfig::default().without_period(), + ); + + for i in 0..10u32 { + inserter + .write(MyRow { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + let q1 = inserter.flush().await.unwrap(); + assert_eq!(q1.rows, 10); + + // Second flush with nothing buffered. + let q2 = inserter.flush().await.unwrap(); + assert_eq!(q2.rows, 0); + + inserter.end().await.unwrap(); + assert_eq!(count_rows(&client).await, 10); +} + +/// max_rows=1 should auto-flush after every single row. +#[tokio::test] +async fn async_inserter_max_rows_one() { + let client = prepare_database!(); + create_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test", + AsyncInserterConfig::default() + .with_max_rows(1) + .without_period(), + ); + + for i in 0..5u32 { + inserter + .write(MyRow { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + assert_eq!(count_rows(&client).await, 5); +} + +/// Rows with large string data still round-trip correctly. +#[tokio::test] +async fn async_inserter_large_strings() { + let client = prepare_database!(); + create_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test", + AsyncInserterConfig::default().without_period(), + ); + + let big = "x".repeat(100_000); + for i in 0..3u32 { + inserter + .write(MyRow { + id: i, + data: big.clone(), + }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + + let rows: Vec = client + .query("SELECT id, data FROM test ORDER BY id") + .fetch_all() + .await + .unwrap(); + assert_eq!(rows.len(), 3); + assert_eq!(rows[0].data.len(), 100_000); +} + +/// max_bytes flush threshold triggers when accumulated serialised data is large. +#[tokio::test] +async fn async_inserter_max_bytes_trigger() { + let client = prepare_database!(); + create_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test", + AsyncInserterConfig::default() + .with_max_rows(u64::MAX) + .with_max_bytes(100) // very small — should trigger after a few rows + .without_period(), + ); + + for i in 0..20u32 { + inserter + .write(MyRow { + id: i, + data: "some payload data here".into(), + }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + assert_eq!(count_rows(&client).await, 20); +} + +/// Small channel capacity (1) forces extreme backpressure but should still work. +#[tokio::test] +async fn async_inserter_tiny_channel() { + let client = prepare_database!(); + create_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test", + AsyncInserterConfig::default() + .with_channel_capacity(1) + .without_period(), + ); + + for i in 0..20u32 { + inserter + .write(MyRow { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + assert_eq!(count_rows(&client).await, 20); +} + +// ═══════════════════════════════════════════════════════════════════════════ +// Failure / error propagation tests +// ═══════════════════════════════════════════════════════════════════════════ + +/// Writing to a non-existent table should propagate the ClickHouse error back +/// to the caller (on flush/end, not on write — writes only serialize). +#[tokio::test] +async fn async_inserter_bad_table() { + let client = prepare_database!(); + // Intentionally do NOT create the table. + + let inserter = AsyncInserter::::new( + &client, + "this_table_does_not_exist", + AsyncInserterConfig::default() + .with_max_rows(1) // force flush after one row + .without_period(), + ); + + // write() serialises into the buffer — the error surfaces on the commit + // triggered by max_rows=1. + let result = inserter + .write(MyRow { + id: 1, + data: "x".into(), + }) + .await; + + // The error might surface on write (if commit happens inline) or on end. + if result.is_ok() { + let end_result = inserter.end().await; + // At least end() should report the error. + assert!( + end_result.is_err() || end_result.unwrap().rows == 0, + "expected error or zero rows for non-existent table" + ); + } +} + +/// Handle becomes inert after the inserter is ended — writes should fail. +#[tokio::test] +async fn async_inserter_handle_after_end() { + let client = prepare_database!(); + create_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test", + AsyncInserterConfig::default().without_period(), + ); + + let handle = inserter.handle(); + + inserter.end().await.unwrap(); + + // The background task has stopped — write via handle should fail. + let result = handle + .write(MyRow { + id: 1, + data: "late".into(), + }) + .await; + assert!(result.is_err(), "write after end() should fail"); +} + +/// flush() via handle after inserter is ended should fail. +#[tokio::test] +async fn async_inserter_flush_after_end() { + let client = prepare_database!(); + create_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test", + AsyncInserterConfig::default().without_period(), + ); + + let handle = inserter.handle(); + + inserter.end().await.unwrap(); + + let result = handle.flush().await; + assert!(result.is_err(), "flush after end() should fail"); +} + +// ═══════════════════════════════════════════════════════════════════════════ +// Stress / concurrency tests +// ═══════════════════════════════════════════════════════════════════════════ + +/// Many concurrent writers with small max_rows to stress the flush path. +#[tokio::test] +async fn async_inserter_stress_concurrent() { + let client = prepare_database!(); + create_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test", + AsyncInserterConfig::default() + .with_max_rows(7) // prime number to create odd batch boundaries + .without_period(), + ); + + let mut tasks = Vec::new(); + for task_id in 0..20u32 { + let handle = inserter.handle(); + tasks.push(tokio::spawn(async move { + for j in 0..50u32 { + let id = task_id * 50 + j; + handle + .write(MyRow { + id, + data: format!("task{task_id}_row{j}"), + }) + .await + .unwrap(); + } + })); + } + + for task in tasks { + task.await.unwrap(); + } + + inserter.end().await.unwrap(); + assert_eq!(count_rows(&client).await, 1000); +} + +/// Interleaved writes and flushes from multiple handles. +#[tokio::test] +async fn async_inserter_interleaved_flush() { + let client = prepare_database!(); + create_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test", + AsyncInserterConfig::default().without_period(), + ); + + let h1 = inserter.handle(); + let h2 = inserter.handle(); + + // Writer 1: write 10 rows, then flush. + for i in 0..10u32 { + h1.write(MyRow { id: i, data: "a".into() }).await.unwrap(); + } + h1.flush().await.unwrap(); + + // Writer 2: write 10 rows, then flush. + for i in 10..20u32 { + h2.write(MyRow { id: i, data: "b".into() }).await.unwrap(); + } + h2.flush().await.unwrap(); + + assert_eq!(count_rows(&client).await, 20); + + drop(h1); + drop(h2); + inserter.end().await.unwrap(); + assert_eq!(count_rows(&client).await, 20); +} + +// ═══════════════════════════════════════════════════════════════════════════ +// Large ugly JSON source tests — Filebeat / Winlogbeat payloads +// ═══════════════════════════════════════════════════════════════════════════ +// +// These tests exercise the full insert round-trip with realistic, deeply +// nested JSON blobs that match what Elastic Beat agents produce in the wild. +// They stress: large String values, Unicode, Windows backslash paths, +// embedded newlines, null-heavy payloads, arrays of objects, and mixed types. + +#[derive(Debug, Clone, PartialEq, Eq, Row, Serialize, Deserialize)] +struct LogRow { + ts: u64, + source: String, + json_data: String, +} + +fn filebeat_nginx_json() -> String { + r#"{ + "@timestamp": "2026-03-12T08:14:22.337Z", + "@metadata": { + "beat": "filebeat", + "type": "_doc", + "version": "8.17.0", + "pipeline": "filebeat-8.17.0-nginx-access-pipeline" + }, + "agent": { + "name": "web-prod-03.dc1.example.com", + "type": "filebeat", + "version": "8.17.0", + "ephemeral_id": "a1b2c3d4-e5f6-7890-abcd-ef1234567890", + "id": "deadbeef-cafe-babe-f00d-123456789abc", + "hostname": "web-prod-03.dc1.example.com" + }, + "log": { + "file": { "path": "/var/log/nginx/access.log", "inode": "1234567" }, + "offset": 9823741, + "flags": ["utf-8", "multiline"] + }, + "message": "192.168.1.100 - jean-françois [12/Mar/2026:08:14:22 +0000] \"GET /api/v2/données/résultat?q=名前&page=1&size=50 HTTP/2.0\" 200 13847 \"https://app.example.com/dashboard/über-ansicht\" \"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36\" \"-\" rt=0.042 uct=0.001 uht=0.040 urt=0.041", + "source": { "address": "192.168.1.100", "ip": "192.168.1.100", "geo": null }, + "http": { + "request": { + "method": "GET", + "referrer": "https://app.example.com/dashboard/über-ansicht", + "headers": { + "Accept": "application/json, text/plain, */*", + "Accept-Language": "de-DE,de;q=0.9,en-US;q=0.8,en;q=0.7,ja;q=0.6", + "X-Request-ID": "req_7f8a9b0c-1d2e-3f4a-5b6c-7d8e9f0a1b2c", + "X-Forwarded-For": "10.0.0.1, 172.16.0.1, 192.168.1.100", + "Cookie": "session=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkrDqWFuLUZyYW7Dp29pcyIsImlhdCI6MTUxNjIzOTAyMn0.fake_sig" + } + }, + "response": { + "status_code": 200, + "body": { "bytes": 13847 }, + "headers": { + "Content-Type": "application/json; charset=utf-8", + "X-Cache": "MISS", + "X-Served-By": "backend-pool-2a" + } + }, + "version": "2.0" + }, + "url": { + "original": "/api/v2/données/résultat?q=名前&page=1&size=50", + "path": "/api/v2/données/résultat", + "query": "q=名前&page=1&size=50", + "domain": "app.example.com", + "scheme": "https", + "port": 443 + }, + "nginx": { + "access": { + "upstream": { + "response_time": 0.041, + "connect_time": 0.001, + "header_time": 0.040, + "addr": ["10.0.2.15:8080", "10.0.2.16:8080"], + "status": [200] + }, + "geoip": { + "country_iso_code": "DE", + "city_name": "München", + "location": { "lat": 48.1351, "lon": 11.5820 } + } + } + }, + "ecs": { "version": "8.0.0" }, + "tags": ["nginx", "web", "production", "dc1"], + "fields": { + "environment": "production", + "team": "platform-engineering", + "cost_center": "CC-4242" + }, + "event": { + "dataset": "nginx.access", + "module": "nginx", + "category": ["web"], + "type": ["access"], + "outcome": "success", + "duration": 42000000, + "created": "2026-03-12T08:14:22.380Z", + "ingested": "2026-03-12T08:14:23.001Z" + } +}"#.to_string() +} + +fn winlogbeat_security_json() -> String { + r#"{ + "@timestamp": "2026-03-12T03:47:11.892Z", + "@metadata": { + "beat": "winlogbeat", + "type": "_doc", + "version": "8.17.0" + }, + "agent": { + "name": "DC01.corp.contoso.com", + "type": "winlogbeat", + "version": "8.17.0", + "ephemeral_id": "f1e2d3c4-b5a6-9780-fedc-ba0987654321", + "id": "01234567-89ab-cdef-0123-456789abcdef" + }, + "winlog": { + "channel": "Security", + "provider_name": "Microsoft-Windows-Security-Auditing", + "provider_guid": "{54849625-5478-4994-A5BA-3E3B0328C30D}", + "event_id": 4625, + "version": 0, + "task": "Logon", + "opcode": "Info", + "keywords": ["Audit Failure"], + "record_id": 987654321, + "computer_name": "DC01.corp.contoso.com", + "process": { "pid": 788, "thread": { "id": 4892 } }, + "api": "wineventlog", + "activity_id": "{A1B2C3D4-E5F6-7890-ABCD-EF1234567890}", + "event_data": { + "SubjectUserSid": "S-1-5-18", + "SubjectUserName": "DC01$", + "SubjectDomainName": "CORP", + "SubjectLogonId": "0x3e7", + "TargetUserSid": "S-1-0-0", + "TargetUserName": "администратор", + "TargetDomainName": "CORP", + "Status": "0xc000006d", + "FailureReason": "%%2313", + "SubStatus": "0xc0000064", + "LogonType": "10", + "LogonProcessName": "User32 ", + "AuthenticationPackageName": "Negotiate", + "WorkstationName": "АТАКУЮЩИЙ-ПК", + "TransmittedServices": "-", + "LmPackageName": "-", + "KeyLength": "0", + "ProcessId": "0x0", + "ProcessName": "-", + "IpAddress": "198.51.100.23", + "IpPort": "49832" + } + }, + "event": { + "code": "4625", + "kind": "event", + "provider": "Microsoft-Windows-Security-Auditing", + "action": "logon-failed", + "category": ["authentication"], + "type": ["start"], + "outcome": "failure", + "created": "2026-03-12T03:47:12.100Z", + "ingested": "2026-03-12T03:47:13.250Z", + "severity": 0 + }, + "host": { + "name": "DC01", + "hostname": "DC01.corp.contoso.com", + "os": { + "family": "windows", + "name": "Windows Server 2022", + "version": "10.0.20348.2340", + "build": "20348.2340", + "platform": "windows", + "type": "windows", + "kernel": "10.0.20348.2340 (WinBuild.160101.0800)" + }, + "ip": ["10.0.0.5", "fe80::1234:5678:abcd:ef01"], + "mac": ["00-15-5D-01-02-03"], + "architecture": "x86_64", + "domain": "corp.contoso.com" + }, + "source": { + "ip": "198.51.100.23", + "port": 49832, + "geo": { + "country_iso_code": "RU", + "city_name": "Москва", + "region_name": "Москва", + "location": { "lat": 55.7558, "lon": 37.6173 }, + "timezone": "Europe/Moscow" + } + }, + "user": { + "name": "администратор", + "domain": "CORP", + "id": "S-1-0-0", + "target": { + "name": "администратор", + "domain": "CORP" + } + }, + "message": "An account failed to log on.\n\nSubject:\n\tSecurity ID:\t\tS-1-5-18\n\tAccount Name:\t\tDC01$\n\tAccount Domain:\t\tCORP\n\tLogon ID:\t\t0x3E7\n\nLogon Information:\n\tLogon Type:\t\t10\n\tRestricted Admin Mode:\t-\n\tVirtual Account:\t\tNo\n\tElevated Token:\t\tNo\n\nFailure Information:\n\tFailure Reason:\t\tUnknown user name or bad password.\n\tStatus:\t\t\t0xC000006D\n\tSub Status:\t\t0xC0000064\n\nNew Logon:\n\tSecurity ID:\t\tS-1-0-0\n\tAccount Name:\t\tадминистратор\n\tAccount Domain:\t\tCORP\n\nProcess Information:\n\tCaller Process ID:\t0x0\n\tCaller Process Name:\t-\n\nNetwork Information:\n\tWorkstation Name:\tАТАКУЮЩИЙ-ПК\n\tSource Network Address:\t198.51.100.23\n\tSource Port:\t\t49832", + "related": { + "ip": ["198.51.100.23", "10.0.0.5"], + "user": ["DC01$", "администратор"] + }, + "ecs": { "version": "8.0.0" }, + "tags": ["security", "authentication", "failed-logon", "brute-force-candidate"] +}"#.to_string() +} + +fn filebeat_multiline_java_json() -> String { + r#"{ + "@timestamp": "2026-03-12T14:22:03.001Z", + "@metadata": { "beat": "filebeat", "version": "8.17.0" }, + "agent": { "name": "app-srv-07", "type": "filebeat", "version": "8.17.0" }, + "log": { + "file": { + "path": "C:\\Program Files\\MyApp\\logs\\application-2026-03-12.log", + "inode": "0" + }, + "offset": 482716, + "flags": ["utf-8", "multiline"] + }, + "message": "2026-03-12 14:22:02,999 ERROR [http-nio-8443-exec-42] com.example.api.UserController - Failed to process request for user_id=café-résumé-42\njava.lang.NullPointerException: Cannot invoke \"com.example.model.UserProfile.getDisplayName()\" because the return value of \"com.example.service.UserService.findById(String)\" is null\n\tat com.example.api.UserController.getUserProfile(UserController.java:142)\n\tat com.example.api.UserController$$FastClassBySpringCGLIB$$abc123.invoke()\n\tat org.springframework.cglib.proxy.MethodProxy.invoke(MethodProxy.java:218)\n\tat org.springframework.aop.framework.CglibAopProxy$CglibMethodInvocation.invokeJoinpoint(CglibAopProxy.java:793)\n\tat org.springframework.aop.framework.ReflectiveMethodInvocation.proceed(ReflectiveMethodInvocation.java:163)\n\tat org.springframework.aop.framework.CglibAopProxy$DynamicAdvisedInterceptor.intercept(CglibAopProxy.java:723)\n\tat com.example.api.UserController$$EnhancerBySpringCGLIB$$def456.getUserProfile()\n\tat sun.reflect.NativeMethodAccessorImpl.invoke0(Native Method)\n\tat sun.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:62)\n\tat sun.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43)\n\tat java.lang.reflect.Method.invoke(Method.java:498)\n\tat org.apache.tomcat.util.threads.TaskThread$WrappingRunnable.run(TaskThread.java:61)\n\tat java.lang.Thread.run(Thread.java:750)\nCaused by: org.hibernate.exception.JDBCConnectionException: Unable to acquire JDBC Connection\n\tat org.hibernate.exception.internal.SQLExceptionTypeDelegate.convert(SQLExceptionTypeDelegate.java:48)\n\tat com.zaxxer.hikari.pool.HikariPool.getConnection(HikariPool.java:163)\n\tat com.zaxxer.hikari.pool.HikariPool.getConnection(HikariPool.java:128)\nCaused by: java.sql.SQLTransientConnectionException: HikariPool-1 - Connection is not available, request timed out after 30000ms.\n\tat com.zaxxer.hikari.pool.HikariPool.createTimeoutException(HikariPool.java:695)\n\t... 42 more", + "error": { + "type": "java.lang.NullPointerException", + "message": "Cannot invoke \"com.example.model.UserProfile.getDisplayName()\"", + "stack_trace": "... (see message field for full trace)" + }, + "host": { + "name": "app-srv-07", + "os": { + "family": "windows", + "name": "Windows Server 2019", + "version": "10.0.17763.5329" + }, + "ip": ["10.10.20.7"] + }, + "service": { + "name": "user-api", + "version": "3.14.159-SNAPSHOT", + "environment": "staging", + "node": { "name": "app-srv-07:8443" } + }, + "labels": { + "deployment_id": "deploy-2026-03-12-r42", + "git_sha": "a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2", + "jira_ticket": "PLAT-9876" + }, + "ecs": { "version": "8.0.0" }, + "tags": ["java", "error", "staging", "connection-pool-exhaustion"] +}"#.to_string() +} + +fn winlogbeat_powershell_json() -> String { + r#"{ + "@timestamp": "2026-03-12T01:15:44.203Z", + "@metadata": { "beat": "winlogbeat", "version": "8.17.0" }, + "agent": { "name": "WS-FINANCE-12", "type": "winlogbeat" }, + "winlog": { + "channel": "Microsoft-Windows-PowerShell/Operational", + "provider_name": "Microsoft-Windows-PowerShell", + "event_id": 4104, + "task": "Execute a Remote Command", + "opcode": "On create calls", + "record_id": 55432, + "computer_name": "WS-FINANCE-12.corp.contoso.com", + "process": { "pid": 6328, "thread": { "id": 7204 } }, + "event_data": { + "MessageNumber": "1", + "MessageTotal": "1", + "ScriptBlockText": "function Invoke-Çömpléx_Tàsk {\n param(\n [Parameter(Mandatory=$true)]\n [string]$Tärget,\n [ValidateSet('Réad','Wríte','Éxecute')]\n [string]$Möde = 'Réad'\n )\n \n $encodedCmd = [Convert]::ToBase64String([Text.Encoding]::Unicode.GetBytes($Tärget))\n $résult = @{\n 'Tïmestamp' = (Get-Date -Format 'yyyy-MM-ddTHH:mm:ss.fffZ')\n 'Üser' = $env:USERNAME\n 'Dömain' = $env:USERDOMAIN\n 'Pàth' = \"C:\\Users\\$env:USERNAME\\AppData\\Local\\Temp\\öutput_$(Get-Random).tmp\"\n 'Àrgs' = @($Tärget, $Möde, $encodedCmd)\n 'Nësted' = @{\n 'Dëep1' = @{\n 'Dëep2' = @{\n 'Dëep3' = @{\n 'value' = 'We\\'re testing deep nesting with spëcial chars: <>&\\\"\\'/'\n }\n }\n }\n }\n }\n \n $résult | ConvertTo-Json -Depth 10 | Out-File -FilePath $résult['Pàth'] -Encoding UTF8\n return $résult\n}", + "ScriptBlockId": "b7c8d9e0-f1a2-3b4c-5d6e-7f8a9b0c1d2e", + "Path": "C:\\Users\\jëan-pierré\\Documents\\Scrïpts\\Ïnvoke-Task.ps1" + } + }, + "event": { + "code": "4104", + "kind": "event", + "provider": "Microsoft-Windows-PowerShell", + "category": ["process"], + "type": ["info"], + "outcome": "success" + }, + "host": { + "name": "WS-FINANCE-12", + "hostname": "WS-FINANCE-12.corp.contoso.com", + "os": { + "family": "windows", + "name": "Windows 11 Enterprise", + "version": "10.0.22631.3155", + "build": "22631.3155" + }, + "ip": ["10.20.30.12", "fe80::abcd:ef01:2345:6789"], + "mac": ["00-50-56-AB-CD-EF"] + }, + "user": { + "name": "jëan-pierré", + "domain": "CORP", + "id": "S-1-5-21-1234567890-1234567890-1234567890-5678" + }, + "process": { + "pid": 6328, + "executable": "C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.exe", + "command_line": "powershell.exe -NoProfile -ExecutionPolicy Bypass -File \"C:\\Users\\jëan-pierré\\Documents\\Scrïpts\\Ïnvoke-Task.ps1\"", + "parent": { + "pid": 4120, + "executable": "C:\\Windows\\explorer.exe" + } + }, + "message": "Creating Scriptblock text (1 of 1):\nfunction Invoke-Çömpléx_Tàsk { ... (see ScriptBlockText for full content)", + "related": { + "user": ["jëan-pierré"] + }, + "ecs": { "version": "8.0.0" }, + "tags": ["powershell", "scriptblock", "finance-dept"] +}"#.to_string() +} + +fn filebeat_kubernetes_json() -> String { + r#"{ + "@timestamp": "2026-03-12T19:33:07.445Z", + "@metadata": { "beat": "filebeat", "version": "8.17.0" }, + "agent": { "name": "k8s-node-pool-a-2", "type": "filebeat" }, + "kubernetes": { + "pod": { + "name": "payment-svc-7b8c9d-xq2f4", + "uid": "12345678-abcd-ef01-2345-67890abcdef0", + "ip": "10.244.3.17", + "labels": { + "app_kubernetes_io/name": "payment-svc", + "app_kubernetes_io/version": "2.71.828", + "app_kubernetes_io/component": "api", + "helm_sh/chart": "payment-svc-2.71.828", + "pod-template-hash": "7b8c9d" + }, + "annotations": { + "prometheus_io/scrape": "true", + "prometheus_io/port": "9090", + "vault_hashicorp_com/agent-inject": "true", + "vault_hashicorp_com/role": "payment-svc-prod" + } + }, + "node": { + "name": "k8s-node-pool-a-2", + "hostname": "k8s-node-pool-a-2.cluster.local", + "labels": { + "kubernetes_io/arch": "amd64", + "node_kubernetes_io/instance-type": "m5.2xlarge", + "topology_kubernetes_io/zone": "ap-southeast-2a" + } + }, + "namespace": "payment-prod", + "replicaset": { "name": "payment-svc-7b8c9d" }, + "deployment": { "name": "payment-svc" }, + "container": { + "name": "payment-api", + "image": "harbor.internal/payment/api:2.71.828-deadbeef", + "id": "containerd://abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789" + } + }, + "container": { + "id": "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789", + "image": { "name": "harbor.internal/payment/api:2.71.828-deadbeef" }, + "runtime": "containerd" + }, + "log": { + "file": { + "path": "/var/log/pods/payment-prod_payment-svc-7b8c9d-xq2f4_12345678-abcd-ef01-2345-67890abcdef0/payment-api/0.log" + } + }, + "message": "{\"level\":\"error\",\"ts\":1741804387.445,\"caller\":\"handler/payment.go:287\",\"msg\":\"payment processing failed\",\"trace_id\":\"abc123def456\",\"span_id\":\"789012\",\"request_id\":\"req-ñoño-42\",\"customer_id\":\"cust_Ωmega_∆lpha\",\"amount\":\"¥123,456.78\",\"currency\":\"JPY\",\"gateway_response\":{\"code\":\"DECLINED_INSUFFICIENT_FUNDS\",\"raw\":\"カード残高不足です。別のお支払い方法をお試しください。\",\"retry_after_ms\":null,\"metadata\":{\"issuer_country\":\"JP\",\"card_brand\":\"JCB\",\"last4\":\"4242\",\"3ds_enrolled\":true,\"risk_score\":0.73}},\"stack\":\"goroutine 847 [running]:\\nruntime/debug.Stack()\\n\\t/usr/local/go/src/runtime/debug/stack.go:24 +0x5e\\ngithub.com/example/payment-svc/internal/handler.(*PaymentHandler).ProcessPayment(...)\\n\\t/app/internal/handler/payment.go:287 +0x1a3\\ngithub.com/example/payment-svc/internal/handler.(*PaymentHandler).HandleRequest(...)\\n\\t/app/internal/handler/payment.go:142 +0x892\"}", + "stream": "stderr", + "event": { + "dataset": "kubernetes.container_logs", + "module": "kubernetes" + }, + "ecs": { "version": "8.0.0" }, + "tags": ["kubernetes", "payment", "production", "pci-zone"] +}"#.to_string() +} + +async fn create_log_table(client: &Client) { + client + .query( + "CREATE TABLE test_logs(ts UInt64, source String, json_data String) \ + ENGINE = MergeTree ORDER BY ts", + ) + .execute() + .await + .unwrap(); +} + +async fn count_log_rows(client: &Client) -> u64 { + client + .query("SELECT count() FROM test_logs") + .fetch_one::() + .await + .unwrap() +} + +/// Filebeat nginx access log — Unicode URL params, geo data, nested headers. +#[tokio::test] +async fn async_inserter_filebeat_nginx() { + let client = prepare_database!(); + create_log_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test_logs", + AsyncInserterConfig::default().without_period(), + ); + + let json = filebeat_nginx_json(); + for i in 0..10u64 { + inserter + .write(LogRow { + ts: 1741760062000 + i, + source: "filebeat-nginx".into(), + json_data: json.clone(), + }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + + let rows: Vec = client + .query("SELECT ts, source, json_data FROM test_logs ORDER BY ts") + .fetch_all() + .await + .unwrap(); + assert_eq!(rows.len(), 10); + assert!(rows[0].json_data.contains("jean-françois")); + assert!(rows[0].json_data.contains("名前")); + assert!(rows[0].json_data.contains("über-ansicht")); + assert!(rows[0].json_data.contains("München")); +} + +/// Winlogbeat security 4625 — Cyrillic usernames, failed logon, nested event_data. +#[tokio::test] +async fn async_inserter_winlogbeat_security() { + let client = prepare_database!(); + create_log_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test_logs", + AsyncInserterConfig::default().without_period(), + ); + + let json = winlogbeat_security_json(); + for i in 0..10u64 { + inserter + .write(LogRow { + ts: 1741744031000 + i, + source: "winlogbeat-security".into(), + json_data: json.clone(), + }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + + let rows: Vec = client + .query("SELECT ts, source, json_data FROM test_logs ORDER BY ts") + .fetch_all() + .await + .unwrap(); + assert_eq!(rows.len(), 10); + assert!(rows[0].json_data.contains("администратор")); + assert!(rows[0].json_data.contains("АТАКУЮЩИЙ-ПК")); + assert!(rows[0].json_data.contains("Москва")); + assert!(rows[0].json_data.contains("S-1-5-18")); +} + +/// Filebeat multiline Java stack trace — Windows paths, embedded newlines, deep exception chain. +#[tokio::test] +async fn async_inserter_filebeat_java_stacktrace() { + let client = prepare_database!(); + create_log_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test_logs", + AsyncInserterConfig::default().without_period(), + ); + + let json = filebeat_multiline_java_json(); + for i in 0..5u64 { + inserter + .write(LogRow { + ts: 1741781723000 + i, + source: "filebeat-java".into(), + json_data: json.clone(), + }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + + let rows: Vec = client + .query("SELECT ts, source, json_data FROM test_logs ORDER BY ts") + .fetch_all() + .await + .unwrap(); + assert_eq!(rows.len(), 5); + assert!(rows[0].json_data.contains("NullPointerException")); + assert!(rows[0].json_data.contains("café-résumé-42")); + assert!(rows[0].json_data.contains("C:\\\\Program Files\\\\MyApp")); + assert!(rows[0].json_data.contains("HikariPool")); +} + +/// Winlogbeat PowerShell scriptblock — deeply nested diacritics, Base64, special chars. +#[tokio::test] +async fn async_inserter_winlogbeat_powershell() { + let client = prepare_database!(); + create_log_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test_logs", + AsyncInserterConfig::default().without_period(), + ); + + let json = winlogbeat_powershell_json(); + for i in 0..5u64 { + inserter + .write(LogRow { + ts: 1741742144000 + i, + source: "winlogbeat-powershell".into(), + json_data: json.clone(), + }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + + let rows: Vec = client + .query("SELECT ts, source, json_data FROM test_logs ORDER BY ts") + .fetch_all() + .await + .unwrap(); + assert_eq!(rows.len(), 5); + assert!(rows[0].json_data.contains("Invoke-Çömpléx_Tàsk")); + assert!(rows[0].json_data.contains("jëan-pierré")); + assert!(rows[0].json_data.contains("Scrïpts")); +} + +/// Filebeat Kubernetes container log — JSON-in-JSON, CJK payment errors, Go stack trace. +#[tokio::test] +async fn async_inserter_filebeat_kubernetes() { + let client = prepare_database!(); + create_log_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test_logs", + AsyncInserterConfig::default().without_period(), + ); + + let json = filebeat_kubernetes_json(); + for i in 0..5u64 { + inserter + .write(LogRow { + ts: 1741804387000 + i, + source: "filebeat-k8s".into(), + json_data: json.clone(), + }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + + let rows: Vec = client + .query("SELECT ts, source, json_data FROM test_logs ORDER BY ts") + .fetch_all() + .await + .unwrap(); + assert_eq!(rows.len(), 5); + assert!(rows[0].json_data.contains("payment-svc-7b8c9d-xq2f4")); + assert!(rows[0].json_data.contains("カード残高不足")); + assert!(rows[0].json_data.contains("req-ñoño-42")); + assert!(rows[0].json_data.contains("cust_Ωmega_∆lpha")); + assert!(rows[0].json_data.contains("¥123,456.78")); +} + +/// Mixed Beat sources in a single batch — concurrent handles, one source per handle. +#[tokio::test] +async fn async_inserter_mixed_beats_concurrent() { + let client = prepare_database!(); + create_log_table(&client).await; + + let inserter = AsyncInserter::::new( + &client, + "test_logs", + AsyncInserterConfig::default() + .with_max_rows(15) // force multiple flushes mid-batch + .without_period(), + ); + + let sources: Vec<(&str, String)> = vec![ + ("filebeat-nginx", filebeat_nginx_json()), + ("winlogbeat-security", winlogbeat_security_json()), + ("filebeat-java", filebeat_multiline_java_json()), + ("winlogbeat-powershell", winlogbeat_powershell_json()), + ("filebeat-k8s", filebeat_kubernetes_json()), + ]; + + let mut tasks = Vec::new(); + for (idx, (source, json)) in sources.into_iter().enumerate() { + let handle = inserter.handle(); + let source = source.to_string(); + tasks.push(tokio::spawn(async move { + for j in 0..20u64 { + let ts = (idx as u64) * 1_000_000 + j; + handle + .write(LogRow { + ts, + source: source.clone(), + json_data: json.clone(), + }) + .await + .unwrap(); + } + })); + } + + for task in tasks { + task.await.unwrap(); + } + + inserter.end().await.unwrap(); + + // 5 sources × 20 rows = 100 + assert_eq!(count_log_rows(&client).await, 100); + + // Verify each source is present. + let nginx_count: u64 = client + .query("SELECT count() FROM test_logs WHERE source = 'filebeat-nginx'") + .fetch_one() + .await + .unwrap(); + assert_eq!(nginx_count, 20); + + let security_count: u64 = client + .query("SELECT count() FROM test_logs WHERE source = 'winlogbeat-security'") + .fetch_one() + .await + .unwrap(); + assert_eq!(security_count, 20); +} diff --git a/tests/it/batcher.rs b/tests/it/batcher.rs new file mode 100644 index 00000000..4f466a7d --- /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/cursor_reborrow.rs b/tests/it/cursor_reborrow.rs new file mode 100644 index 00000000..40b73024 --- /dev/null +++ b/tests/it/cursor_reborrow.rs @@ -0,0 +1,252 @@ +// Tests for the unsafe reborrow in RowCursor::poll_next and Next::poll. +// +// These specifically exercise the code paths that previously used +// polonius-the-crab and now use a manual unsafe reborrow. The key +// scenarios are the get-or-retry loop (NotEnoughData -> extend -> retry) +// and borrowed deserialization (T::Value<'_> borrowing from the buffer). + +#![cfg(feature = "test-util")] + +use clickhouse::{Client, Row, test}; +use serde::{Deserialize, Serialize}; + +// -- Mock-based tests (no ClickHouse needed) -------------------------------- + +#[tokio::test] +async fn cursor_single_row() { + #[derive(Debug, Clone, Row, Serialize, Deserialize, PartialEq)] + struct R { + x: u32, + } + + let mock = test::Mock::new(); + let client = Client::default().with_mock(&mock); + mock.add(test::handlers::provide([R { x: 42 }])); + + let mut cursor = client.query("SELECT x").fetch::().unwrap(); + assert_eq!(cursor.next().await.unwrap(), Some(R { x: 42 })); + assert_eq!(cursor.next().await.unwrap(), None); +} + +#[tokio::test] +async fn cursor_multiple_rows() { + // The loop in poll_next is the bit that needs the reborrow. Multiple + // rows means the loop iterates, which exercises extend() after a + // successful deserialisation on the previous iteration. + #[derive(Debug, Clone, Row, Serialize, Deserialize, PartialEq)] + struct R { + id: u64, + data: String, + } + + let mock = test::Mock::new(); + let client = Client::default().with_mock(&mock); + let rows: Vec = (0..100) + .map(|i| R { + id: i, + data: format!("row-{i}"), + }) + .collect(); + mock.add(test::handlers::provide(rows.clone())); + + let mut cursor = client.query("SELECT id, data").fetch::().unwrap(); + let mut got = Vec::new(); + while let Some(row) = cursor.next().await.unwrap() { + got.push(row); + } + assert_eq!(got, rows); +} + +#[tokio::test] +async fn cursor_empty_result() { + #[derive(Debug, Clone, Row, Serialize, Deserialize, PartialEq)] + struct R { + x: u32, + } + + let mock = test::Mock::new(); + let client = Client::default().with_mock(&mock); + mock.add(test::handlers::provide(Vec::::new())); + + let mut cursor = client.query("SELECT x").fetch::().unwrap(); + assert_eq!(cursor.next().await.unwrap(), None); +} + +#[tokio::test] +async fn cursor_fetch_all_and_fetch_one() { + // fetch_all and fetch_one both go through poll_next internally. + #[derive(Debug, Clone, Row, Serialize, Deserialize, PartialEq)] + struct R { + v: String, + } + + let mock = test::Mock::new(); + let client = Client::default().with_mock(&mock); + + let rows = vec![ + R { + v: "aaa".to_string(), + }, + R { + v: "bbb".to_string(), + }, + R { + v: "ccc".to_string(), + }, + ]; + mock.add(test::handlers::provide(rows.clone())); + let got = client + .query("SELECT v") + .fetch_all::() + .await + .unwrap(); + assert_eq!(got, rows); + + mock.add(test::handlers::provide([R { + v: "one".to_string(), + }])); + let got = client + .query("SELECT v") + .fetch_one::() + .await + .unwrap(); + assert_eq!(got, R { + v: "one".to_string(), + }); +} + +// -- Integration tests (need a real ClickHouse) ----------------------------- + +#[tokio::test] +async fn cursor_large_result_spanning_chunks() { + // Large enough to span multiple HTTP response chunks, exercising the + // NotEnoughData -> raw.poll_next -> extend -> retry path in the loop. + // This is the core path the unsafe reborrow protects. + #[derive(Debug, Clone, Row, Serialize, Deserialize, PartialEq)] + struct R { + id: u64, + payload: String, + } + + let client = prepare_database!(); + client + .query( + "CREATE TABLE test (id UInt64, payload String) \ + ENGINE = MergeTree ORDER BY id", + ) + .execute() + .await + .unwrap(); + + // 500 rows with ~200 bytes each = ~100KB, enough to span chunks. + let expected: Vec = (0..500) + .map(|i| R { + id: i, + payload: format!("{i:0>200}"), + }) + .collect(); + + let mut insert = client.insert::("test").await.unwrap(); + for row in &expected { + insert.write(row).await.unwrap(); + } + insert.end().await.unwrap(); + + let mut cursor = client + .query("SELECT id, payload FROM test ORDER BY id") + .fetch::() + .unwrap(); + + let mut got = Vec::new(); + while let Some(row) = cursor.next().await.unwrap() { + got.push(row); + } + assert_eq!(got.len(), expected.len()); + assert_eq!(got, expected); +} + +#[tokio::test] +async fn cursor_borrowed_rows() { + // Borrowed deserialization is the reason the unsafe exists — the + // returned T::Value<'_> borrows from the cursor's internal buffer. + #[derive(Debug, Row, Serialize, Deserialize, PartialEq)] + struct Borrowed<'a> { + id: u64, + data: &'a str, + } + + let client = prepare_database!(); + crate::create_simple_table(&client, "test").await; + + let mut insert = client.insert::>("test").await.unwrap(); + insert + .write(&Borrowed { id: 1, data: "one" }) + .await + .unwrap(); + insert + .write(&Borrowed { + id: 2, + data: "two", + }) + .await + .unwrap(); + insert + .write(&Borrowed { + id: 3, + data: "three", + }) + .await + .unwrap(); + insert.end().await.unwrap(); + + let mut cursor = client + .query("SELECT id, data FROM test ORDER BY id") + .fetch::>() + .unwrap(); + + let row = cursor.next().await.unwrap().unwrap(); + assert_eq!(row, Borrowed { id: 1, data: "one" }); + let row = cursor.next().await.unwrap().unwrap(); + assert_eq!(row, Borrowed { + id: 2, + data: "two", + }); + let row = cursor.next().await.unwrap().unwrap(); + assert_eq!( + row, + Borrowed { + id: 3, + data: "three" + } + ); + assert_eq!(cursor.next().await.unwrap(), None); +} + +#[tokio::test] +async fn cursor_small_block_size() { + // Force ClickHouse to send one row per chunk. This maximises the + // number of extend() calls per row, hammering the reborrow path. + let client = prepare_database!(); + crate::create_simple_table(&client, "test").await; + + let mut insert = client.insert::("test").await.unwrap(); + for i in 0..50 { + insert + .write(&crate::SimpleRow::new(i, format!("val-{i}"))) + .await + .unwrap(); + } + insert.end().await.unwrap(); + + let mut cursor = client + .with_option("max_block_size", "1") + .query("SELECT ?fields FROM test ORDER BY id") + .fetch::() + .unwrap(); + + let mut count = 0u64; + while cursor.next().await.unwrap().is_some() { + count += 1; + } + assert_eq!(count, 50); +} diff --git a/tests/it/dynamic.rs b/tests/it/dynamic.rs new file mode 100644 index 00000000..6c3f5afe --- /dev/null +++ b/tests/it/dynamic.rs @@ -0,0 +1,170 @@ +//! Integration tests for dynamic (schema-driven) inserts. +//! +//! These tests require a running ClickHouse instance. + +use clickhouse::sql::Identifier; +use serde_json::json; + +#[tokio::test] +async fn inserts_simple_types() { + let client = prepare_database!(); + let table = "dynamic_simple"; + + client + .query( + "CREATE TABLE ?(id UInt64, name String, score Float64, active UInt8) \ + ENGINE = MergeTree ORDER BY id", + ) + .with_option("wait_end_of_query", "1") + .bind(Identifier(table)) + .execute() + .await + .unwrap(); + + let mut insert = client.dynamic_insert(&test_database_name!(), table); + insert + .write_map(json!({"id": 1, "name": "alice", "score": 9.5, "active": 1}).as_object().unwrap()) + .await + .unwrap(); + insert + .write_map(json!({"id": 2, "name": "bob", "score": 7.2, "active": 0}).as_object().unwrap()) + .await + .unwrap(); + let rows = insert.end().await.unwrap(); + assert_eq!(rows, 2); + + // Query back + let result = client + .query(&format!("SELECT id, name, score FROM {table} ORDER BY id")) + .fetch_all::<(u64, String, f64)>() + .await + .unwrap(); + + assert_eq!(result.len(), 2); + assert_eq!(result[0], (1, "alice".to_string(), 9.5)); + assert_eq!(result[1], (2, "bob".to_string(), 7.2)); +} + +#[tokio::test] +async fn inserts_nullable_columns() { + let client = prepare_database!(); + let table = "dynamic_nullable"; + + client + .query( + "CREATE TABLE ?(id UInt64, label Nullable(String)) \ + ENGINE = MergeTree ORDER BY id", + ) + .with_option("wait_end_of_query", "1") + .bind(Identifier(table)) + .execute() + .await + .unwrap(); + + let mut insert = client.dynamic_insert(&test_database_name!(), table); + insert + .write_map(json!({"id": 1, "label": "present"}).as_object().unwrap()) + .await + .unwrap(); + insert + .write_map(json!({"id": 2, "label": null}).as_object().unwrap()) + .await + .unwrap(); + insert.end().await.unwrap(); + + let result = client + .query(&format!( + "SELECT id, label FROM {table} ORDER BY id" + )) + .fetch_all::<(u64, Option)>() + .await + .unwrap(); + + assert_eq!(result.len(), 2); + assert_eq!(result[0], (1, Some("present".to_string()))); + assert_eq!(result[1], (2, None)); +} + +#[tokio::test] +async fn skips_columns_with_defaults() { + let client = prepare_database!(); + let table = "dynamic_defaults"; + + client + .query( + "CREATE TABLE ?(\ + id UInt64, \ + name String, \ + created_at DateTime64(3) DEFAULT now64(3)\ + ) ENGINE = MergeTree ORDER BY id", + ) + .with_option("wait_end_of_query", "1") + .bind(Identifier(table)) + .execute() + .await + .unwrap(); + + // Insert without created_at — should use server default + let mut insert = client.dynamic_insert(&test_database_name!(), table); + insert + .write_map(json!({"id": 1, "name": "auto-ts"}).as_object().unwrap()) + .await + .unwrap(); + insert.end().await.unwrap(); + + // Verify row exists and created_at was populated by server + let result = client + .query(&format!( + "SELECT id, name, created_at > 0 as has_ts FROM {table}" + )) + .fetch_all::<(u64, String, u8)>() + .await + .unwrap(); + + assert_eq!(result.len(), 1); + assert_eq!(result[0].0, 1); + assert_eq!(result[0].1, "auto-ts"); + assert_eq!(result[0].2, 1); // created_at was filled +} + +#[tokio::test] +async fn batcher_flushes_on_end() { + let client = prepare_database!(); + let table = "dynamic_batcher"; + + client + .query( + "CREATE TABLE ?(id UInt64, value String) ENGINE = MergeTree ORDER BY id", + ) + .with_option("wait_end_of_query", "1") + .bind(Identifier(table)) + .execute() + .await + .unwrap(); + + let batcher = client.dynamic_batcher( + &test_database_name!(), + table, + clickhouse::dynamic::DynamicBatchConfig { + max_rows: 100, + ..Default::default() + }, + ); + + for i in 0..10u64 { + batcher + .write_map(json!({"id": i, "value": format!("row-{i}")}).as_object().unwrap().clone()) + .await + .unwrap(); + } + + let total = batcher.end().await.unwrap(); + assert_eq!(total, 10); + + let count = client + .query(&format!("SELECT count() FROM {table}")) + .fetch_one::() + .await + .unwrap(); + assert_eq!(count, 10); +} diff --git a/tests/it/main.rs b/tests/it/main.rs index a137381d..4d1c3ec2 100644 --- a/tests/it/main.rs +++ b/tests/it/main.rs @@ -249,7 +249,10 @@ pub(crate) mod decimals { mod chrono; mod cloud_jwt; mod compression; +#[cfg(feature = "native-transport")] +mod native; mod cursor_error; +mod cursor_reborrow; mod cursor_stats; mod fetch_bytes; mod https_errors; @@ -257,6 +260,10 @@ mod insert; mod insert_formatted; #[cfg(feature = "inserter")] mod inserter; +#[cfg(feature = "batcher")] +mod batcher; +#[cfg(feature = "async-inserter")] +mod async_inserter; mod int128; mod int256; mod ip; @@ -272,6 +279,7 @@ mod time; mod user_agent; mod uuid; mod variant; +mod dynamic; #[derive(Clone, Copy, PartialEq, Eq)] enum TestEnv { diff --git a/tests/it/native.rs b/tests/it/native.rs new file mode 100644 index 00000000..6b63c1b0 --- /dev/null +++ b/tests/it/native.rs @@ -0,0 +1,3934 @@ +//! 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"); +} + +/// 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; + + 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 }); +} + +// --------------------------------------------------------------------------- +// 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; + + 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() }); +} + +/// 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; + + 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"); +} + +// --------------------------------------------------------------------------- +// 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"); +} + +// ═══════════════════════════════════════════════════════════════════════════ +// AsyncNativeInserter tests +// ═══════════════════════════════════════════════════════════════════════════ + +#[tokio::test] +async fn native_async_inserter_basic() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct R { + id: u32, + data: String, + } + + let client = prepare_native_database("async_inserter_basic").await; + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, data String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default().without_period(), + ); + + for i in 0..100u32 { + inserter + .write(R { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + + let rows: Vec = client + .query("SELECT id, data FROM t ORDER BY id") + .fetch_all() + .await + .unwrap(); + assert_eq!(rows.len(), 100); + assert_eq!(rows[0].id, 0); + assert_eq!(rows[99].id, 99); +} + +#[tokio::test] +async fn native_async_inserter_flush() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct R { + id: u32, + data: String, + } + + let client = prepare_native_database("async_inserter_flush").await; + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, data String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default().without_period(), + ); + + for i in 0..50u32 { + inserter + .write(R { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + let q = inserter.flush().await.unwrap(); + assert_eq!(q.rows, 50); + + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one() + .await + .unwrap(); + assert_eq!(count, 50); + + for i in 50..100u32 { + inserter + .write(R { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one() + .await + .unwrap(); + assert_eq!(count, 100); +} + +#[tokio::test] +async fn native_async_inserter_concurrent_handles() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct R { + id: u32, + data: String, + } + + let client = prepare_native_database("async_inserter_concurrent").await; + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, data String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default().without_period(), + ); + + let mut tasks = Vec::new(); + for chunk_start in (0..100u32).step_by(10) { + let handle = inserter.handle(); + tasks.push(tokio::spawn(async move { + for i in chunk_start..chunk_start + 10 { + handle + .write(R { id: i, data: i.to_string() }) + .await + .unwrap(); + } + })); + } + + for task in tasks { + task.await.unwrap(); + } + + inserter.end().await.unwrap(); + + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one() + .await + .unwrap(); + assert_eq!(count, 100); +} + +#[tokio::test] +async fn native_async_inserter_empty_end() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct R { + id: u32, + data: String, + } + + let client = prepare_native_database("async_inserter_empty").await; + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, data String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default().without_period(), + ); + + inserter.end().await.unwrap(); + + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one() + .await + .unwrap(); + assert_eq!(count, 0); +} + +// ── Edge cases ─────────────────────────────────────────────────────────── + +/// Writing a single row should work. +#[tokio::test] +async fn native_async_inserter_single_row() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct R { + id: u32, + data: String, + } + + let client = prepare_native_database("async_inserter_single").await; + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, data String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default().without_period(), + ); + + inserter + .write(R { + id: 42, + data: "hello".into(), + }) + .await + .unwrap(); + + inserter.end().await.unwrap(); + + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one() + .await + .unwrap(); + assert_eq!(count, 1); +} + +/// Flush on empty buffer returns zero quantities (not an error). +#[tokio::test] +async fn native_async_inserter_flush_empty() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct R { + id: u32, + data: String, + } + + let client = prepare_native_database("async_inserter_flush_empty").await; + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, data String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default().without_period(), + ); + + let q = inserter.flush().await.unwrap(); + assert_eq!(q.rows, 0); + + inserter.end().await.unwrap(); +} + +/// Multiple flushes in a row without writes in between. +#[tokio::test] +async fn native_async_inserter_double_flush() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct R { + id: u32, + data: String, + } + + let client = prepare_native_database("async_inserter_double_flush").await; + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, data String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default().without_period(), + ); + + for i in 0..10u32 { + inserter + .write(R { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + let q1 = inserter.flush().await.unwrap(); + assert_eq!(q1.rows, 10); + + let q2 = inserter.flush().await.unwrap(); + assert_eq!(q2.rows, 0); + + inserter.end().await.unwrap(); + + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one() + .await + .unwrap(); + assert_eq!(count, 10); +} + +/// max_rows=1 should auto-flush after every single row. +#[tokio::test] +async fn native_async_inserter_max_rows_one() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct R { + id: u32, + data: String, + } + + let client = prepare_native_database("async_inserter_max1").await; + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, data String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default() + .with_max_rows(1) + .without_period(), + ); + + for i in 0..5u32 { + inserter + .write(R { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one() + .await + .unwrap(); + assert_eq!(count, 5); +} + +/// Large string values round-trip correctly. +#[tokio::test] +async fn native_async_inserter_large_strings() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct R { + id: u32, + data: String, + } + + let client = prepare_native_database("async_inserter_large_str").await; + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, data String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default().without_period(), + ); + + let big = "x".repeat(100_000); + for i in 0..3u32 { + inserter + .write(R { + id: i, + data: big.clone(), + }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one() + .await + .unwrap(); + assert_eq!(count, 3); +} + +/// Small channel capacity (1) forces extreme backpressure. +#[tokio::test] +async fn native_async_inserter_tiny_channel() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct R { + id: u32, + data: String, + } + + let client = prepare_native_database("async_inserter_tiny_ch").await; + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, data String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default() + .with_channel_capacity(1) + .without_period(), + ); + + for i in 0..20u32 { + inserter + .write(R { id: i, data: i.to_string() }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one() + .await + .unwrap(); + assert_eq!(count, 20); +} + +// ── Failure / error propagation ────────────────────────────────────────── + +/// Handle becomes inert after the inserter is ended — writes should fail. +#[tokio::test] +async fn native_async_inserter_handle_after_end() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct R { + id: u32, + data: String, + } + + let client = prepare_native_database("async_inserter_after_end").await; + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, data String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default().without_period(), + ); + + let handle = inserter.handle(); + + inserter.end().await.unwrap(); + + let result = handle + .write(R { + id: 1, + data: "late".into(), + }) + .await; + assert!(result.is_err(), "write after end() should fail"); +} + +/// flush() via handle after inserter is ended should fail. +#[tokio::test] +async fn native_async_inserter_flush_after_end() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct R { + id: u32, + data: String, + } + + let client = prepare_native_database("async_inserter_flush_end").await; + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, data String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default().without_period(), + ); + + let handle = inserter.handle(); + + inserter.end().await.unwrap(); + + let result = handle.flush().await; + assert!(result.is_err(), "flush after end() should fail"); +} + +// ── Stress / concurrency ───────────────────────────────────────────────── + +/// Many concurrent writers with small max_rows to stress the flush path. +#[tokio::test] +async fn native_async_inserter_stress_concurrent() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct R { + id: u32, + data: String, + } + + let client = prepare_native_database("async_inserter_stress").await; + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, data String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default() + .with_max_rows(7) // prime number for odd batch boundaries + .without_period(), + ); + + let mut tasks = Vec::new(); + for task_id in 0..20u32 { + let handle = inserter.handle(); + tasks.push(tokio::spawn(async move { + for j in 0..50u32 { + let id = task_id * 50 + j; + handle + .write(R { + id, + data: format!("task{task_id}_row{j}"), + }) + .await + .unwrap(); + } + })); + } + + for task in tasks { + task.await.unwrap(); + } + + inserter.end().await.unwrap(); + + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one() + .await + .unwrap(); + assert_eq!(count, 1000); +} + +/// Interleaved writes and flushes from multiple handles. +#[tokio::test] +async fn native_async_inserter_interleaved_flush() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct R { + id: u32, + data: String, + } + + let client = prepare_native_database("async_inserter_interleave").await; + client + .query(&format!( + "CREATE TABLE t{} (id UInt32, data String) {}", + on_cluster(), + test_engine("id"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default().without_period(), + ); + + let h1 = inserter.handle(); + let h2 = inserter.handle(); + + for i in 0..10u32 { + h1.write(R { id: i, data: "a".into() }).await.unwrap(); + } + h1.flush().await.unwrap(); + + for i in 10..20u32 { + h2.write(R { id: i, data: "b".into() }).await.unwrap(); + } + h2.flush().await.unwrap(); + + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one() + .await + .unwrap(); + assert_eq!(count, 20); + + drop(h1); + drop(h2); + inserter.end().await.unwrap(); +} + +// ═══════════════════════════════════════════════════════════════════════════ +// Large ugly JSON source tests — Filebeat / Winlogbeat payloads (native TCP) +// ═══════════════════════════════════════════════════════════════════════════ +// +// Realistic, deeply nested JSON blobs matching Elastic Beat agent output. +// Stresses: large String values, Unicode (CJK, Cyrillic, diacritics), +// Windows backslash paths, embedded newlines/tabs, null fields, arrays of +// objects, JSON-in-JSON (Kubernetes container logs), and mixed types. + +fn filebeat_nginx_json() -> String { + r#"{ + "@timestamp": "2026-03-12T08:14:22.337Z", + "@metadata": { + "beat": "filebeat", + "type": "_doc", + "version": "8.17.0", + "pipeline": "filebeat-8.17.0-nginx-access-pipeline" + }, + "agent": { + "name": "web-prod-03.dc1.example.com", + "type": "filebeat", + "version": "8.17.0", + "ephemeral_id": "a1b2c3d4-e5f6-7890-abcd-ef1234567890", + "id": "deadbeef-cafe-babe-f00d-123456789abc", + "hostname": "web-prod-03.dc1.example.com" + }, + "log": { + "file": { "path": "/var/log/nginx/access.log", "inode": "1234567" }, + "offset": 9823741, + "flags": ["utf-8", "multiline"] + }, + "message": "192.168.1.100 - jean-françois [12/Mar/2026:08:14:22 +0000] \"GET /api/v2/données/résultat?q=名前&page=1&size=50 HTTP/2.0\" 200 13847 \"https://app.example.com/dashboard/über-ansicht\" \"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36\" \"-\" rt=0.042 uct=0.001 uht=0.040 urt=0.041", + "source": { "address": "192.168.1.100", "ip": "192.168.1.100", "geo": null }, + "http": { + "request": { + "method": "GET", + "referrer": "https://app.example.com/dashboard/über-ansicht", + "headers": { + "Accept": "application/json, text/plain, */*", + "Accept-Language": "de-DE,de;q=0.9,en-US;q=0.8,en;q=0.7,ja;q=0.6", + "X-Request-ID": "req_7f8a9b0c-1d2e-3f4a-5b6c-7d8e9f0a1b2c", + "X-Forwarded-For": "10.0.0.1, 172.16.0.1, 192.168.1.100", + "Cookie": "session=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkrDqWFuLUZyYW7Dp29pcyIsImlhdCI6MTUxNjIzOTAyMn0.fake_sig" + } + }, + "response": { + "status_code": 200, + "body": { "bytes": 13847 }, + "headers": { + "Content-Type": "application/json; charset=utf-8", + "X-Cache": "MISS", + "X-Served-By": "backend-pool-2a" + } + }, + "version": "2.0" + }, + "url": { + "original": "/api/v2/données/résultat?q=名前&page=1&size=50", + "path": "/api/v2/données/résultat", + "query": "q=名前&page=1&size=50", + "domain": "app.example.com", + "scheme": "https", + "port": 443 + }, + "nginx": { + "access": { + "upstream": { + "response_time": 0.041, + "connect_time": 0.001, + "header_time": 0.040, + "addr": ["10.0.2.15:8080", "10.0.2.16:8080"], + "status": [200] + }, + "geoip": { + "country_iso_code": "DE", + "city_name": "München", + "location": { "lat": 48.1351, "lon": 11.5820 } + } + } + }, + "ecs": { "version": "8.0.0" }, + "tags": ["nginx", "web", "production", "dc1"], + "fields": { + "environment": "production", + "team": "platform-engineering", + "cost_center": "CC-4242" + }, + "event": { + "dataset": "nginx.access", + "module": "nginx", + "category": ["web"], + "type": ["access"], + "outcome": "success", + "duration": 42000000, + "created": "2026-03-12T08:14:22.380Z", + "ingested": "2026-03-12T08:14:23.001Z" + } +}"#.to_string() +} + +fn winlogbeat_security_json() -> String { + r#"{ + "@timestamp": "2026-03-12T03:47:11.892Z", + "@metadata": { + "beat": "winlogbeat", + "type": "_doc", + "version": "8.17.0" + }, + "agent": { + "name": "DC01.corp.contoso.com", + "type": "winlogbeat", + "version": "8.17.0", + "ephemeral_id": "f1e2d3c4-b5a6-9780-fedc-ba0987654321", + "id": "01234567-89ab-cdef-0123-456789abcdef" + }, + "winlog": { + "channel": "Security", + "provider_name": "Microsoft-Windows-Security-Auditing", + "provider_guid": "{54849625-5478-4994-A5BA-3E3B0328C30D}", + "event_id": 4625, + "version": 0, + "task": "Logon", + "opcode": "Info", + "keywords": ["Audit Failure"], + "record_id": 987654321, + "computer_name": "DC01.corp.contoso.com", + "process": { "pid": 788, "thread": { "id": 4892 } }, + "api": "wineventlog", + "activity_id": "{A1B2C3D4-E5F6-7890-ABCD-EF1234567890}", + "event_data": { + "SubjectUserSid": "S-1-5-18", + "SubjectUserName": "DC01$", + "SubjectDomainName": "CORP", + "SubjectLogonId": "0x3e7", + "TargetUserSid": "S-1-0-0", + "TargetUserName": "администратор", + "TargetDomainName": "CORP", + "Status": "0xc000006d", + "FailureReason": "%%2313", + "SubStatus": "0xc0000064", + "LogonType": "10", + "LogonProcessName": "User32 ", + "AuthenticationPackageName": "Negotiate", + "WorkstationName": "АТАКУЮЩИЙ-ПК", + "TransmittedServices": "-", + "LmPackageName": "-", + "KeyLength": "0", + "ProcessId": "0x0", + "ProcessName": "-", + "IpAddress": "198.51.100.23", + "IpPort": "49832" + } + }, + "event": { + "code": "4625", + "kind": "event", + "provider": "Microsoft-Windows-Security-Auditing", + "action": "logon-failed", + "category": ["authentication"], + "type": ["start"], + "outcome": "failure", + "created": "2026-03-12T03:47:12.100Z", + "ingested": "2026-03-12T03:47:13.250Z", + "severity": 0 + }, + "host": { + "name": "DC01", + "hostname": "DC01.corp.contoso.com", + "os": { + "family": "windows", + "name": "Windows Server 2022", + "version": "10.0.20348.2340", + "build": "20348.2340", + "platform": "windows", + "type": "windows", + "kernel": "10.0.20348.2340 (WinBuild.160101.0800)" + }, + "ip": ["10.0.0.5", "fe80::1234:5678:abcd:ef01"], + "mac": ["00-15-5D-01-02-03"], + "architecture": "x86_64", + "domain": "corp.contoso.com" + }, + "source": { + "ip": "198.51.100.23", + "port": 49832, + "geo": { + "country_iso_code": "RU", + "city_name": "Москва", + "region_name": "Москва", + "location": { "lat": 55.7558, "lon": 37.6173 }, + "timezone": "Europe/Moscow" + } + }, + "user": { + "name": "администратор", + "domain": "CORP", + "id": "S-1-0-0", + "target": { + "name": "администратор", + "domain": "CORP" + } + }, + "message": "An account failed to log on.\n\nSubject:\n\tSecurity ID:\t\tS-1-5-18\n\tAccount Name:\t\tDC01$\n\tAccount Domain:\t\tCORP\n\tLogon ID:\t\t0x3E7\n\nLogon Information:\n\tLogon Type:\t\t10\n\tRestricted Admin Mode:\t-\n\tVirtual Account:\t\tNo\n\tElevated Token:\t\tNo\n\nFailure Information:\n\tFailure Reason:\t\tUnknown user name or bad password.\n\tStatus:\t\t\t0xC000006D\n\tSub Status:\t\t0xC0000064\n\nNew Logon:\n\tSecurity ID:\t\tS-1-0-0\n\tAccount Name:\t\tадминистратор\n\tAccount Domain:\t\tCORP\n\nProcess Information:\n\tCaller Process ID:\t0x0\n\tCaller Process Name:\t-\n\nNetwork Information:\n\tWorkstation Name:\tАТАКУЮЩИЙ-ПК\n\tSource Network Address:\t198.51.100.23\n\tSource Port:\t\t49832", + "related": { + "ip": ["198.51.100.23", "10.0.0.5"], + "user": ["DC01$", "администратор"] + }, + "ecs": { "version": "8.0.0" }, + "tags": ["security", "authentication", "failed-logon", "brute-force-candidate"] +}"#.to_string() +} + +fn filebeat_multiline_java_json() -> String { + r#"{ + "@timestamp": "2026-03-12T14:22:03.001Z", + "@metadata": { "beat": "filebeat", "version": "8.17.0" }, + "agent": { "name": "app-srv-07", "type": "filebeat", "version": "8.17.0" }, + "log": { + "file": { + "path": "C:\\Program Files\\MyApp\\logs\\application-2026-03-12.log", + "inode": "0" + }, + "offset": 482716, + "flags": ["utf-8", "multiline"] + }, + "message": "2026-03-12 14:22:02,999 ERROR [http-nio-8443-exec-42] com.example.api.UserController - Failed to process request for user_id=café-résumé-42\njava.lang.NullPointerException: Cannot invoke \"com.example.model.UserProfile.getDisplayName()\" because the return value of \"com.example.service.UserService.findById(String)\" is null\n\tat com.example.api.UserController.getUserProfile(UserController.java:142)\n\tat com.example.api.UserController$$FastClassBySpringCGLIB$$abc123.invoke()\n\tat org.springframework.cglib.proxy.MethodProxy.invoke(MethodProxy.java:218)\n\tat org.springframework.aop.framework.CglibAopProxy$CglibMethodInvocation.invokeJoinpoint(CglibAopProxy.java:793)\n\tat org.springframework.aop.framework.ReflectiveMethodInvocation.proceed(ReflectiveMethodInvocation.java:163)\n\tat org.springframework.aop.framework.CglibAopProxy$DynamicAdvisedInterceptor.intercept(CglibAopProxy.java:723)\n\tat com.example.api.UserController$$EnhancerBySpringCGLIB$$def456.getUserProfile()\n\tat sun.reflect.NativeMethodAccessorImpl.invoke0(Native Method)\n\tat sun.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:62)\n\tat sun.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43)\n\tat java.lang.reflect.Method.invoke(Method.java:498)\n\tat org.apache.tomcat.util.threads.TaskThread$WrappingRunnable.run(TaskThread.java:61)\n\tat java.lang.Thread.run(Thread.java:750)\nCaused by: org.hibernate.exception.JDBCConnectionException: Unable to acquire JDBC Connection\n\tat org.hibernate.exception.internal.SQLExceptionTypeDelegate.convert(SQLExceptionTypeDelegate.java:48)\n\tat com.zaxxer.hikari.pool.HikariPool.getConnection(HikariPool.java:163)\n\tat com.zaxxer.hikari.pool.HikariPool.getConnection(HikariPool.java:128)\nCaused by: java.sql.SQLTransientConnectionException: HikariPool-1 - Connection is not available, request timed out after 30000ms.\n\tat com.zaxxer.hikari.pool.HikariPool.createTimeoutException(HikariPool.java:695)\n\t... 42 more", + "error": { + "type": "java.lang.NullPointerException", + "message": "Cannot invoke \"com.example.model.UserProfile.getDisplayName()\"", + "stack_trace": "... (see message field for full trace)" + }, + "host": { + "name": "app-srv-07", + "os": { + "family": "windows", + "name": "Windows Server 2019", + "version": "10.0.17763.5329" + }, + "ip": ["10.10.20.7"] + }, + "service": { + "name": "user-api", + "version": "3.14.159-SNAPSHOT", + "environment": "staging", + "node": { "name": "app-srv-07:8443" } + }, + "labels": { + "deployment_id": "deploy-2026-03-12-r42", + "git_sha": "a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2", + "jira_ticket": "PLAT-9876" + }, + "ecs": { "version": "8.0.0" }, + "tags": ["java", "error", "staging", "connection-pool-exhaustion"] +}"#.to_string() +} + +fn winlogbeat_powershell_json() -> String { + r#"{ + "@timestamp": "2026-03-12T01:15:44.203Z", + "@metadata": { "beat": "winlogbeat", "version": "8.17.0" }, + "agent": { "name": "WS-FINANCE-12", "type": "winlogbeat" }, + "winlog": { + "channel": "Microsoft-Windows-PowerShell/Operational", + "provider_name": "Microsoft-Windows-PowerShell", + "event_id": 4104, + "task": "Execute a Remote Command", + "opcode": "On create calls", + "record_id": 55432, + "computer_name": "WS-FINANCE-12.corp.contoso.com", + "process": { "pid": 6328, "thread": { "id": 7204 } }, + "event_data": { + "MessageNumber": "1", + "MessageTotal": "1", + "ScriptBlockText": "function Invoke-Çömpléx_Tàsk {\n param(\n [Parameter(Mandatory=$true)]\n [string]$Tärget,\n [ValidateSet('Réad','Wríte','Éxecute')]\n [string]$Möde = 'Réad'\n )\n \n $encodedCmd = [Convert]::ToBase64String([Text.Encoding]::Unicode.GetBytes($Tärget))\n $résult = @{\n 'Tïmestamp' = (Get-Date -Format 'yyyy-MM-ddTHH:mm:ss.fffZ')\n 'Üser' = $env:USERNAME\n 'Dömain' = $env:USERDOMAIN\n 'Pàth' = \"C:\\Users\\$env:USERNAME\\AppData\\Local\\Temp\\öutput_$(Get-Random).tmp\"\n 'Àrgs' = @($Tärget, $Möde, $encodedCmd)\n 'Nësted' = @{\n 'Dëep1' = @{\n 'Dëep2' = @{\n 'Dëep3' = @{\n 'value' = 'We\\'re testing deep nesting with spëcial chars: <>&\\\"\\'/'\n }\n }\n }\n }\n }\n \n $résult | ConvertTo-Json -Depth 10 | Out-File -FilePath $résult['Pàth'] -Encoding UTF8\n return $résult\n}", + "ScriptBlockId": "b7c8d9e0-f1a2-3b4c-5d6e-7f8a9b0c1d2e", + "Path": "C:\\Users\\jëan-pierré\\Documents\\Scrïpts\\Ïnvoke-Task.ps1" + } + }, + "event": { + "code": "4104", + "kind": "event", + "provider": "Microsoft-Windows-PowerShell", + "category": ["process"], + "type": ["info"], + "outcome": "success" + }, + "host": { + "name": "WS-FINANCE-12", + "hostname": "WS-FINANCE-12.corp.contoso.com", + "os": { + "family": "windows", + "name": "Windows 11 Enterprise", + "version": "10.0.22631.3155", + "build": "22631.3155" + }, + "ip": ["10.20.30.12", "fe80::abcd:ef01:2345:6789"], + "mac": ["00-50-56-AB-CD-EF"] + }, + "user": { + "name": "jëan-pierré", + "domain": "CORP", + "id": "S-1-5-21-1234567890-1234567890-1234567890-5678" + }, + "process": { + "pid": 6328, + "executable": "C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.exe", + "command_line": "powershell.exe -NoProfile -ExecutionPolicy Bypass -File \"C:\\Users\\jëan-pierré\\Documents\\Scrïpts\\Ïnvoke-Task.ps1\"", + "parent": { + "pid": 4120, + "executable": "C:\\Windows\\explorer.exe" + } + }, + "message": "Creating Scriptblock text (1 of 1):\nfunction Invoke-Çömpléx_Tàsk { ... (see ScriptBlockText for full content)", + "related": { + "user": ["jëan-pierré"] + }, + "ecs": { "version": "8.0.0" }, + "tags": ["powershell", "scriptblock", "finance-dept"] +}"#.to_string() +} + +fn filebeat_kubernetes_json() -> String { + r#"{ + "@timestamp": "2026-03-12T19:33:07.445Z", + "@metadata": { "beat": "filebeat", "version": "8.17.0" }, + "agent": { "name": "k8s-node-pool-a-2", "type": "filebeat" }, + "kubernetes": { + "pod": { + "name": "payment-svc-7b8c9d-xq2f4", + "uid": "12345678-abcd-ef01-2345-67890abcdef0", + "ip": "10.244.3.17", + "labels": { + "app_kubernetes_io/name": "payment-svc", + "app_kubernetes_io/version": "2.71.828", + "app_kubernetes_io/component": "api", + "helm_sh/chart": "payment-svc-2.71.828", + "pod-template-hash": "7b8c9d" + }, + "annotations": { + "prometheus_io/scrape": "true", + "prometheus_io/port": "9090", + "vault_hashicorp_com/agent-inject": "true", + "vault_hashicorp_com/role": "payment-svc-prod" + } + }, + "node": { + "name": "k8s-node-pool-a-2", + "hostname": "k8s-node-pool-a-2.cluster.local", + "labels": { + "kubernetes_io/arch": "amd64", + "node_kubernetes_io/instance-type": "m5.2xlarge", + "topology_kubernetes_io/zone": "ap-southeast-2a" + } + }, + "namespace": "payment-prod", + "replicaset": { "name": "payment-svc-7b8c9d" }, + "deployment": { "name": "payment-svc" }, + "container": { + "name": "payment-api", + "image": "harbor.internal/payment/api:2.71.828-deadbeef", + "id": "containerd://abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789" + } + }, + "container": { + "id": "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789", + "image": { "name": "harbor.internal/payment/api:2.71.828-deadbeef" }, + "runtime": "containerd" + }, + "log": { + "file": { + "path": "/var/log/pods/payment-prod_payment-svc-7b8c9d-xq2f4_12345678-abcd-ef01-2345-67890abcdef0/payment-api/0.log" + } + }, + "message": "{\"level\":\"error\",\"ts\":1741804387.445,\"caller\":\"handler/payment.go:287\",\"msg\":\"payment processing failed\",\"trace_id\":\"abc123def456\",\"span_id\":\"789012\",\"request_id\":\"req-ñoño-42\",\"customer_id\":\"cust_Ωmega_∆lpha\",\"amount\":\"¥123,456.78\",\"currency\":\"JPY\",\"gateway_response\":{\"code\":\"DECLINED_INSUFFICIENT_FUNDS\",\"raw\":\"カード残高不足です。別のお支払い方法をお試しください。\",\"retry_after_ms\":null,\"metadata\":{\"issuer_country\":\"JP\",\"card_brand\":\"JCB\",\"last4\":\"4242\",\"3ds_enrolled\":true,\"risk_score\":0.73}},\"stack\":\"goroutine 847 [running]:\\nruntime/debug.Stack()\\n\\t/usr/local/go/src/runtime/debug/stack.go:24 +0x5e\\ngithub.com/example/payment-svc/internal/handler.(*PaymentHandler).ProcessPayment(...)\\n\\t/app/internal/handler/payment.go:287 +0x1a3\\ngithub.com/example/payment-svc/internal/handler.(*PaymentHandler).HandleRequest(...)\\n\\t/app/internal/handler/payment.go:142 +0x892\"}", + "stream": "stderr", + "event": { + "dataset": "kubernetes.container_logs", + "module": "kubernetes" + }, + "ecs": { "version": "8.0.0" }, + "tags": ["kubernetes", "payment", "production", "pci-zone"] +}"#.to_string() +} + +#[tokio::test] +async fn native_async_inserter_filebeat_nginx() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct LogRow { + ts: u64, + source: String, + json_data: String, + } + + let client = prepare_native_database("async_fb_nginx").await; + client + .query(&format!( + "CREATE TABLE t{} (ts UInt64, source String, json_data String) {}", + on_cluster(), + test_engine("ts"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default().without_period(), + ); + + let json = filebeat_nginx_json(); + for i in 0..10u64 { + inserter + .write(LogRow { + ts: 1741760062000 + i, + source: "filebeat-nginx".into(), + json_data: json.clone(), + }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + + let rows: Vec = client + .query("SELECT ts, source, json_data FROM t ORDER BY ts") + .fetch_all() + .await + .unwrap(); + assert_eq!(rows.len(), 10); + assert!(rows[0].json_data.contains("jean-françois")); + assert!(rows[0].json_data.contains("名前")); + assert!(rows[0].json_data.contains("über-ansicht")); + assert!(rows[0].json_data.contains("München")); +} + +#[tokio::test] +async fn native_async_inserter_winlogbeat_security() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct LogRow { + ts: u64, + source: String, + json_data: String, + } + + let client = prepare_native_database("async_wlb_security").await; + client + .query(&format!( + "CREATE TABLE t{} (ts UInt64, source String, json_data String) {}", + on_cluster(), + test_engine("ts"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default().without_period(), + ); + + let json = winlogbeat_security_json(); + for i in 0..10u64 { + inserter + .write(LogRow { + ts: 1741744031000 + i, + source: "winlogbeat-security".into(), + json_data: json.clone(), + }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + + let rows: Vec = client + .query("SELECT ts, source, json_data FROM t ORDER BY ts") + .fetch_all() + .await + .unwrap(); + assert_eq!(rows.len(), 10); + assert!(rows[0].json_data.contains("администратор")); + assert!(rows[0].json_data.contains("АТАКУЮЩИЙ-ПК")); + assert!(rows[0].json_data.contains("Москва")); + assert!(rows[0].json_data.contains("S-1-5-18")); +} + +#[tokio::test] +async fn native_async_inserter_filebeat_java_stacktrace() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct LogRow { + ts: u64, + source: String, + json_data: String, + } + + let client = prepare_native_database("async_fb_java").await; + client + .query(&format!( + "CREATE TABLE t{} (ts UInt64, source String, json_data String) {}", + on_cluster(), + test_engine("ts"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default().without_period(), + ); + + let json = filebeat_multiline_java_json(); + for i in 0..5u64 { + inserter + .write(LogRow { + ts: 1741781723000 + i, + source: "filebeat-java".into(), + json_data: json.clone(), + }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + + let rows: Vec = client + .query("SELECT ts, source, json_data FROM t ORDER BY ts") + .fetch_all() + .await + .unwrap(); + assert_eq!(rows.len(), 5); + assert!(rows[0].json_data.contains("NullPointerException")); + assert!(rows[0].json_data.contains("café-résumé-42")); + assert!(rows[0].json_data.contains("C:\\\\Program Files\\\\MyApp")); + assert!(rows[0].json_data.contains("HikariPool")); +} + +#[tokio::test] +async fn native_async_inserter_winlogbeat_powershell() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct LogRow { + ts: u64, + source: String, + json_data: String, + } + + let client = prepare_native_database("async_wlb_powershell").await; + client + .query(&format!( + "CREATE TABLE t{} (ts UInt64, source String, json_data String) {}", + on_cluster(), + test_engine("ts"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default().without_period(), + ); + + let json = winlogbeat_powershell_json(); + for i in 0..5u64 { + inserter + .write(LogRow { + ts: 1741742144000 + i, + source: "winlogbeat-powershell".into(), + json_data: json.clone(), + }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + + let rows: Vec = client + .query("SELECT ts, source, json_data FROM t ORDER BY ts") + .fetch_all() + .await + .unwrap(); + assert_eq!(rows.len(), 5); + assert!(rows[0].json_data.contains("Invoke-Çömpléx_Tàsk")); + assert!(rows[0].json_data.contains("jëan-pierré")); + assert!(rows[0].json_data.contains("Scrïpts")); +} + +#[tokio::test] +async fn native_async_inserter_filebeat_kubernetes() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct LogRow { + ts: u64, + source: String, + json_data: String, + } + + let client = prepare_native_database("async_fb_k8s").await; + client + .query(&format!( + "CREATE TABLE t{} (ts UInt64, source String, json_data String) {}", + on_cluster(), + test_engine("ts"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default().without_period(), + ); + + let json = filebeat_kubernetes_json(); + for i in 0..5u64 { + inserter + .write(LogRow { + ts: 1741804387000 + i, + source: "filebeat-k8s".into(), + json_data: json.clone(), + }) + .await + .unwrap(); + } + + inserter.end().await.unwrap(); + + let rows: Vec = client + .query("SELECT ts, source, json_data FROM t ORDER BY ts") + .fetch_all() + .await + .unwrap(); + assert_eq!(rows.len(), 5); + assert!(rows[0].json_data.contains("payment-svc-7b8c9d-xq2f4")); + assert!(rows[0].json_data.contains("カード残高不足")); + assert!(rows[0].json_data.contains("req-ñoño-42")); + assert!(rows[0].json_data.contains("cust_Ωmega_∆lpha")); + assert!(rows[0].json_data.contains("¥123,456.78")); +} + +/// Mixed Beat sources in a single batch — concurrent handles, one source per handle. +#[tokio::test] +async fn native_async_inserter_mixed_beats_concurrent() { + use clickhouse::native::{AsyncNativeInserter, AsyncNativeInserterConfig}; + + #[derive(Debug, Clone, PartialEq, Row, Serialize, Deserialize)] + struct LogRow { + ts: u64, + source: String, + json_data: String, + } + + let client = prepare_native_database("async_mixed_beats").await; + client + .query(&format!( + "CREATE TABLE t{} (ts UInt64, source String, json_data String) {}", + on_cluster(), + test_engine("ts"), + )) + .execute() + .await + .unwrap(); + + let inserter = AsyncNativeInserter::::new( + &client, + "t", + AsyncNativeInserterConfig::default() + .with_max_rows(15) // force multiple flushes mid-batch + .without_period(), + ); + + let sources: Vec<(&str, String)> = vec![ + ("filebeat-nginx", filebeat_nginx_json()), + ("winlogbeat-security", winlogbeat_security_json()), + ("filebeat-java", filebeat_multiline_java_json()), + ("winlogbeat-powershell", winlogbeat_powershell_json()), + ("filebeat-k8s", filebeat_kubernetes_json()), + ]; + + let mut tasks = Vec::new(); + for (idx, (source, json)) in sources.into_iter().enumerate() { + let handle = inserter.handle(); + let source = source.to_string(); + tasks.push(tokio::spawn(async move { + for j in 0..20u64 { + let ts = (idx as u64) * 1_000_000 + j; + handle + .write(LogRow { + ts, + source: source.clone(), + json_data: json.clone(), + }) + .await + .unwrap(); + } + })); + } + + for task in tasks { + task.await.unwrap(); + } + + inserter.end().await.unwrap(); + + // 5 sources × 20 rows = 100 + let count: u64 = client + .query("SELECT count() FROM t") + .fetch_one() + .await + .unwrap(); + assert_eq!(count, 100); + + // Verify each source is present. + let nginx_count: u64 = client + .query("SELECT count() FROM t WHERE source = 'filebeat-nginx'") + .fetch_one() + .await + .unwrap(); + assert_eq!(nginx_count, 20); + + let security_count: u64 = client + .query("SELECT count() FROM t WHERE source = 'winlogbeat-security'") + .fetch_one() + .await + .unwrap(); + assert_eq!(security_count, 20); +} diff --git a/tests/it/rbwnat_smoke.rs b/tests/it/rbwnat_smoke.rs index 77ae3005..2bffd153 100644 --- a/tests/it/rbwnat_smoke.rs +++ b/tests/it/rbwnat_smoke.rs @@ -3,9 +3,8 @@ use crate::geo_types::{LineString, MultiLineString, MultiPolygon, Point, Polygon use crate::{SimpleRow, create_simple_table, execute_statements, get_client, insert_and_select}; use clickhouse::Row; use clickhouse::sql::Identifier; -use fxhash::FxHashMap; use indexmap::IndexMap; -use linked_hash_map::LinkedHashMap; +use rustc_hash::FxHashMap; use serde::{Deserialize, Serialize}; use serde_repr::{Deserialize_repr, Serialize_repr}; use std::collections::HashMap; @@ -392,9 +391,9 @@ async fn maps_third_party() { #[derive(Clone, Debug, Row, Serialize, Deserialize, PartialEq)] struct Data { im: IndexMap, - lhm: LinkedHashMap, + lhm: IndexMap, fx: FxHashMap, - weird_but_ok: LinkedHashMap>>>, + weird_but_ok: IndexMap>>>, } let client = prepare_database!(); @@ -417,9 +416,9 @@ async fn maps_third_party() { let rows = vec![Data { im: IndexMap::from_iter(vec![(1, "one".to_string()), (2, "two".to_string())]), - lhm: LinkedHashMap::from_iter(vec![(3, "three".to_string()), (4, "four".to_string())]), + lhm: IndexMap::from_iter(vec![(3, "three".to_string()), (4, "four".to_string())]), fx: FxHashMap::from_iter(vec![(5, "five".to_string()), (6, "six".to_string())]), - weird_but_ok: LinkedHashMap::from_iter(vec![( + weird_but_ok: IndexMap::from_iter(vec![( 7u128, IndexMap::from_iter(vec![( -8i8,