connectorx/sources/clickhouse/
mod.rs

1//! Source implementation for ClickHouse using the native protocol.
2
3mod errors;
4mod typesystem;
5
6pub use self::errors::ClickHouseSourceError;
7pub use self::typesystem::{ClickHouseTypeSystem, TypeMetadata};
8
9use crate::{
10    data_order::DataOrder,
11    errors::ConnectorXError,
12    sources::{
13        clickhouse::typesystem::DataType, PartitionParser, Produce, Source, SourcePartition,
14    },
15    sql::{count_query, limit0_query, CXQuery},
16};
17use anyhow::anyhow;
18use chrono::{DateTime, Duration, NaiveDate, NaiveTime, Utc};
19use chrono_tz::Tz;
20use clickhouse::Client;
21use fehler::{throw, throws};
22use rust_decimal::Decimal;
23use serde::Deserialize;
24use serde_json::Value as JsonValue;
25use sqlparser::dialect::{ClickHouseDialect, GenericDialect};
26use std::io::{Cursor, Read};
27use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
28use std::sync::Arc;
29use tokio::runtime::Runtime;
30use uuid::Uuid;
31
32/// ClickHouse source that uses the HTTP protocol.
33pub struct ClickHouseSource {
34    rt: Arc<Runtime>,
35    pub client: Client,
36    origin_query: Option<String>,
37    queries: Vec<CXQuery<String>>,
38    names: Vec<String>,
39    schema: Vec<ClickHouseTypeSystem>,
40    metadata: Vec<TypeMetadata>,
41}
42
43impl ClickHouseSource {
44    #[throws(ClickHouseSourceError)]
45    pub fn new(rt: Arc<Runtime>, conn: &str) -> Self {
46        let url = url::Url::parse(conn)?;
47
48        let use_https = url
49            .query_pairs()
50            .find(|(k, v)| k == "protocol" && v == "https")
51            .is_some();
52
53        let base_url = format!(
54            "{}://{}:{}",
55            if use_https { "https" } else { "http" },
56            url.host_str().unwrap_or("localhost"),
57            url.port().unwrap_or(8123)
58        );
59
60        // Only the client's `lz4` feature is enabled, so it decodes LZ4 only;
61        // newer servers default to ZSTD.
62        let mut client = Client::default()
63            .with_url(&base_url)
64            .with_setting("network_compression_method", "lz4");
65
66        let database = url.path().trim_start_matches('/');
67        if !database.is_empty() {
68            client = client.with_database(database);
69        }
70
71        let username = url.username();
72        if !username.is_empty() {
73            client = client.with_user(username);
74        }
75
76        let password = url.password().unwrap_or("");
77        if !password.is_empty() {
78            client = client.with_password(password);
79        }
80
81        Self {
82            rt,
83            client,
84            origin_query: None,
85            queries: vec![],
86            names: vec![],
87            schema: vec![],
88            metadata: vec![],
89        }
90    }
91}
92
93impl Source for ClickHouseSource
94where
95    ClickHouseSourcePartition:
96        SourcePartition<TypeSystem = ClickHouseTypeSystem, Error = ClickHouseSourceError>,
97{
98    const DATA_ORDERS: &'static [DataOrder] = &[DataOrder::RowMajor];
99    type Partition = ClickHouseSourcePartition;
100    type TypeSystem = ClickHouseTypeSystem;
101    type Error = ClickHouseSourceError;
102
103    #[throws(ClickHouseSourceError)]
104    fn set_data_order(&mut self, data_order: DataOrder) {
105        if !matches!(data_order, DataOrder::RowMajor) {
106            throw!(ConnectorXError::UnsupportedDataOrder(data_order));
107        }
108    }
109
110    fn set_queries<Q: ToString>(&mut self, queries: &[CXQuery<Q>]) {
111        self.queries = queries.iter().map(|q| q.map(Q::to_string)).collect();
112    }
113
114    fn set_origin_query(&mut self, query: Option<String>) {
115        self.origin_query = query;
116    }
117
118    #[throws(ClickHouseSourceError)]
119    fn fetch_metadata(&mut self) {
120        assert!(!self.queries.is_empty());
121
122        let first_query = &self.queries[0];
123        let l1query = limit0_query(first_query, &ClickHouseDialect {})?;
124
125        let describe_query = format!("DESCRIBE ({})", l1query.as_str());
126
127        let response = self.rt.block_on(async {
128            let mut cursor = self
129                .client
130                .query(&describe_query)
131                .fetch_bytes("JSONCompact")
132                .map_err(|e| anyhow!("ClickHouse error: {}", e))?;
133            let bytes = cursor
134                .collect()
135                .await
136                .map_err(|e| anyhow!("ClickHouse error: {}", e))?;
137            Ok::<_, ClickHouseSourceError>(bytes)
138        })?;
139
140        #[derive(Debug, Deserialize)]
141        struct DescribeResponse {
142            data: Vec<Vec<JsonValue>>,
143        }
144
145        let parsed: DescribeResponse = serde_json::from_slice(&response)
146            .map_err(|e| anyhow!("Failed to parse DESCRIBE response: {}", e))?;
147
148        let mut names = Vec::new();
149        let mut types = Vec::new();
150        let mut metadata = Vec::new();
151
152        for row in parsed.data {
153            if row.len() >= 2 {
154                let name = row[0].as_str().unwrap_or("").to_string();
155                let type_str = row[1].as_str().unwrap_or("String");
156                let (ts, meta) = ClickHouseTypeSystem::from_type_str_with_metadata(type_str);
157                names.push(name);
158                types.push(ts);
159                metadata.push(meta);
160            }
161        }
162
163        self.names = names;
164        self.schema = types;
165        self.metadata = metadata;
166    }
167
168    #[throws(ClickHouseSourceError)]
169    fn result_rows(&mut self) -> Option<usize> {
170        match &self.origin_query {
171            Some(q) => {
172                let cxq = CXQuery::Naked(q.clone());
173                let cquery = count_query(&cxq, &ClickHouseDialect {})?;
174
175                let response = self.rt.block_on(async {
176                    let mut cursor = self
177                        .client
178                        .query(cquery.as_str())
179                        .fetch_bytes("JSONCompact")
180                        .map_err(|e| anyhow!("ClickHouse error: {}", e))?;
181                    let bytes = cursor
182                        .collect()
183                        .await
184                        .map_err(|e| anyhow!("ClickHouse error: {}", e))?;
185                    Ok::<_, ClickHouseSourceError>(bytes)
186                })?;
187
188                #[derive(Debug, Deserialize)]
189                struct CountResponse {
190                    data: Vec<Vec<JsonValue>>,
191                }
192
193                let parsed: CountResponse = serde_json::from_slice(&response)
194                    .map_err(|e| anyhow!("Failed to parse count response: {}", e))?;
195
196                if let Some(row) = parsed.data.first() {
197                    if let Some(count_val) = row.first() {
198                        let count = match count_val {
199                            JsonValue::Number(n) => n.as_u64().unwrap_or(0),
200                            JsonValue::String(s) => s.parse().unwrap_or(0),
201                            _ => 0,
202                        };
203                        return Some(count as usize);
204                    }
205                }
206                None
207            }
208            None => None,
209        }
210    }
211
212    fn names(&self) -> Vec<String> {
213        self.names.clone()
214    }
215
216    fn schema(&self) -> Vec<Self::TypeSystem> {
217        self.schema.clone()
218    }
219
220    #[throws(ClickHouseSourceError)]
221    fn partition(self) -> Vec<Self::Partition> {
222        let mut ret = vec![];
223        for query in self.queries {
224            ret.push(ClickHouseSourcePartition::new(
225                self.rt.clone(),
226                self.client.clone(),
227                &query,
228                &self.schema,
229                &self.metadata,
230            ));
231        }
232        ret
233    }
234}
235
236pub struct ClickHouseSourcePartition {
237    rt: Arc<Runtime>,
238    client: Client,
239    query: CXQuery<String>,
240    schema: Vec<ClickHouseTypeSystem>,
241    metadata: Vec<TypeMetadata>,
242    nrows: usize,
243    ncols: usize,
244}
245
246impl ClickHouseSourcePartition {
247    pub fn new(
248        rt: Arc<Runtime>,
249        client: Client,
250        query: &CXQuery<String>,
251        schema: &[ClickHouseTypeSystem],
252        metadata: &[TypeMetadata],
253    ) -> Self {
254        Self {
255            rt,
256            client,
257            query: query.clone(),
258            schema: schema.to_vec(),
259            metadata: metadata.to_vec(),
260            nrows: 0,
261            ncols: schema.len(),
262        }
263    }
264}
265
266impl SourcePartition for ClickHouseSourcePartition {
267    type TypeSystem = ClickHouseTypeSystem;
268    type Parser<'a> = ClickHouseSourceParser<'a>;
269    type Error = ClickHouseSourceError;
270
271    #[throws(ClickHouseSourceError)]
272    fn result_rows(&mut self) {
273        let cquery = count_query(&self.query, &GenericDialect {})?;
274
275        let response = self.rt.block_on(async {
276            let mut cursor = self
277                .client
278                .query(cquery.as_str())
279                .fetch_bytes("JSONCompact")
280                .map_err(|e| anyhow!("ClickHouse error: {}", e))?;
281            let bytes = cursor
282                .collect()
283                .await
284                .map_err(|e| anyhow!("ClickHouse error: {}", e))?;
285            Ok::<_, ClickHouseSourceError>(bytes)
286        })?;
287
288        #[derive(Debug, Deserialize)]
289        struct CountResponse {
290            data: Vec<Vec<JsonValue>>,
291        }
292
293        let parsed: CountResponse = serde_json::from_slice(&response)
294            .map_err(|e| anyhow!("Failed to parse count response: {}", e))?;
295
296        if let Some(row) = parsed.data.first() {
297            if let Some(count_val) = row.first() {
298                let count = match count_val {
299                    JsonValue::Number(n) => n.as_u64().unwrap_or(0),
300                    JsonValue::String(s) => s.parse().unwrap_or(0),
301                    _ => 0,
302                };
303                self.nrows = count as usize;
304            }
305        }
306    }
307
308    #[throws(ClickHouseSourceError)]
309    fn parser(&mut self) -> Self::Parser<'_> {
310        ClickHouseSourceParser::new(
311            self.rt.clone(),
312            self.client.clone(),
313            self.query.clone(),
314            &self.schema,
315            &self.metadata,
316        )?
317    }
318
319    fn nrows(&self) -> usize {
320        self.nrows
321    }
322
323    fn ncols(&self) -> usize {
324        self.ncols
325    }
326}
327
328struct BinaryReader<'a> {
329    cursor: Cursor<&'a [u8]>,
330}
331
332impl<'a> BinaryReader<'a> {
333    fn new(data: &'a [u8]) -> Self {
334        Self {
335            cursor: Cursor::new(data),
336        }
337    }
338
339    fn read_bytes<const N: usize>(&mut self) -> Result<[u8; N], ClickHouseSourceError> {
340        let mut buf = [0u8; N];
341        self.cursor
342            .read_exact(&mut buf)
343            .map_err(|e| anyhow!("Failed to read {} bytes: {}", N, e))?;
344        Ok(buf)
345    }
346
347    fn read_u8(&mut self) -> Result<u8, ClickHouseSourceError> {
348        Ok(self.read_bytes::<1>()?[0])
349    }
350
351    fn read_i8(&mut self) -> Result<i8, ClickHouseSourceError> {
352        Ok(self.read_u8()? as i8)
353    }
354
355    fn read_u16(&mut self) -> Result<u16, ClickHouseSourceError> {
356        Ok(u16::from_le_bytes(self.read_bytes()?))
357    }
358
359    fn read_i16(&mut self) -> Result<i16, ClickHouseSourceError> {
360        Ok(i16::from_le_bytes(self.read_bytes()?))
361    }
362
363    fn read_u32(&mut self) -> Result<u32, ClickHouseSourceError> {
364        Ok(u32::from_le_bytes(self.read_bytes()?))
365    }
366
367    fn read_i32(&mut self) -> Result<i32, ClickHouseSourceError> {
368        Ok(i32::from_le_bytes(self.read_bytes()?))
369    }
370
371    fn read_u64(&mut self) -> Result<u64, ClickHouseSourceError> {
372        Ok(u64::from_le_bytes(self.read_bytes()?))
373    }
374
375    fn read_i64(&mut self) -> Result<i64, ClickHouseSourceError> {
376        Ok(i64::from_le_bytes(self.read_bytes()?))
377    }
378
379    fn read_f32(&mut self) -> Result<f32, ClickHouseSourceError> {
380        Ok(f32::from_le_bytes(self.read_bytes()?))
381    }
382
383    fn read_f64(&mut self) -> Result<f64, ClickHouseSourceError> {
384        Ok(f64::from_le_bytes(self.read_bytes()?))
385    }
386
387    fn read_varint(&mut self) -> Result<u64, ClickHouseSourceError> {
388        let mut result: u64 = 0;
389        let mut shift = 0;
390        loop {
391            let byte = self.read_u8()?;
392            result |= ((byte & 0x7f) as u64) << shift;
393            if byte & 0x80 == 0 {
394                break;
395            }
396            shift += 7;
397            if shift >= 64 {
398                return Err(anyhow!("Varint too long").into());
399            }
400        }
401        Ok(result)
402    }
403
404    fn read_fixed_string(&mut self, len: usize) -> Result<Vec<u8>, ClickHouseSourceError> {
405        let mut buf = vec![0u8; len];
406        self.cursor
407            .read_exact(&mut buf)
408            .map_err(|e| anyhow!("Failed to read FixedString: {}", e))?;
409        Ok(buf)
410    }
411
412    fn read_string(&mut self) -> Result<String, ClickHouseSourceError> {
413        let len = self.read_varint()? as usize;
414        let mut buf = vec![0u8; len];
415        self.cursor
416            .read_exact(&mut buf)
417            .map_err(|e| anyhow!("Failed to read string: {}", e))?;
418        String::from_utf8(buf).map_err(|e| anyhow!("Invalid UTF-8: {}", e).into())
419    }
420
421    fn read_uuid(&mut self) -> Result<Uuid, ClickHouseSourceError> {
422        // ClickHouse stores UUID as two UInt64 in big-endian order
423        let high = self.read_u64()?;
424        let low = self.read_u64()?;
425
426        Ok(Uuid::from_u64_pair(high, low))
427    }
428
429    fn read_bool(&mut self) -> Result<bool, ClickHouseSourceError> {
430        Ok(self.read_u8()? != 0)
431    }
432
433    fn read_date(&mut self) -> Result<NaiveDate, ClickHouseSourceError> {
434        // ClickHouse stores Date as UInt16 representing days since 1970-01-01
435        let days = self.read_u16()? as i64;
436        let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
437        epoch
438            .checked_add_signed(Duration::days(days))
439            .ok_or_else(|| anyhow!("Invalid date value: {} days since epoch", days).into())
440    }
441
442    fn read_date32(&mut self) -> Result<NaiveDate, ClickHouseSourceError> {
443        // ClickHouse stores Date32 as Int32 representing days since 1970-01-01
444        // negative values represent dates before 1970-01-01
445        let days = self.read_i32()? as i64;
446        let epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
447        epoch
448            .checked_add_signed(Duration::days(days))
449            .ok_or_else(|| anyhow!("Invalid date32 value: {} days since epoch", days).into())
450    }
451
452    fn read_datetime(&mut self, tz: Option<&Tz>) -> Result<DateTime<Utc>, ClickHouseSourceError> {
453        let seconds = self.read_u32()? as i64;
454        chrono::DateTime::from_timestamp(seconds, 0)
455            .map(|dt| dt.with_timezone(tz.unwrap_or(&Tz::UTC)).to_utc())
456            .ok_or_else(|| {
457                anyhow!("Invalid datetime value: {} seconds since epoch", seconds).into()
458            })
459    }
460
461    fn read_datetime64(
462        &mut self,
463        precision: u8,
464        tz: Option<&Tz>,
465    ) -> Result<DateTime<Utc>, ClickHouseSourceError> {
466        if precision > 9 {
467            return Err(anyhow!("Unsupported DateTime64 precision: {}", precision).into());
468        }
469        let ticks = self.read_i64()?;
470        let scale = 10_i64.pow(precision as u32);
471        let seconds = ticks.div_euclid(scale);
472        let fractional_ticks = ticks.rem_euclid(scale) as u32;
473        let nanos = fractional_ticks * 10_u32.pow(9 - precision as u32);
474
475        DateTime::from_timestamp(seconds, nanos)
476            .map(|dt| dt.with_timezone(tz.unwrap_or(&Tz::UTC)).to_utc())
477            .ok_or_else(|| anyhow!("Invalid datetime64 value: {} ticks", ticks).into())
478    }
479
480    fn read_decimal(&mut self, precision: u8, scale: u8) -> Result<Decimal, ClickHouseSourceError> {
481        match precision {
482            1..=9 => {
483                let value = self.read_i32()?;
484                Ok(Decimal::new(value as i64, scale as u32))
485            }
486            10..=18 => {
487                let value = self.read_i64()?;
488                Ok(Decimal::new(value, scale as u32))
489            }
490            _ => Err(anyhow!("Unsupported Decimal precision: {}", precision).into()),
491        }
492    }
493
494    fn read_time(&mut self) -> Result<NaiveTime, ClickHouseSourceError> {
495        let seconds = self.read_u32()? as i64;
496        Ok(
497            NaiveTime::from_num_seconds_from_midnight_opt(seconds as u32, 0)
498                .ok_or_else(|| anyhow!("Invalid time value: {} seconds since midnight", seconds))?,
499        )
500    }
501
502    fn read_time64(&mut self, precision: u8) -> Result<NaiveTime, ClickHouseSourceError> {
503        if precision > 9 {
504            return Err(anyhow!("Unsupported Time64 precision: {}", precision).into());
505        }
506        let ticks = self.read_i64()?;
507        let nanos = ticks * 10_i64.pow(9 - precision as u32);
508        Ok(NaiveTime::from_num_seconds_from_midnight_opt(
509            (nanos / 1_000_000_000) as u32,
510            (nanos % 1_000_000_000) as u32,
511        )
512        .ok_or_else(|| anyhow!("Invalid time64 value: {} ticks since midnight", ticks))?)
513    }
514
515    fn read_ipv4(&mut self) -> Result<IpAddr, ClickHouseSourceError> {
516        let bytes = self.read_u32()?;
517
518        Ok(IpAddr::V4(Ipv4Addr::from_bits(bytes)))
519    }
520
521    fn read_ipv6(&mut self) -> Result<IpAddr, ClickHouseSourceError> {
522        let seg1 = u16::from_be(self.read_u16()?);
523        let seg2 = u16::from_be(self.read_u16()?);
524        let seg3 = u16::from_be(self.read_u16()?);
525        let seg4 = u16::from_be(self.read_u16()?);
526        let seg5 = u16::from_be(self.read_u16()?);
527        let seg6 = u16::from_be(self.read_u16()?);
528        let seg7 = u16::from_be(self.read_u16()?);
529        let seg8 = u16::from_be(self.read_u16()?);
530
531        Ok(IpAddr::V6(Ipv6Addr::from_segments([
532            seg1, seg2, seg3, seg4, seg5, seg6, seg7, seg8,
533        ])))
534    }
535
536    fn read_enum8(&mut self) -> Result<i8, ClickHouseSourceError> {
537        self.read_i8()
538    }
539
540    fn read_enum16(&mut self) -> Result<i16, ClickHouseSourceError> {
541        self.read_i16()
542    }
543
544    fn read_array<T, F>(&mut self, read_elem: F) -> Result<Vec<Option<T>>, ClickHouseSourceError>
545    where
546        F: Fn(&mut Self) -> Result<T, ClickHouseSourceError>,
547    {
548        let len = self.read_varint()? as usize;
549        let mut result = Vec::with_capacity(len);
550        for _ in 0..len {
551            result.push(Some(read_elem(self)?));
552        }
553        Ok(result)
554    }
555
556    fn is_empty(&self) -> bool {
557        self.cursor.position() as usize >= self.cursor.get_ref().len()
558    }
559}
560
561pub struct ClickHouseSourceParser<'a> {
562    rt: Arc<Runtime>,
563    client: Client,
564    query: CXQuery<String>,
565    schema: Vec<ClickHouseTypeSystem>,
566    metadata: Vec<TypeMetadata>,
567    rowbuf: Vec<Vec<DataType>>,
568    ncols: usize,
569    current_row: usize,
570    current_col: usize,
571    is_finished: bool,
572    _phantom: std::marker::PhantomData<&'a ()>,
573}
574
575impl<'a> ClickHouseSourceParser<'a> {
576    #[throws(ClickHouseSourceError)]
577    pub fn new(
578        rt: Arc<Runtime>,
579        client: Client,
580        query: CXQuery<String>,
581        schema: &[ClickHouseTypeSystem],
582        metadata: &[TypeMetadata],
583    ) -> Self {
584        Self {
585            rt,
586            client,
587            query,
588            schema: schema.to_vec(),
589            metadata: metadata.to_vec(),
590            rowbuf: Vec::new(),
591            ncols: schema.len(),
592            current_row: 0,
593            current_col: 0,
594            is_finished: false,
595            _phantom: std::marker::PhantomData,
596        }
597    }
598
599    fn parse_row_binary(
600        &self,
601        reader: &mut BinaryReader,
602    ) -> Result<Vec<DataType>, ClickHouseSourceError> {
603        let mut row = Vec::with_capacity(self.ncols);
604
605        for (col_idx, col_type) in self.schema.iter().enumerate() {
606            let is_nullable = col_type.is_nullable();
607            let meta = &self.metadata[col_idx];
608
609            if is_nullable {
610                let null_flag = reader.read_u8()?;
611                if null_flag == 1 {
612                    row.push(DataType::Null);
613                    continue;
614                }
615            }
616
617            let value = match col_type {
618                ClickHouseTypeSystem::Int8(_) => DataType::Int8(reader.read_i8()?),
619                ClickHouseTypeSystem::Int16(_) => DataType::Int16(reader.read_i16()?),
620                ClickHouseTypeSystem::Int32(_) => DataType::Int32(reader.read_i32()?),
621                ClickHouseTypeSystem::Int64(_) => DataType::Int64(reader.read_i64()?),
622                ClickHouseTypeSystem::UInt8(_) => DataType::UInt8(reader.read_u8()?),
623                ClickHouseTypeSystem::UInt16(_) => DataType::UInt16(reader.read_u16()?),
624                ClickHouseTypeSystem::UInt32(_) => DataType::UInt32(reader.read_u32()?),
625                ClickHouseTypeSystem::UInt64(_) => DataType::UInt64(reader.read_u64()?),
626
627                ClickHouseTypeSystem::Float32(_) => DataType::Float32(reader.read_f32()?),
628                ClickHouseTypeSystem::Float64(_) => DataType::Float64(reader.read_f64()?),
629
630                ClickHouseTypeSystem::Decimal(_) => {
631                    DataType::Decimal(reader.read_decimal(meta.precision, meta.scale)?)
632                }
633
634                ClickHouseTypeSystem::String(_) => DataType::String(reader.read_string()?),
635
636                ClickHouseTypeSystem::FixedString(_) => {
637                    DataType::FixedString(reader.read_fixed_string(meta.length)?)
638                }
639
640                ClickHouseTypeSystem::Date(_) => DataType::Date(reader.read_date()?),
641                ClickHouseTypeSystem::Date32(_) => DataType::Date32(reader.read_date32()?),
642
643                ClickHouseTypeSystem::DateTime(_) => {
644                    DataType::DateTime(reader.read_datetime(meta.timezone.as_ref())?)
645                }
646                ClickHouseTypeSystem::DateTime64(_) => DataType::DateTime64(
647                    reader.read_datetime64(meta.precision, meta.timezone.as_ref())?,
648                ),
649
650                ClickHouseTypeSystem::Time(_) => DataType::Time(reader.read_time()?),
651                ClickHouseTypeSystem::Time64(_) => {
652                    DataType::Time64(reader.read_time64(meta.precision)?)
653                }
654
655                ClickHouseTypeSystem::Enum8(_) => {
656                    let enum_value = reader.read_enum8()?;
657                    let enum_str = meta
658                        .named_values
659                        .as_ref()
660                        .and_then(|h| h.get(&(enum_value as i16)));
661                    DataType::Enum8(enum_str.cloned().unwrap_or_default())
662                }
663                ClickHouseTypeSystem::Enum16(_) => {
664                    let enum_value = reader.read_enum16()?;
665                    let enum_str = meta.named_values.as_ref().and_then(|h| h.get(&enum_value));
666                    DataType::Enum16(enum_str.cloned().unwrap_or_default())
667                }
668
669                ClickHouseTypeSystem::UUID(_) => DataType::UUID(reader.read_uuid()?),
670
671                ClickHouseTypeSystem::IPv4(_) => DataType::IPv4(reader.read_ipv4()?),
672                ClickHouseTypeSystem::IPv6(_) => DataType::IPv6(reader.read_ipv6()?),
673
674                ClickHouseTypeSystem::Bool(_) => DataType::Bool(reader.read_bool()?),
675
676                ClickHouseTypeSystem::ArrayBool(_) => {
677                    DataType::ArrayBool(reader.read_array(BinaryReader::read_bool)?)
678                }
679                ClickHouseTypeSystem::ArrayString(_) => {
680                    DataType::ArrayString(reader.read_array(BinaryReader::read_string)?)
681                }
682                ClickHouseTypeSystem::ArrayInt8(_) => {
683                    DataType::ArrayInt8(reader.read_array(BinaryReader::read_i8)?)
684                }
685                ClickHouseTypeSystem::ArrayInt16(_) => {
686                    DataType::ArrayInt16(reader.read_array(BinaryReader::read_i16)?)
687                }
688                ClickHouseTypeSystem::ArrayInt32(_) => {
689                    DataType::ArrayInt32(reader.read_array(BinaryReader::read_i32)?)
690                }
691                ClickHouseTypeSystem::ArrayInt64(_) => {
692                    DataType::ArrayInt64(reader.read_array(BinaryReader::read_i64)?)
693                }
694                ClickHouseTypeSystem::ArrayUInt8(_) => {
695                    DataType::ArrayUInt8(reader.read_array(BinaryReader::read_u8)?)
696                }
697                ClickHouseTypeSystem::ArrayUInt16(_) => {
698                    DataType::ArrayUInt16(reader.read_array(BinaryReader::read_u16)?)
699                }
700                ClickHouseTypeSystem::ArrayUInt32(_) => {
701                    DataType::ArrayUInt32(reader.read_array(BinaryReader::read_u32)?)
702                }
703                ClickHouseTypeSystem::ArrayUInt64(_) => {
704                    DataType::ArrayUInt64(reader.read_array(BinaryReader::read_u64)?)
705                }
706                ClickHouseTypeSystem::ArrayFloat32(_) => {
707                    DataType::ArrayFloat32(reader.read_array(BinaryReader::read_f32)?)
708                }
709                ClickHouseTypeSystem::ArrayFloat64(_) => {
710                    DataType::ArrayFloat64(reader.read_array(BinaryReader::read_f64)?)
711                }
712                ClickHouseTypeSystem::ArrayDecimal(_) => {
713                    let precision = meta.precision;
714                    let scale = meta.scale;
715                    DataType::ArrayDecimal(reader.read_array(|r| r.read_decimal(precision, scale))?)
716                }
717            };
718
719            row.push(value);
720        }
721
722        Ok(row)
723    }
724
725    #[throws(ClickHouseSourceError)]
726    fn next_loc(&mut self) -> (usize, usize) {
727        let ret = (self.current_row, self.current_col);
728        self.current_row += (self.current_col + 1) / self.ncols;
729        self.current_col = (self.current_col + 1) % self.ncols;
730        ret
731    }
732}
733
734impl<'a> PartitionParser<'a> for ClickHouseSourceParser<'a> {
735    type TypeSystem = ClickHouseTypeSystem;
736    type Error = ClickHouseSourceError;
737
738    #[throws(ClickHouseSourceError)]
739    fn fetch_next(&mut self) -> (usize, bool) {
740        assert!(self.current_col == 0);
741
742        if self.is_finished {
743            return (0, true);
744        }
745
746        let response = self.rt.block_on(async {
747            let mut cursor = self
748                .client
749                .query(self.query.as_str())
750                .fetch_bytes("RowBinary")
751                .map_err(|e| anyhow!("ClickHouse error: {}", e))?;
752            let bytes = cursor
753                .collect()
754                .await
755                .map_err(|e| anyhow!("ClickHouse error: {}", e))?;
756            Ok::<_, ClickHouseSourceError>(bytes)
757        })?;
758        let mut reader = BinaryReader::new(&response);
759        let mut rows = Vec::new();
760
761        while !reader.is_empty() {
762            match self.parse_row_binary(&mut reader) {
763                Ok(row) => rows.push(row),
764                Err(_) => break,
765            }
766        }
767
768        self.rowbuf = rows;
769        self.current_row = 0;
770        self.is_finished = true;
771
772        (self.rowbuf.len(), true)
773    }
774}
775
776macro_rules! impl_produce {
777    ($rust_type:ty, [$($variant:ident),+]) => {
778        impl<'r, 'a> Produce<'r, $rust_type> for ClickHouseSourceParser<'a> {
779            type Error = ClickHouseSourceError;
780
781            #[throws(ClickHouseSourceError)]
782            fn produce(&'r mut self) -> $rust_type {
783                let (ridx, cidx) = self.next_loc()?;
784                let value = &self.rowbuf[ridx][cidx];
785
786                match value {
787                    $(DataType::$variant(v) => *v as $rust_type,)+
788                    _ => throw!(ConnectorXError::cannot_produce::<$rust_type>(Some(
789                        format!("{:?}", value)
790                    ))),
791                }
792            }
793        }
794
795        impl<'r, 'a> Produce<'r, Option<$rust_type>> for ClickHouseSourceParser<'a> {
796            type Error = ClickHouseSourceError;
797
798            #[throws(ClickHouseSourceError)]
799            fn produce(&'r mut self) -> Option<$rust_type> {
800                let (ridx, cidx) = self.next_loc()?;
801                let value = &self.rowbuf[ridx][cidx];
802
803                match value {
804                    DataType::Null => None,
805                    $(DataType::$variant(v) => Some(*v as $rust_type),)+
806                    _ => throw!(ConnectorXError::cannot_produce::<$rust_type>(Some(
807                        format!("{:?}", value)
808                    ))),
809                }
810            }
811        }
812    };
813}
814
815macro_rules! impl_produce_with_clone {
816    ($rust_type:ty, [$($variant:ident),+]) => {
817        impl<'r, 'a> Produce<'r, $rust_type> for ClickHouseSourceParser<'a> {
818            type Error = ClickHouseSourceError;
819
820            #[throws(ClickHouseSourceError)]
821            fn produce(&'r mut self) -> $rust_type {
822                let (ridx, cidx) = self.next_loc()?;
823                let value = &self.rowbuf[ridx][cidx];
824
825                match value {
826                    $(DataType::$variant(v) => v.clone() as $rust_type,)+
827                    _ => throw!(ConnectorXError::cannot_produce::<$rust_type>(Some(
828                        format!("{:?}", value)
829                    ))),
830                }
831            }
832        }
833
834        impl<'r, 'a> Produce<'r, Option<$rust_type>> for ClickHouseSourceParser<'a> {
835            type Error = ClickHouseSourceError;
836
837            #[throws(ClickHouseSourceError)]
838            fn produce(&'r mut self) -> Option<$rust_type> {
839                let (ridx, cidx) = self.next_loc()?;
840                let value = &self.rowbuf[ridx][cidx];
841
842                match value {
843                    DataType::Null => None,
844                    $(DataType::$variant(v) => Some(v.clone() as $rust_type),)+
845                    _ => throw!(ConnectorXError::cannot_produce::<$rust_type>(Some(
846                        format!("{:?}", value)
847                    ))),
848                }
849            }
850        }
851    };
852}
853
854macro_rules! impl_produce_vec {
855    ($rust_type:ty, [$($variant:ident),+]) => {
856        impl<'r, 'a> Produce<'r, Vec<Option<$rust_type>>> for ClickHouseSourceParser<'a> {
857            type Error = ClickHouseSourceError;
858
859            #[throws(ClickHouseSourceError)]
860            fn produce(&'r mut self) -> Vec<Option<$rust_type>> {
861                let (ridx, cidx) = self.next_loc()?;
862                let value = &self.rowbuf[ridx][cidx];
863
864                match value {
865                    $(DataType::$variant(v) => v.clone(),)+
866                    _ => throw!(ConnectorXError::cannot_produce::<Vec<Option<$rust_type>>>(Some(
867                        format!("{:?}", value)
868                    ))),
869                }
870            }
871        }
872
873        impl<'r, 'a> Produce<'r, Option<Vec<Option<$rust_type>>>> for ClickHouseSourceParser<'a> {
874            type Error = ClickHouseSourceError;
875
876            #[throws(ClickHouseSourceError)]
877            fn produce(&'r mut self) -> Option<Vec<Option<$rust_type>>> {
878                let (ridx, cidx) = self.next_loc()?;
879                let value = &self.rowbuf[ridx][cidx];
880
881                match value {
882                    DataType::Null => None,
883                    $(DataType::$variant(v) => Some(v.clone()),)+
884                    _ => throw!(ConnectorXError::cannot_produce::<Option<Vec<Option<$rust_type>>>>(Some(
885                        format!("{:?}", value)
886                    ))),
887                }
888            }
889        }
890    };
891}
892
893impl_produce!(i8, [Int8]);
894impl_produce!(i16, [Int16]);
895impl_produce!(i32, [Int32]);
896impl_produce!(i64, [Int64]);
897impl_produce!(u8, [UInt8]);
898impl_produce!(u16, [UInt16]);
899impl_produce!(u32, [UInt32]);
900impl_produce!(u64, [UInt64]);
901impl_produce!(f32, [Float32]);
902impl_produce!(f64, [Float64]);
903impl_produce!(Decimal, [Decimal]);
904impl_produce_with_clone!(String, [String, Enum8, Enum16]);
905impl_produce_with_clone!(Vec<u8>, [FixedString]);
906impl_produce!(NaiveDate, [Date, Date32]);
907impl_produce!(DateTime<Utc>, [DateTime, DateTime64]);
908impl_produce!(NaiveTime, [Time, Time64]);
909impl_produce!(Uuid, [UUID]);
910impl_produce!(IpAddr, [IPv4, IPv6]);
911impl_produce!(bool, [Bool]);
912impl_produce_vec!(bool, [ArrayBool]);
913impl_produce_vec!(String, [ArrayString]);
914impl_produce_vec!(i8, [ArrayInt8]);
915impl_produce_vec!(i16, [ArrayInt16]);
916impl_produce_vec!(i32, [ArrayInt32]);
917impl_produce_vec!(i64, [ArrayInt64]);
918impl_produce_vec!(u8, [ArrayUInt8]);
919impl_produce_vec!(u16, [ArrayUInt16]);
920impl_produce_vec!(u32, [ArrayUInt32]);
921impl_produce_vec!(u64, [ArrayUInt64]);
922impl_produce_vec!(f32, [ArrayFloat32]);
923impl_produce_vec!(f64, [ArrayFloat64]);
924impl_produce_vec!(Decimal, [ArrayDecimal]);
connectorx/sources/clickhouse/mod.rs

connectorx/sources/clickhouse/
mod.rs
