This is an automated email from the ASF dual-hosted git repository. CritasWang pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/iotdb-client-rust.git
commit 658fe28e2def97d3d15506fcb04ea9fe760cf859 Author: CritasWang <[email protected]> AuthorDate: Fri Jul 10 17:35:06 2026 +0800 Phase 4-5: SessionDataSet iterator, TableSession, SessionPool/TableSessionPool; live integration verified against IoTDB 2.0.6 --- docker-compose-1c1d.yml | 66 ++++++ docker-compose.yml | 36 +++ src/client/dataset.rs | 383 ++++++++++++++++++++++++++++++++ src/client/mod.rs | 11 +- src/client/pool.rs | 528 ++++++++++++++++++++++++++++++++++++++++++++ src/client/session.rs | 62 +++++- src/client/table_session.rs | 241 ++++++++++++++++++++ src/data/tablet.rs | 21 ++ src/data/tsblock.rs | 43 +++- src/lib.rs | 4 + 10 files changed, 1376 insertions(+), 19 deletions(-) diff --git a/docker-compose-1c1d.yml b/docker-compose-1c1d.yml new file mode 100644 index 0000000..847caae --- /dev/null +++ b/docker-compose-1c1d.yml @@ -0,0 +1,66 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Single-node IoTDB cluster (1 ConfigNode + 1 DataNode) for integration tests. +# Usage: docker compose -f docker-compose-1c1d.yml up -d +# cargo test -- --ignored # or: cargo test (integration tests detect the server) + +services: + iotdb-datanode: + image: apache/iotdb:2.0.6-datanode + container_name: iotdb-rust-datanode + restart: always + depends_on: + iotdb-confignode: + condition: service_healthy + healthcheck: + test: ["CMD", "ls", "/iotdb/data/datanode/system"] + interval: 10s + timeout: 60s + retries: 30 + start_period: 30s + ports: + - "6667:6667" + networks: + - iotdb-network + environment: + - dn_rpc_address=iotdb-datanode + - dn_internal_address=iotdb-datanode + - dn_seed_config_node=iotdb-confignode:10710 + - IOTDB_JMX_OPTS=-Xms1G -Xmx1G -XX:MaxDirectMemorySize=384M + + iotdb-confignode: + image: apache/iotdb:2.0.6-confignode + container_name: iotdb-rust-confignode + restart: always + healthcheck: + test: ["CMD", "ls", "/iotdb/data"] + interval: 3s + timeout: 5s + retries: 30 + start_period: 30s + networks: + - iotdb-network + environment: + - cn_internal_address=iotdb-confignode + - cn_internal_port=10710 + - cn_seed_config_node=iotdb-confignode:10710 + - CONFIGNODE_JMX_OPTS=-Xms384M -Xmx384M -XX:MaxDirectMemorySize=192M + +networks: + iotdb-network: + driver: bridge diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..0aaef9e --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,36 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Standalone IoTDB for integration tests (simplest topology). +# Usage: docker compose up -d && cargo test +# For a 1C1D cluster topology, see docker-compose-1c1d.yml. + +services: + iotdb: + image: apache/iotdb:2.0.6-standalone + container_name: iotdb-rust-standalone + restart: unless-stopped + ports: + - "6667:6667" + environment: + - IOTDB_JMX_OPTS=-Xms1G -Xmx1G -XX:MaxDirectMemorySize=384M + healthcheck: + test: ["CMD", "ls", "/iotdb/data/datanode/system"] + interval: 5s + timeout: 10s + retries: 30 + start_period: 20s diff --git a/src/client/dataset.rs b/src/client/dataset.rs new file mode 100644 index 0000000..457b213 --- /dev/null +++ b/src/client/dataset.rs @@ -0,0 +1,383 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Query result iteration — the SessionDataSet state machine (protocol +//! spec §4.3): drain the current TsBlock, decode the next cached block, +//! fetch more pages while `moreData`, then auto-close. + +use std::collections::VecDeque; + +use crate::client::session::{QueryHandle, Session}; +use crate::data::tsblock::TsBlock; +use crate::data::value::Value; +use crate::error::{Error, Result}; + +/// One result row: the timestamp (`None` when the server set +/// `ignoreTimeStamp`, e.g. aggregations and SHOW statements) and one +/// [`Value`] per output column. +#[derive(Debug, Clone, PartialEq)] +pub struct Row { + pub timestamp: Option<i64>, + pub values: Vec<Value>, +} + +/// An iterable query result set. +/// +/// Borrows the [`Session`] mutably for its whole lifetime: `fetchResultsV2` +/// and `closeOperation` must reach the node that owns the query id, so the +/// connection stays pinned to the result set until it is closed (spec +/// gotcha #13). The borrow also means no other statement can run on the +/// session while a result set is open — close (or drop) it first. +/// +/// Not a `std::iter::Iterator` — [`SessionDataSet::next_row`] returns +/// `Result<Option<Row>>` so decode and fetch errors surface cleanly. +/// The query is closed automatically on exhaustion, on [`SessionDataSet::close`], +/// or on drop (best-effort, matching the other SDKs). +pub struct SessionDataSet<'a> { + session: &'a mut Session, + query_id: i64, + statement: String, + columns: Vec<String>, + data_type_list: Vec<String>, + ignore_time_stamp: bool, + /// Output column ordinal → physical TsBlock column index; `-1` = time + /// column; identity mapping when `None` (spec §4.2 field 17). + column_index_map: Option<Vec<i32>>, + /// Cached serialized TsBlocks not yet decoded. + pending_blocks: VecDeque<Vec<u8>>, + /// Currently decoded block and the next row to yield from it. + current: Option<TsBlock>, + row_index: usize, + more_data: bool, + closed: bool, +} + +impl<'a> SessionDataSet<'a> { + pub(crate) fn new(session: &'a mut Session, handle: QueryHandle) -> SessionDataSet<'a> { + SessionDataSet { + session, + query_id: handle.query_id, + statement: handle.statement, + columns: handle.columns, + data_type_list: handle.data_type_list, + ignore_time_stamp: handle.ignore_time_stamp, + column_index_map: handle.column_index2_ts_block_column_index_list, + pending_blocks: handle.query_result.into(), + current: None, + row_index: 0, + more_data: handle.more_data, + closed: false, + } + } + + /// Output column names. + pub fn columns(&self) -> &[String] { + &self.columns + } + + /// Output column type names as reported by the server + /// (`"INT64"`, `"TEXT"`, …), parallel to [`SessionDataSet::columns`]. + pub fn column_types(&self) -> &[String] { + &self.data_type_list + } + + /// Whether the server flagged this result as having no meaningful + /// timestamps (row timestamps are then `None`). + pub fn ignore_time_stamp(&self) -> bool { + self.ignore_time_stamp + } + + /// Advance to the next row, fetching further pages as needed. + /// Returns `Ok(None)` when exhausted (the query is then auto-closed). + pub fn next_row(&mut self) -> Result<Option<Row>> { + if self.closed { + return Ok(None); + } + loop { + // 1. Unread rows in the current block → yield the next one. + if let Some(block) = &self.current { + if self.row_index < block.position_count { + let row = self.assemble_row()?; + self.row_index += 1; + return Ok(Some(row)); + } + self.current = None; + } + // 2. More cached blocks → decode the next non-empty one. + if let Some(bytes) = self.pending_blocks.pop_front() { + self.current = Some(TsBlock::decode(&bytes)?); + self.row_index = 0; + continue; + } + // 3. Server has more pages → fetch and refill. + if self.more_data { + let (blocks, more) = self.session.fetch_results(self.query_id, &self.statement)?; + self.pending_blocks = blocks.into(); + self.more_data = more; + if !self.pending_blocks.is_empty() { + continue; + } + if self.more_data { + // hasResultSet with an empty page and moreData still + // set — loop and fetch again rather than spin here. + continue; + } + } + // 4. Exhausted → auto-close and report the end. + self.close(); + return Ok(None); + } + } + + /// Build the output row at `self.row_index` of the current block, + /// mapping output columns through `columnIndex2TsBlockColumnIndexList` + /// (`-1` = time column; identity when absent — spec §5.4). + fn assemble_row(&self) -> Result<Row> { + let block = self.current.as_ref().expect("current block"); + let i = self.row_index; + let mut values = Vec::with_capacity(self.columns.len()); + for ordinal in 0..self.columns.len() { + let physical = match &self.column_index_map { + Some(map) => *map.get(ordinal).ok_or_else(|| { + Error::Decode(format!( + "column index map has {} entries, need output column {ordinal}", + map.len() + )) + })?, + None => ordinal as i32, + }; + if physical == -1 { + values.push(Value::Timestamp(block.timestamps[i])); + continue; + } + let column = usize::try_from(physical) + .ok() + .and_then(|p| block.columns.get(p)) + .ok_or_else(|| { + Error::Decode(format!( + "column index map points at TsBlock column {physical}, block has {}", + block.columns.len() + )) + })?; + values.push(column[i].clone()); + } + let timestamp = (!self.ignore_time_stamp).then(|| block.timestamps[i]); + Ok(Row { timestamp, values }) + } + + /// Close the query (best-effort `closeOperation`, errors swallowed) and + /// release the session borrow's obligations. Idempotent; also invoked + /// automatically on exhaustion and on drop. + pub fn close(&mut self) { + if self.closed { + return; + } + self.closed = true; + self.pending_blocks.clear(); + self.current = None; + self.session.close_query(self.query_id); + } +} + +impl Drop for SessionDataSet<'_> { + fn drop(&mut self) { + self.close(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::client::session::SessionConfig; + use crate::data::tsblock::test_util::{header, int32_block, time_column}; + use crate::data::tsblock::{ENCODING_BINARY_ARRAY, ENCODING_INT32_ARRAY}; + use crate::data::TSDataType; + + /// A never-opened session: fetch_results errors, close_query is a no-op — + /// exactly what the offline state-machine tests need. + fn offline_session() -> Session { + Session::new(SessionConfig::default()) + } + + fn handle(blocks: Vec<Vec<u8>>, more_data: bool) -> QueryHandle { + QueryHandle { + query_id: 1, + statement: "SELECT s1 FROM root.sg.d1".into(), + columns: vec!["root.sg.d1.s1".into()], + data_type_list: vec!["INT32".into()], + ignore_time_stamp: false, + query_result: blocks, + more_data, + column_index2_ts_block_column_index_list: None, + } + } + + #[test] + fn drains_rows_across_multiple_cached_blocks() { + let mut session = offline_session(); + let blocks = vec![int32_block(&[1, 2], &[10, 20]), int32_block(&[3], &[30])]; + let mut ds = SessionDataSet::new(&mut session, handle(blocks, false)); + + let mut rows = Vec::new(); + while let Some(row) = ds.next_row().unwrap() { + rows.push(row); + } + assert_eq!( + rows, + vec![ + Row { + timestamp: Some(1), + values: vec![Value::Int32(10)] + }, + Row { + timestamp: Some(2), + values: vec![Value::Int32(20)] + }, + Row { + timestamp: Some(3), + values: vec![Value::Int32(30)] + }, + ] + ); + // Exhausted → auto-closed; further calls keep returning None. + assert_eq!(ds.next_row().unwrap(), None); + } + + #[test] + fn empty_result_set_yields_none_immediately() { + let mut session = offline_session(); + let mut ds = SessionDataSet::new(&mut session, handle(vec![], false)); + assert_eq!(ds.next_row().unwrap(), None); + } + + #[test] + fn empty_blocks_are_skipped() { + let mut session = offline_session(); + let blocks = vec![int32_block(&[], &[]), int32_block(&[7], &[70])]; + let mut ds = SessionDataSet::new(&mut session, handle(blocks, false)); + let row = ds.next_row().unwrap().unwrap(); + assert_eq!(row.timestamp, Some(7)); + assert_eq!(ds.next_row().unwrap(), None); + } + + #[test] + fn more_data_on_dead_session_surfaces_fetch_error() { + // moreData=true forces a fetch_results call, which fails on a + // never-opened session — the error must propagate, not panic. + let mut session = offline_session(); + let mut ds = SessionDataSet::new(&mut session, handle(vec![], true)); + assert!(ds.next_row().is_err()); + } + + #[test] + fn ignore_time_stamp_omits_timestamps() { + let mut session = offline_session(); + let mut h = handle(vec![int32_block(&[5], &[50])], false); + h.ignore_time_stamp = true; + let mut ds = SessionDataSet::new(&mut session, h); + assert!(ds.ignore_time_stamp()); + let row = ds.next_row().unwrap().unwrap(); + assert_eq!(row.timestamp, None); + assert_eq!(row.values, vec![Value::Int32(50)]); + } + + #[test] + fn column_index_map_with_time_column_and_duplicates() { + // Physical block: one TEXT column. Output: [time, text, text again]. + let mut b = header(&[TSDataType::Text], 2, &[ENCODING_BINARY_ARRAY]); + b.extend_from_slice(&time_column(&[100, 200])); + b.push(0); // mayHaveNull + for s in ["a", "b"] { + b.extend_from_slice(&(s.len() as i32).to_be_bytes()); + b.extend_from_slice(s.as_bytes()); + } + + let mut session = offline_session(); + let h = QueryHandle { + query_id: 9, + statement: "SELECT time, tag, tag FROM t".into(), + columns: vec!["time".into(), "tag".into(), "tag".into()], + data_type_list: vec!["TIMESTAMP".into(), "TEXT".into(), "TEXT".into()], + ignore_time_stamp: true, + query_result: vec![b], + more_data: false, + column_index2_ts_block_column_index_list: Some(vec![-1, 0, 0]), + }; + let mut ds = SessionDataSet::new(&mut session, h); + assert_eq!(ds.columns(), ["time", "tag", "tag"]); + assert_eq!(ds.column_types(), ["TIMESTAMP", "TEXT", "TEXT"]); + + let row = ds.next_row().unwrap().unwrap(); + assert_eq!( + row.values, + vec![ + Value::Timestamp(100), + Value::Text("a".into()), + Value::Text("a".into()), + ] + ); + let row = ds.next_row().unwrap().unwrap(); + assert_eq!(row.values[0], Value::Timestamp(200)); + assert_eq!(ds.next_row().unwrap(), None); + } + + #[test] + fn identity_mapping_when_map_absent() { + let types = [TSDataType::Int32, TSDataType::Int32]; + let mut b = header(&types, 1, &[ENCODING_INT32_ARRAY, ENCODING_INT32_ARRAY]); + b.extend_from_slice(&time_column(&[1])); + for v in [11i32, 22] { + b.push(0); + b.extend_from_slice(&v.to_be_bytes()); + } + + let mut session = offline_session(); + let mut h = handle(vec![b], false); + h.columns = vec!["s1".into(), "s2".into()]; + h.data_type_list = vec!["INT32".into(), "INT32".into()]; + let mut ds = SessionDataSet::new(&mut session, h); + let row = ds.next_row().unwrap().unwrap(); + assert_eq!(row.values, vec![Value::Int32(11), Value::Int32(22)]); + } + + #[test] + fn out_of_range_map_entry_is_decode_error() { + let mut session = offline_session(); + let mut h = handle(vec![int32_block(&[1], &[10])], false); + h.column_index2_ts_block_column_index_list = Some(vec![5]); + let mut ds = SessionDataSet::new(&mut session, h); + assert!(matches!(ds.next_row(), Err(Error::Decode(_)))); + } + + #[test] + fn corrupt_block_is_decode_error() { + let mut session = offline_session(); + let mut ds = SessionDataSet::new(&mut session, handle(vec![vec![0, 1]], false)); + assert!(matches!(ds.next_row(), Err(Error::Decode(_)))); + } + + #[test] + fn explicit_close_stops_iteration() { + let mut session = offline_session(); + let blocks = vec![int32_block(&[1, 2], &[10, 20])]; + let mut ds = SessionDataSet::new(&mut session, handle(blocks, false)); + assert!(ds.next_row().unwrap().is_some()); + ds.close(); + ds.close(); // idempotent + assert_eq!(ds.next_row().unwrap(), None); + } +} diff --git a/src/client/mod.rs b/src/client/mod.rs index 34b96db..917cd13 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -15,9 +15,16 @@ // specific language governing permissions and limitations // under the License. -//! Session-layer API: tree model (`Session`) and, later, table model -//! (`TableSession`) plus pooled variants — mirroring the Node.js and C# SDKs. +//! Session-layer API: tree model ([`Session`]), table model +//! ([`TableSession`]), pooled variants and query result iteration — +//! mirroring the Node.js and C# SDKs. +pub mod dataset; +pub mod pool; pub mod session; +pub mod table_session; +pub use dataset::{Row, SessionDataSet}; +pub use pool::{PooledSession, SessionPool, SessionPoolConfig, TableSessionPool}; pub use session::{QueryHandle, Session, SessionConfig}; +pub use table_session::{TableSession, TableSessionBuilder}; diff --git a/src/client/pool.rs b/src/client/pool.rs new file mode 100644 index 0000000..2e6a864 --- /dev/null +++ b/src/client/pool.rs @@ -0,0 +1,528 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Session pools (protocol spec §7): a bounded set of open sessions, +//! eagerly grown to `min_size` and lazily to `max_size`, handed out as +//! RAII guards. Round-robin endpoint spread happens naturally via the +//! rotating start index in [`Session::open`]. + +use std::collections::VecDeque; +use std::ops::{Deref, DerefMut}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Condvar, Mutex}; +use std::time::{Duration, Instant}; + +use crate::client::session::{Session, SessionConfig}; +use crate::error::{Error, Result}; + +/// Configuration for a [`SessionPool`] (or [`TableSessionPool`]). +/// +/// `session` carries the per-connection settings — node urls, +/// user/password, dialect/database, connect timeout — shared by every +/// pooled session. +#[derive(Debug, Clone)] +pub struct SessionPoolConfig { + pub session: SessionConfig, + /// Upper bound on live sessions (idle + handed out). Default 8, + /// matching the C# SDK's pool size. + pub max_size: usize, + /// Sessions opened eagerly when the pool is created. Default 0. + pub min_size: usize, + /// How long [`SessionPool::acquire`] waits for an idle session once the + /// pool is at `max_size`. Default 60 s (Node.js `waitTimeout`). + pub acquire_timeout: Duration, +} + +impl Default for SessionPoolConfig { + fn default() -> Self { + Self { + session: SessionConfig::default(), + max_size: 8, + min_size: 0, + acquire_timeout: Duration::from_secs(60), + } + } +} + +impl SessionPoolConfig { + /// Set endpoints from `"host:port"` node-url strings. + pub fn with_node_urls<S: AsRef<str>>(mut self, node_urls: &[S]) -> Result<Self> { + self.session = self.session.with_node_urls(node_urls)?; + Ok(self) + } +} + +/// Idle sessions plus the pool lifecycle flag, guarded by one mutex so +/// [`Condvar`] waiters observe both consistently. +struct PoolState { + idle: VecDeque<Session>, + closed: bool, +} + +/// A pool of open tree-model [`Session`]s. +/// +/// Sessions are created lazily up to `max_size` (after the eager +/// `min_size`); [`SessionPool::acquire`] blocks up to `acquire_timeout` +/// when the pool is exhausted. Dead sessions are discarded on acquire and +/// on release. The pool tracks the most recent `USE <db>` seen on any +/// released session and replays it on acquire so every handed-out session +/// is in the pool's current database (spec §6.2). +pub struct SessionPool { + config: SessionPoolConfig, + state: Mutex<PoolState>, + available: Condvar, + /// Sessions alive right now: idle + handed out. Only mutated while + /// holding `state`, so `live < max_size` checks cannot overshoot. + live: AtomicUsize, + /// Pool-level current database, updated from released sessions. + database: Mutex<Option<String>>, +} + +impl SessionPool { + /// Create the pool and eagerly open `min_size` sessions. + pub fn new(config: SessionPoolConfig) -> Result<SessionPool> { + if config.min_size > config.max_size { + return Err(Error::Client(format!( + "pool min_size ({}) > max_size ({})", + config.min_size, config.max_size + ))); + } + let database = config.session.database.clone(); + let pool = SessionPool { + config, + state: Mutex::new(PoolState { + idle: VecDeque::new(), + closed: false, + }), + available: Condvar::new(), + live: AtomicUsize::new(0), + database: Mutex::new(database), + }; + for _ in 0..pool.config.min_size { + let session = pool.open_session()?; + let mut state = pool.state.lock().expect("pool lock poisoned"); + state.idle.push_back(session); + pool.live.fetch_add(1, Ordering::Relaxed); + } + Ok(pool) + } + + /// Sessions alive right now (idle + handed out). + pub fn live_count(&self) -> usize { + self.live.load(Ordering::Relaxed) + } + + /// Acquire a session, blocking up to `acquire_timeout` when the pool is + /// at capacity with nothing idle. Dead idle sessions are discarded and + /// the acquire retried; a fresh session is opened while under + /// `max_size`. + pub fn acquire(&self) -> Result<PooledSession<'_>> { + let deadline = Instant::now() + self.config.acquire_timeout; + let mut state = self.state.lock().expect("pool lock poisoned"); + loop { + if state.closed { + return Err(Error::Client("session pool is closed".into())); + } + // Idle session available → validate liveness, evict the dead. + while let Some(session) = state.idle.pop_front() { + if session.is_open() { + drop(state); + return self.hand_out(session); + } + self.live.fetch_sub(1, Ordering::Relaxed); + } + // Below capacity → grow lazily. Count the slot while still + // holding the lock so concurrent acquires cannot overshoot. + if self.live.load(Ordering::Relaxed) < self.config.max_size { + self.live.fetch_add(1, Ordering::Relaxed); + drop(state); + match self.open_session() { + Ok(session) => return self.hand_out(session), + Err(e) => { + self.live.fetch_sub(1, Ordering::Relaxed); + self.available.notify_one(); + return Err(e); + } + } + } + // At capacity → wait for a release, bounded by the deadline. + let now = Instant::now(); + if now >= deadline { + return Err(Error::Client(format!( + "pool exhausted: no session available within {:?} ({} live, max {})", + self.config.acquire_timeout, + self.live.load(Ordering::Relaxed), + self.config.max_size + ))); + } + let (guard, _) = self + .available + .wait_timeout(state, deadline - now) + .expect("pool lock poisoned"); + state = guard; + } + } + + /// Convenience: acquire a session, run one non-query statement, release. + /// `USE <db>` propagates to the whole pool via the database tracking. + pub fn execute_non_query(&self, sql: &str) -> Result<()> { + self.acquire()?.execute_non_query(sql) + } + + /// Close the pool: no further acquires; drain and close all idle + /// sessions. Sessions currently handed out are closed when their guards + /// drop. + pub fn close(&self) { + let drained = { + let mut state = self.state.lock().expect("pool lock poisoned"); + state.closed = true; + std::mem::take(&mut state.idle) + }; + self.live.fetch_sub(drained.len(), Ordering::Relaxed); + for mut session in drained { + let _ = session.close(); + } + self.available.notify_all(); + } + + fn open_session(&self) -> Result<Session> { + let mut config = self.config.session.clone(); + // New sessions start in the pool's current database (config key + // "db"), so no catch-up USE is needed on first hand-out. + config.database = self.database.lock().expect("pool lock poisoned").clone(); + let mut session = Session::new(config); + session.open()?; + Ok(session) + } + + /// Final step of acquire: sync the session onto the pool's current + /// database before handing it out. On USE failure the session is + /// discarded, not returned to the pool. + fn hand_out(&self, mut session: Session) -> Result<PooledSession<'_>> { + let pool_db = self.database.lock().expect("pool lock poisoned").clone(); + if let Some(db) = pool_db { + if session.database() != Some(db.as_str()) { + if let Err(e) = session.execute_non_query(&format!("USE {db}")) { + let _ = session.close(); + self.live.fetch_sub(1, Ordering::Relaxed); + self.available.notify_one(); + return Err(e); + } + } + } + Ok(PooledSession { + pool: self, + session: Some(session), + }) + } + + /// Return a session from a dropped guard. Dead sessions are discarded; + /// live ones update the pool database and go back to the idle queue. + fn release(&self, session: Session) { + if !session.is_open() { + self.live.fetch_sub(1, Ordering::Relaxed); + self.available.notify_one(); + return; + } + if let Some(db) = session.database() { + let mut pool_db = self.database.lock().expect("pool lock poisoned"); + if pool_db.as_deref() != Some(db) { + *pool_db = Some(db.to_string()); + } + } + let mut state = self.state.lock().expect("pool lock poisoned"); + if state.closed { + drop(state); + self.live.fetch_sub(1, Ordering::Relaxed); + let mut session = session; + let _ = session.close(); + } else { + state.idle.push_back(session); + drop(state); + } + self.available.notify_one(); + } + + /// Test hook: push a pre-built session (possibly dead) into the idle + /// queue, counting it as live. + #[cfg(test)] + fn inject_idle(&self, session: Session) { + let mut state = self.state.lock().expect("pool lock poisoned"); + state.idle.push_back(session); + self.live.fetch_add(1, Ordering::Relaxed); + } +} + +impl Drop for SessionPool { + fn drop(&mut self) { + self.close(); + } +} + +/// RAII guard for a pooled [`Session`]. Derefs to the session; returns it +/// to the pool on drop (dead sessions are discarded instead). +pub struct PooledSession<'a> { + pool: &'a SessionPool, + session: Option<Session>, +} + +impl std::fmt::Debug for PooledSession<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PooledSession") + .field("open", &self.session.as_ref().is_some_and(Session::is_open)) + .finish_non_exhaustive() + } +} + +impl Deref for PooledSession<'_> { + type Target = Session; + + fn deref(&self) -> &Session { + self.session.as_ref().expect("session taken") + } +} + +impl DerefMut for PooledSession<'_> { + fn deref_mut(&mut self) -> &mut Session { + self.session.as_mut().expect("session taken") + } +} + +impl Drop for PooledSession<'_> { + fn drop(&mut self) { + if let Some(session) = self.session.take() { + self.pool.release(session); + } + } +} + +/// A pool of table-dialect sessions — a [`SessionPool`] whose sessions are +/// opened with `sql_dialect="table"` (and optionally a database), per +/// protocol spec §6. `USE <db>` on any pooled session propagates pool-wide +/// on release. +pub struct TableSessionPool { + pool: SessionPool, +} + +impl TableSessionPool { + /// Create the pool, forcing the table dialect on the session config. + pub fn new(mut config: SessionPoolConfig) -> Result<TableSessionPool> { + config.session.sql_dialect = "table".into(); + Ok(TableSessionPool { + pool: SessionPool::new(config)?, + }) + } + + /// Acquire a table-dialect session guard. + pub fn acquire(&self) -> Result<PooledSession<'_>> { + self.pool.acquire() + } + + /// Convenience: acquire, run one non-query statement, release. + pub fn execute_non_query(&self, sql: &str) -> Result<()> { + self.pool.execute_non_query(sql) + } + + pub fn live_count(&self) -> usize { + self.pool.live_count() + } + + pub fn close(&self) { + self.pool.close() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::connection::Endpoint; + use std::net::TcpStream; + + /// Endpoint nothing listens on: connects are refused ~immediately. + fn dead_endpoint_config() -> SessionConfig { + SessionConfig { + endpoints: vec![Endpoint::new("127.0.0.1", 1)], + connect_timeout: Duration::from_millis(200), + ..SessionConfig::default() + } + } + + fn live_server_available() -> bool { + TcpStream::connect_timeout( + &"127.0.0.1:6667".parse().unwrap(), + Duration::from_millis(300), + ) + .is_ok() + } + + #[test] + fn config_defaults() { + let cfg = SessionPoolConfig::default(); + assert_eq!(cfg.max_size, 8); + assert_eq!(cfg.min_size, 0); + assert_eq!(cfg.acquire_timeout, Duration::from_secs(60)); + } + + #[test] + fn min_greater_than_max_is_rejected() { + let cfg = SessionPoolConfig { + min_size: 9, + max_size: 8, + ..Default::default() + }; + assert!(SessionPool::new(cfg).is_err()); + } + + #[test] + fn exhausted_pool_times_out_with_typed_error() { + // max_size 0: nothing idle, no growth allowed → deadline error. + let cfg = SessionPoolConfig { + max_size: 0, + acquire_timeout: Duration::from_millis(50), + session: dead_endpoint_config(), + ..Default::default() + }; + let pool = SessionPool::new(cfg).unwrap(); + let start = Instant::now(); + match pool.acquire() { + Err(Error::Client(msg)) => assert!(msg.contains("pool exhausted"), "{msg}"), + other => panic!("expected pool-exhausted error, got {other:?}"), + } + assert!(start.elapsed() >= Duration::from_millis(50)); + } + + #[test] + fn dead_idle_sessions_are_evicted_on_acquire() { + let cfg = SessionPoolConfig { + max_size: 1, + acquire_timeout: Duration::from_millis(50), + session: dead_endpoint_config(), + ..Default::default() + }; + let pool = SessionPool::new(cfg).unwrap(); + // A never-opened session is dead (is_open() == false). + pool.inject_idle(Session::new(dead_endpoint_config())); + assert_eq!(pool.live_count(), 1); + + // Acquire evicts the dead session, then tries to open a fresh one, + // which fails against the dead endpoint — a connect error, NOT + // "pool exhausted". + match pool.acquire() { + Err(Error::Thrift(_)) => {} + other => panic!("expected thrift connect error, got {other:?}"), + } + // The dead session and the failed growth slot are both released. + assert_eq!(pool.live_count(), 0); + } + + #[test] + fn acquire_on_closed_pool_fails() { + let cfg = SessionPoolConfig { + session: dead_endpoint_config(), + ..Default::default() + }; + let pool = SessionPool::new(cfg).unwrap(); + pool.close(); + match pool.acquire() { + Err(Error::Client(msg)) => assert!(msg.contains("closed"), "{msg}"), + other => panic!("expected closed-pool error, got {other:?}"), + }; + } + + #[test] + fn releasing_dead_session_shrinks_live_count() { + let cfg = SessionPoolConfig { + max_size: 1, + session: dead_endpoint_config(), + ..Default::default() + }; + let pool = SessionPool::new(cfg).unwrap(); + // Simulate a handed-out session dying before release. + pool.live.fetch_add(1, Ordering::Relaxed); + pool.release(Session::new(dead_endpoint_config())); + assert_eq!(pool.live_count(), 0); + let state = pool.state.lock().unwrap(); + assert!(state.idle.is_empty()); + } + + #[test] + fn table_pool_forces_table_dialect() { + let cfg = SessionPoolConfig { + session: SessionConfig { + sql_dialect: "tree".into(), // deliberately wrong + ..dead_endpoint_config() + }, + ..Default::default() + }; + let pool = TableSessionPool::new(cfg).unwrap(); + assert_eq!(pool.pool.config.session.sql_dialect, "table"); + } + + /// Live-server tests: acquire/release round-trip, reuse, and blocking + /// hand-off. Skipped when no IoTDB server is reachable. + #[test] + fn live_pool_acquire_release_reuse() { + if !live_server_available() { + eprintln!("skipping live_pool_acquire_release_reuse: no server on 127.0.0.1:6667"); + return; + } + let cfg = SessionPoolConfig { + max_size: 2, + acquire_timeout: Duration::from_secs(5), + ..Default::default() + }; + let pool = SessionPool::new(cfg).unwrap(); + + { + let mut s1 = pool.acquire().unwrap(); + assert!(s1.is_open()); + let mut ds = s1.execute_query("SHOW DATABASES").unwrap(); + while ds.next_row().unwrap().is_some() {} + drop(ds); + let _s2 = pool.acquire().unwrap(); + assert_eq!(pool.live_count(), 2); + } + // Both released; a re-acquire reuses an idle session (no growth). + let _s3 = pool.acquire().unwrap(); + assert_eq!(pool.live_count(), 2); + drop(_s3); + pool.close(); + assert_eq!(pool.live_count(), 0); + } + + #[test] + fn live_pool_waiter_wakes_on_release() { + if !live_server_available() { + eprintln!("skipping live_pool_waiter_wakes_on_release: no server on 127.0.0.1:6667"); + return; + } + let cfg = SessionPoolConfig { + max_size: 1, + acquire_timeout: Duration::from_secs(5), + ..Default::default() + }; + let pool = std::sync::Arc::new(SessionPool::new(cfg).unwrap()); + let guard = pool.acquire().unwrap(); + + let p2 = pool.clone(); + let waiter = std::thread::spawn(move || p2.acquire().map(|s| s.is_open())); + std::thread::sleep(Duration::from_millis(100)); + drop(guard); // wakes the waiter + assert!(waiter.join().unwrap().unwrap()); + assert_eq!(pool.live_count(), 1); + } +} diff --git a/src/client/session.rs b/src/client/session.rs index a0c0a68..5909ea9 100644 --- a/src/client/session.rs +++ b/src/client/session.rs @@ -21,7 +21,9 @@ use std::collections::BTreeMap; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; +use crate::client::dataset::SessionDataSet; use crate::connection::{Connection, Endpoint}; +use crate::data::Tablet; use crate::error::{Error, Result}; use crate::protocol::client::{ TIClientRPCServiceSyncClient, TSCloseOperationReq, TSCloseSessionReq, TSExecuteStatementReq, @@ -217,9 +219,19 @@ impl Session { Ok(()) } + /// Execute a query statement, returning a [`SessionDataSet`] that + /// borrows this session until it is dropped or closed — fetches and + /// closeOperation must reach the node that owns the query id (spec + /// gotcha #13), so the connection stays pinned to the result set. + pub fn execute_query(&mut self, sql: &str) -> Result<SessionDataSet<'_>> { + let handle = self.execute_query_raw(sql)?; + Ok(SessionDataSet::new(self, handle)) + } + /// Execute a query statement, returning raw TsBlock bytes plus metadata. - /// Decoding happens in the data layer. - pub fn execute_query(&mut self, sql: &str) -> Result<QueryHandle> { + /// Low-level path: decoding and pagination are the caller's problem — + /// prefer [`Session::execute_query`]. + pub fn execute_query_raw(&mut self, sql: &str) -> Result<QueryHandle> { let req = self.statement_req(sql); let resp = self .connection_mut()? @@ -276,12 +288,42 @@ impl Session { } } + /// Insert a [`Tablet`] (tree or table model). Serializes per protocol + /// spec §3 — column-major values with trailing null bitmaps, i64-BE + /// timestamps — sorting rows by timestamp first (spec §3.5). Table-model + /// tablets are sent with `writeToTable=true` plus their column + /// categories, and are never aligned (spec §6). + pub fn insert_tablet(&mut self, tablet: &Tablet) -> Result<()> { + // Serialization sorts in place; clone so the caller's tablet order + // is untouched (the clone is cheap relative to the RPC). + let mut tablet = tablet.clone(); + let values = tablet.serialize_values(); + let timestamps = tablet.serialize_timestamps(); + let (write_to_table, column_categories) = match tablet.column_categories() { + Some(categories) => ( + Some(true), + Some(categories.iter().map(|c| c.code()).collect()), + ), + None => (None, None), + }; + self.insert_tablet_raw( + tablet.target(), + tablet.measurements().to_vec(), + tablet.types().iter().map(|t| t.code()).collect(), + values, + timestamps, + tablet.row_count() as i32, + tablet.is_aligned(), + write_to_table, + column_categories, + ) + } + /// Insert a tablet from pre-serialized buffers (see protocol spec §3). /// /// `values` is the column-major value buffer with trailing null bitmaps; - /// `timestamps` is the `size × 8-byte i64 BE` buffer. Serialization from a - /// `Tablet` lives in the data layer; a typed `insert_tablet(&Tablet)` is - /// wired in the next phase. + /// `timestamps` is the `size × 8-byte i64 BE` buffer. Prefer the typed + /// [`Session::insert_tablet`]; this is the low-level escape hatch. #[allow(clippy::too_many_arguments)] pub fn insert_tablet_raw( &mut self, @@ -488,9 +530,13 @@ mod tests { session.open().expect("open session"); assert!(session.is_open()); - let handle = session.execute_query("SHOW DATABASES").expect("query"); - assert!(!handle.columns.is_empty()); - session.close_query(handle.query_id); + { + let mut dataset = session.execute_query("SHOW DATABASES").expect("query"); + assert!(!dataset.columns().is_empty()); + while let Some(row) = dataset.next_row().expect("next_row") { + assert_eq!(row.values.len(), dataset.columns().len()); + } + } // dataset drop closes the query and releases the session borrow session.close().expect("close session"); assert!(!session.is_open()); diff --git a/src/client/table_session.rs b/src/client/table_session.rs new file mode 100644 index 0000000..10f740e --- /dev/null +++ b/src/client/table_session.rs @@ -0,0 +1,241 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Table-model (relational) session — a thin wrapper over [`Session`] with +//! `sql_dialect="table"` enforced at open time (protocol spec §6). + +use std::time::Duration; + +use crate::client::dataset::SessionDataSet; +use crate::client::session::{Session, SessionConfig}; +use crate::connection::Endpoint; +use crate::data::Tablet; +use crate::error::{Error, Result}; + +/// Builder for [`TableSession`]. Defaults mirror [`SessionConfig`] except +/// the SQL dialect, which is always `"table"`. +/// +/// ```no_run +/// use iotdb_client::TableSession; +/// +/// let mut session = TableSession::builder() +/// .node_urls(&["127.0.0.1:6667"])? +/// .database("mydb") +/// .build()?; +/// session.execute_non_query("SHOW TABLES")?; +/// # Ok::<(), iotdb_client::Error>(()) +/// ``` +#[derive(Debug, Clone)] +pub struct TableSessionBuilder { + config: SessionConfig, +} + +impl Default for TableSessionBuilder { + fn default() -> Self { + Self { + config: SessionConfig { + sql_dialect: "table".into(), + ..SessionConfig::default() + }, + } + } +} + +impl TableSessionBuilder { + pub fn new() -> Self { + Self::default() + } + + /// Endpoints from `"host:port"` node-url strings. + pub fn node_urls<S: AsRef<str>>(mut self, node_urls: &[S]) -> Result<Self> { + self.config = self.config.with_node_urls(node_urls)?; + Ok(self) + } + + pub fn endpoints(mut self, endpoints: Vec<Endpoint>) -> Self { + self.config.endpoints = endpoints; + self + } + + pub fn username(mut self, username: impl Into<String>) -> Self { + self.config.username = username.into(); + self + } + + pub fn password(mut self, password: impl Into<String>) -> Self { + self.config.password = password.into(); + self + } + + /// Database to select at open time (sent as config key `"db"`). + pub fn database(mut self, database: impl Into<String>) -> Self { + self.config.database = Some(database.into()); + self + } + + pub fn fetch_size(mut self, fetch_size: i32) -> Self { + self.config.fetch_size = fetch_size; + self + } + + pub fn zone_id(mut self, zone_id: impl Into<String>) -> Self { + self.config.zone_id = zone_id.into(); + self + } + + pub fn connect_timeout(mut self, timeout: Duration) -> Self { + self.config.connect_timeout = timeout; + self + } + + pub fn query_timeout_ms(mut self, timeout_ms: i64) -> Self { + self.config.query_timeout_ms = timeout_ms; + self + } + + /// The [`SessionConfig`] this builder resolves to (dialect always + /// `"table"`). + pub fn config(&self) -> &SessionConfig { + &self.config + } + + /// Open the session against the configured endpoints. + pub fn build(mut self) -> Result<TableSession> { + // Enforce the dialect even if the config was mutated externally. + self.config.sql_dialect = "table".into(); + let mut session = Session::new(self.config); + session.open()?; + Ok(TableSession { session }) + } +} + +/// A table-model (relational) session. All statements run in the `"table"` +/// SQL dialect; inserts require table-model tablets +/// ([`Tablet::new_table`]). +pub struct TableSession { + session: Session, +} + +impl TableSession { + pub fn builder() -> TableSessionBuilder { + TableSessionBuilder::new() + } + + /// The database currently selected on this session, if any — tracked + /// from the open-time `"db"` config key and `USE <db>` responses. + pub fn database(&self) -> Option<&str> { + self.session.database() + } + + pub fn is_open(&self) -> bool { + self.session.is_open() + } + + /// Insert a table-model tablet (`writeToTable=true` + column + /// categories on the wire). Rejects tree-model tablets. + pub fn insert(&mut self, tablet: &Tablet) -> Result<()> { + if !tablet.is_table_model() { + return Err(Error::Client( + "TableSession::insert requires a table-model tablet (Tablet::new_table)".into(), + )); + } + self.session.insert_tablet(tablet) + } + + /// Execute a query; the returned [`SessionDataSet`] borrows this + /// session until closed or dropped. + pub fn execute_query(&mut self, sql: &str) -> Result<SessionDataSet<'_>> { + self.session.execute_query(sql) + } + + /// Execute a non-query statement (DDL/DML), tracking `USE <db>`. + pub fn execute_non_query(&mut self, sql: &str) -> Result<()> { + self.session.execute_non_query(sql) + } + + pub fn close(&mut self) -> Result<()> { + self.session.close() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::data::TSDataType; + + #[test] + fn builder_defaults_to_table_dialect() { + let b = TableSessionBuilder::new() + .node_urls(&["10.0.0.1:6667"]) + .unwrap() + .username("u") + .password("p") + .database("db1") + .fetch_size(500) + .zone_id("UTC") + .connect_timeout(Duration::from_secs(3)) + .query_timeout_ms(1_000); + let cfg = b.config(); + assert_eq!(cfg.sql_dialect, "table"); + assert_eq!(cfg.endpoints, vec![Endpoint::new("10.0.0.1", 6667)]); + assert_eq!(cfg.username, "u"); + assert_eq!(cfg.password, "p"); + assert_eq!(cfg.database.as_deref(), Some("db1")); + assert_eq!(cfg.fetch_size, 500); + assert_eq!(cfg.zone_id, "UTC"); + assert_eq!(cfg.connect_timeout, Duration::from_secs(3)); + assert_eq!(cfg.query_timeout_ms, 1_000); + } + + #[test] + fn insert_rejects_tree_model_tablets() { + let mut session = TableSession { + session: Session::new(SessionConfig::default()), + }; + let tablet = Tablet::new("root.sg.d1", vec!["s1".into()], vec![TSDataType::Int32]).unwrap(); + match session.insert(&tablet) { + Err(Error::Client(msg)) => assert!(msg.contains("table-model")), + other => panic!("expected client error, got {other:?}"), + } + } + + /// Live-server test; skipped when no IoTDB instance is reachable. + #[test] + fn live_table_session_roundtrip() { + use std::net::TcpStream; + if TcpStream::connect_timeout( + &"127.0.0.1:6667".parse().unwrap(), + Duration::from_millis(300), + ) + .is_err() + { + eprintln!("skipping live_table_session_roundtrip: no IoTDB server on 127.0.0.1:6667"); + return; + } + + let mut session = TableSession::builder().build().expect("open"); + assert!(session.is_open()); + session + .execute_non_query("CREATE DATABASE IF NOT EXISTS rust_client_test") + .expect("create db"); + session + .execute_non_query("USE rust_client_test") + .expect("use db"); + assert_eq!(session.database(), Some("rust_client_test")); + session.close().expect("close"); + } +} diff --git a/src/data/tablet.rs b/src/data/tablet.rs index 66ab5ba..8cffaf8 100644 --- a/src/data/tablet.rs +++ b/src/data/tablet.rs @@ -44,6 +44,9 @@ pub struct Tablet { values: Vec<Vec<Option<Value>>>, /// Table model only: one category per column. column_categories: Option<Vec<ColumnCategory>>, + /// Tree model only: write to an aligned device. Always `false` on the + /// wire for table-model tablets (spec §6). + aligned: bool, } impl Tablet { @@ -68,9 +71,21 @@ impl Tablet { timestamps: Vec::new(), values: vec![Vec::new(); columns], column_categories: None, + aligned: false, }) } + /// Creates an empty tree-model tablet for an **aligned** device. + pub fn new_aligned( + device_id: impl Into<String>, + measurements: Vec<String>, + types: Vec<TSDataType>, + ) -> Result<Tablet> { + let mut tablet = Tablet::new(device_id, measurements, types)?; + tablet.aligned = true; + Ok(tablet) + } + /// Creates an empty table-model tablet for `table_name` with one /// [`ColumnCategory`] per column. pub fn new_table( @@ -118,6 +133,12 @@ impl Tablet { self.column_categories.as_deref() } + /// True for aligned-device tree-model tablets (spec §3.1 field 8). + /// Always `false` for table-model tablets. + pub fn is_aligned(&self) -> bool { + self.aligned + } + pub fn row_count(&self) -> usize { self.timestamps.len() } diff --git a/src/data/tsblock.rs b/src/data/tsblock.rs index a393cf2..f9265e8 100644 --- a/src/data/tsblock.rs +++ b/src/data/tsblock.rs @@ -23,12 +23,13 @@ use super::value::Value; use super::TSDataType; use crate::error::{Error, Result}; -/// TsBlock column encodings (spec §5.2). -const ENCODING_BYTE_ARRAY: u8 = 0; -const ENCODING_INT32_ARRAY: u8 = 1; -const ENCODING_INT64_ARRAY: u8 = 2; -const ENCODING_BINARY_ARRAY: u8 = 3; -const ENCODING_RLE: u8 = 4; +/// TsBlock column encodings (spec §5.2). `pub(crate)` so the dataset tests +/// can craft synthetic blocks with the same constants. +pub(crate) const ENCODING_BYTE_ARRAY: u8 = 0; +pub(crate) const ENCODING_INT32_ARRAY: u8 = 1; +pub(crate) const ENCODING_INT64_ARRAY: u8 = 2; +pub(crate) const ENCODING_BINARY_ARRAY: u8 = 3; +pub(crate) const ENCODING_RLE: u8 = 4; /// One decoded TsBlock: `position_count` rows of a time column plus /// `columns.len()` value columns. Null cells are [`Value::Null`]. @@ -294,13 +295,15 @@ impl<'a> Reader<'a> { } } +/// Encoding helpers for crafting synthetic TsBlocks in unit tests (shared +/// with the dataset state-machine tests). #[cfg(test)] -mod tests { +pub(crate) mod test_util { use super::*; /// Builds a TsBlock header: valueColumnCount, type bytes, positionCount, /// time encoding (Int64Array), value encodings. - fn header(types: &[TSDataType], position_count: i32, encodings: &[u8]) -> Vec<u8> { + pub(crate) fn header(types: &[TSDataType], position_count: i32, encodings: &[u8]) -> Vec<u8> { let mut b = Vec::new(); b.extend_from_slice(&(types.len() as i32).to_be_bytes()); for &t in types { @@ -313,7 +316,7 @@ mod tests { } /// Dense Int64Array time column with no nulls. - fn time_column(ts: &[i64]) -> Vec<u8> { + pub(crate) fn time_column(ts: &[i64]) -> Vec<u8> { let mut b = vec![0u8]; // mayHaveNull = 0 for t in ts { b.extend_from_slice(&t.to_be_bytes()); @@ -321,6 +324,28 @@ mod tests { b } + /// One-Int32-column block: timestamps `ts`, dense non-null values `vals`. + pub(crate) fn int32_block(ts: &[i64], vals: &[i32]) -> Vec<u8> { + assert_eq!(ts.len(), vals.len()); + let mut b = header( + &[TSDataType::Int32], + ts.len() as i32, + &[ENCODING_INT32_ARRAY], + ); + b.extend_from_slice(&time_column(ts)); + b.push(0); // mayHaveNull + for v in vals { + b.extend_from_slice(&v.to_be_bytes()); + } + b + } +} + +#[cfg(test)] +mod tests { + use super::test_util::{header, time_column}; + use super::*; + #[test] fn int32_column_no_nulls() { let mut b = header(&[TSDataType::Int32], 2, &[ENCODING_INT32_ARRAY]); diff --git a/src/lib.rs b/src/lib.rs index 130d7a8..5cc766f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -13,6 +13,10 @@ pub mod data; pub mod error; pub mod protocol; +pub use client::dataset::{Row, SessionDataSet}; +pub use client::pool::{PooledSession, SessionPool, SessionPoolConfig, TableSessionPool}; pub use client::session::{QueryHandle, Session, SessionConfig}; +pub use client::table_session::{TableSession, TableSessionBuilder}; pub use connection::Endpoint; +pub use data::{ColumnCategory, TSDataType, Tablet, TsBlock, Value}; pub use error::{Error, Result};
