├── .github └── workflows │ └── test.yml ├── .gitignore ├── Cargo.toml ├── LICENSE ├── README.md ├── clickhouse-rs-cityhash-sys ├── Cargo.toml ├── build.rs └── src │ ├── cc │ ├── BUCK │ ├── CMakeLists.txt │ ├── COPYING │ ├── city.cc │ ├── city.h │ ├── citycrc.h │ └── config.h │ └── lib.rs ├── docker-compose.yml ├── examples └── simple.rs ├── extras └── ci │ ├── generate_certs.sh │ └── overrides.xml ├── src ├── binary │ ├── encoder.rs │ ├── mod.rs │ ├── parser.rs │ ├── protocol.rs │ ├── read_ex.rs │ └── uvarint.rs ├── client_info.rs ├── connecting_stream.rs ├── errors │ ├── codes.rs │ └── mod.rs ├── io │ ├── mod.rs │ ├── read_to_end.rs │ ├── stream.rs │ └── transport.rs ├── lib.rs ├── pool │ ├── futures │ │ ├── get_handle.rs │ │ └── mod.rs │ └── mod.rs ├── retry_guard.rs └── types │ ├── block │ ├── block_info.rs │ ├── builder.rs │ ├── chunk_iterator.rs │ ├── compressed.rs │ ├── mod.rs │ └── row.rs │ ├── cmd.rs │ ├── column │ ├── array.rs │ ├── chrono_datetime.rs │ ├── chunk.rs │ ├── column_data.rs │ ├── concat.rs │ ├── date.rs │ ├── datetime64.rs │ ├── decimal.rs │ ├── enums.rs │ ├── factory.rs │ ├── fixed_string.rs │ ├── ip.rs │ ├── iter │ │ └── mod.rs │ ├── list.rs │ ├── low_cardinality.rs │ ├── map.rs │ ├── mod.rs │ ├── nullable.rs │ ├── numeric.rs │ ├── simple_agg_func.rs │ ├── string.rs │ ├── string_pool.rs │ └── util.rs │ ├── date_converter.rs │ ├── decimal.rs │ ├── enums.rs │ ├── from_sql.rs │ ├── marshal.rs │ ├── mod.rs │ ├── options.rs │ ├── query.rs │ ├── query_result │ ├── mod.rs │ └── stream_blocks.rs │ ├── stat_buffer.rs │ ├── unmarshal.rs │ ├── value.rs │ └── value_ref.rs └── tests └── clickhouse.rs /.github/workflows/test.yml: -------------------------------------------------------------------------------- 1 | name: Test 2 | on: 3 | workflow_dispatch: {} 4 | push: 5 | branches: 6 | - async-await 7 | pull_request: 8 | branches: 9 | - async-await 10 | 11 | env: 12 | CARGO_TERM_COLOR: always 13 | RUST_BACKTRACE: 1 14 | 15 | jobs: 16 | build: 17 | runs-on: ubuntu-latest 18 | services: 19 | clickhouse: 20 | image: clickhouse/clickhouse-server 21 | ports: 22 | - 9000:9000 23 | steps: 24 | - uses: actions/checkout@v3 25 | - name: Build 26 | run: cargo build --verbose 27 | - name: Run tests 28 | run: cargo test --verbose 29 | 30 | build-native-tls: 31 | runs-on: ubuntu-latest 32 | env: 33 | # NOTE: not all tests "secure" aware, so let's define DATABASE_URL explicitly 34 | # NOTE: sometimes for native-tls default connection_timeout (500ms) is not enough, interestingly that for rustls it is OK. 35 | DATABASE_URL: "tcp://localhost:9440?compression=lz4&ping_timeout=2s&retry_timeout=3s&secure=true&skip_verify=true&connection_timeout=5s" 36 | steps: 37 | - uses: actions/checkout@v3 38 | # NOTE: 39 | # - we cannot use "services" because they are executed before the steps, i.e. repository checkout. 40 | # - "job.container.network" is empty, hence "host" 41 | # - github actions does not support YAML anchors (sigh) 42 | - name: Run clickhouse-server 43 | run: docker run 44 | -v ./extras/ci/generate_certs.sh:/docker-entrypoint-initdb.d/generate_certs.sh 45 | -v ./extras/ci/overrides.xml:/etc/clickhouse-server/config.d/overrides.xml 46 | -e CH_SSL_CERTIFICATE=/etc/clickhouse-server/config.d/server.crt 47 | -e CH_SSL_PRIVATE_KEY=/etc/clickhouse-server/config.d/server.key 48 | --network host 49 | --rm 50 | --detach 51 | --publish 9440:9440 52 | clickhouse/clickhouse-server 53 | - name: Build 54 | run: cargo build --features tls-native-tls --verbose 55 | - name: Run tests 56 | run: cargo test --features tls-native-tls --verbose 57 | 58 | build-rustls: 59 | runs-on: ubuntu-latest 60 | env: 61 | # NOTE: not all tests "secure" aware, so let's define DATABASE_URL explicitly 62 | DATABASE_URL: "tcp://localhost:9440?compression=lz4&ping_timeout=2s&retry_timeout=3s&secure=true&skip_verify=true" 63 | steps: 64 | - uses: actions/checkout@v3 65 | # NOTE: 66 | # - we cannot use "services" because they are executed before the steps, i.e. repository checkout. 67 | # - "job.container.network" is empty, hence "host" 68 | # - github actions does not support YAML anchors (sigh) 69 | - name: Run clickhouse-server 70 | run: docker run 71 | -v ./extras/ci/generate_certs.sh:/docker-entrypoint-initdb.d/generate_certs.sh 72 | -v ./extras/ci/overrides.xml:/etc/clickhouse-server/config.d/overrides.xml 73 | -e CH_SSL_CERTIFICATE=/etc/clickhouse-server/config.d/server.crt 74 | -e CH_SSL_PRIVATE_KEY=/etc/clickhouse-server/config.d/server.key 75 | --network host 76 | --rm 77 | --detach 78 | --publish 9440:9440 79 | clickhouse/clickhouse-server 80 | - name: Build 81 | run: cargo build --features tls-rustls --verbose 82 | - name: Run tests 83 | run: cargo test --features tls-rustls --verbose 84 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | target 2 | Cargo.lock 3 | .idea/ 4 | *.iml 5 | .DS_Store 6 | **/*.rs.bk 7 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "clickhouse-rs" 3 | version = "1.1.0-alpha.1" 4 | authors = ["Mikhail Sukharev "] 5 | license = "MIT" 6 | homepage = "https://github.com/suharev7/clickhouse-rs" 7 | repository = "https://github.com/suharev7/clickhouse-rs" 8 | description = "Asynchronous Yandex ClickHouse client library." 9 | documentation = "https://docs.rs/clickhouse-rs/" 10 | readme = "README.md" 11 | keywords = ["tokio", "database", "clickhouse"] 12 | categories = ["database"] 13 | edition = "2021" 14 | exclude = ["tests/*", "examples/*"] 15 | 16 | [features] 17 | default = ["tokio_io"] 18 | _tls = [] # meta feature for the clickhouse-rs generic TLS code 19 | tls = ["tls-native-tls"] # backward compatibility 20 | tls-native-tls = ["tokio-native-tls", "native-tls", "_tls"] 21 | tls-rustls = ["tokio-rustls", "rustls", "rustls-pemfile", "webpki-roots", "_tls"] 22 | async_std = ["async-std"] 23 | tokio_io = ["tokio"] 24 | 25 | [dependencies] 26 | byteorder = "^1.4" 27 | chrono-tz = "^0.8" 28 | crossbeam = "^0.8" 29 | thiserror = "^1.0" 30 | futures-core = "^0.3" 31 | futures-sink = "^0.3" 32 | hostname = "^0.3" 33 | lazy_static = "1.4.0" 34 | lz4 = "^1.24" 35 | pin-project = "^1.1" 36 | url = "^2" 37 | uuid = "^1.4" 38 | combine = "^4.6" 39 | percent-encoding = "^2.3" 40 | either = "^1.6" 41 | cfg-if = "1.0.0" 42 | 43 | [dependencies.futures-util] 44 | version = "^0.3" 45 | features = ["sink"] 46 | 47 | [dependencies.tokio] 48 | version = "^1.32" 49 | default-features = false 50 | features = ["io-util", "time", "net", "sync", "rt-multi-thread"] 51 | optional = true 52 | 53 | [dependencies.async-std] 54 | version = "1.6" 55 | optional = true 56 | 57 | [dependencies.clickhouse-rs-cityhash-sys] 58 | path = "clickhouse-rs-cityhash-sys" 59 | version = "0.1.2" 60 | 61 | [dependencies.log] 62 | version = "0.4.8" 63 | features = ["std", "serde"] 64 | 65 | [dependencies.native-tls] 66 | version = "0.2" 67 | optional = true 68 | 69 | [dependencies.tokio-native-tls] 70 | version = "^0.3" 71 | optional = true 72 | 73 | [dependencies.rustls] 74 | version = "0.22.1" 75 | optional = true 76 | 77 | [dependencies.rustls-pemfile] 78 | version = "2.0" 79 | optional = true 80 | 81 | [dependencies.tokio-rustls] 82 | version = "0.25.0" 83 | optional = true 84 | 85 | [dependencies.webpki-roots] 86 | version = "*" 87 | optional = true 88 | 89 | [dependencies.chrono] 90 | version = "^0.4" 91 | default-features = false 92 | features = ["std"] 93 | 94 | [dev-dependencies] 95 | env_logger = "^0.10" 96 | pretty_assertions = "1.3.0" 97 | rand = "^0.8" 98 | 99 | [dev-dependencies.tokio] 100 | version = "^1.32" 101 | default-features = false 102 | features = ["macros"] 103 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Mikhail Sukharev 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Async ClickHouse Client 2 | 3 | [![Build Status](https://travis-ci.com/suharev7/clickhouse-rs.svg?branch=master)](https://travis-ci.com/suharev7/clickhouse-rs) 4 | [![Crate info](https://img.shields.io/crates/v/clickhouse-rs.svg)](https://crates.io/crates/clickhouse-rs) 5 | [![Documentation](https://docs.rs/clickhouse-rs/badge.svg)](https://docs.rs/clickhouse-rs) 6 | [![dependency status](https://deps.rs/repo/github/suharev7/clickhouse-rs/status.svg)](https://deps.rs/repo/github/suharev7/clickhouse-rs) 7 | [![Coverage Status](https://coveralls.io/repos/github/suharev7/clickhouse-rs/badge.svg)](https://coveralls.io/github/suharev7/clickhouse-rs) 8 | 9 | Asynchronous [Yandex ClickHouse](https://clickhouse.yandex/) client library for rust programming language. 10 | 11 | ## Installation 12 | Library hosted on [crates.io](https://crates.io/crates/clickhouse-rs/). 13 | ```toml 14 | [dependencies] 15 | clickhouse-rs = "*" 16 | ``` 17 | 18 | ## Supported data types 19 | 20 | * Date 21 | * DateTime 22 | * Decimal(P, S) 23 | * Float32, Float64 24 | * String, FixedString(N) 25 | * UInt8, UInt16, UInt32, UInt64, UInt128, Int8, Int16, Int32, Int64, Int128 26 | * Nullable(T) 27 | * Array(UInt/Int/Float/String/Date/DateTime) 28 | * SimpleAggregateFunction(F, T) 29 | * IPv4/IPv6 30 | * UUID 31 | * Bool 32 | 33 | ## DNS 34 | 35 | ```url 36 | schema://user:password@host[:port]/database?param1=value1&...¶mN=valueN 37 | ``` 38 | 39 | parameters: 40 | 41 | - `compression` - Whether or not use compression (defaults to `none`). Possible choices: 42 | * `none` 43 | * `lz4` 44 | 45 | - `connection_timeout` - Timeout for connection (defaults to `500 ms`) 46 | - `query_timeout` - Timeout for queries (defaults to `180 sec`). 47 | - `insert_timeout` - Timeout for inserts (defaults to `180 sec`). 48 | - `execute_timeout` - Timeout for execute (defaults to `180 sec`). 49 | - `keepalive` - TCP keep alive timeout in milliseconds. 50 | - `nodelay` - Whether to enable `TCP_NODELAY` (defaults to `true`). 51 | 52 | - `pool_min` - Lower bound of opened connections for `Pool` (defaults to `10`). 53 | - `pool_max` - Upper bound of opened connections for `Pool` (defaults to `20`). 54 | 55 | - `ping_before_query` - Ping server every time before execute any query. (defaults to `true`). 56 | - `send_retries` - Count of retry to send request to server. (defaults to `3`). 57 | - `retry_timeout` - Amount of time to wait before next retry. (defaults to `5 sec`). 58 | - `ping_timeout` - Timeout for ping (defaults to `500 ms`). 59 | 60 | 61 | - `alt_hosts` - Comma separated list of single address host for load-balancing. 62 | 63 | example: 64 | ```url 65 | tcp://user:password@host:9000/clicks?compression=lz4&ping_timeout=42ms 66 | ``` 67 | 68 | ## Optional features 69 | 70 | `clickhouse-rs` puts some functionality behind optional features to optimize compile time 71 | for the most common use cases. The following features are available. 72 | 73 | - `tokio_io` *(enabled by default)* — I/O based on [Tokio](https://tokio.rs/). 74 | - `async_std` — I/O based on [async-std](https://async.rs/) (doesn't work together with `tokio_io`). 75 | - `tls` — TLS support (allowed only with `tokio_io`). 76 | 77 | ## Example 78 | 79 | ```rust 80 | use clickhouse_rs::{Block, Pool}; 81 | use std::error::Error; 82 | 83 | #[tokio::main] 84 | async fn main() -> Result<(), Box> { 85 | let ddl = r" 86 | CREATE TABLE IF NOT EXISTS payment ( 87 | customer_id UInt32, 88 | amount UInt32, 89 | account_name Nullable(FixedString(3)) 90 | ) Engine=Memory"; 91 | 92 | let block = Block::new() 93 | .column("customer_id", vec![1_u32, 3, 5, 7, 9]) 94 | .column("amount", vec![2_u32, 4, 6, 8, 10]) 95 | .column("account_name", vec![Some("foo"), None, None, None, Some("bar")]); 96 | 97 | let pool = Pool::new(database_url); 98 | 99 | let mut client = pool.get_handle().await?; 100 | client.execute(ddl).await?; 101 | client.insert("payment", block).await?; 102 | let block = client.query("SELECT * FROM payment").fetch_all().await?; 103 | 104 | for row in block.rows() { 105 | let id: u32 = row.get("customer_id")?; 106 | let amount: u32 = row.get("amount")?; 107 | let name: Option<&str> = row.get("account_name")?; 108 | println!("Found payment {}: {} {:?}", id, amount, name); 109 | } 110 | 111 | Ok(()) 112 | } 113 | ``` 114 | -------------------------------------------------------------------------------- /clickhouse-rs-cityhash-sys/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "clickhouse-rs-cityhash-sys" 3 | version = "0.1.2" 4 | authors = ["mikhail "] 5 | license = "MIT" 6 | links = "clickhouse-rs" 7 | build = "build.rs" 8 | description = "Rust CityHash sys package for clickhouse-rs." 9 | 10 | [build-dependencies] 11 | cc = "1.0.28" 12 | -------------------------------------------------------------------------------- /clickhouse-rs-cityhash-sys/build.rs: -------------------------------------------------------------------------------- 1 | extern crate cc; 2 | 3 | fn main() { 4 | let mut compiler = cc::Build::new(); 5 | compiler.file("src/cc/city.cc").cpp(true).opt_level(3); 6 | 7 | compiler.compile("libchcityhash.a"); 8 | } 9 | -------------------------------------------------------------------------------- /clickhouse-rs-cityhash-sys/src/cc/BUCK: -------------------------------------------------------------------------------- 1 | cxx_library( 2 | name = 'cityhash', 3 | header_namespace = 'cityhash', 4 | exported_headers = subdir_glob([ 5 | ('', '*.h'), 6 | ]), 7 | srcs = glob([ 8 | '*.cc', 9 | ]), 10 | visibility = [ 11 | '//...', 12 | ], 13 | ) 14 | -------------------------------------------------------------------------------- /clickhouse-rs-cityhash-sys/src/cc/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | ADD_LIBRARY (cityhash-lib STATIC 2 | city.cc 3 | ) 4 | 5 | set_property(TARGET cityhash-lib PROPERTY POSITION_INDEPENDENT_CODE ON) 6 | -------------------------------------------------------------------------------- /clickhouse-rs-cityhash-sys/src/cc/COPYING: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2011 Google, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | -------------------------------------------------------------------------------- /clickhouse-rs-cityhash-sys/src/cc/city.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2011 Google, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | // 21 | // CityHash, by Geoff Pike and Jyrki Alakuijala 22 | // 23 | // This file provides a few functions for hashing strings. On x86-64 24 | // hardware in 2011, CityHash64() is faster than other high-quality 25 | // hash functions, such as Murmur. This is largely due to higher 26 | // instruction-level parallelism. CityHash64() and CityHash128() also perform 27 | // well on hash-quality tests. 28 | // 29 | // CityHash128() is optimized for relatively long strings and returns 30 | // a 128-bit hash. For strings more than about 2000 bytes it can be 31 | // faster than CityHash64(). 32 | // 33 | // Functions in the CityHash family are not suitable for cryptography. 34 | // 35 | // WARNING: This code has not been tested on big-endian platforms! 36 | // It is known to work well on little-endian platforms that have a small penalty 37 | // for unaligned reads, such as current Intel and AMD moderate-to-high-end CPUs. 38 | // 39 | // By the way, for some hash functions, given strings a and b, the hash 40 | // of a+b is easily derived from the hashes of a and b. This property 41 | // doesn't hold for any hash functions in this file. 42 | 43 | #ifndef CITY_HASH_H_ 44 | #define CITY_HASH_H_ 45 | 46 | #include // for size_t. 47 | #include 48 | #include 49 | 50 | typedef uint8_t uint8; 51 | typedef uint32_t uint32; 52 | typedef uint64_t uint64; 53 | typedef std::pair uint128; 54 | 55 | inline uint64 Uint128Low64(const uint128& x) { return x.first; } 56 | inline uint64 Uint128High64(const uint128& x) { return x.second; } 57 | 58 | // Hash function for a byte array. 59 | uint64 CityHash64(const char *buf, size_t len); 60 | 61 | // Hash function for a byte array. For convenience, a 64-bit seed is also 62 | // hashed into the result. 63 | uint64 CityHash64WithSeed(const char *buf, size_t len, uint64 seed); 64 | 65 | // Hash function for a byte array. For convenience, two seeds are also 66 | // hashed into the result. 67 | uint64 CityHash64WithSeeds(const char *buf, size_t len, 68 | uint64 seed0, uint64 seed1); 69 | 70 | // Hash function for a byte array. 71 | uint128 CityHash128(const char *s, size_t len); 72 | 73 | // Hash function for a byte array. For convenience, a 128-bit seed is also 74 | // hashed into the result. 75 | uint128 CityHash128WithSeed(const char *s, size_t len, uint128 seed); 76 | 77 | // Hash 128 input bits down to 64 bits of output. 78 | // This is intended to be a reasonably good hash function. 79 | inline uint64 Hash128to64(const uint128& x) { 80 | // Murmur-inspired hashing. 81 | const uint64 kMul = 0x9ddfea08eb382d69ULL; 82 | uint64 a = (Uint128Low64(x) ^ Uint128High64(x)) * kMul; 83 | a ^= (a >> 47); 84 | uint64 b = (Uint128High64(x) ^ a) * kMul; 85 | b ^= (b >> 47); 86 | b *= kMul; 87 | return b; 88 | } 89 | 90 | #endif // CITY_HASH_H_ 91 | -------------------------------------------------------------------------------- /clickhouse-rs-cityhash-sys/src/cc/citycrc.h: -------------------------------------------------------------------------------- 1 | // Copyright (c) 2011 Google, Inc. 2 | // 3 | // Permission is hereby granted, free of charge, to any person obtaining a copy 4 | // of this software and associated documentation files (the "Software"), to deal 5 | // in the Software without restriction, including without limitation the rights 6 | // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | // copies of the Software, and to permit persons to whom the Software is 8 | // furnished to do so, subject to the following conditions: 9 | // 10 | // The above copyright notice and this permission notice shall be included in 11 | // all copies or substantial portions of the Software. 12 | // 13 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | // THE SOFTWARE. 20 | // 21 | // CityHash, by Geoff Pike and Jyrki Alakuijala 22 | // 23 | // This file declares the subset of the CityHash functions that require 24 | // _mm_crc32_u64(). See the CityHash README for details. 25 | // 26 | // Functions in the CityHash family are not suitable for cryptography. 27 | 28 | #ifndef CITY_HASH_CRC_H_ 29 | #define CITY_HASH_CRC_H_ 30 | 31 | #include "city.h" 32 | 33 | // Hash function for a byte array. 34 | uint128 CityHashCrc128(const char *s, size_t len); 35 | 36 | // Hash function for a byte array. For convenience, a 128-bit seed is also 37 | // hashed into the result. 38 | uint128 CityHashCrc128WithSeed(const char *s, size_t len, uint128 seed); 39 | 40 | // Hash function for a byte array. Sets result[0] ... result[3]. 41 | void CityHashCrc256(const char *s, size_t len, uint64 *result); 42 | 43 | #endif // CITY_HASH_CRC_H_ 44 | -------------------------------------------------------------------------------- /clickhouse-rs-cityhash-sys/src/cc/config.h: -------------------------------------------------------------------------------- 1 | /* config.h. Generated from config.h.in by configure. */ 2 | /* config.h.in. Generated from configure.ac by autoheader. */ 3 | 4 | /* Define if building universal (internal helper macro) */ 5 | /* #undef AC_APPLE_UNIVERSAL_BUILD */ 6 | 7 | /* Define to 1 if the compiler supports __builtin_expect. */ 8 | #if _MSC_VER 9 | # define HAVE_BUILTIN_EXPECT 0 10 | #else 11 | # define HAVE_BUILTIN_EXPECT 1 12 | #endif 13 | 14 | /* Define to 1 if you have the header file. */ 15 | #define HAVE_DLFCN_H 1 16 | 17 | /* Define to 1 if you have the header file. */ 18 | #define HAVE_INTTYPES_H 1 19 | 20 | /* Define to 1 if you have the header file. */ 21 | #define HAVE_MEMORY_H 1 22 | 23 | /* Define to 1 if you have the header file. */ 24 | #define HAVE_STDINT_H 1 25 | 26 | /* Define to 1 if you have the header file. */ 27 | #define HAVE_STDLIB_H 1 28 | 29 | /* Define to 1 if you have the header file. */ 30 | #define HAVE_STRINGS_H 1 31 | 32 | /* Define to 1 if you have the header file. */ 33 | #define HAVE_STRING_H 1 34 | 35 | /* Define to 1 if you have the header file. */ 36 | #define HAVE_SYS_STAT_H 1 37 | 38 | /* Define to 1 if you have the header file. */ 39 | #define HAVE_SYS_TYPES_H 1 40 | 41 | /* Define to 1 if you have the header file. */ 42 | #define HAVE_UNISTD_H 1 43 | 44 | /* Define to the sub-directory in which libtool stores uninstalled libraries. 45 | */ 46 | #define LT_OBJDIR ".libs/" 47 | 48 | /* Define to the address where bug reports for this package should be sent. */ 49 | #define PACKAGE_BUGREPORT "cityhash-discuss@googlegroups.com" 50 | 51 | /* Define to the full name of this package. */ 52 | #define PACKAGE_NAME "CityHash" 53 | 54 | /* Define to the full name and version of this package. */ 55 | #define PACKAGE_STRING "CityHash 1.0.2" 56 | 57 | /* Define to the one symbol short name of this package. */ 58 | #define PACKAGE_TARNAME "cityhash" 59 | 60 | /* Define to the home page for this package. */ 61 | #define PACKAGE_URL "" 62 | 63 | /* Define to the version of this package. */ 64 | #define PACKAGE_VERSION "1.0.2" 65 | 66 | /* Define to 1 if you have the ANSI C header files. */ 67 | #define STDC_HEADERS 1 68 | 69 | /* Define WORDS_BIGENDIAN to 1 if your processor stores words with the most 70 | significant byte first (like Motorola and SPARC, unlike Intel). */ 71 | #if defined AC_APPLE_UNIVERSAL_BUILD 72 | # if defined __BIG_ENDIAN__ 73 | # define WORDS_BIGENDIAN 1 74 | # endif 75 | #else 76 | # ifndef WORDS_BIGENDIAN 77 | /* # undef WORDS_BIGENDIAN */ 78 | # endif 79 | #endif 80 | 81 | /* Define for Solaris 2.5.1 so the uint32_t typedef from , 82 | , or is not used. If the typedef were allowed, the 83 | #define below would cause a syntax error. */ 84 | /* #undef _UINT32_T */ 85 | 86 | /* Define for Solaris 2.5.1 so the uint64_t typedef from , 87 | , or is not used. If the typedef were allowed, the 88 | #define below would cause a syntax error. */ 89 | /* #undef _UINT64_T */ 90 | 91 | /* Define for Solaris 2.5.1 so the uint8_t typedef from , 92 | , or is not used. If the typedef were allowed, the 93 | #define below would cause a syntax error. */ 94 | /* #undef _UINT8_T */ 95 | 96 | /* Define to `__inline__' or `__inline' if that's what the C compiler 97 | calls it, or to nothing if 'inline' is not supported under any name. */ 98 | #ifndef __cplusplus 99 | /* #undef inline */ 100 | #endif 101 | 102 | /* Define to `unsigned int' if does not define. */ 103 | /* #undef size_t */ 104 | 105 | /* Define to `int' if does not define. */ 106 | /* #undef ssize_t */ 107 | 108 | /* Define to the type of an unsigned integer type of width exactly 32 bits if 109 | such a type exists and the standard includes do not define it. */ 110 | /* #undef uint32_t */ 111 | 112 | /* Define to the type of an unsigned integer type of width exactly 64 bits if 113 | such a type exists and the standard includes do not define it. */ 114 | /* #undef uint64_t */ 115 | 116 | /* Define to the type of an unsigned integer type of width exactly 8 bits if 117 | such a type exists and the standard includes do not define it. */ 118 | /* #undef uint8_t */ 119 | 120 | #ifdef _MSC_VER 121 | #include 122 | typedef SSIZE_T ssize_t; 123 | #else 124 | #include 125 | #endif 126 | -------------------------------------------------------------------------------- /clickhouse-rs-cityhash-sys/src/lib.rs: -------------------------------------------------------------------------------- 1 | #![allow(non_upper_case_globals)] 2 | #![allow(non_camel_case_types)] 3 | #![allow(non_snake_case)] 4 | 5 | use std::os::raw::c_char; 6 | 7 | #[repr(C)] 8 | #[derive(Debug, Copy, Clone, Eq, PartialEq)] 9 | pub struct UInt128 { 10 | pub lo: u64, 11 | pub hi: u64, 12 | } 13 | 14 | extern "C" { 15 | fn CityHash128(s: *const c_char, len: usize) -> UInt128; 16 | } 17 | 18 | pub fn city_hash_128>(source: B) -> UInt128 { 19 | let buffer = source.as_ref(); 20 | unsafe { CityHash128(buffer.as_ptr() as *const c_char, buffer.len()) } 21 | } 22 | 23 | #[test] 24 | fn test_city_hash_128() { 25 | let expected = UInt128 { 26 | lo: 0x900ff195577748fe, 27 | hi: 0x13a9176355b20d7e, 28 | }; 29 | let actual = city_hash_128("abc"); 30 | assert_eq!(expected, actual) 31 | } 32 | -------------------------------------------------------------------------------- /docker-compose.yml: -------------------------------------------------------------------------------- 1 | version: '3' 2 | services: 3 | clickhouse: 4 | image: yandex/clickhouse-server 5 | ports: 6 | - 8123:8123 7 | - 9000:9000 -------------------------------------------------------------------------------- /examples/simple.rs: -------------------------------------------------------------------------------- 1 | use clickhouse_rs::{row, types::Block, Pool}; 2 | use futures_util::StreamExt; 3 | use std::{env, error::Error}; 4 | 5 | async fn execute(database_url: String) -> Result<(), Box> { 6 | env::set_var("RUST_LOG", "clickhouse_rs=debug"); 7 | env_logger::init(); 8 | 9 | let ddl = r" 10 | CREATE TABLE IF NOT EXISTS payment ( 11 | customer_id UInt32, 12 | amount UInt32, 13 | account_name Nullable(FixedString(3)) 14 | ) Engine=Memory"; 15 | 16 | let mut block = Block::with_capacity(5); 17 | block.push(row! { customer_id: 1_u32, amount: 2_u32, account_name: Some("foo") })?; 18 | block.push(row! { customer_id: 3_u32, amount: 4_u32, account_name: None::<&str> })?; 19 | block.push(row! { customer_id: 5_u32, amount: 6_u32, account_name: None::<&str> })?; 20 | block.push(row! { customer_id: 7_u32, amount: 8_u32, account_name: None::<&str> })?; 21 | block.push(row! { customer_id: 9_u32, amount: 10_u32, account_name: Some("bar") })?; 22 | 23 | let pool = Pool::new(database_url); 24 | 25 | let mut client = pool.get_handle().await?; 26 | client.execute(ddl).await?; 27 | client.insert("payment", block).await?; 28 | let mut stream = client.query("SELECT * FROM payment").stream(); 29 | 30 | while let Some(row) = stream.next().await { 31 | let row = row?; 32 | let id: u32 = row.get("customer_id")?; 33 | let amount: u32 = row.get("amount")?; 34 | let name: Option<&str> = row.get("account_name")?; 35 | println!("Found payment {id}: {amount} {name:?}"); 36 | } 37 | 38 | Ok(()) 39 | } 40 | 41 | #[cfg(all(feature = "tokio_io", not(feature = "_tls")))] 42 | #[tokio::main] 43 | async fn main() -> Result<(), Box> { 44 | let database_url = 45 | env::var("DATABASE_URL").unwrap_or_else(|_| "tcp://localhost:9000?compression=lz4".into()); 46 | execute(database_url).await 47 | } 48 | 49 | #[cfg(all(feature = "tokio_io", feature = "_tls"))] 50 | #[tokio::main] 51 | async fn main() -> Result<(), Box> { 52 | let database_url = env::var("DATABASE_URL") 53 | .unwrap_or_else(|_| "tcp://localhost:9440?secure=true&skip_verify=true".into()); 54 | execute(database_url).await 55 | } 56 | 57 | #[cfg(feature = "async_std")] 58 | fn main() { 59 | use async_std::task; 60 | let database_url = 61 | env::var("DATABASE_URL").unwrap_or_else(|_| "tcp://localhost:9000?compression=lz4".into()); 62 | task::block_on(execute(database_url)).unwrap(); 63 | } 64 | -------------------------------------------------------------------------------- /extras/ci/generate_certs.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | crt=$CH_SSL_CERTIFICATE 4 | key=$CH_SSL_PRIVATE_KEY 5 | 6 | openssl req -subj "/CN=localhost" -new -newkey rsa:2048 -days 365 -nodes -x509 -keyout "$key" -out "$crt" 7 | chown clickhouse:clickhouse "$crt" "$key" 8 | -------------------------------------------------------------------------------- /extras/ci/overrides.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | none 7 | true 8 | true 9 | sslv2,sslv3 10 | true 11 | 12 | 13 | 9440 14 | 15 | 16 | 1 17 | 18 | 19 | -------------------------------------------------------------------------------- /src/binary/encoder.rs: -------------------------------------------------------------------------------- 1 | use crate::{ 2 | binary, 3 | types::{Marshal, StatBuffer}, 4 | }; 5 | 6 | const MAX_VARINT_LEN64: usize = 10; 7 | 8 | #[derive(Default)] 9 | pub struct Encoder { 10 | buffer: Vec, 11 | } 12 | 13 | impl Encoder { 14 | pub fn new() -> Self { 15 | Encoder { buffer: Vec::new() } 16 | } 17 | 18 | pub fn uvarint(&mut self, v: u64) { 19 | let mut scratch = [0u8; MAX_VARINT_LEN64]; 20 | let ln = binary::put_uvarint(&mut scratch[..], v); 21 | self.write_bytes(&scratch[..ln]); 22 | } 23 | 24 | pub fn string(&mut self, text: impl AsRef) { 25 | let bytes = text.as_ref().as_bytes(); 26 | self.byte_string(bytes); 27 | } 28 | 29 | pub fn byte_string(&mut self, source: impl AsRef<[u8]>) { 30 | self.uvarint(source.as_ref().len() as u64); 31 | self.write_bytes(source.as_ref()); 32 | } 33 | 34 | pub fn write(&mut self, value: T) 35 | where 36 | T: Copy + Marshal + StatBuffer, 37 | { 38 | let mut buffer = T::buffer(); 39 | value.marshal(buffer.as_mut()); 40 | self.write_bytes(buffer.as_ref()); 41 | } 42 | 43 | pub fn write_bytes(&mut self, b: &[u8]) { 44 | self.buffer.extend_from_slice(b); 45 | } 46 | 47 | pub fn get_buffer(self) -> Vec { 48 | self.buffer 49 | } 50 | 51 | pub fn get_buffer_ref(&self) -> &[u8] { 52 | self.buffer.as_ref() 53 | } 54 | } 55 | -------------------------------------------------------------------------------- /src/binary/mod.rs: -------------------------------------------------------------------------------- 1 | pub(crate) use self::{encoder::Encoder, parser::Parser, read_ex::ReadEx, uvarint::put_uvarint}; 2 | 3 | mod encoder; 4 | mod parser; 5 | pub mod protocol; 6 | mod read_ex; 7 | mod uvarint; 8 | -------------------------------------------------------------------------------- /src/binary/parser.rs: -------------------------------------------------------------------------------- 1 | use std::io::{self, Read}; 2 | 3 | use chrono_tz::Tz; 4 | use log::{trace, warn}; 5 | 6 | use crate::{ 7 | binary::{protocol, ReadEx}, 8 | errors::{DriverError, Error, Result, ServerError}, 9 | io::transport::TransportInfo, 10 | types::{Block, Packet, ProfileInfo, Progress, ServerInfo, TableColumns}, 11 | }; 12 | 13 | /// The internal clickhouse response parser. 14 | pub(crate) struct Parser<'i, T> { 15 | reader: T, 16 | info: &'i TransportInfo, 17 | } 18 | 19 | /// The parser can be used to parse clickhouse responses into values. Generally 20 | /// you normally do not use this directly as it's already done for you by 21 | /// the client but in some more complex situations it might be useful to be 22 | /// able to parse the clickhouse responses. 23 | impl<'i, T: Read> Parser<'i, T> { 24 | /// Creates a new parser that parses the data behind the reader. More 25 | /// than one value can be behind the reader in which case the parser can 26 | /// be invoked multiple times. In other words: the stream does not have 27 | /// to be terminated. 28 | pub(crate) fn new(reader: T, info: &'i TransportInfo) -> Parser { 29 | Self { reader, info } 30 | } 31 | 32 | /// Parses a single value out of the stream. If there are multiple 33 | /// values you can call this multiple times. If the reader is not yet 34 | /// ready this will block. 35 | pub(crate) fn parse_packet(&mut self, revision: u64) -> Result> { 36 | let packet = self.reader.read_uvarint()?; 37 | match packet { 38 | protocol::SERVER_HELLO => Ok(self.parse_server_info()?), 39 | protocol::SERVER_PONG => Ok(self.parse_pong()), 40 | protocol::SERVER_PROGRESS => Ok(self.parse_progress(revision)?), 41 | protocol::SERVER_PROFILE_INFO => Ok(self.parse_profile_info()?), 42 | protocol::SERVER_TABLE_COLUMNS => Ok(self.parse_table_columns()?), 43 | protocol::SERVER_EXCEPTION => Ok(self.parse_exception()?), 44 | protocol::SERVER_DATA | protocol::SERVER_TOTALS | protocol::SERVER_EXTREMES => { 45 | Ok(self.parse_block()?) 46 | } 47 | protocol::SERVER_END_OF_STREAM => Ok(Packet::Eof(())), 48 | _ => Err(Error::Driver(DriverError::UnknownPacket { packet })), 49 | } 50 | } 51 | 52 | fn parse_block(&mut self) -> Result> { 53 | match self.info.timezone { 54 | None => Err(Error::Driver(DriverError::UnexpectedPacket)), 55 | Some(tz) => { 56 | self.reader.skip_string()?; 57 | let block = Block::load(&mut self.reader, tz, self.info.compress)?; 58 | Ok(Packet::Block(block)) 59 | } 60 | } 61 | } 62 | 63 | fn parse_server_info(&mut self) -> Result> { 64 | let name = self.reader.read_string()?; 65 | let major_version = self.reader.read_uvarint()?; 66 | let minor_version = self.reader.read_uvarint()?; 67 | let revision = self.reader.read_uvarint()?; 68 | 69 | let timezone = if revision >= protocol::DBMS_MIN_REVISION_WITH_SERVER_TIMEZONE { 70 | match self.reader.read_string()?.parse() { 71 | Ok(tz) => tz, 72 | Err(err) => return Err(io::Error::new(io::ErrorKind::Other, err).into()), 73 | } 74 | } else { 75 | Tz::UTC 76 | }; 77 | 78 | let display_name = if revision >= protocol::DBMS_MIN_REVISION_WITH_SERVER_DISPLAY_NAME { 79 | self.reader.read_string()? 80 | } else { 81 | "".into() 82 | }; 83 | 84 | let patch_version = if revision >= protocol::DBMS_MIN_REVISION_WITH_VERSION_PATCH { 85 | self.reader.read_uvarint()? 86 | } else { 87 | 0 88 | }; 89 | 90 | let server_info = ServerInfo { 91 | name, 92 | major_version, 93 | minor_version, 94 | revision, 95 | timezone, 96 | display_name, 97 | patch_version, 98 | }; 99 | 100 | trace!("[hello] <- {:?}", &server_info); 101 | Ok(Packet::Hello((), server_info)) 102 | } 103 | 104 | fn parse_progress(&mut self, revision: u64) -> Result> { 105 | let rows = self.reader.read_uvarint()?; 106 | let bytes = self.reader.read_uvarint()?; 107 | 108 | let total_rows = if revision >= protocol::DBMS_MIN_REVISION_WITH_TOTAL_ROWS_IN_PROGRESS { 109 | self.reader.read_uvarint()? 110 | } else { 111 | 0 112 | }; 113 | 114 | let (written_rows, written_bytes) = 115 | if revision >= protocol::DBMS_MIN_REVISION_WITH_CLIENT_WRITE_INFO { 116 | (self.reader.read_uvarint()?, self.reader.read_uvarint()?) 117 | } else { 118 | (0, 0) 119 | }; 120 | 121 | let progress = Progress { 122 | rows, 123 | bytes, 124 | total_rows, 125 | written_rows, 126 | written_bytes, 127 | }; 128 | 129 | trace!( 130 | "[process] <- Progress: rows={}, bytes={}, total rows={}, written_rows={}, write_bytes={}", 131 | progress.rows, 132 | progress.bytes, 133 | progress.total_rows, 134 | progress.written_rows, 135 | progress.written_bytes, 136 | ); 137 | 138 | Ok(Packet::Progress(progress)) 139 | } 140 | 141 | fn parse_profile_info(&mut self) -> Result> { 142 | let info = Packet::ProfileInfo(ProfileInfo { 143 | rows: self.reader.read_uvarint()?, 144 | blocks: self.reader.read_uvarint()?, 145 | bytes: self.reader.read_uvarint()?, 146 | applied_limit: self.reader.read_scalar()?, 147 | rows_before_limit: self.reader.read_uvarint()?, 148 | calculated_rows_before_limit: self.reader.read_scalar()?, 149 | }); 150 | 151 | trace!("profile_info: {:?}", info); 152 | Ok(info) 153 | } 154 | 155 | fn parse_table_columns(&mut self) -> Result> { 156 | let table_columns = Packet::TableColumns(TableColumns { 157 | table_name: self.reader.read_string()?, 158 | columns: self.reader.read_string()?, 159 | }); 160 | 161 | trace!("table_columns: {:?}", table_columns); 162 | Ok(table_columns) 163 | } 164 | 165 | fn parse_exception(&mut self) -> Result> { 166 | let exception = ServerError { 167 | code: self.reader.read_scalar()?, 168 | name: self.reader.read_string()?, 169 | message: self.reader.read_string()?, 170 | stack_trace: self.reader.read_string()?, 171 | }; 172 | 173 | warn!("server exception: {:?}", exception); 174 | Ok(Packet::Exception(exception)) 175 | } 176 | 177 | fn parse_pong(&self) -> Packet<()> { 178 | trace!("[process] <- pong"); 179 | Packet::Pong(()) 180 | } 181 | } 182 | -------------------------------------------------------------------------------- /src/binary/protocol.rs: -------------------------------------------------------------------------------- 1 | pub const DBMS_MIN_REVISION_WITH_TOTAL_ROWS_IN_PROGRESS: u64 = 51554; 2 | pub const DBMS_MIN_REVISION_WITH_SERVER_TIMEZONE: u64 = 54058; 3 | pub const DBMS_MIN_REVISION_WITH_QUOTA_KEY_IN_CLIENT_INFO: u64 = 54060; 4 | pub const DBMS_MIN_REVISION_WITH_SERVER_DISPLAY_NAME: u64 = 54372; 5 | pub const DBMS_MIN_REVISION_WITH_VERSION_PATCH: u64 = 54401; 6 | pub const _DBMS_MIN_REVISION_WITH_SERVER_LOGS: u64 = 54406; 7 | pub const DBMS_MIN_REVISION_WITH_CLIENT_WRITE_INFO: u64 = 54420; 8 | pub const DBMS_MIN_REVISION_WITH_SETTINGS_SERIALIZED_AS_STRINGS: u64 = 54429; 9 | 10 | pub const CLIENT_HELLO: u64 = 0; 11 | pub const CLIENT_QUERY: u64 = 1; 12 | pub const CLIENT_DATA: u64 = 2; 13 | pub const CLIENT_CANCEL: u64 = 3; 14 | pub const CLIENT_PING: u64 = 4; 15 | 16 | pub const COMPRESS_ENABLE: u64 = 1; 17 | pub const COMPRESS_DISABLE: u64 = 0; 18 | 19 | pub const STATE_COMPLETE: u64 = 2; 20 | 21 | pub const SERVER_HELLO: u64 = 0; 22 | pub const SERVER_DATA: u64 = 1; 23 | pub const SERVER_EXCEPTION: u64 = 2; 24 | pub const SERVER_PROGRESS: u64 = 3; 25 | pub const SERVER_PONG: u64 = 4; 26 | pub const SERVER_END_OF_STREAM: u64 = 5; 27 | pub const SERVER_PROFILE_INFO: u64 = 6; 28 | pub const SERVER_TOTALS: u64 = 7; 29 | pub const SERVER_EXTREMES: u64 = 8; 30 | pub const _SERVER_TABLES_STATUS_RESPONSE: u64 = 9; 31 | pub const _SERVER_LOG: u64 = 10; 32 | pub const SERVER_TABLE_COLUMNS: u64 = 11; 33 | -------------------------------------------------------------------------------- /src/binary/read_ex.rs: -------------------------------------------------------------------------------- 1 | use std::{io, mem::MaybeUninit}; 2 | 3 | use crate::{ 4 | errors::{DriverError, Error, Result}, 5 | types::{column::StringPool, StatBuffer, Unmarshal}, 6 | }; 7 | 8 | pub(crate) trait ReadEx { 9 | fn read_bytes(&mut self, rv: &mut [u8]) -> Result<()>; 10 | fn read_scalar(&mut self) -> Result 11 | where 12 | V: Copy + Unmarshal + StatBuffer; 13 | fn read_string(&mut self) -> Result; 14 | fn skip_string(&mut self) -> Result<()>; 15 | fn read_uvarint(&mut self) -> Result; 16 | fn read_str_into_buffer(&mut self, pool: &mut StringPool) -> Result<()>; 17 | } 18 | 19 | const MAX_STACK_BUFFER_LEN: usize = 1024; 20 | 21 | impl ReadEx for T 22 | where 23 | T: io::Read, 24 | { 25 | fn read_bytes(&mut self, rv: &mut [u8]) -> Result<()> { 26 | let mut i = 0; 27 | while i < rv.len() { 28 | let res_nread = { 29 | let buf = &mut rv[i..]; 30 | self.read(buf) 31 | }; 32 | match res_nread { 33 | Ok(0) => { 34 | let ret = io::Error::new(io::ErrorKind::WouldBlock, "would block"); 35 | return Err(ret.into()); 36 | } 37 | Ok(nread) => i += nread, 38 | Err(e) => return Err(From::from(e)), 39 | } 40 | } 41 | Ok(()) 42 | } 43 | 44 | fn read_scalar(&mut self) -> Result 45 | where 46 | V: Copy + Unmarshal + StatBuffer, 47 | { 48 | let mut buffer = V::buffer(); 49 | self.read_bytes(buffer.as_mut())?; 50 | Ok(V::unmarshal(buffer.as_ref())) 51 | } 52 | 53 | fn read_string(&mut self) -> Result { 54 | let str_len = self.read_uvarint()? as usize; 55 | let mut buffer = vec![0_u8; str_len]; 56 | self.read_bytes(buffer.as_mut())?; 57 | Ok(String::from_utf8(buffer)?) 58 | } 59 | 60 | fn skip_string(&mut self) -> Result<()> { 61 | let str_len = self.read_uvarint()? as usize; 62 | 63 | if str_len <= MAX_STACK_BUFFER_LEN { 64 | unsafe { 65 | let mut buffer: [MaybeUninit; MAX_STACK_BUFFER_LEN] = 66 | MaybeUninit::uninit().assume_init(); 67 | self.read_bytes( 68 | &mut *(&mut buffer[..str_len] as *mut [MaybeUninit] as *mut [u8]), 69 | )?; 70 | } 71 | } else { 72 | let mut buffer = vec![0_u8; str_len]; 73 | self.read_bytes(buffer.as_mut())?; 74 | } 75 | 76 | Ok(()) 77 | } 78 | 79 | fn read_uvarint(&mut self) -> Result { 80 | let mut x = 0_u64; 81 | let mut s = 0_u32; 82 | let mut i = 0_usize; 83 | loop { 84 | let b: u8 = self.read_scalar()?; 85 | 86 | if b < 0x80 { 87 | if i > 9 || i == 9 && b > 1 { 88 | return Err(Error::Driver(DriverError::Overflow)); 89 | } 90 | return Ok(x | (u64::from(b) << s)); 91 | } 92 | 93 | x |= u64::from(b & 0x7f) << s; 94 | s += 7; 95 | 96 | i += 1; 97 | } 98 | } 99 | 100 | fn read_str_into_buffer(&mut self, pool: &mut StringPool) -> Result<()> { 101 | let str_len = self.read_uvarint()? as usize; 102 | let buffer = pool.allocate(str_len); 103 | self.read_bytes(buffer)?; 104 | Ok(()) 105 | } 106 | } 107 | 108 | #[test] 109 | fn test_read_uvarint() { 110 | use super::ReadEx; 111 | use std::io::Cursor; 112 | 113 | let bytes = [194_u8, 10]; 114 | let mut cursor = Cursor::new(bytes); 115 | 116 | let actual = cursor.read_uvarint().unwrap(); 117 | 118 | assert_eq!(actual, 1346) 119 | } 120 | -------------------------------------------------------------------------------- /src/binary/uvarint.rs: -------------------------------------------------------------------------------- 1 | // put_uvarint encodes a uint64 into buf and returns the number of bytes written. 2 | // If the buffer is too small, put_uvarint will panic. 3 | pub fn put_uvarint(mut buffer: impl AsMut<[u8]>, x: u64) -> usize { 4 | let mut i = 0; 5 | let mut mx = x; 6 | let buf = buffer.as_mut(); 7 | while mx >= 0x80 { 8 | buf[i] = mx as u8 | 0x80; 9 | mx >>= 7; 10 | i += 1; 11 | } 12 | buf[i] = mx as u8; 13 | i + 1 14 | } 15 | 16 | #[cfg(test)] 17 | mod test { 18 | #[test] 19 | fn test_put_uvarint() { 20 | let expected = [148u8, 145, 6, 0, 0, 0, 0, 0, 0, 0]; 21 | let mut buffer = [0u8; 10]; 22 | 23 | let actual = super::put_uvarint(&mut buffer[..], 100_500); 24 | 25 | assert_eq!(actual, 3); 26 | assert_eq!(buffer, expected); 27 | } 28 | } 29 | -------------------------------------------------------------------------------- /src/client_info.rs: -------------------------------------------------------------------------------- 1 | use crate::binary::Encoder; 2 | 3 | pub static CLIENT_NAME: &str = "Rust SQLDriver"; 4 | 5 | pub const CLICK_HOUSE_REVISION: u64 = 54429; // DBMS_MIN_REVISION_WITH_SETTINGS_SERIALIZED_AS_STRINGS 6 | pub const CLICK_HOUSE_DBMSVERSION_MAJOR: u64 = 1; 7 | pub const CLICK_HOUSE_DBMSVERSION_MINOR: u64 = 1; 8 | 9 | pub fn write(encoder: &mut Encoder) { 10 | encoder.string(CLIENT_NAME); 11 | encoder.uvarint(CLICK_HOUSE_DBMSVERSION_MAJOR); 12 | encoder.uvarint(CLICK_HOUSE_DBMSVERSION_MINOR); 13 | encoder.uvarint(CLICK_HOUSE_REVISION); 14 | } 15 | 16 | pub fn description() -> String { 17 | format!( 18 | "{CLIENT_NAME} {CLICK_HOUSE_DBMSVERSION_MAJOR}.{CLICK_HOUSE_DBMSVERSION_MINOR}.{CLICK_HOUSE_REVISION}", 19 | ) 20 | } 21 | 22 | #[test] 23 | fn test_description() { 24 | assert_eq!( 25 | description(), 26 | format!( 27 | "Rust SQLDriver {}.{}.{}", 28 | CLICK_HOUSE_DBMSVERSION_MAJOR, CLICK_HOUSE_DBMSVERSION_MINOR, CLICK_HOUSE_REVISION 29 | ) 30 | ) 31 | } 32 | -------------------------------------------------------------------------------- /src/connecting_stream.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | future::Future, 3 | io, 4 | pin::Pin, 5 | task::{Context, Poll}, 6 | }; 7 | 8 | use futures_util::future::{select_ok, BoxFuture, SelectOk, TryFutureExt}; 9 | #[cfg(feature = "_tls")] 10 | use futures_util::FutureExt; 11 | 12 | #[cfg(feature = "async_std")] 13 | use async_std::net::TcpStream; 14 | #[cfg(feature = "tls-native-tls")] 15 | use native_tls::TlsConnector; 16 | #[cfg(feature = "tokio_io")] 17 | use tokio::net::TcpStream; 18 | #[cfg(feature = "tls-rustls")] 19 | use { 20 | rustls::{ 21 | client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier}, 22 | crypto::{verify_tls12_signature, verify_tls13_signature}, 23 | pki_types::{CertificateDer, ServerName, UnixTime}, 24 | ClientConfig, DigitallySignedStruct, Error as TlsError, RootCertStore, 25 | }, 26 | std::sync::Arc, 27 | tokio_rustls::TlsConnector, 28 | }; 29 | 30 | use pin_project::pin_project; 31 | use url::Url; 32 | 33 | use crate::{errors::ConnectionError, io::Stream as InnerStream, Options}; 34 | #[cfg(feature = "tls-native-tls")] 35 | use tokio_native_tls::TlsStream; 36 | #[cfg(feature = "tls-rustls")] 37 | use tokio_rustls::client::TlsStream; 38 | 39 | type Result = std::result::Result; 40 | 41 | type ConnectingFuture = BoxFuture<'static, Result>; 42 | 43 | #[pin_project(project = TcpStateProj)] 44 | enum TcpState { 45 | Wait(#[pin] SelectOk>), 46 | Fail(Option), 47 | } 48 | 49 | #[cfg(feature = "_tls")] 50 | #[pin_project(project = TlsStateProj)] 51 | enum TlsState { 52 | Wait(#[pin] ConnectingFuture>), 53 | Fail(Option), 54 | } 55 | 56 | #[pin_project(project = StateProj)] 57 | enum State { 58 | Tcp(#[pin] TcpState), 59 | #[cfg(feature = "_tls")] 60 | Tls(#[pin] TlsState), 61 | } 62 | 63 | impl TcpState { 64 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 65 | match self.project() { 66 | TcpStateProj::Wait(inner) => match inner.poll(cx) { 67 | Poll::Ready(Ok((tcp, _))) => Poll::Ready(Ok(InnerStream::Plain(tcp))), 68 | Poll::Ready(Err(err)) => Poll::Ready(Err(err)), 69 | Poll::Pending => Poll::Pending, 70 | }, 71 | TcpStateProj::Fail(ref mut err) => Poll::Ready(Err(err.take().unwrap())), 72 | } 73 | } 74 | } 75 | 76 | #[cfg(feature = "_tls")] 77 | impl TlsState { 78 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 79 | match self.project() { 80 | TlsStateProj::Wait(ref mut inner) => match inner.poll_unpin(cx) { 81 | Poll::Ready(Ok(tls)) => Poll::Ready(Ok(InnerStream::Secure(tls))), 82 | Poll::Ready(Err(err)) => Poll::Ready(Err(err)), 83 | Poll::Pending => Poll::Pending, 84 | }, 85 | TlsStateProj::Fail(ref mut err) => { 86 | let e = err.take().unwrap(); 87 | Poll::Ready(Err(e)) 88 | } 89 | } 90 | } 91 | } 92 | 93 | impl State { 94 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 95 | match self.project() { 96 | StateProj::Tcp(inner) => inner.poll(cx), 97 | #[cfg(feature = "_tls")] 98 | StateProj::Tls(inner) => inner.poll(cx), 99 | } 100 | } 101 | 102 | fn tcp_err(e: io::Error) -> Self { 103 | let conn_error = ConnectionError::IoError(e); 104 | State::Tcp(TcpState::Fail(Some(conn_error))) 105 | } 106 | 107 | #[cfg(feature = "_tls")] 108 | fn tls_host_err() -> Self { 109 | State::Tls(TlsState::Fail(Some(ConnectionError::TlsHostNotProvided))) 110 | } 111 | 112 | fn tcp_wait(socket: SelectOk>) -> Self { 113 | State::Tcp(TcpState::Wait(socket)) 114 | } 115 | 116 | #[cfg(feature = "_tls")] 117 | fn tls_wait(s: ConnectingFuture>) -> Self { 118 | State::Tls(TlsState::Wait(s)) 119 | } 120 | } 121 | 122 | #[pin_project] 123 | pub(crate) struct ConnectingStream { 124 | #[pin] 125 | state: State, 126 | } 127 | 128 | #[derive(Debug)] 129 | struct DummyTlsVerifier; 130 | 131 | #[cfg(feature = "tls-rustls")] 132 | impl ServerCertVerifier for DummyTlsVerifier { 133 | fn verify_server_cert( 134 | &self, 135 | _end_entity: &CertificateDer<'_>, 136 | _intermediates: &[CertificateDer<'_>], 137 | _server_name: &ServerName<'_>, 138 | _ocsp_response: &[u8], 139 | _now: UnixTime, 140 | ) -> std::result::Result { 141 | Ok(ServerCertVerified::assertion()) 142 | } 143 | 144 | fn verify_tls12_signature( 145 | &self, 146 | message: &[u8], 147 | cert: &CertificateDer<'_>, 148 | dss: &DigitallySignedStruct, 149 | ) -> std::result::Result { 150 | verify_tls12_signature( 151 | message, 152 | cert, 153 | dss, 154 | &rustls::crypto::ring::default_provider().signature_verification_algorithms, 155 | ) 156 | } 157 | 158 | fn verify_tls13_signature( 159 | &self, 160 | message: &[u8], 161 | cert: &CertificateDer<'_>, 162 | dss: &DigitallySignedStruct, 163 | ) -> std::result::Result { 164 | verify_tls13_signature( 165 | message, 166 | cert, 167 | dss, 168 | &rustls::crypto::ring::default_provider().signature_verification_algorithms, 169 | ) 170 | } 171 | 172 | fn supported_verify_schemes(&self) -> Vec { 173 | rustls::crypto::ring::default_provider() 174 | .signature_verification_algorithms 175 | .supported_schemes() 176 | } 177 | } 178 | 179 | impl ConnectingStream { 180 | #[allow(unused_variables)] 181 | pub(crate) fn new(addr: &Url, options: &Options) -> Self { 182 | match addr.socket_addrs(|| None) { 183 | Ok(addresses) => { 184 | let streams: Vec<_> = addresses 185 | .iter() 186 | .copied() 187 | .map(|address| -> ConnectingFuture { 188 | Box::pin(TcpStream::connect(address).map_err(ConnectionError::IoError)) 189 | }) 190 | .collect(); 191 | 192 | if streams.is_empty() { 193 | let err = io::Error::new( 194 | io::ErrorKind::InvalidInput, 195 | "Could not resolve to any address.", 196 | ); 197 | return Self { 198 | state: State::tcp_err(err), 199 | }; 200 | } 201 | 202 | let socket = select_ok(streams); 203 | 204 | #[cfg(feature = "_tls")] 205 | { 206 | if options.secure { 207 | return ConnectingStream::new_tls_connection(addr, socket, options); 208 | } 209 | } 210 | 211 | Self { 212 | state: State::tcp_wait(socket), 213 | } 214 | } 215 | Err(err) => Self { 216 | state: State::tcp_err(err), 217 | }, 218 | } 219 | } 220 | 221 | #[cfg(feature = "tls-native-tls")] 222 | fn new_tls_connection( 223 | addr: &Url, 224 | socket: SelectOk>, 225 | options: &Options, 226 | ) -> Self { 227 | match addr.host_str().map(|host| host.to_owned()) { 228 | None => Self { 229 | state: State::tls_host_err(), 230 | }, 231 | Some(host) => { 232 | let mut builder = TlsConnector::builder(); 233 | builder.danger_accept_invalid_certs(options.skip_verify); 234 | if let Some(certificate) = options.certificate.clone() { 235 | let native_cert = native_tls::Certificate::from(certificate); 236 | builder.add_root_certificate(native_cert); 237 | } 238 | 239 | Self { 240 | state: State::tls_wait(Box::pin(async move { 241 | let (s, _) = socket.await?; 242 | 243 | let cx = builder.build()?; 244 | let cx = tokio_native_tls::TlsConnector::from(cx); 245 | 246 | Ok(cx.connect(&host, s).await?) 247 | })), 248 | } 249 | } 250 | } 251 | } 252 | 253 | #[cfg(feature = "tls-rustls")] 254 | fn new_tls_connection( 255 | addr: &Url, 256 | socket: SelectOk>, 257 | options: &Options, 258 | ) -> Self { 259 | match addr.host_str().map(|host| host.to_owned()) { 260 | None => Self { 261 | state: State::tls_host_err(), 262 | }, 263 | Some(host) => { 264 | let config = if options.skip_verify { 265 | ClientConfig::builder() 266 | .dangerous() 267 | .with_custom_certificate_verifier(Arc::new(DummyTlsVerifier)) 268 | .with_no_client_auth() 269 | } else { 270 | let mut cert_store = RootCertStore::empty(); 271 | cert_store.extend( 272 | webpki_roots::TLS_SERVER_ROOTS 273 | .iter() 274 | .cloned() 275 | ); 276 | if let Some(certificates) = options.certificate.clone() { 277 | for certificate in 278 | Into::>>::into( 279 | certificates, 280 | ) 281 | { 282 | match cert_store.add(certificate) { 283 | Ok(_) => {}, 284 | Err(err) => { 285 | let err = io::Error::new( 286 | io::ErrorKind::InvalidInput, 287 | format!("Could not load certificate: {}.", err), 288 | ); 289 | return Self { state: State::tcp_err(err) }; 290 | }, 291 | } 292 | } 293 | } 294 | ClientConfig::builder() 295 | .with_root_certificates(cert_store) 296 | .with_no_client_auth() 297 | }; 298 | Self { 299 | state: State::tls_wait(Box::pin(async move { 300 | let (s, _) = socket.await?; 301 | let cx = TlsConnector::from(Arc::new(config)); 302 | let host = ServerName::try_from(host) 303 | .map_err(|_| ConnectionError::TlsHostNotProvided)?; 304 | Ok(cx 305 | .connect(host, s) 306 | .await 307 | .map_err(|e| ConnectionError::IoError(e))?) 308 | })), 309 | } 310 | } 311 | } 312 | } 313 | } 314 | 315 | impl Future for ConnectingStream { 316 | type Output = Result; 317 | 318 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 319 | self.project().state.poll(cx) 320 | } 321 | } 322 | -------------------------------------------------------------------------------- /src/errors/mod.rs: -------------------------------------------------------------------------------- 1 | use std::{borrow::Cow, io, result, str::Utf8Error, string::FromUtf8Error}; 2 | 3 | use thiserror::Error; 4 | #[cfg(feature = "tokio_io")] 5 | use tokio::time::error::Elapsed; 6 | use url::ParseError; 7 | 8 | #[cfg(feature = "tls-native-tls")] 9 | use native_tls::Error as TlsError; 10 | #[cfg(feature = "tls-rustls")] 11 | use rustls::Error as TlsError; 12 | 13 | /// Clickhouse error codes 14 | pub mod codes; 15 | 16 | /// Result type alias for this library. 17 | pub type Result = result::Result; 18 | 19 | /// This type enumerates library errors. 20 | #[derive(Debug, Error)] 21 | pub enum Error { 22 | #[error("Driver error: `{}`", _0)] 23 | Driver(#[source] DriverError), 24 | 25 | #[error("Input/output error: `{}`", _0)] 26 | Io(#[source] io::Error), 27 | 28 | #[error("Connections error: `{}`", _0)] 29 | Connection(#[source] ConnectionError), 30 | 31 | #[error("Other error: `{}`", _0)] 32 | Other(Cow<'static, str>), 33 | 34 | #[error("Server error: `{}`", _0)] 35 | Server(#[source] ServerError), 36 | 37 | #[error("URL error: `{}`", _0)] 38 | Url(#[source] UrlError), 39 | 40 | #[error("From SQL error: `{}`", _0)] 41 | FromSql(#[source] FromSqlError), 42 | } 43 | 44 | /// This type represents Clickhouse server error. 45 | #[derive(Debug, Error, Clone)] 46 | #[error("ERROR {} ({}): {}", name, code, message)] 47 | pub struct ServerError { 48 | pub code: u32, 49 | pub name: String, 50 | pub message: String, 51 | pub stack_trace: String, 52 | } 53 | 54 | /// This type enumerates connection errors. 55 | #[derive(Debug, Error)] 56 | pub enum ConnectionError { 57 | #[error("TLS connection requires hostname to be provided")] 58 | TlsHostNotProvided, 59 | 60 | #[error("Input/output error: `{}`", _0)] 61 | IoError(#[source] io::Error), 62 | 63 | #[cfg(feature = "_tls")] 64 | #[error("TLS connection error: `{}`", _0)] 65 | TlsError(#[source] TlsError), 66 | 67 | #[error("Connection broken")] 68 | Broken, 69 | } 70 | 71 | /// This type enumerates connection URL errors. 72 | #[derive(Debug, Error)] 73 | pub enum UrlError { 74 | #[error("Invalid or incomplete connection URL")] 75 | Invalid, 76 | 77 | #[error("Invalid value `{}' for connection URL parameter `{}'", value, param)] 78 | InvalidParamValue { param: String, value: String }, 79 | 80 | #[error("URL parse error: {}", _0)] 81 | Parse(#[source] ParseError), 82 | 83 | #[error("Unknown connection URL parameter `{}'", param)] 84 | UnknownParameter { param: String }, 85 | 86 | #[error("Unsupported connection URL scheme `{}'", scheme)] 87 | UnsupportedScheme { scheme: String }, 88 | } 89 | 90 | /// This type enumerates driver errors. 91 | #[derive(Debug, Error)] 92 | pub enum DriverError { 93 | #[error("Varint overflows a 64-bit integer.")] 94 | Overflow, 95 | 96 | #[error("Unknown packet 0x{:x}.", packet)] 97 | UnknownPacket { packet: u64 }, 98 | 99 | #[error("Unexpected packet.")] 100 | UnexpectedPacket, 101 | 102 | #[error("Timeout error.")] 103 | Timeout, 104 | 105 | #[error("Invalid utf-8 sequence.")] 106 | Utf8Error(Utf8Error), 107 | 108 | #[error("Deserialize error: `{}`", _0)] 109 | Deserialize(Cow<'static, str>), 110 | } 111 | 112 | /// This type enumerates cast from sql type errors. 113 | #[derive(Debug, Error)] 114 | pub enum FromSqlError { 115 | #[error("SqlType::{} cannot be cast to {}.", src, dst)] 116 | InvalidType { 117 | src: Cow<'static, str>, 118 | dst: Cow<'static, str>, 119 | }, 120 | 121 | #[error("Out of range.")] 122 | OutOfRange, 123 | 124 | #[error("Unsupported operation.")] 125 | UnsupportedOperation, 126 | } 127 | 128 | impl Error { 129 | pub(crate) fn is_would_block(&self) -> bool { 130 | if let Error::Io(ref e) = self { 131 | if e.kind() == io::ErrorKind::WouldBlock { 132 | return true; 133 | } 134 | } 135 | false 136 | } 137 | } 138 | 139 | impl From for Error { 140 | fn from(error: ConnectionError) -> Self { 141 | Error::Connection(error) 142 | } 143 | } 144 | 145 | #[cfg(feature = "_tls")] 146 | impl From for ConnectionError { 147 | fn from(error: TlsError) -> Self { 148 | ConnectionError::TlsError(error) 149 | } 150 | } 151 | 152 | impl From for Error { 153 | fn from(err: DriverError) -> Self { 154 | Error::Driver(err) 155 | } 156 | } 157 | 158 | impl From for Error { 159 | fn from(err: io::Error) -> Self { 160 | Error::Io(err) 161 | } 162 | } 163 | 164 | impl From for Error { 165 | fn from(err: ServerError) -> Self { 166 | Error::Server(err) 167 | } 168 | } 169 | 170 | impl From for Error { 171 | fn from(err: UrlError) -> Self { 172 | Error::Url(err) 173 | } 174 | } 175 | 176 | impl From for Error { 177 | fn from(err: String) -> Self { 178 | Error::Other(Cow::from(err)) 179 | } 180 | } 181 | 182 | impl From<&str> for Error { 183 | fn from(err: &str) -> Self { 184 | Error::Other(err.to_string().into()) 185 | } 186 | } 187 | 188 | impl From for Error { 189 | fn from(err: FromUtf8Error) -> Self { 190 | Error::Other(err.to_string().into()) 191 | } 192 | } 193 | 194 | #[cfg(feature = "tokio_io")] 195 | impl From for Error { 196 | fn from(_err: Elapsed) -> Self { 197 | Error::Driver(DriverError::Timeout) 198 | } 199 | } 200 | 201 | impl From for Error { 202 | fn from(err: ParseError) -> Self { 203 | Error::Url(UrlError::Parse(err)) 204 | } 205 | } 206 | 207 | impl From for io::Error { 208 | fn from(err: Error) -> Self { 209 | match err { 210 | Error::Io(error) => error, 211 | e => io::Error::new(io::ErrorKind::Other, e.to_string()), 212 | } 213 | } 214 | } 215 | 216 | impl From for Error { 217 | fn from(err: Utf8Error) -> Self { 218 | Error::Driver(DriverError::Utf8Error(err)) 219 | } 220 | } 221 | 222 | #[cfg(test)] 223 | mod tests { 224 | #[test] 225 | fn to_std_error_without_recursion() { 226 | let src_err: super::Error = From::from("Somth went wrong."); 227 | let dst_err: Box = src_err.into(); 228 | assert_eq!(dst_err.to_string(), "Other error: `Somth went wrong.`"); 229 | } 230 | 231 | #[test] 232 | fn to_io_error_without_recursion() { 233 | let src_err: super::Error = From::from("Somth went wrong."); 234 | let dst_err: std::io::Error = src_err.into(); 235 | assert_eq!(dst_err.to_string(), "Other error: `Somth went wrong.`"); 236 | } 237 | } 238 | -------------------------------------------------------------------------------- /src/io/mod.rs: -------------------------------------------------------------------------------- 1 | pub(crate) use self::{stream::Stream, transport::ClickhouseTransport}; 2 | 3 | mod read_to_end; 4 | pub(crate) mod stream; 5 | pub(crate) mod transport; 6 | -------------------------------------------------------------------------------- /src/io/read_to_end.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | io, 3 | pin::Pin, 4 | task::{Context, Poll}, 5 | }; 6 | 7 | use futures_util::ready; 8 | 9 | use crate::io::Stream as InnerStream; 10 | 11 | struct Guard<'a> { 12 | buf: &'a mut Vec, 13 | len: usize, 14 | } 15 | 16 | impl Drop for Guard<'_> { 17 | fn drop(&mut self) { 18 | unsafe { 19 | self.buf.set_len(self.len); 20 | } 21 | } 22 | } 23 | 24 | pub(crate) fn read_to_end( 25 | mut rd: Pin<&mut InnerStream>, 26 | cx: &mut Context<'_>, 27 | buf: &mut Vec, 28 | ) -> Poll> { 29 | let start_len = buf.len(); 30 | let mut g = Guard { 31 | len: buf.len(), 32 | buf, 33 | }; 34 | let ret; 35 | loop { 36 | if g.len == g.buf.len() { 37 | unsafe { 38 | g.buf.reserve(32); 39 | let capacity = g.buf.capacity(); 40 | g.buf.set_len(capacity); 41 | } 42 | } 43 | 44 | match ready!(rd.as_mut().poll_read(cx, &mut g.buf[g.len..])) { 45 | Ok(0) => { 46 | ret = Poll::Ready(Ok(g.len - start_len)); 47 | break; 48 | } 49 | Ok(n) => g.len += n, 50 | Err(e) => { 51 | ret = Poll::Ready(Err(e)); 52 | break; 53 | } 54 | } 55 | } 56 | 57 | ret 58 | } 59 | -------------------------------------------------------------------------------- /src/io/stream.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | io, 3 | pin::Pin, 4 | task::{Context, Poll}, 5 | time::Duration, 6 | }; 7 | 8 | #[cfg(feature = "tokio_io")] 9 | use tokio::{io::ReadBuf, net::TcpStream}; 10 | #[cfg(feature = "tls-native-tls")] 11 | use tokio_native_tls::TlsStream; 12 | #[cfg(feature = "tls-rustls")] 13 | use tokio_rustls::client::TlsStream; 14 | 15 | #[cfg(feature = "async_std")] 16 | use async_std::io::prelude::*; 17 | #[cfg(feature = "async_std")] 18 | use async_std::net::TcpStream; 19 | use pin_project::pin_project; 20 | #[cfg(feature = "tokio_io")] 21 | use tokio::io::{AsyncRead, AsyncWrite}; 22 | 23 | #[cfg(all(feature = "_tls", feature = "tokio_io"))] 24 | type SecureTcpStream = TlsStream; 25 | 26 | #[pin_project(project = StreamProj)] 27 | pub(crate) enum Stream { 28 | Plain(#[pin] TcpStream), 29 | #[cfg(feature = "_tls")] 30 | Secure(#[pin] SecureTcpStream), 31 | } 32 | 33 | impl From for Stream { 34 | fn from(stream: TcpStream) -> Self { 35 | Self::Plain(stream) 36 | } 37 | } 38 | 39 | #[cfg(feature = "_tls")] 40 | impl From for Stream { 41 | fn from(stream: SecureTcpStream) -> Stream { 42 | Self::Secure(stream) 43 | } 44 | } 45 | 46 | impl Stream { 47 | #[cfg(feature = "async_std")] 48 | pub(crate) fn set_keepalive(&mut self, keepalive: Option) -> io::Result<()> { 49 | if keepalive.is_some() { 50 | log::warn!("`std_async` doesn't support `set_keepalive`.") 51 | } 52 | Ok(()) 53 | } 54 | 55 | #[cfg(not(feature = "async_std"))] 56 | #[allow(clippy::all)] 57 | pub(crate) fn set_keepalive(&mut self, keepalive: Option) -> io::Result<()> { 58 | // match *self { 59 | // Self::Plain(ref mut stream) => stream.set_keepalive(keepalive), 60 | // #[cfg(feature = "_tls")] 61 | // Self::Secure(ref mut stream) => stream.get_mut().set_keepalive(keepalive), 62 | // }.map_err(|err| io::Error::new(err.kind(), format!("set_keepalive error: {}", err))) 63 | if keepalive.is_some() { 64 | // https://github.com/tokio-rs/tokio/issues/3082 65 | log::warn!("`tokio` dropped `set_keepalive` in v0.3 and is currently working on getting it back") 66 | } 67 | Ok(()) 68 | } 69 | 70 | pub(crate) fn set_nodelay(&mut self, nodelay: bool) -> io::Result<()> { 71 | match *self { 72 | Self::Plain(ref mut stream) => stream.set_nodelay(nodelay), 73 | #[cfg(feature = "tls-native-tls")] 74 | Self::Secure(ref mut stream) => { 75 | stream.get_mut().get_mut().get_mut().set_nodelay(nodelay) 76 | } 77 | #[cfg(feature = "tls-rustls")] 78 | Self::Secure(ref mut stream) => stream.get_mut().0.set_nodelay(nodelay), 79 | } 80 | .map_err(|err| io::Error::new(err.kind(), format!("set_nodelay error: {err}"))) 81 | } 82 | 83 | #[cfg(feature = "async_std")] 84 | pub(crate) fn poll_read( 85 | self: Pin<&mut Self>, 86 | cx: &mut Context<'_>, 87 | buf: &mut [u8], 88 | ) -> Poll> { 89 | match self.project() { 90 | StreamProj::Plain(stream) => stream.poll_read(cx, buf), 91 | #[cfg(feature = "_tls")] 92 | StreamProj::Secure(stream) => stream.poll_read(cx, buf), 93 | } 94 | } 95 | 96 | #[cfg(not(feature = "async_std"))] 97 | pub(crate) fn poll_read( 98 | self: Pin<&mut Self>, 99 | cx: &mut Context<'_>, 100 | buf: &mut [u8], 101 | ) -> Poll> { 102 | let mut read_buf = ReadBuf::new(buf); 103 | 104 | let result = match self.project() { 105 | StreamProj::Plain(stream) => stream.poll_read(cx, &mut read_buf), 106 | #[cfg(feature = "_tls")] 107 | StreamProj::Secure(stream) => stream.poll_read(cx, &mut read_buf), 108 | }; 109 | 110 | match result { 111 | Poll::Ready(Ok(())) => Poll::Ready(Ok(read_buf.filled().len())), 112 | Poll::Ready(Err(x)) => Poll::Ready(Err(x)), 113 | Poll::Pending => Poll::Pending, 114 | } 115 | } 116 | 117 | pub(crate) fn poll_write( 118 | self: Pin<&mut Self>, 119 | cx: &mut Context<'_>, 120 | buf: &[u8], 121 | ) -> Poll> { 122 | match self.project() { 123 | StreamProj::Plain(stream) => stream.poll_write(cx, buf), 124 | #[cfg(feature = "_tls")] 125 | StreamProj::Secure(stream) => stream.poll_write(cx, buf), 126 | } 127 | } 128 | } 129 | -------------------------------------------------------------------------------- /src/pool/futures/get_handle.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | future::Future, 3 | pin::Pin, 4 | task::{Context, Poll}, 5 | }; 6 | 7 | use pin_project::pin_project; 8 | 9 | use crate::{errors::Result, pool::Pool, ClientHandle}; 10 | 11 | /// Future that resolves to a `ClientHandle`. 12 | #[pin_project] 13 | pub struct GetHandle { 14 | #[pin] 15 | pool: Pool, 16 | } 17 | 18 | impl GetHandle { 19 | pub(crate) fn new(pool: &Pool) -> Self { 20 | Self { pool: pool.clone() } 21 | } 22 | } 23 | 24 | impl Future for GetHandle { 25 | type Output = Result; 26 | 27 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 28 | self.project().pool.poll(cx) 29 | } 30 | } 31 | -------------------------------------------------------------------------------- /src/pool/futures/mod.rs: -------------------------------------------------------------------------------- 1 | mod get_handle; 2 | 3 | pub use self::get_handle::GetHandle; 4 | -------------------------------------------------------------------------------- /src/retry_guard.rs: -------------------------------------------------------------------------------- 1 | use std::{mem, time::Duration}; 2 | 3 | use log::warn; 4 | 5 | use crate::{errors::Result, types::OptionsSource, Client, ClientHandle, Pool}; 6 | 7 | pub(crate) async fn retry_guard( 8 | handle: &mut ClientHandle, 9 | source: &OptionsSource, 10 | pool: Option, 11 | max_attempt: usize, 12 | duration: Duration, 13 | ) -> Result<()> { 14 | let mut attempt = 0; 15 | let mut skip_check = false; 16 | 17 | loop { 18 | if skip_check { 19 | skip_check = false; 20 | } else { 21 | match check(handle).await { 22 | Ok(()) => return Ok(()), 23 | Err(err) => { 24 | if attempt >= max_attempt { 25 | return Err(err); 26 | } 27 | } 28 | } 29 | } 30 | 31 | match reconnect(handle, source, pool.clone()).await { 32 | Ok(()) => continue, 33 | Err(err) => { 34 | skip_check = true; 35 | if attempt >= max_attempt { 36 | return Err(err); 37 | } 38 | 39 | #[cfg(feature = "async_std")] 40 | { 41 | use async_std::task; 42 | task::sleep(duration).await; 43 | } 44 | 45 | #[cfg(not(feature = "async_std"))] 46 | { 47 | tokio::time::sleep(duration).await; 48 | } 49 | } 50 | } 51 | 52 | attempt += 1; 53 | } 54 | } 55 | 56 | async fn check(c: &mut ClientHandle) -> Result<()> { 57 | c.ping().await 58 | } 59 | 60 | async fn reconnect(c: &mut ClientHandle, source: &OptionsSource, pool: Option) -> Result<()> { 61 | warn!("[reconnect]"); 62 | let mut nc = match pool { 63 | None => Client::open(source.clone(), pool).await?, 64 | Some(p) => p.get_handle().await?, 65 | }; 66 | mem::swap(c, &mut nc); 67 | Ok(()) 68 | } 69 | -------------------------------------------------------------------------------- /src/types/block/block_info.rs: -------------------------------------------------------------------------------- 1 | use crate::{ 2 | binary::{Encoder, ReadEx}, 3 | errors::{DriverError, Error, Result}, 4 | }; 5 | 6 | const BLOCK_INFO_OVERFLOWS: u64 = 1; 7 | const BLOCK_INFO_BUCKET_NUM: u64 = 2; 8 | const END_FIELD: u64 = 0; 9 | 10 | #[allow(dead_code)] 11 | #[derive(Copy, Clone)] 12 | pub struct BlockInfo { 13 | is_overflows: bool, 14 | bucket_num: i32, 15 | } 16 | 17 | impl Default for BlockInfo { 18 | fn default() -> Self { 19 | Self { 20 | is_overflows: false, 21 | bucket_num: -1, 22 | } 23 | } 24 | } 25 | 26 | impl BlockInfo { 27 | pub(crate) fn read(reader: &mut R) -> Result { 28 | let mut block_info = BlockInfo::default(); 29 | 30 | loop { 31 | match reader.read_uvarint()? { 32 | BLOCK_INFO_OVERFLOWS => block_info.is_overflows = reader.read_scalar()?, 33 | BLOCK_INFO_BUCKET_NUM => block_info.bucket_num = reader.read_scalar()?, 34 | END_FIELD => break, 35 | i => { 36 | return Err(Error::Driver(DriverError::UnknownPacket { packet: i })); 37 | } 38 | } 39 | } 40 | 41 | Ok(block_info) 42 | } 43 | 44 | pub fn write(&self, encoder: &mut Encoder) { 45 | encoder.uvarint(BLOCK_INFO_OVERFLOWS); 46 | encoder.write(self.is_overflows); 47 | encoder.uvarint(BLOCK_INFO_BUCKET_NUM); 48 | encoder.write(self.bucket_num); 49 | encoder.uvarint(END_FIELD); 50 | } 51 | } 52 | 53 | #[cfg(test)] 54 | mod test { 55 | use std::io::Cursor; 56 | 57 | use super::*; 58 | 59 | #[test] 60 | fn test_read_write_blockinfo() { 61 | let b1 = BlockInfo { 62 | is_overflows: true, 63 | bucket_num: 3, 64 | }; 65 | 66 | let mut enc = Encoder::new(); 67 | b1.write(&mut enc); 68 | 69 | let mut buf = Cursor::new(enc.get_buffer()); 70 | let b2 = BlockInfo::read(&mut buf).expect("Failed to deserialize BlockInfo"); 71 | 72 | assert_eq!(b1.is_overflows, b2.is_overflows); 73 | assert_eq!(b1.bucket_num, b2.bucket_num); 74 | } 75 | } 76 | -------------------------------------------------------------------------------- /src/types/block/builder.rs: -------------------------------------------------------------------------------- 1 | use std::{borrow::Cow, marker}; 2 | 3 | use chrono_tz::Tz; 4 | use either::Either; 5 | 6 | use crate::{ 7 | errors::{Error, FromSqlError, Result}, 8 | types::{ 9 | block::ColumnIdx, 10 | column::{datetime64::DEFAULT_TZ, ArcColumnWrapper, ColumnData}, 11 | Column, ColumnType, Value, 12 | }, 13 | Block, 14 | }; 15 | 16 | pub trait RowBuilder { 17 | fn apply(self, block: &mut Block) -> Result<()>; 18 | } 19 | 20 | pub struct RNil; 21 | 22 | pub struct RCons 23 | where 24 | T: RowBuilder, 25 | { 26 | key: Cow<'static, str>, 27 | value: Value, 28 | tail: T, 29 | } 30 | 31 | impl RNil { 32 | pub fn put(self, key: Cow<'static, str>, value: Value) -> RCons { 33 | RCons { 34 | key, 35 | value, 36 | tail: RNil, 37 | } 38 | } 39 | } 40 | 41 | impl RCons 42 | where 43 | T: RowBuilder, 44 | { 45 | pub fn put(self, key: Cow<'static, str>, value: Value) -> RCons { 46 | RCons { 47 | key, 48 | value, 49 | tail: self, 50 | } 51 | } 52 | } 53 | 54 | impl RowBuilder for RNil { 55 | #[inline(always)] 56 | fn apply(self, _block: &mut Block) -> Result<()> { 57 | Ok(()) 58 | } 59 | } 60 | 61 | impl RowBuilder for RCons 62 | where 63 | T: RowBuilder, 64 | { 65 | #[inline(always)] 66 | fn apply(self, block: &mut Block) -> Result<()> { 67 | put_param(self.key, self.value, block)?; 68 | self.tail.apply(block) 69 | } 70 | } 71 | 72 | impl RowBuilder for Vec<(String, Value)> { 73 | fn apply(self, block: &mut Block) -> Result<()> { 74 | for (k, v) in self { 75 | put_param(k.into(), v, block)?; 76 | } 77 | Ok(()) 78 | } 79 | } 80 | 81 | fn put_param( 82 | key: Cow<'static, str>, 83 | value: Value, 84 | block: &mut Block, 85 | ) -> Result<()> { 86 | let col_index = match key.as_ref().get_index(&block.columns) { 87 | Ok(col_index) => col_index, 88 | Err(Error::FromSql(FromSqlError::OutOfRange)) => { 89 | if block.row_count() <= 1 { 90 | let sql_type = From::from(value.clone()); 91 | 92 | let timezone = extract_timezone(&value); 93 | 94 | let column = Column { 95 | name: key.clone().into(), 96 | data: ::from_type::( 97 | sql_type, 98 | timezone, 99 | block.capacity, 100 | )?, 101 | _marker: marker::PhantomData, 102 | }; 103 | 104 | block.columns.push(column); 105 | return put_param(key, value, block); 106 | } else { 107 | return Err(Error::FromSql(FromSqlError::OutOfRange)); 108 | } 109 | } 110 | Err(err) => return Err(err), 111 | }; 112 | 113 | block.columns[col_index].push(value); 114 | Ok(()) 115 | } 116 | 117 | fn extract_timezone(value: &Value) -> Tz { 118 | match value { 119 | Value::Date(_) => *DEFAULT_TZ, 120 | Value::DateTime(_, tz) => *tz, 121 | Value::Nullable(Either::Right(d)) => extract_timezone(d), 122 | Value::Array(_, data) => { 123 | if let Some(v) = data.first() { 124 | extract_timezone(v) 125 | } else { 126 | *DEFAULT_TZ 127 | } 128 | } 129 | _ => *DEFAULT_TZ, 130 | } 131 | } 132 | 133 | #[cfg(test)] 134 | mod test { 135 | use chrono::prelude::*; 136 | use chrono_tz::Tz::{self, UTC}; 137 | 138 | use crate::{ 139 | row, 140 | types::{DateTimeType, Decimal, Simple, SqlType}, 141 | }; 142 | 143 | use super::*; 144 | 145 | #[test] 146 | fn test_push_row() { 147 | let date_value: NaiveDate = NaiveDate::from_ymd_opt(2016, 10, 22).unwrap(); 148 | let date_time_value: DateTime = UTC.with_ymd_and_hms(2014, 7, 8, 14, 0, 0).unwrap(); 149 | 150 | let decimal = Decimal::of(2.0_f64, 4); 151 | 152 | let mut block = Block::::new(); 153 | block 154 | .push(row! { 155 | i8_field: 1_i8, 156 | i16_field: 1_i16, 157 | i32_field: 1_i32, 158 | i64_field: 1_i64, 159 | 160 | u8_field: 1_u8, 161 | u16_field: 1_u16, 162 | u32_field: 1_u32, 163 | u64_field: 1_u64, 164 | 165 | f32_field: 4.66_f32, 166 | f64_field: 2.71_f64, 167 | 168 | str_field: "text", 169 | opt_filed: Some("text"), 170 | nil_filed: Option::<&str>::None, 171 | 172 | date_field: date_value, 173 | date_time_field: date_time_value, 174 | 175 | decimal_field: decimal 176 | }) 177 | .unwrap(); 178 | 179 | assert_eq!(block.row_count(), 1); 180 | 181 | assert_eq!(block.columns[0].sql_type(), SqlType::Int8); 182 | assert_eq!(block.columns[1].sql_type(), SqlType::Int16); 183 | assert_eq!(block.columns[2].sql_type(), SqlType::Int32); 184 | assert_eq!(block.columns[3].sql_type(), SqlType::Int64); 185 | 186 | assert_eq!(block.columns[4].sql_type(), SqlType::UInt8); 187 | assert_eq!(block.columns[5].sql_type(), SqlType::UInt16); 188 | assert_eq!(block.columns[6].sql_type(), SqlType::UInt32); 189 | assert_eq!(block.columns[7].sql_type(), SqlType::UInt64); 190 | 191 | assert_eq!(block.columns[8].sql_type(), SqlType::Float32); 192 | assert_eq!(block.columns[9].sql_type(), SqlType::Float64); 193 | 194 | assert_eq!(block.columns[10].sql_type(), SqlType::String); 195 | assert_eq!( 196 | block.columns[11].sql_type(), 197 | SqlType::Nullable(SqlType::String.into()) 198 | ); 199 | assert_eq!( 200 | block.columns[12].sql_type(), 201 | SqlType::Nullable(SqlType::String.into()) 202 | ); 203 | 204 | assert_eq!(block.columns[13].sql_type(), SqlType::Date); 205 | assert_eq!( 206 | block.columns[14].sql_type(), 207 | SqlType::DateTime(DateTimeType::Chrono) 208 | ); 209 | assert_eq!(block.columns[15].sql_type(), SqlType::Decimal(18, 4)); 210 | } 211 | } 212 | -------------------------------------------------------------------------------- /src/types/block/chunk_iterator.rs: -------------------------------------------------------------------------------- 1 | use std::cmp; 2 | 3 | use crate::types::{Block, ColumnType}; 4 | 5 | pub(crate) struct ChunkIterator { 6 | position: usize, 7 | size: usize, 8 | block: Block, 9 | } 10 | 11 | impl Iterator for ChunkIterator { 12 | type Item = Block; 13 | 14 | fn next(&mut self) -> Option { 15 | let m = self.block.row_count(); 16 | 17 | if m == 0 && self.position == 0 { 18 | self.position += 1; 19 | return Some(Block::default()); 20 | } 21 | 22 | if self.position >= m { 23 | return None; 24 | } 25 | 26 | let mut result = Block::new(); 27 | let size = cmp::min(self.size, m - self.position); 28 | 29 | for column in self.block.columns().iter() { 30 | let range = self.position..self.position + size; 31 | let data = column.slice(range); 32 | result = result.column(column.name(), data); 33 | } 34 | 35 | self.position += size; 36 | Some(result) 37 | } 38 | } 39 | 40 | impl ChunkIterator { 41 | pub fn new(size: usize, block: Block) -> ChunkIterator { 42 | ChunkIterator { 43 | position: 0, 44 | size, 45 | block, 46 | } 47 | } 48 | } 49 | -------------------------------------------------------------------------------- /src/types/block/compressed.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | io, 3 | io::Read, 4 | mem, 5 | os::raw::{c_char, c_int}, 6 | }; 7 | 8 | use byteorder::{LittleEndian, WriteBytesExt}; 9 | use clickhouse_rs_cityhash_sys::{city_hash_128, UInt128}; 10 | use lz4::liblz4::LZ4_decompress_safe; 11 | 12 | use crate::{ 13 | binary::ReadEx, 14 | errors::{Error, Result}, 15 | }; 16 | 17 | const DBMS_MAX_COMPRESSED_SIZE: u32 = 0x4000_0000; // 1GB 18 | 19 | pub(crate) struct CompressedReader<'a, R> { 20 | reader: &'a mut R, 21 | cursor: io::Cursor>, 22 | } 23 | 24 | pub(crate) fn make(reader: &mut R) -> CompressedReader { 25 | CompressedReader { 26 | reader, 27 | cursor: io::Cursor::new(Vec::new()), 28 | } 29 | } 30 | 31 | impl<'a, R> CompressedReader<'a, R> 32 | where 33 | R: Read + ReadEx, 34 | { 35 | fn is_empty(&self) -> bool { 36 | let len = self.cursor.get_ref().len(); 37 | let pos = self.cursor.position() as usize; 38 | len == pos 39 | } 40 | 41 | fn fill(&mut self) -> Result<()> { 42 | let cursor = mem::replace(&mut self.cursor, io::Cursor::new(Vec::new())); 43 | let buffer = cursor.into_inner(); 44 | 45 | let tmp = decompress_buffer(&mut self.reader, buffer)?; 46 | self.cursor = io::Cursor::new(tmp); 47 | Ok(()) 48 | } 49 | } 50 | 51 | impl<'a, R> Read for CompressedReader<'a, R> 52 | where 53 | R: Read + ReadEx, 54 | { 55 | fn read(&mut self, buf: &mut [u8]) -> io::Result { 56 | if self.is_empty() { 57 | self.fill()?; 58 | } 59 | 60 | self.cursor.read(buf) 61 | } 62 | } 63 | 64 | fn decompress_buffer(reader: &mut R, mut buffer: Vec) -> Result> 65 | where 66 | R: ReadEx, 67 | { 68 | let h = UInt128 { 69 | lo: reader.read_scalar()?, 70 | hi: reader.read_scalar()?, 71 | }; 72 | 73 | let method: u8 = reader.read_scalar()?; 74 | if method != 0x82 { 75 | let message: String = format!("unsupported compression method {method}"); 76 | return Err(raise_error(message)); 77 | } 78 | 79 | let compressed: u32 = reader.read_scalar()?; 80 | let original: u32 = reader.read_scalar()?; 81 | 82 | if compressed > DBMS_MAX_COMPRESSED_SIZE { 83 | return Err(raise_error("compressed data too big".to_string())); 84 | } 85 | 86 | buffer.resize(compressed as usize, 0_u8); 87 | { 88 | let mut cursor = io::Cursor::new(&mut buffer); 89 | cursor.write_u8(0x82)?; 90 | cursor.write_u32::(compressed)?; 91 | cursor.write_u32::(original)?; 92 | } 93 | reader.read_bytes(&mut buffer[9..])?; 94 | 95 | if h != city_hash_128(&buffer) { 96 | return Err(raise_error("data was corrupted".to_string())); 97 | } 98 | 99 | let data = vec![0_u8; original as usize]; 100 | let status = unsafe { 101 | LZ4_decompress_safe( 102 | (buffer.as_mut_ptr() as *const c_char).add(9), 103 | data.as_ptr() as *mut c_char, 104 | (compressed - 9) as c_int, 105 | original as c_int, 106 | ) 107 | }; 108 | 109 | if status < 0 { 110 | return Err(raise_error("can't decompress data".to_string())); 111 | } 112 | 113 | Ok(data) 114 | } 115 | 116 | fn raise_error(message: String) -> Error { 117 | message.into() 118 | } 119 | 120 | #[cfg(test)] 121 | mod test { 122 | use super::*; 123 | 124 | #[test] 125 | fn test_decompress() { 126 | let expected = [ 127 | 1u8, 0, 2, 255, 255, 255, 255, 0, 1, 1, 1, 115, 6, 83, 116, 114, 105, 110, 103, 3, 97, 128 | 98, 99, 129 | ]; 130 | 131 | let source = [ 132 | 245_u8, 5, 222, 235, 225, 158, 59, 108, 225, 31, 65, 215, 66, 66, 36, 92, 130, 34, 0, 133 | 0, 0, 23, 0, 0, 0, 240, 8, 1, 0, 2, 255, 255, 255, 255, 0, 1, 1, 1, 115, 6, 83, 116, 134 | 114, 105, 110, 103, 3, 97, 98, 99, 135 | ]; 136 | 137 | let mut cursor = io::Cursor::new(&source[..]); 138 | let actual = decompress_buffer(&mut cursor, Vec::new()).unwrap(); 139 | 140 | assert_eq!(actual, expected); 141 | } 142 | } 143 | -------------------------------------------------------------------------------- /src/types/block/row.rs: -------------------------------------------------------------------------------- 1 | use std::{marker, sync::Arc}; 2 | 3 | use crate::{ 4 | errors::Result, 5 | types::{block::ColumnIdx, Block, Column, ColumnType, FromSql, SqlType}, 6 | }; 7 | 8 | /// A row from Clickhouse 9 | pub struct Row<'a, K: ColumnType> { 10 | pub(crate) row: usize, 11 | pub(crate) block_ref: BlockRef<'a, K>, 12 | pub(crate) kind: marker::PhantomData, 13 | } 14 | 15 | impl<'a, K: ColumnType> Row<'a, K> { 16 | /// Get the value of a particular cell of the row. 17 | pub fn get(&'a self, col: I) -> Result 18 | where 19 | T: FromSql<'a>, 20 | I: ColumnIdx + Copy, 21 | { 22 | self.block_ref.get(self.row, col) 23 | } 24 | 25 | /// Return the number of cells in the current row. 26 | pub fn len(&self) -> usize { 27 | self.block_ref.column_count() 28 | } 29 | 30 | /// Returns `true` if the row contains no cells. 31 | pub fn is_empty(&self) -> bool { 32 | self.len() == 0 33 | } 34 | 35 | /// Get the name of a particular cell of the row. 36 | pub fn name(&self, col: I) -> Result<&str> { 37 | Ok(self.block_ref.get_column(col)?.name()) 38 | } 39 | 40 | /// Get the type of a particular cell of the row. 41 | pub fn sql_type(&self, col: I) -> Result { 42 | Ok(self.block_ref.get_column(col)?.sql_type()) 43 | } 44 | } 45 | 46 | pub(crate) enum BlockRef<'a, K: ColumnType> { 47 | Borrowed(&'a Block), 48 | Owned(Arc>), 49 | } 50 | 51 | impl<'a, K: ColumnType> Clone for BlockRef<'a, K> { 52 | fn clone(&self) -> Self { 53 | match self { 54 | BlockRef::Borrowed(block_ref) => BlockRef::Borrowed(*block_ref), 55 | BlockRef::Owned(block_ref) => BlockRef::Owned(block_ref.clone()), 56 | } 57 | } 58 | } 59 | 60 | impl<'a, K: ColumnType> BlockRef<'a, K> { 61 | fn row_count(&self) -> usize { 62 | match self { 63 | BlockRef::Borrowed(block) => block.row_count(), 64 | BlockRef::Owned(block) => block.row_count(), 65 | } 66 | } 67 | 68 | fn column_count(&self) -> usize { 69 | match self { 70 | BlockRef::Borrowed(block) => block.column_count(), 71 | BlockRef::Owned(block) => block.column_count(), 72 | } 73 | } 74 | 75 | fn get<'s, T, I>(&'s self, row: usize, col: I) -> Result 76 | where 77 | T: FromSql<'s>, 78 | I: ColumnIdx + Copy, 79 | { 80 | match self { 81 | BlockRef::Borrowed(block) => block.get(row, col), 82 | BlockRef::Owned(block) => block.get(row, col), 83 | } 84 | } 85 | 86 | fn get_column(&self, col: I) -> Result<&Column> { 87 | match self { 88 | BlockRef::Borrowed(block) => { 89 | let column_index = col.get_index(block.columns())?; 90 | Ok(&block.columns[column_index]) 91 | } 92 | BlockRef::Owned(block) => { 93 | let column_index = col.get_index(block.columns())?; 94 | Ok(&block.columns[column_index]) 95 | } 96 | } 97 | } 98 | } 99 | 100 | /// Immutable rows iterator 101 | pub struct Rows<'a, K: ColumnType> { 102 | pub(crate) row: usize, 103 | pub(crate) block_ref: BlockRef<'a, K>, 104 | pub(crate) kind: marker::PhantomData, 105 | } 106 | 107 | impl<'a, K: ColumnType> Iterator for Rows<'a, K> { 108 | type Item = Row<'a, K>; 109 | 110 | fn next(&mut self) -> Option { 111 | if self.row >= self.block_ref.row_count() { 112 | return None; 113 | } 114 | let result = Some(Row { 115 | row: self.row, 116 | block_ref: self.block_ref.clone(), 117 | kind: marker::PhantomData, 118 | }); 119 | self.row += 1; 120 | result 121 | } 122 | } 123 | -------------------------------------------------------------------------------- /src/types/cmd.rs: -------------------------------------------------------------------------------- 1 | use log::trace; 2 | 3 | use crate::{ 4 | binary::{protocol, Encoder}, 5 | client_info, 6 | errors::Result, 7 | types::{Context, Options, Query, SettingType, Simple}, 8 | Block, 9 | }; 10 | 11 | /// Represents Clickhouse commands. 12 | pub(crate) enum Cmd { 13 | Hello(Context), 14 | Ping, 15 | SendQuery(Query, Context), 16 | SendData(Block, Context), 17 | Union(Box, Box), 18 | Cancel, 19 | } 20 | 21 | impl Cmd { 22 | /// Returns the packed command as a byte vector. 23 | #[inline(always)] 24 | pub(crate) fn get_packed_command(&self) -> Result> { 25 | encode_command(self) 26 | } 27 | } 28 | 29 | #[derive(Debug, PartialOrd, PartialEq)] 30 | enum SettingsBinaryFormat { 31 | Old, 32 | Strings, 33 | } 34 | 35 | fn encode_command(cmd: &Cmd) -> Result> { 36 | match cmd { 37 | Cmd::Hello(context) => encode_hello(context), 38 | Cmd::Ping => Ok(encode_ping()), 39 | Cmd::SendQuery(query, context) => encode_query(query, context), 40 | Cmd::SendData(block, context) => encode_data(block, context), 41 | Cmd::Union(first, second) => encode_union(first.as_ref(), second.as_ref()), 42 | Cmd::Cancel => Ok(encode_cancel()), 43 | } 44 | } 45 | 46 | fn encode_hello(context: &Context) -> Result> { 47 | trace!("[hello] -> {}", client_info::description()); 48 | 49 | let mut encoder = Encoder::new(); 50 | encoder.uvarint(protocol::CLIENT_HELLO); 51 | client_info::write(&mut encoder); 52 | 53 | let options = context.options.get()?; 54 | 55 | encoder.string(&options.database); 56 | encoder.string(&options.username); 57 | encoder.string(&options.password); 58 | 59 | Ok(encoder.get_buffer()) 60 | } 61 | 62 | fn encode_ping() -> Vec { 63 | trace!("[ping] -> ping"); 64 | 65 | let mut encoder = Encoder::new(); 66 | encoder.uvarint(protocol::CLIENT_PING); 67 | encoder.get_buffer() 68 | } 69 | 70 | fn encode_cancel() -> Vec { 71 | trace!("[cancel]"); 72 | 73 | let mut encoder = Encoder::new(); 74 | encoder.uvarint(protocol::CLIENT_CANCEL); 75 | encoder.get_buffer() 76 | } 77 | 78 | fn encode_query(query: &Query, context: &Context) -> Result> { 79 | trace!("[send query] {}", query.get_sql()); 80 | 81 | // DBMS_MIN_REVISION_WITH_CLIENT_INFO 82 | let mut encoder = Encoder::new(); 83 | encoder.uvarint(protocol::CLIENT_QUERY); 84 | encoder.string(""); 85 | 86 | { 87 | let hostname = &context.hostname; 88 | encoder.uvarint(1); 89 | encoder.string(""); 90 | encoder.string(query.get_id()); // initial_query_id; 91 | encoder.string("[::ffff:127.0.0.1]:0"); 92 | encoder.uvarint(1); // iface type TCP; 93 | encoder.string(hostname); 94 | encoder.string(hostname); 95 | } 96 | client_info::write(&mut encoder); 97 | 98 | if context.server_info.revision >= protocol::DBMS_MIN_REVISION_WITH_QUOTA_KEY_IN_CLIENT_INFO { 99 | encoder.string(""); 100 | } 101 | 102 | if context.server_info.revision >= protocol::DBMS_MIN_REVISION_WITH_VERSION_PATCH { 103 | encoder.uvarint(0); 104 | } 105 | 106 | let options = context.options.get()?; 107 | 108 | let settings_format = if context.server_info.revision 109 | >= protocol::DBMS_MIN_REVISION_WITH_SETTINGS_SERIALIZED_AS_STRINGS 110 | { 111 | SettingsBinaryFormat::Strings 112 | } else { 113 | SettingsBinaryFormat::Old 114 | }; 115 | 116 | serialize_settings(&mut encoder, &options, settings_format); 117 | 118 | encoder.uvarint(protocol::STATE_COMPLETE); 119 | 120 | encoder.uvarint(if options.compression { 121 | protocol::COMPRESS_ENABLE 122 | } else { 123 | protocol::COMPRESS_DISABLE 124 | }); 125 | 126 | let options = context.options.get()?; 127 | 128 | encoder.string(query.get_sql()); 129 | Block::::default().send_data(&mut encoder, options.compression); 130 | 131 | Ok(encoder.get_buffer()) 132 | } 133 | 134 | fn serialize_settings(encoder: &mut Encoder, options: &Options, format: SettingsBinaryFormat) { 135 | if format < SettingsBinaryFormat::Strings { 136 | for (name, value) in &options.settings { 137 | encoder.string(name); 138 | match &value.value { 139 | SettingType::String(val) => encoder.string(val), 140 | // Bool and UInt64 is the same 141 | SettingType::Bool(val) => encoder.uvarint((val == &true) as u64), 142 | // Float is written in string representation 143 | SettingType::Float64(val) => encoder.string(val.to_string()), 144 | SettingType::UInt64(val) => encoder.uvarint(*val), 145 | } 146 | } 147 | } else { 148 | for (name, value) in &options.settings { 149 | encoder.string(name); 150 | encoder.write(value.is_important); 151 | encoder.string(value.to_string()); 152 | } 153 | } 154 | 155 | encoder.string(""); // end of settings marker 156 | } 157 | 158 | fn encode_data(block: &Block, context: &Context) -> Result> { 159 | let mut encoder = Encoder::new(); 160 | let options = context.options.get()?; 161 | block.send_data(&mut encoder, options.compression); 162 | Ok(encoder.get_buffer()) 163 | } 164 | 165 | fn encode_union(first: &Cmd, second: &Cmd) -> Result> { 166 | let mut result = encode_command(first)?; 167 | result.extend((encode_command(second)?).iter()); 168 | Ok(result) 169 | } 170 | -------------------------------------------------------------------------------- /src/types/column/array.rs: -------------------------------------------------------------------------------- 1 | use std::sync::Arc; 2 | 3 | use chrono_tz::Tz; 4 | 5 | use crate::{ 6 | binary::{Encoder, ReadEx}, 7 | errors::Result, 8 | types::{ 9 | column::{ 10 | column_data::{ArcColumnData, BoxColumnData}, 11 | list::List, 12 | ArcColumnWrapper, ColumnData, 13 | }, 14 | SqlType, Value, ValueRef, 15 | }, 16 | }; 17 | 18 | pub(crate) struct ArrayColumnData { 19 | pub(crate) inner: ArcColumnData, 20 | pub(crate) offsets: List, 21 | } 22 | 23 | impl ArrayColumnData { 24 | pub(crate) fn load( 25 | reader: &mut R, 26 | type_name: &str, 27 | rows: usize, 28 | tz: Tz, 29 | ) -> Result { 30 | let mut offsets = List::with_capacity(rows); 31 | offsets.resize(rows, 0_u64); 32 | reader.read_bytes(offsets.as_mut())?; 33 | 34 | let size = match rows { 35 | 0 => 0, 36 | _ => offsets.at(rows - 1) as usize, 37 | }; 38 | let inner = 39 | ::load_data::(reader, type_name, size, tz)?; 40 | 41 | Ok(ArrayColumnData { inner, offsets }) 42 | } 43 | } 44 | 45 | impl ColumnData for ArrayColumnData { 46 | fn sql_type(&self) -> SqlType { 47 | let inner_type = self.inner.sql_type(); 48 | SqlType::Array(inner_type.into()) 49 | } 50 | 51 | fn save(&self, encoder: &mut Encoder, start: usize, end: usize) { 52 | let mut offset = 0_u64; 53 | 54 | for i in start..end { 55 | offset = self.offsets.at(i); 56 | encoder.write(offset); 57 | } 58 | 59 | self.inner.save(encoder, 0, offset as usize); 60 | } 61 | 62 | fn len(&self) -> usize { 63 | self.offsets.len() 64 | } 65 | 66 | fn push(&mut self, value: Value) { 67 | if let Value::Array(_, vs) = value { 68 | let offsets_len = self.offsets.len(); 69 | let prev = if offsets_len == 0 { 70 | 0_usize 71 | } else { 72 | self.offsets.at(offsets_len - 1) as usize 73 | }; 74 | 75 | let inner_column = Arc::get_mut(&mut self.inner).unwrap(); 76 | self.offsets.push((prev + vs.len()) as u64); 77 | for v in vs.iter() { 78 | inner_column.push(v.clone()); 79 | } 80 | } else { 81 | panic!("value should be an array") 82 | } 83 | } 84 | 85 | fn at(&self, index: usize) -> ValueRef { 86 | let sql_type = self.inner.sql_type(); 87 | 88 | let start = if index > 0 { 89 | self.offsets.at(index - 1) as usize 90 | } else { 91 | 0_usize 92 | }; 93 | let end = self.offsets.at(index) as usize; 94 | let mut vs = Vec::with_capacity(end - start); 95 | for i in start..end { 96 | let v = self.inner.at(i); 97 | vs.push(v); 98 | } 99 | ValueRef::Array(sql_type.into(), Arc::new(vs)) 100 | } 101 | 102 | fn clone_instance(&self) -> BoxColumnData { 103 | Box::new(Self { 104 | inner: self.inner.clone(), 105 | offsets: self.offsets.clone(), 106 | }) 107 | } 108 | 109 | unsafe fn get_internal( 110 | &self, 111 | pointers: &[*mut *const u8], 112 | level: u8, 113 | props: u32, 114 | ) -> Result<()> { 115 | if level == self.sql_type().level() { 116 | *pointers[0] = self.offsets.as_ptr() as *const u8; 117 | *(pointers[1] as *mut usize) = self.offsets.len(); 118 | Ok(()) 119 | } else { 120 | self.inner.get_internal(pointers, level, props) 121 | } 122 | } 123 | 124 | fn cast_to(&self, _this: &ArcColumnData, target: &SqlType) -> Option { 125 | if let SqlType::Array(inner_target) = target { 126 | if let Some(inner) = self.inner.cast_to(&self.inner, inner_target) { 127 | return Some(Arc::new(ArrayColumnData { 128 | inner, 129 | offsets: self.offsets.clone(), 130 | })); 131 | } 132 | } 133 | None 134 | } 135 | 136 | fn get_timezone(&self) -> Option { 137 | self.inner.get_timezone() 138 | } 139 | } 140 | 141 | #[cfg(test)] 142 | mod test { 143 | use std::io::Cursor; 144 | 145 | use super::*; 146 | use crate::types::{column::datetime64::DEFAULT_TZ, Block, Simple}; 147 | 148 | #[test] 149 | fn test_write_and_read() { 150 | let block = Block::::new().column( 151 | "vals", 152 | vec![vec![7_u32, 8], vec![9, 1, 2], vec![3, 4, 5, 6]], 153 | ); 154 | 155 | let mut encoder = Encoder::new(); 156 | block.write(&mut encoder, false); 157 | 158 | let mut reader = Cursor::new(encoder.get_buffer_ref()); 159 | let rblock = Block::load(&mut reader, *DEFAULT_TZ, false).unwrap(); 160 | 161 | assert_eq!(block, rblock); 162 | } 163 | } 164 | -------------------------------------------------------------------------------- /src/types/column/chrono_datetime.rs: -------------------------------------------------------------------------------- 1 | use chrono::prelude::*; 2 | use chrono_tz::Tz; 3 | use std::{ptr, slice, sync::Arc}; 4 | 5 | use crate::{ 6 | binary::Encoder, 7 | errors::Result, 8 | types::{ 9 | column::{ 10 | column_data::{ArcColumnData, BoxColumnData, ColumnData}, 11 | datetime64::{from_datetime, DEFAULT_TZ}, 12 | nullable::NullableColumnData, 13 | ArcColumnWrapper, ColumnFrom, ColumnWrapper, Value, ValueRef, 14 | }, 15 | DateTimeType, SqlType, 16 | }, 17 | }; 18 | 19 | pub struct ChronoDateTimeColumnData { 20 | data: Vec>, 21 | tz: Tz, 22 | } 23 | 24 | impl ChronoDateTimeColumnData { 25 | pub(crate) fn with_capacity(capacity: usize, tz: Tz) -> ChronoDateTimeColumnData { 26 | let data = Vec::with_capacity(capacity); 27 | ChronoDateTimeColumnData { data, tz } 28 | } 29 | } 30 | 31 | pub(crate) struct ChronoDateTimeAdapter { 32 | pub(crate) column: ArcColumnData, 33 | dst_type: SqlType, 34 | } 35 | 36 | impl ChronoDateTimeAdapter { 37 | pub(crate) fn new(column: ArcColumnData, dst_type: SqlType) -> ChronoDateTimeAdapter { 38 | ChronoDateTimeAdapter { column, dst_type } 39 | } 40 | } 41 | 42 | impl ColumnFrom for Vec> { 43 | fn column_from(data: Self) -> W::Wrapper { 44 | let tz = if data.is_empty() { 45 | *DEFAULT_TZ 46 | } else { 47 | data[0].timezone() 48 | }; 49 | W::wrap(ChronoDateTimeColumnData { data, tz }) 50 | } 51 | } 52 | 53 | impl ColumnFrom for Vec>> { 54 | fn column_from(source: Self) -> ::Wrapper { 55 | let n = source.len(); 56 | 57 | let tz = source 58 | .iter() 59 | .find_map(|u| u.map(|v| v.timezone())) 60 | .unwrap_or(*DEFAULT_TZ); 61 | 62 | let mut values: Vec> = Vec::with_capacity(n); 63 | let mut nulls = Vec::with_capacity(n); 64 | 65 | for time in source { 66 | match time { 67 | None => { 68 | nulls.push(1); 69 | values.push(tz.timestamp_opt(0, 0).unwrap()) 70 | } 71 | Some(time) => { 72 | nulls.push(0); 73 | values.push(time) 74 | } 75 | } 76 | } 77 | 78 | W::wrap(NullableColumnData { 79 | inner: Vec::column_from::(values), 80 | nulls, 81 | }) 82 | } 83 | } 84 | 85 | impl ColumnData for ChronoDateTimeColumnData { 86 | fn sql_type(&self) -> SqlType { 87 | SqlType::DateTime(DateTimeType::Chrono) 88 | } 89 | 90 | fn save(&self, _encoder: &mut Encoder, _start: usize, _end: usize) { 91 | unimplemented!() 92 | } 93 | 94 | fn len(&self) -> usize { 95 | self.data.len() 96 | } 97 | 98 | fn push(&mut self, value: Value) { 99 | let time: DateTime = value.into(); 100 | self.data.push(time); 101 | } 102 | 103 | fn at(&self, index: usize) -> ValueRef { 104 | let v = &self.data[index]; 105 | ValueRef::DateTime(v.timestamp() as u32, v.timezone()) 106 | } 107 | 108 | fn clone_instance(&self) -> BoxColumnData { 109 | Box::new(Self { 110 | data: self.data.clone(), 111 | tz: self.tz, 112 | }) 113 | } 114 | 115 | unsafe fn get_internal( 116 | &self, 117 | pointers: &[*mut *const u8], 118 | level: u8, 119 | _props: u32, 120 | ) -> Result<()> { 121 | assert_eq!(level, 0); 122 | *pointers[0] = self.data.as_ptr() as *const u8; 123 | *pointers[1] = &self.tz as *const Tz as *const u8; 124 | *(pointers[2] as *mut usize) = self.len(); 125 | Ok(()) 126 | } 127 | 128 | fn cast_to(&self, this: &ArcColumnData, target: &SqlType) -> Option { 129 | if target.is_datetime() { 130 | let clone = this.clone(); 131 | let adapter = ChronoDateTimeAdapter::new(clone, target.clone()); 132 | Some(Arc::new(adapter)) 133 | } else { 134 | None 135 | } 136 | } 137 | 138 | fn get_timezone(&self) -> Option { 139 | Some(self.tz) 140 | } 141 | } 142 | 143 | #[inline(always)] 144 | fn is_chrono_datetime(column: &dyn ColumnData) -> bool { 145 | column.sql_type() == SqlType::DateTime(DateTimeType::Chrono) 146 | } 147 | 148 | pub(crate) fn get_date_slice<'a>(column: &dyn ColumnData) -> Result<&'a [DateTime]> { 149 | unsafe { 150 | let mut data: *const DateTime = ptr::null(); 151 | let mut tz: *const Tz = ptr::null(); 152 | let mut len: usize = 0; 153 | column.get_internal( 154 | &[ 155 | &mut data as *mut *const DateTime as *mut *const u8, 156 | &mut tz as *mut *const Tz as *mut *const u8, 157 | &mut len as *mut usize as *mut *const u8, 158 | ], 159 | 0, 160 | 0, 161 | )?; 162 | assert_ne!(data, ptr::null()); 163 | assert_ne!(tz, ptr::null()); 164 | Ok(slice::from_raw_parts(data, len)) 165 | } 166 | } 167 | 168 | impl ColumnData for ChronoDateTimeAdapter { 169 | fn sql_type(&self) -> SqlType { 170 | self.dst_type.clone() 171 | } 172 | 173 | fn save(&self, encoder: &mut Encoder, start: usize, end: usize) { 174 | if !is_chrono_datetime(self.column.as_ref()) { 175 | panic!("Invalid column type {}.", self.column.sql_type()); 176 | } 177 | 178 | let dates = get_date_slice(self.column.as_ref()).unwrap(); 179 | 180 | match self.dst_type { 181 | SqlType::DateTime(DateTimeType::DateTime64(precision, tz)) => { 182 | for date in &dates[start..end] { 183 | let value = from_datetime(date.with_timezone(&tz), precision); 184 | encoder.write(value); 185 | } 186 | } 187 | SqlType::DateTime(DateTimeType::DateTime32) => { 188 | for date in &dates[start..end] { 189 | let value = date.timestamp() as u32; 190 | encoder.write(value); 191 | } 192 | } 193 | _ => unimplemented!(), 194 | } 195 | } 196 | 197 | fn len(&self) -> usize { 198 | self.column.len() 199 | } 200 | 201 | fn push(&mut self, _value: Value) { 202 | unimplemented!() 203 | } 204 | 205 | fn at(&self, _index: usize) -> ValueRef { 206 | unimplemented!() 207 | } 208 | 209 | fn clone_instance(&self) -> BoxColumnData { 210 | unimplemented!() 211 | } 212 | 213 | unsafe fn get_internal( 214 | &self, 215 | _pointers: &[*mut *const u8], 216 | _level: u8, 217 | _props: u32, 218 | ) -> Result<()> { 219 | unimplemented!() 220 | } 221 | 222 | fn get_timezone(&self) -> Option { 223 | self.column.get_timezone() 224 | } 225 | } 226 | -------------------------------------------------------------------------------- /src/types/column/chunk.rs: -------------------------------------------------------------------------------- 1 | use chrono_tz::Tz; 2 | use std::{cmp, ops}; 3 | 4 | use crate::{ 5 | binary::Encoder, 6 | types::{ 7 | column::column_data::{ArcColumnData, BoxColumnData}, 8 | SqlType, Value, ValueRef, 9 | }, 10 | }; 11 | 12 | use super::ColumnData; 13 | 14 | pub struct ChunkColumnData { 15 | data: ArcColumnData, 16 | range: ops::Range, 17 | } 18 | 19 | impl ChunkColumnData { 20 | pub(crate) fn new(data: ArcColumnData, range: ops::Range) -> Self { 21 | Self { data, range } 22 | } 23 | } 24 | 25 | impl ColumnData for ChunkColumnData { 26 | fn sql_type(&self) -> SqlType { 27 | self.data.sql_type() 28 | } 29 | 30 | fn save(&self, encoder: &mut Encoder, start: usize, end: usize) { 31 | self.data.save( 32 | encoder, 33 | self.range.start + start, 34 | cmp::min(self.range.end, self.range.start + end), 35 | ) 36 | } 37 | 38 | fn len(&self) -> usize { 39 | self.range.len() 40 | } 41 | 42 | fn push(&mut self, _value: Value) { 43 | unimplemented!() 44 | } 45 | 46 | fn at(&self, index: usize) -> ValueRef { 47 | if index >= self.range.len() { 48 | panic!("out of range"); 49 | } 50 | 51 | self.data.at(index + self.range.start) 52 | } 53 | 54 | fn clone_instance(&self) -> BoxColumnData { 55 | unimplemented!() 56 | } 57 | 58 | fn get_timezone(&self) -> Option { 59 | self.data.get_timezone() 60 | } 61 | } 62 | -------------------------------------------------------------------------------- /src/types/column/column_data.rs: -------------------------------------------------------------------------------- 1 | use chrono_tz::Tz; 2 | use std::sync::Arc; 3 | 4 | use crate::{ 5 | binary::Encoder, 6 | errors::{Error, FromSqlError, Result}, 7 | types::{SqlType, Value, ValueRef}, 8 | }; 9 | 10 | pub(crate) type ArcColumnData = Arc; 11 | 12 | pub(crate) type BoxColumnData = Box; 13 | 14 | pub trait LowCardinalityAccessor { 15 | fn get_string(&self, _index: usize) -> &[u8] { 16 | unimplemented!() 17 | } 18 | } 19 | 20 | pub trait ColumnData { 21 | fn sql_type(&self) -> SqlType; 22 | fn save(&self, encoder: &mut Encoder, start: usize, end: usize); 23 | fn len(&self) -> usize; 24 | fn push(&mut self, value: Value); 25 | fn at(&self, index: usize) -> ValueRef; 26 | 27 | fn clone_instance(&self) -> BoxColumnData; 28 | 29 | unsafe fn get_internal( 30 | &self, 31 | _pointers: &[*mut *const u8], 32 | _level: u8, 33 | _props: u32, 34 | ) -> Result<()> { 35 | Err(Error::FromSql(FromSqlError::UnsupportedOperation)) 36 | } 37 | 38 | unsafe fn get_internals(&self, _data: *mut (), _level: u8, _props: u32) -> Result<()> { 39 | Err(Error::FromSql(FromSqlError::UnsupportedOperation)) 40 | } 41 | 42 | fn cast_to(&self, _this: &ArcColumnData, _target: &SqlType) -> Option { 43 | None 44 | } 45 | 46 | fn get_timezone(&self) -> Option; 47 | 48 | fn get_low_cardinality_accessor(&self) -> Option<&dyn LowCardinalityAccessor> { 49 | None 50 | } 51 | } 52 | 53 | pub(crate) trait ColumnDataExt { 54 | fn append>(&mut self, value: T); 55 | } 56 | 57 | impl ColumnDataExt for C { 58 | fn append>(&mut self, value: T) { 59 | self.push(value.into()); 60 | } 61 | } 62 | -------------------------------------------------------------------------------- /src/types/column/concat.rs: -------------------------------------------------------------------------------- 1 | use chrono_tz::Tz; 2 | 3 | use crate::{ 4 | binary::Encoder, 5 | errors::{Error, FromSqlError, Result}, 6 | types::{SqlType, Value, ValueRef}, 7 | }; 8 | 9 | use super::column_data::{ArcColumnData, BoxColumnData, ColumnData}; 10 | 11 | pub struct ConcatColumnData { 12 | data: Vec, 13 | index: Vec, 14 | } 15 | 16 | impl ConcatColumnData { 17 | pub(crate) fn concat(data: Vec) -> Self { 18 | Self::check_columns(&data); 19 | 20 | let index = build_index(data.iter().map(|x| x.len())); 21 | Self { data, index } 22 | } 23 | 24 | fn check_columns(data: &[ArcColumnData]) { 25 | match data.first() { 26 | None => panic!("data should not be empty."), 27 | Some(first) => { 28 | for column in data.iter().skip(1) { 29 | if first.sql_type() != column.sql_type() { 30 | panic!( 31 | "all columns should have the same type ({:?} != {:?}).", 32 | first.sql_type(), 33 | column.sql_type() 34 | ); 35 | } 36 | } 37 | } 38 | } 39 | } 40 | } 41 | 42 | impl ColumnData for ConcatColumnData { 43 | fn sql_type(&self) -> SqlType { 44 | self.data[0].sql_type() 45 | } 46 | 47 | fn save(&self, _: &mut Encoder, _: usize, _: usize) { 48 | unimplemented!() 49 | } 50 | 51 | fn len(&self) -> usize { 52 | *self.index.last().unwrap() 53 | } 54 | 55 | fn push(&mut self, _value: Value) { 56 | unimplemented!() 57 | } 58 | 59 | fn at(&self, index: usize) -> ValueRef { 60 | let chunk_index = find_chunk(&self.index, index); 61 | let chunk = &self.data[chunk_index]; 62 | chunk.at(index - self.index[chunk_index]) 63 | } 64 | 65 | fn clone_instance(&self) -> BoxColumnData { 66 | unimplemented!() 67 | } 68 | 69 | unsafe fn get_internal( 70 | &self, 71 | pointers: &[*mut *const u8], 72 | level: u8, 73 | _props: u32, 74 | ) -> Result<()> { 75 | if level == 0xff { 76 | *pointers[0] = &self.data as *const Vec as *mut u8; 77 | Ok(()) 78 | } else { 79 | Err(Error::FromSql(FromSqlError::UnsupportedOperation)) 80 | } 81 | } 82 | 83 | fn get_timezone(&self) -> Option { 84 | self.data.first().and_then(|chunk| chunk.get_timezone()) 85 | } 86 | } 87 | 88 | fn build_index<'a, I>(sizes: I) -> Vec 89 | where 90 | I: Iterator + 'a, 91 | { 92 | let mut acc = 0; 93 | let mut index = vec![acc]; 94 | 95 | for size in sizes { 96 | acc += size; 97 | index.push(acc); 98 | } 99 | 100 | index 101 | } 102 | 103 | fn find_chunk(index: &[usize], ix: usize) -> usize { 104 | let mut lo = 0_usize; 105 | let mut hi = index.len() - 1; 106 | 107 | while lo < hi { 108 | let mid = lo + (hi - lo) / 2; 109 | 110 | if index[lo] == index[lo + 1] { 111 | lo += 1; 112 | continue; 113 | } 114 | 115 | if ix < index[mid] { 116 | hi = mid; 117 | } else if ix >= index[mid + 1] { 118 | lo = mid + 1; 119 | } else { 120 | return mid; 121 | } 122 | } 123 | 124 | 0 125 | } 126 | 127 | #[cfg(test)] 128 | mod test { 129 | use std::sync::Arc; 130 | 131 | use crate::types::column::{ 132 | column_data::ColumnDataExt, numeric::VectorColumnData, string::StringColumnData, 133 | }; 134 | 135 | use super::*; 136 | 137 | #[test] 138 | fn test_build_index() { 139 | let sizes = [2_usize, 3, 4]; 140 | let index = build_index(sizes.iter().cloned()); 141 | assert_eq!(index, vec![0, 2, 5, 9]) 142 | } 143 | 144 | #[test] 145 | fn test_find_chunk() { 146 | let index = vec![0_usize, 2, 5, 9]; 147 | assert_eq!(find_chunk(&index, 0), 0); 148 | assert_eq!(find_chunk(&index, 1), 0); 149 | assert_eq!(find_chunk(&index, 2), 1); 150 | assert_eq!(find_chunk(&index, 3), 1); 151 | assert_eq!(find_chunk(&index, 4), 1); 152 | assert_eq!(find_chunk(&index, 5), 2); 153 | assert_eq!(find_chunk(&index, 6), 2); 154 | 155 | assert_eq!(find_chunk(&index, 7), 2); 156 | assert_eq!(find_chunk(&[0], 7), 0); 157 | } 158 | 159 | #[test] 160 | fn test_find_chunk2() { 161 | let index = vec![0_usize, 0, 5]; 162 | assert_eq!(find_chunk(&index, 0), 1); 163 | assert_eq!(find_chunk(&index, 1), 1); 164 | assert_eq!(find_chunk(&index, 2), 1); 165 | assert_eq!(find_chunk(&index, 3), 1); 166 | assert_eq!(find_chunk(&index, 4), 1); 167 | assert_eq!(find_chunk(&index, 5), 0); 168 | } 169 | 170 | #[test] 171 | fn test_find_chunk5() { 172 | let index = vec![ 173 | 0_usize, 0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48, 51, 54, 57, 174 | 60, 63, 66, 69, 175 | ]; 176 | for i in 0..69 { 177 | assert_eq!(find_chunk(&index, i), 1 + i / 3); 178 | } 179 | } 180 | 181 | #[test] 182 | fn test_concat_column() { 183 | let xs = vec![make_string_column(), make_string_column()]; 184 | let actual = ConcatColumnData::concat(xs); 185 | 186 | assert_eq!( 187 | actual.at(0).as_str().unwrap(), 188 | "13298a5f-6a10-4fbe-9644-807f7ebf82cc" 189 | ); 190 | assert_eq!( 191 | actual.at(1).as_str().unwrap(), 192 | "df0e62bb-c0db-4728-a558-821f8e8da38c" 193 | ); 194 | assert_eq!( 195 | actual.at(2).as_str().unwrap(), 196 | "13298a5f-6a10-4fbe-9644-807f7ebf82cc" 197 | ); 198 | assert_eq!( 199 | actual.at(3).as_str().unwrap(), 200 | "df0e62bb-c0db-4728-a558-821f8e8da38c" 201 | ); 202 | 203 | assert_eq!(actual.len(), 4); 204 | } 205 | 206 | #[test] 207 | fn test_concat_num_column() { 208 | let xs = vec![make_num_column(), make_num_column()]; 209 | let actual = ConcatColumnData::concat(xs); 210 | 211 | assert_eq!(u32::from(actual.at(0)), 1_u32); 212 | assert_eq!(u32::from(actual.at(1)), 2_u32); 213 | assert_eq!(u32::from(actual.at(2)), 1_u32); 214 | assert_eq!(u32::from(actual.at(3)), 2_u32); 215 | 216 | assert_eq!(actual.len(), 4); 217 | } 218 | 219 | fn make_string_column() -> ArcColumnData { 220 | let mut data = StringColumnData::with_capacity(1); 221 | data.append("13298a5f-6a10-4fbe-9644-807f7ebf82cc".to_string()); 222 | data.append("df0e62bb-c0db-4728-a558-821f8e8da38c".to_string()); 223 | Arc::new(data) 224 | } 225 | 226 | fn make_num_column() -> ArcColumnData { 227 | let mut data = VectorColumnData::::with_capacity(1); 228 | data.append(1_u32); 229 | data.append(2_u32); 230 | Arc::new(data) 231 | } 232 | } 233 | -------------------------------------------------------------------------------- /src/types/column/date.rs: -------------------------------------------------------------------------------- 1 | use std::{fmt, ptr, sync::Arc}; 2 | 3 | use chrono::prelude::*; 4 | use chrono_tz::Tz; 5 | use either::Either; 6 | 7 | use crate::{ 8 | binary::{Encoder, ReadEx}, 9 | errors::Result, 10 | types::{ 11 | column::{ 12 | array::ArrayColumnData, 13 | column_data::{BoxColumnData, ColumnData, LowCardinalityAccessor}, 14 | datetime64::DEFAULT_TZ, 15 | list::List, 16 | nullable::NullableColumnData, 17 | numeric::save_data, 18 | ArcColumnWrapper, ColumnFrom, ColumnWrapper, 19 | }, 20 | DateConverter, Marshal, SqlType, StatBuffer, Unmarshal, Value, ValueRef, 21 | }, 22 | }; 23 | 24 | pub struct DateColumnData 25 | where 26 | T: StatBuffer 27 | + Unmarshal 28 | + Marshal 29 | + Copy 30 | + Into 31 | + From 32 | + fmt::Display 33 | + Sync 34 | + Default 35 | + 'static, 36 | { 37 | data: List, 38 | tz: Tz, 39 | } 40 | 41 | impl DateColumnData 42 | where 43 | T: StatBuffer 44 | + Unmarshal 45 | + Marshal 46 | + Copy 47 | + Into 48 | + From 49 | + fmt::Display 50 | + Sync 51 | + Default 52 | + 'static, 53 | { 54 | pub(crate) fn with_capacity(capacity: usize, timezone: Tz) -> DateColumnData { 55 | DateColumnData { 56 | data: List::with_capacity(capacity), 57 | tz: timezone, 58 | } 59 | } 60 | 61 | pub(crate) fn load( 62 | reader: &mut R, 63 | size: usize, 64 | tz: Tz, 65 | ) -> Result> { 66 | let mut data = List::with_capacity(size); 67 | unsafe { 68 | data.set_len(size); 69 | } 70 | reader.read_bytes(data.as_mut())?; 71 | Ok(DateColumnData { data, tz }) 72 | } 73 | } 74 | 75 | impl ColumnFrom for Vec { 76 | fn column_from(source: Self) -> W::Wrapper { 77 | let mut data = List::::with_capacity(source.len()); 78 | for s in source { 79 | data.push(u16::get_days(s)); 80 | } 81 | 82 | let column: DateColumnData = DateColumnData { 83 | data, 84 | tz: *DEFAULT_TZ, 85 | }; 86 | W::wrap(column) 87 | } 88 | } 89 | 90 | impl ColumnFrom for Vec> { 91 | fn column_from(source: Self) -> W::Wrapper { 92 | let fake: Vec = Vec::with_capacity(source.len()); 93 | let inner = Vec::column_from::(fake); 94 | let sql_type = inner.sql_type(); 95 | 96 | let mut data = ArrayColumnData { 97 | inner, 98 | offsets: List::with_capacity(source.len()), 99 | }; 100 | 101 | for vs in source { 102 | let mut inner = Vec::with_capacity(vs.len()); 103 | for v in vs { 104 | let days = u16::get_days(v); 105 | let value: Value = Value::Date(days); 106 | inner.push(value); 107 | } 108 | data.push(Value::Array(sql_type.clone().into(), Arc::new(inner))); 109 | } 110 | 111 | W::wrap(data) 112 | } 113 | } 114 | 115 | impl ColumnFrom for Vec>> { 116 | fn column_from(source: Self) -> W::Wrapper { 117 | let fake: Vec> = Vec::with_capacity(source.len()); 118 | let inner = Vec::column_from::(fake); 119 | let sql_type = inner.sql_type(); 120 | 121 | let mut data = ArrayColumnData { 122 | inner, 123 | offsets: List::with_capacity(source.len()), 124 | }; 125 | 126 | for vs in source { 127 | let mut inner = Vec::with_capacity(vs.len()); 128 | for v in vs { 129 | let value: Value = Value::DateTime(v.timestamp() as u32, v.timezone()); 130 | inner.push(value); 131 | } 132 | data.push(Value::Array(sql_type.clone().into(), Arc::new(inner))); 133 | } 134 | 135 | W::wrap(data) 136 | } 137 | } 138 | 139 | impl ColumnFrom for Vec> { 140 | fn column_from(source: Self) -> ::Wrapper { 141 | let fake: Vec = Vec::with_capacity(source.len()); 142 | let inner = Vec::column_from::(fake); 143 | 144 | let mut data = NullableColumnData { 145 | inner, 146 | nulls: Vec::with_capacity(source.len()), 147 | }; 148 | 149 | for value in source { 150 | match value { 151 | None => data.push(Value::Nullable(Either::Left(SqlType::Date.into()))), 152 | Some(d) => { 153 | let days = u16::get_days(d); 154 | let value = Value::Date(days); 155 | data.push(Value::Nullable(Either::Right(Box::new(value)))) 156 | } 157 | } 158 | } 159 | 160 | W::wrap(data) 161 | } 162 | } 163 | 164 | impl LowCardinalityAccessor for DateColumnData where 165 | T: StatBuffer 166 | + Unmarshal 167 | + Marshal 168 | + Copy 169 | + Into 170 | + From 171 | + fmt::Display 172 | + Sync 173 | + Send 174 | + DateConverter 175 | + Default 176 | + 'static 177 | { 178 | } 179 | 180 | impl ColumnData for DateColumnData 181 | where 182 | T: StatBuffer 183 | + Unmarshal 184 | + Marshal 185 | + Copy 186 | + Into 187 | + From 188 | + fmt::Display 189 | + Sync 190 | + Send 191 | + DateConverter 192 | + Default 193 | + 'static, 194 | { 195 | fn sql_type(&self) -> SqlType { 196 | T::date_type() 197 | } 198 | 199 | fn save(&self, encoder: &mut Encoder, start: usize, end: usize) { 200 | save_data::(self.data.as_ref(), encoder, start, end); 201 | } 202 | 203 | fn len(&self) -> usize { 204 | self.data.len() 205 | } 206 | 207 | fn push(&mut self, value: Value) { 208 | self.data.push(T::get_stamp(value)); 209 | } 210 | 211 | fn at(&self, index: usize) -> ValueRef { 212 | self.data.at(index).to_date(self.tz) 213 | } 214 | 215 | fn clone_instance(&self) -> BoxColumnData { 216 | Box::new(Self { 217 | data: self.data.clone(), 218 | tz: self.tz, 219 | }) 220 | } 221 | 222 | unsafe fn get_internal( 223 | &self, 224 | pointers: &[*mut *const u8], 225 | level: u8, 226 | _props: u32, 227 | ) -> Result<()> { 228 | assert_eq!(level, 0); 229 | *pointers[0] = self.data.as_ptr() as *const u8; 230 | *pointers[1] = &self.tz as *const Tz as *const u8; 231 | *(pointers[2] as *mut usize) = self.len(); 232 | Ok(()) 233 | } 234 | 235 | unsafe fn get_internals(&self, data_ptr: *mut (), _level: u8, _props: u32) -> Result<()> { 236 | unsafe { 237 | let data_ref = &mut *(data_ptr as *mut DateTimeInternals); 238 | data_ref.begin = self.data.as_ptr().cast(); 239 | data_ref.len = self.data.len(); 240 | data_ref.tz = self.tz; 241 | Ok(()) 242 | } 243 | } 244 | 245 | fn get_timezone(&self) -> Option { 246 | Some(self.tz) 247 | } 248 | 249 | fn get_low_cardinality_accessor(&self) -> Option<&dyn LowCardinalityAccessor> { 250 | Some(self) 251 | } 252 | } 253 | 254 | #[derive(Debug)] 255 | pub(crate) struct DateTimeInternals { 256 | pub(crate) begin: *const (), 257 | pub(crate) len: usize, 258 | pub(crate) tz: Tz, 259 | pub(crate) precision: Option, 260 | } 261 | 262 | impl Default for DateTimeInternals { 263 | fn default() -> Self { 264 | Self { 265 | begin: ptr::null(), 266 | len: 0, 267 | tz: Tz::Zulu, 268 | precision: None, 269 | } 270 | } 271 | } 272 | 273 | #[cfg(test)] 274 | mod test { 275 | use chrono::TimeZone; 276 | 277 | use crate::types::column::{datetime64::DEFAULT_TZ, ArcColumnWrapper}; 278 | 279 | use super::*; 280 | 281 | #[test] 282 | fn test_create_date() { 283 | let column = 284 | Vec::column_from::(vec![ 285 | NaiveDate::from_ymd_opt(2016, 10, 22).unwrap() 286 | ]); 287 | assert_eq!("2016-10-22", format!("{:#}", column.at(0))); 288 | assert_eq!(SqlType::Date, column.sql_type()); 289 | } 290 | 291 | #[test] 292 | fn test_create_date_time() { 293 | let tz = *DEFAULT_TZ; 294 | let column = Vec::column_from::(vec![tz 295 | .with_ymd_and_hms(2016, 10, 22, 12, 0, 0) 296 | .unwrap()]); 297 | assert_eq!(format!("{}", column.at(0)), "2016-10-22 12:00:00"); 298 | } 299 | } 300 | -------------------------------------------------------------------------------- /src/types/column/datetime64.rs: -------------------------------------------------------------------------------- 1 | use chrono::{prelude::*, LocalResult}; 2 | use chrono_tz::Tz; 3 | use lazy_static::lazy_static; 4 | 5 | use crate::{ 6 | binary::{Encoder, ReadEx}, 7 | errors::Result, 8 | types::{ 9 | column::{ 10 | column_data::{BoxColumnData, ColumnData}, 11 | date::DateTimeInternals, 12 | list::List, 13 | }, 14 | DateTimeType, SqlType, Value, ValueRef, 15 | }, 16 | }; 17 | 18 | lazy_static! { 19 | pub(crate) static ref DEFAULT_TZ: Tz = Tz::Zulu; 20 | } 21 | 22 | pub struct DateTime64ColumnData { 23 | data: List, 24 | params: (u32, Tz), 25 | } 26 | 27 | impl DateTime64ColumnData { 28 | pub(crate) fn load( 29 | reader: &mut R, 30 | size: usize, 31 | precision: u32, 32 | tz: Tz, 33 | ) -> Result { 34 | let mut data = List::with_capacity(size); 35 | unsafe { 36 | data.set_len(size); 37 | } 38 | reader.read_bytes(data.as_mut())?; 39 | Ok(DateTime64ColumnData { 40 | data, 41 | params: (precision, tz), 42 | }) 43 | } 44 | 45 | pub(crate) fn with_capacity(capacity: usize, precision: u32, timezone: Tz) -> Self { 46 | DateTime64ColumnData { 47 | data: List::with_capacity(capacity), 48 | params: (precision, timezone), 49 | } 50 | } 51 | } 52 | 53 | impl ColumnData for DateTime64ColumnData { 54 | fn sql_type(&self) -> SqlType { 55 | let (precision, tz) = self.params; 56 | SqlType::DateTime(DateTimeType::DateTime64(precision, tz)) 57 | } 58 | 59 | fn save(&self, _encoder: &mut Encoder, _start: usize, _end: usize) { 60 | unimplemented!() 61 | } 62 | 63 | fn len(&self) -> usize { 64 | self.data.len() 65 | } 66 | 67 | fn push(&mut self, value: Value) { 68 | let (precision, tz) = &self.params; 69 | let time = DateTime::::from(value); 70 | let stamp = from_datetime(time.with_timezone(tz), *precision); 71 | self.data.push(stamp) 72 | } 73 | 74 | fn at(&self, index: usize) -> ValueRef { 75 | let value = self.data.at(index); 76 | ValueRef::DateTime64(value, &self.params) 77 | } 78 | 79 | fn clone_instance(&self) -> BoxColumnData { 80 | Box::new(Self { 81 | data: self.data.clone(), 82 | params: self.params, 83 | }) 84 | } 85 | 86 | unsafe fn get_internal( 87 | &self, 88 | pointers: &[*mut *const u8], 89 | level: u8, 90 | _props: u32, 91 | ) -> Result<()> { 92 | assert_eq!(level, 0); 93 | let (precision, tz) = &self.params; 94 | *pointers[0] = self.data.as_ptr() as *const u8; 95 | *pointers[1] = tz as *const Tz as *const u8; 96 | *(pointers[2] as *mut usize) = self.len(); 97 | *(pointers[3] as *mut Option) = Some(*precision); 98 | Ok(()) 99 | } 100 | 101 | unsafe fn get_internals(&self, data_ptr: *mut (), _level: u8, _props: u32) -> Result<()> { 102 | let (precision, tz) = &self.params; 103 | unsafe { 104 | let data_ref = &mut *(data_ptr as *mut DateTimeInternals); 105 | data_ref.begin = self.data.as_ptr().cast(); 106 | data_ref.len = self.data.len(); 107 | data_ref.tz = *tz; 108 | data_ref.precision = Some(*precision); 109 | Ok(()) 110 | } 111 | } 112 | 113 | fn get_timezone(&self) -> Option { 114 | let (_, tz) = self.params; 115 | Some(tz) 116 | } 117 | } 118 | 119 | pub(crate) fn from_datetime(time: DateTime, precision: u32) -> i64 { 120 | let base10: i64 = 10; 121 | let timestamp = time.timestamp_nanos_opt().unwrap(); 122 | timestamp / base10.pow(9 - precision) 123 | } 124 | 125 | #[inline(always)] 126 | pub(crate) fn to_datetime(value: i64, precision: u32, tz: Tz) -> DateTime { 127 | to_datetime_opt(value, precision, tz).unwrap() 128 | } 129 | 130 | #[inline(always)] 131 | pub(crate) fn to_datetime_opt(value: i64, precision: u32, tz: Tz) -> LocalResult> { 132 | let base10: i64 = 10; 133 | 134 | let nano = if precision < 19 { 135 | value * base10.pow(9 - precision) 136 | } else { 137 | 0_i64 138 | }; 139 | 140 | let sec = nano / 1_000_000_000; 141 | let nsec = nano - sec * 1_000_000_000; 142 | 143 | tz.timestamp_opt(sec, nsec as u32) 144 | } 145 | 146 | #[inline(always)] 147 | pub(crate) fn to_native_datetime_opt(value: i64, precision: u32) -> Option { 148 | let base10: i64 = 10; 149 | 150 | let nano = if precision < 19 { 151 | value * base10.pow(9 - precision) 152 | } else { 153 | 0_i64 154 | }; 155 | 156 | let sec = nano / 1_000_000_000; 157 | let nsec = nano - sec * 1_000_000_000; 158 | 159 | NaiveDateTime::from_timestamp_opt(sec, nsec as u32) 160 | } 161 | 162 | #[cfg(test)] 163 | mod test { 164 | use super::*; 165 | 166 | #[test] 167 | fn test_to_datetime() { 168 | let expected = DateTime::parse_from_rfc3339("2019-01-01T00:00:00-00:00").unwrap(); 169 | let actual = to_datetime(1_546_300_800_000, 3, Tz::UTC); 170 | assert_eq!(actual, expected) 171 | } 172 | 173 | #[test] 174 | fn test_from_datetime() { 175 | let origin = DateTime::parse_from_rfc3339("2019-01-01T00:00:00-00:00").unwrap(); 176 | let actual = from_datetime(origin, 3); 177 | assert_eq!(actual, 1_546_300_800_000) 178 | } 179 | } 180 | -------------------------------------------------------------------------------- /src/types/column/decimal.rs: -------------------------------------------------------------------------------- 1 | use chrono_tz::Tz; 2 | use either::Either; 3 | use std::sync::Arc; 4 | 5 | use crate::{ 6 | binary::{Encoder, ReadEx}, 7 | errors::Result, 8 | types::{ 9 | column::{ 10 | column_data::{BoxColumnData, ColumnData}, 11 | list::List, 12 | nullable::NullableColumnData, 13 | util::extract_nulls_and_values, 14 | BoxColumnWrapper, ColumnFrom, ColumnWrapper, VectorColumnData, 15 | }, 16 | decimal::{Decimal, NoBits}, 17 | from_sql::FromSql, 18 | Column, ColumnType, SqlType, Value, ValueRef, 19 | }, 20 | }; 21 | 22 | pub(crate) struct DecimalColumnData { 23 | pub(crate) inner: Box, 24 | pub(crate) precision: u8, 25 | pub(crate) scale: u8, 26 | pub(crate) nobits: NoBits, 27 | } 28 | 29 | pub(crate) struct DecimalAdapter { 30 | pub(crate) column: Column, 31 | pub(crate) precision: u8, 32 | pub(crate) scale: u8, 33 | pub(crate) nobits: NoBits, 34 | } 35 | 36 | pub(crate) struct NullableDecimalAdapter { 37 | pub(crate) column: Column, 38 | pub(crate) precision: u8, 39 | pub(crate) scale: u8, 40 | pub(crate) nobits: NoBits, 41 | } 42 | 43 | impl DecimalColumnData { 44 | pub(crate) fn load( 45 | reader: &mut T, 46 | precision: u8, 47 | scale: u8, 48 | nobits: NoBits, 49 | size: usize, 50 | tz: Tz, 51 | ) -> Result { 52 | let type_name = match nobits { 53 | NoBits::N32 => "Int32", 54 | NoBits::N64 => "Int64", 55 | }; 56 | let inner = 57 | ::load_data::(reader, type_name, size, tz)?; 58 | 59 | Ok(DecimalColumnData { 60 | inner, 61 | precision, 62 | scale, 63 | nobits, 64 | }) 65 | } 66 | } 67 | 68 | impl ColumnFrom for Vec { 69 | fn column_from(source: Self) -> W::Wrapper { 70 | let mut data = List::::with_capacity(source.len()); 71 | let mut precision = 18; 72 | let mut opt_scale = None; 73 | for s in source { 74 | precision = std::cmp::max(precision, s.precision); 75 | if let Some(old_scale) = opt_scale { 76 | if old_scale != s.scale { 77 | panic!("scale != scale") 78 | } 79 | } else { 80 | opt_scale = Some(s.scale); 81 | } 82 | 83 | data.push(s.internal()); 84 | } 85 | let scale = opt_scale.unwrap_or(4); 86 | let inner = Box::new(VectorColumnData { data }); 87 | 88 | let column = DecimalColumnData { 89 | inner, 90 | precision, 91 | scale, 92 | nobits: NoBits::N64, 93 | }; 94 | 95 | W::wrap(column) 96 | } 97 | } 98 | 99 | impl ColumnFrom for Vec> { 100 | fn column_from(source: Self) -> W::Wrapper { 101 | let mut nulls: Vec = Vec::with_capacity(source.len()); 102 | let mut data = List::::with_capacity(source.len()); 103 | let mut precision = 18; 104 | let mut opt_scale = None; 105 | for os in source { 106 | if let Some(s) = os { 107 | precision = std::cmp::max(precision, s.precision); 108 | if let Some(old_scale) = opt_scale { 109 | if old_scale != s.scale { 110 | panic!("scale != scale") 111 | } 112 | } else { 113 | opt_scale = Some(s.scale); 114 | } 115 | 116 | data.push(s.internal()); 117 | nulls.push(0); 118 | } else { 119 | data.push(0); 120 | nulls.push(1); 121 | } 122 | } 123 | let scale = opt_scale.unwrap_or(4); 124 | let inner = Box::new(VectorColumnData { data }); 125 | 126 | let inner = DecimalColumnData { 127 | inner, 128 | precision, 129 | scale, 130 | nobits: NoBits::N64, 131 | }; 132 | 133 | W::wrap(NullableColumnData { 134 | nulls, 135 | inner: Arc::new(inner), 136 | }) 137 | } 138 | } 139 | 140 | impl ColumnData for DecimalColumnData { 141 | fn sql_type(&self) -> SqlType { 142 | SqlType::Decimal(self.precision, self.scale) 143 | } 144 | 145 | fn save(&self, encoder: &mut Encoder, start: usize, end: usize) { 146 | self.inner.save(encoder, start, end) 147 | } 148 | 149 | fn len(&self) -> usize { 150 | self.inner.len() 151 | } 152 | 153 | fn push(&mut self, value: Value) { 154 | if let Value::Decimal(decimal) = value { 155 | match self.nobits { 156 | NoBits::N32 => { 157 | let internal: i32 = decimal.internal(); 158 | self.inner.push(internal.into()) 159 | } 160 | NoBits::N64 => { 161 | let internal: i64 = decimal.internal(); 162 | self.inner.push(internal.into()) 163 | } 164 | } 165 | } else { 166 | panic!("value should be decimal ({value:?})"); 167 | } 168 | } 169 | 170 | fn at(&self, index: usize) -> ValueRef { 171 | let underlying: i64 = match self.nobits { 172 | NoBits::N32 => i64::from(i32::from(self.inner.at(index))), 173 | NoBits::N64 => i64::from(self.inner.at(index)), 174 | }; 175 | 176 | ValueRef::Decimal(Decimal { 177 | underlying, 178 | precision: self.precision, 179 | scale: self.scale, 180 | nobits: self.nobits, 181 | }) 182 | } 183 | 184 | fn clone_instance(&self) -> BoxColumnData { 185 | Box::new(Self { 186 | inner: self.inner.clone_instance(), 187 | precision: self.precision, 188 | scale: self.scale, 189 | nobits: self.nobits, 190 | }) 191 | } 192 | 193 | unsafe fn get_internal( 194 | &self, 195 | pointers: &[*mut *const u8], 196 | level: u8, 197 | _props: u32, 198 | ) -> Result<()> { 199 | assert_eq!(level, 0); 200 | self.inner.get_internal(pointers, 0, 0)?; 201 | *(pointers[2] as *mut NoBits) = self.nobits; 202 | Ok(()) 203 | } 204 | 205 | fn get_timezone(&self) -> Option { 206 | None 207 | } 208 | } 209 | 210 | impl ColumnData for DecimalAdapter { 211 | fn sql_type(&self) -> SqlType { 212 | SqlType::Decimal(self.precision, self.scale) 213 | } 214 | 215 | fn save(&self, encoder: &mut Encoder, start: usize, end: usize) { 216 | for i in start..end { 217 | if let ValueRef::Decimal(decimal) = self.at(i) { 218 | match self.nobits { 219 | NoBits::N32 => { 220 | let internal: i32 = decimal.internal(); 221 | encoder.write(internal); 222 | } 223 | NoBits::N64 => { 224 | let internal: i64 = decimal.internal(); 225 | encoder.write(internal); 226 | } 227 | } 228 | } else { 229 | panic!("should be decimal"); 230 | } 231 | } 232 | } 233 | 234 | fn len(&self) -> usize { 235 | self.column.len() 236 | } 237 | 238 | fn push(&mut self, _: Value) { 239 | unimplemented!() 240 | } 241 | 242 | fn at(&self, index: usize) -> ValueRef { 243 | if let ValueRef::Decimal(decimal) = self.column.at(index) { 244 | let mut d = decimal.set_scale(self.scale); 245 | d.precision = self.precision; 246 | d.nobits = self.nobits; 247 | ValueRef::Decimal(d) 248 | } else { 249 | panic!("should be decimal"); 250 | } 251 | } 252 | 253 | fn clone_instance(&self) -> BoxColumnData { 254 | unimplemented!() 255 | } 256 | 257 | fn get_timezone(&self) -> Option { 258 | None 259 | } 260 | } 261 | 262 | impl ColumnData for NullableDecimalAdapter { 263 | fn sql_type(&self) -> SqlType { 264 | SqlType::Nullable(SqlType::Decimal(self.precision, self.scale).into()) 265 | } 266 | 267 | fn save(&self, encoder: &mut Encoder, start: usize, end: usize) { 268 | let (nulls, values) = extract_nulls_and_values::(self, start, end).unwrap(); 269 | encoder.write_bytes(nulls.as_ref()); 270 | 271 | for value in values { 272 | let underlying = if let Some(v) = value { v.underlying } else { 0 }; 273 | 274 | match self.nobits { 275 | NoBits::N32 => { 276 | encoder.write(underlying as i32); 277 | } 278 | NoBits::N64 => { 279 | encoder.write(underlying); 280 | } 281 | } 282 | } 283 | } 284 | 285 | fn len(&self) -> usize { 286 | self.column.len() 287 | } 288 | 289 | fn push(&mut self, _: Value) { 290 | unimplemented!() 291 | } 292 | 293 | fn at(&self, index: usize) -> ValueRef { 294 | let value: Option = Option::from_sql(self.column.at(index)).unwrap(); 295 | match value { 296 | None => ValueRef::Nullable(Either::Left(self.sql_type().into())), 297 | Some(mut v) => { 298 | v = v.set_scale(self.scale); 299 | v.precision = self.precision; 300 | v.nobits = self.nobits; 301 | let inner = ValueRef::Decimal(v); 302 | ValueRef::Nullable(Either::Right(Box::new(inner))) 303 | } 304 | } 305 | } 306 | 307 | fn clone_instance(&self) -> BoxColumnData { 308 | unimplemented!() 309 | } 310 | 311 | fn get_timezone(&self) -> Option { 312 | self.column.data.get_timezone() 313 | } 314 | } 315 | -------------------------------------------------------------------------------- /src/types/column/fixed_string.rs: -------------------------------------------------------------------------------- 1 | use chrono_tz::Tz; 2 | use std::cmp; 3 | 4 | use crate::{ 5 | binary::{Encoder, ReadEx}, 6 | errors::Result, 7 | types::{ 8 | column::{ 9 | column_data::{BoxColumnData, LowCardinalityAccessor}, 10 | util::extract_nulls_and_values, 11 | }, 12 | value::get_str_buffer, 13 | Column, ColumnType, SqlType, Value, ValueRef, 14 | }, 15 | }; 16 | 17 | use super::column_data::ColumnData; 18 | 19 | pub(crate) struct FixedStringColumnData { 20 | buffer: Vec, 21 | str_len: usize, 22 | } 23 | 24 | pub(crate) struct FixedStringAdapter { 25 | pub(crate) column: Column, 26 | pub(crate) str_len: usize, 27 | } 28 | 29 | pub(crate) struct NullableFixedStringAdapter { 30 | pub(crate) column: Column, 31 | pub(crate) str_len: usize, 32 | } 33 | 34 | impl FixedStringColumnData { 35 | pub fn with_capacity(capacity: usize, str_len: usize) -> Self { 36 | Self { 37 | buffer: Vec::with_capacity(capacity * str_len), 38 | str_len, 39 | } 40 | } 41 | 42 | pub(crate) fn load(reader: &mut T, size: usize, str_len: usize) -> Result { 43 | let mut instance = Self::with_capacity(size, str_len); 44 | 45 | for _ in 0..size { 46 | let old_len = instance.buffer.len(); 47 | instance.buffer.resize(old_len + str_len, 0_u8); 48 | reader.read_bytes(&mut instance.buffer[old_len..old_len + str_len])?; 49 | } 50 | 51 | Ok(instance) 52 | } 53 | } 54 | 55 | impl LowCardinalityAccessor for FixedStringColumnData { 56 | fn get_string(&self, index: usize) -> &[u8] { 57 | let shift = index * self.str_len; 58 | &self.buffer[shift..shift + self.str_len] 59 | } 60 | } 61 | 62 | impl ColumnData for FixedStringColumnData { 63 | fn sql_type(&self) -> SqlType { 64 | SqlType::FixedString(self.str_len) 65 | } 66 | 67 | fn save(&self, encoder: &mut Encoder, start: usize, end: usize) { 68 | let start_index = start * self.str_len; 69 | let end_index = end * self.str_len; 70 | encoder.write_bytes(&self.buffer[start_index..end_index]); 71 | } 72 | 73 | fn len(&self) -> usize { 74 | self.buffer.len() / self.str_len 75 | } 76 | 77 | fn push(&mut self, value: Value) { 78 | let bs = get_str_buffer(&value); 79 | let l = cmp::min(bs.len(), self.str_len); 80 | let old_len = self.buffer.len(); 81 | self.buffer.extend(bs.iter().take(l)); 82 | self.buffer.resize(old_len + self.str_len, 0_u8); 83 | } 84 | 85 | fn at(&self, index: usize) -> ValueRef { 86 | let shift = index * self.str_len; 87 | let str_ref = &self.buffer[shift..shift + self.str_len]; 88 | ValueRef::String(str_ref) 89 | } 90 | 91 | fn clone_instance(&self) -> BoxColumnData { 92 | Box::new(Self { 93 | buffer: self.buffer.clone(), 94 | str_len: self.str_len, 95 | }) 96 | } 97 | 98 | unsafe fn get_internal( 99 | &self, 100 | pointers: &[*mut *const u8], 101 | level: u8, 102 | _props: u32, 103 | ) -> Result<()> { 104 | assert_eq!(level, 0); 105 | *pointers[0] = self.buffer.as_ptr(); 106 | *(pointers[1] as *mut usize) = self.len(); 107 | Ok(()) 108 | } 109 | 110 | fn get_timezone(&self) -> Option { 111 | None 112 | } 113 | 114 | fn get_low_cardinality_accessor(&self) -> Option<&dyn LowCardinalityAccessor> { 115 | Some(self) 116 | } 117 | } 118 | 119 | impl ColumnData for FixedStringAdapter { 120 | fn sql_type(&self) -> SqlType { 121 | SqlType::FixedString(self.str_len) 122 | } 123 | 124 | fn save(&self, encoder: &mut Encoder, start: usize, end: usize) { 125 | let mut buffer = Vec::with_capacity(self.str_len); 126 | for index in start..end { 127 | buffer.clear(); 128 | match self.column.at(index) { 129 | ValueRef::String(_) => { 130 | let string_ref = self.column.at(index).as_bytes().unwrap(); 131 | buffer.extend(string_ref.as_ref()); 132 | } 133 | ValueRef::Array(SqlType::UInt8, vs) => { 134 | let mut string_val: Vec = Vec::with_capacity(vs.len()); 135 | for v in vs.iter() { 136 | let byte: u8 = v.clone().into(); 137 | string_val.push(byte); 138 | } 139 | let string_ref: &[u8] = string_val.as_ref(); 140 | buffer.extend(string_ref); 141 | } 142 | _ => unimplemented!(), 143 | } 144 | buffer.resize(self.str_len, 0); 145 | encoder.write_bytes(&buffer[..]); 146 | } 147 | } 148 | 149 | fn len(&self) -> usize { 150 | self.column.len() 151 | } 152 | 153 | fn push(&mut self, _value: Value) { 154 | unimplemented!() 155 | } 156 | 157 | fn at(&self, index: usize) -> ValueRef { 158 | self.column.at(index) 159 | } 160 | 161 | fn clone_instance(&self) -> BoxColumnData { 162 | unimplemented!() 163 | } 164 | 165 | fn get_timezone(&self) -> Option { 166 | None 167 | } 168 | } 169 | 170 | impl ColumnData for NullableFixedStringAdapter { 171 | fn sql_type(&self) -> SqlType { 172 | SqlType::Nullable(SqlType::FixedString(self.str_len).into()) 173 | } 174 | 175 | fn save(&self, encoder: &mut Encoder, start: usize, end: usize) { 176 | let (nulls, values) = extract_nulls_and_values::<&[u8], Self>(self, start, end).unwrap(); 177 | encoder.write_bytes(nulls.as_ref()); 178 | 179 | let mut buffer = Vec::with_capacity(self.str_len); 180 | for value in values { 181 | buffer.clear(); 182 | if let Some(string_ref) = value { 183 | buffer.extend(string_ref); 184 | } 185 | buffer.resize(self.str_len, 0); 186 | encoder.write_bytes(buffer.as_ref()); 187 | } 188 | } 189 | 190 | fn len(&self) -> usize { 191 | self.column.len() 192 | } 193 | 194 | fn push(&mut self, _value: Value) { 195 | unimplemented!() 196 | } 197 | 198 | fn at(&self, index: usize) -> ValueRef { 199 | self.column.at(index) 200 | } 201 | 202 | fn clone_instance(&self) -> BoxColumnData { 203 | unimplemented!() 204 | } 205 | 206 | fn get_timezone(&self) -> Option { 207 | None 208 | } 209 | } 210 | -------------------------------------------------------------------------------- /src/types/column/ip.rs: -------------------------------------------------------------------------------- 1 | use chrono_tz::Tz; 2 | use std::{ 3 | marker::PhantomData, 4 | net::{Ipv4Addr, Ipv6Addr}, 5 | sync::Arc, 6 | }; 7 | 8 | use crate::{ 9 | binary::{Encoder, ReadEx}, 10 | errors::Result, 11 | types::{ 12 | column::{column_data::BoxColumnData, nullable::NullableColumnData, ColumnWrapper}, 13 | SqlType, Value, ValueRef, 14 | }, 15 | }; 16 | 17 | use super::{column_data::ColumnData, ColumnFrom}; 18 | 19 | pub(crate) trait IpVersion: Copy + Sync + Send + 'static { 20 | fn sql_type() -> SqlType; 21 | fn size() -> usize; 22 | fn push(inner: &mut Vec, value: Value); 23 | fn get(inner: &[u8], index: usize) -> ValueRef; 24 | } 25 | 26 | #[derive(Copy, Clone)] 27 | pub(crate) struct Ipv4; 28 | 29 | #[derive(Copy, Clone)] 30 | pub(crate) struct Ipv6; 31 | 32 | #[derive(Copy, Clone)] 33 | pub(crate) struct Uuid; 34 | 35 | impl IpVersion for Ipv4 { 36 | #[inline(always)] 37 | fn sql_type() -> SqlType { 38 | SqlType::Ipv4 39 | } 40 | 41 | #[inline(always)] 42 | fn size() -> usize { 43 | 4 44 | } 45 | 46 | #[inline(always)] 47 | fn push(inner: &mut Vec, value: Value) { 48 | if let Value::Ipv4(v) = value { 49 | inner.extend(&v); 50 | } else { 51 | panic!(); 52 | } 53 | } 54 | 55 | #[inline(always)] 56 | fn get(inner: &[u8], index: usize) -> ValueRef { 57 | let mut v: [u8; 4] = Default::default(); 58 | v.copy_from_slice(&inner[index * 4..(index + 1) * 4]); 59 | ValueRef::Ipv4(v) 60 | } 61 | } 62 | 63 | impl IpVersion for Ipv6 { 64 | #[inline(always)] 65 | fn sql_type() -> SqlType { 66 | SqlType::Ipv6 67 | } 68 | 69 | #[inline(always)] 70 | fn size() -> usize { 71 | 16 72 | } 73 | 74 | #[inline(always)] 75 | fn push(inner: &mut Vec, value: Value) { 76 | if let Value::Ipv6(v) = value { 77 | inner.extend(&v); 78 | } else { 79 | panic!(); 80 | } 81 | } 82 | 83 | #[inline(always)] 84 | fn get(inner: &[u8], index: usize) -> ValueRef { 85 | let mut v: [u8; 16] = Default::default(); 86 | v.copy_from_slice(&inner[index * 16..(index + 1) * 16]); 87 | ValueRef::Ipv6(v) 88 | } 89 | } 90 | 91 | impl IpVersion for Uuid { 92 | #[inline(always)] 93 | fn sql_type() -> SqlType { 94 | SqlType::Uuid 95 | } 96 | 97 | #[inline(always)] 98 | fn size() -> usize { 99 | 16 100 | } 101 | 102 | #[inline(always)] 103 | fn push(inner: &mut Vec, value: Value) { 104 | if let Value::Uuid(v) = value { 105 | inner.extend(&v); 106 | } else { 107 | panic!(); 108 | } 109 | } 110 | 111 | #[inline(always)] 112 | fn get(inner: &[u8], index: usize) -> ValueRef { 113 | let mut v: [u8; 16] = Default::default(); 114 | v.copy_from_slice(&inner[index * 16..(index + 1) * 16]); 115 | ValueRef::Uuid(v) 116 | } 117 | } 118 | 119 | impl ColumnFrom for Vec { 120 | fn column_from(data: Self) -> W::Wrapper { 121 | let mut inner = Vec::with_capacity(data.len()); 122 | for ip in data { 123 | let mut buffer = ip.octets(); 124 | buffer.reverse(); 125 | inner.extend(&buffer); 126 | } 127 | 128 | W::wrap(IpColumnData:: { 129 | inner, 130 | phantom: PhantomData, 131 | }) 132 | } 133 | } 134 | 135 | impl ColumnFrom for Vec { 136 | fn column_from(data: Self) -> W::Wrapper { 137 | let mut inner = Vec::with_capacity(data.len()); 138 | for ip in data { 139 | inner.extend(&ip.octets()); 140 | } 141 | 142 | W::wrap(IpColumnData:: { 143 | inner, 144 | phantom: PhantomData, 145 | }) 146 | } 147 | } 148 | 149 | impl ColumnFrom for Vec { 150 | fn column_from(data: Self) -> W::Wrapper { 151 | let mut inner = Vec::with_capacity(data.len()); 152 | for uuid in data { 153 | let mut buffer = *uuid.as_bytes(); 154 | buffer[..8].reverse(); 155 | buffer[8..].reverse(); 156 | inner.extend(&buffer); 157 | } 158 | 159 | W::wrap(IpColumnData:: { 160 | inner, 161 | phantom: PhantomData, 162 | }) 163 | } 164 | } 165 | 166 | impl ColumnFrom for Vec> { 167 | fn column_from(source: Self) -> ::Wrapper { 168 | let n = source.len(); 169 | let mut inner: Vec = Vec::with_capacity(n * 4); 170 | let mut nulls: Vec = Vec::with_capacity(n); 171 | 172 | for ip in source { 173 | match ip { 174 | None => { 175 | inner.extend(&[0; 4]); 176 | nulls.push(1); 177 | } 178 | Some(ip) => { 179 | let mut buffer = ip.octets(); 180 | buffer.reverse(); 181 | inner.extend(&buffer); 182 | nulls.push(0); 183 | } 184 | } 185 | } 186 | 187 | let inner = Arc::new(IpColumnData:: { 188 | inner, 189 | phantom: PhantomData, 190 | }); 191 | 192 | let data = NullableColumnData { inner, nulls }; 193 | 194 | W::wrap(data) 195 | } 196 | } 197 | 198 | impl ColumnFrom for Vec> { 199 | fn column_from(source: Self) -> ::Wrapper { 200 | let n = source.len(); 201 | let mut inner: Vec = Vec::with_capacity(n * 16); 202 | let mut nulls: Vec = Vec::with_capacity(n); 203 | 204 | for ip in source { 205 | match ip { 206 | None => { 207 | inner.extend(&[0; 16]); 208 | nulls.push(1); 209 | } 210 | Some(ip) => { 211 | inner.extend(&ip.octets()); 212 | nulls.push(0); 213 | } 214 | } 215 | } 216 | 217 | let inner = Arc::new(IpColumnData:: { 218 | inner, 219 | phantom: PhantomData, 220 | }); 221 | 222 | let data = NullableColumnData { inner, nulls }; 223 | 224 | W::wrap(data) 225 | } 226 | } 227 | 228 | impl ColumnFrom for Vec> { 229 | fn column_from(source: Self) -> ::Wrapper { 230 | let n = source.len(); 231 | let mut inner: Vec = Vec::with_capacity(n * 16); 232 | let mut nulls: Vec = Vec::with_capacity(n); 233 | 234 | for uuid in source { 235 | match uuid { 236 | None => { 237 | inner.extend(&[0; 16]); 238 | nulls.push(1); 239 | } 240 | Some(uuid) => { 241 | let mut buffer = *uuid.as_bytes(); 242 | buffer[..8].reverse(); 243 | buffer[8..].reverse(); 244 | inner.extend(&buffer); 245 | nulls.push(0); 246 | } 247 | } 248 | } 249 | 250 | let inner = Arc::new(IpColumnData:: { 251 | inner, 252 | phantom: PhantomData, 253 | }); 254 | 255 | let data = NullableColumnData { inner, nulls }; 256 | 257 | W::wrap(data) 258 | } 259 | } 260 | 261 | pub(crate) struct IpColumnData { 262 | pub(crate) inner: Vec, 263 | pub(crate) phantom: PhantomData, 264 | } 265 | 266 | impl IpColumnData { 267 | pub fn with_capacity(capacity: usize) -> Self { 268 | Self { 269 | inner: Vec::with_capacity(capacity * V::size()), 270 | phantom: PhantomData, 271 | } 272 | } 273 | 274 | pub(crate) fn load(reader: &mut R, size: usize) -> Result { 275 | let mut inner = vec![0; size * V::size()]; 276 | reader.read_bytes(inner.as_mut())?; 277 | 278 | Ok(Self { 279 | inner, 280 | phantom: PhantomData, 281 | }) 282 | } 283 | } 284 | 285 | impl ColumnData for IpColumnData { 286 | fn sql_type(&self) -> SqlType { 287 | V::sql_type() 288 | } 289 | 290 | fn save(&self, encoder: &mut Encoder, start: usize, end: usize) { 291 | let start_index = start * V::size(); 292 | let end_index = end * V::size(); 293 | 294 | let slice: &[u8] = self.inner.as_ref(); 295 | encoder.write_bytes(&slice[start_index..end_index]); 296 | } 297 | 298 | fn len(&self) -> usize { 299 | self.inner.len() / V::size() 300 | } 301 | 302 | fn push(&mut self, value: Value) { 303 | V::push(&mut self.inner, value) 304 | } 305 | 306 | fn at(&self, index: usize) -> ValueRef { 307 | V::get(&self.inner, index) 308 | } 309 | 310 | fn clone_instance(&self) -> BoxColumnData { 311 | Box::new(Self { 312 | inner: self.inner.clone(), 313 | phantom: PhantomData, 314 | }) 315 | } 316 | 317 | unsafe fn get_internal( 318 | &self, 319 | pointers: &[*mut *const u8], 320 | level: u8, 321 | _props: u32, 322 | ) -> Result<()> { 323 | assert_eq!(level, 0); 324 | *pointers[0] = &self.inner as *const Vec as *const u8; 325 | Ok(()) 326 | } 327 | 328 | fn get_timezone(&self) -> Option { 329 | None 330 | } 331 | } 332 | -------------------------------------------------------------------------------- /src/types/column/list.rs: -------------------------------------------------------------------------------- 1 | use std::{fmt, mem, slice}; 2 | 3 | use crate::types::{Marshal, StatBuffer, Unmarshal}; 4 | 5 | #[derive(Clone)] 6 | pub struct List 7 | where 8 | T: StatBuffer + Unmarshal + Marshal + Copy + Sync + 'static, 9 | { 10 | data: Vec, 11 | } 12 | 13 | impl List 14 | where 15 | T: StatBuffer + Unmarshal + Marshal + Copy + Sync + 'static, 16 | { 17 | pub fn len(&self) -> usize { 18 | self.data.len() 19 | } 20 | 21 | pub fn at(&self, index: usize) -> T { 22 | self.data[index] 23 | } 24 | 25 | pub fn push(&mut self, value: T) { 26 | self.data.push(value); 27 | } 28 | 29 | #[cfg(test)] 30 | pub fn new() -> List { 31 | List { data: Vec::new() } 32 | } 33 | 34 | pub fn with_capacity(capacity: usize) -> List { 35 | Self { 36 | data: Vec::with_capacity(capacity), 37 | } 38 | } 39 | 40 | pub fn resize(&mut self, new_len: usize, value: T) { 41 | self.data.resize(new_len, value); 42 | } 43 | 44 | pub(super) unsafe fn set_len(&mut self, new_len: usize) { 45 | self.data.set_len(new_len); 46 | } 47 | 48 | pub(super) unsafe fn as_ptr(&self) -> *const T { 49 | self.data.as_ptr() 50 | } 51 | 52 | pub fn map(&self, f: F) -> List 53 | where 54 | N: StatBuffer + Unmarshal + Marshal + Copy + Sync + 'static, 55 | F: Fn(T) -> N, 56 | { 57 | let new_capacity = std::cmp::max(self.data.capacity(), self.len() + 1); 58 | let mut new_list = List::::with_capacity(new_capacity); 59 | for i in 0..self.len() { 60 | new_list.push(f(self.at(i))); 61 | } 62 | new_list 63 | } 64 | } 65 | 66 | impl fmt::Debug for List 67 | where 68 | T: StatBuffer + Unmarshal + Marshal + Copy + Sync + 'static + fmt::Debug, 69 | { 70 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 71 | write!(f, "{:?}", self.data) 72 | } 73 | } 74 | 75 | impl AsRef<[u8]> for List 76 | where 77 | T: StatBuffer + Unmarshal + Marshal + Copy + Sync + 'static, 78 | { 79 | fn as_ref(&self) -> &[u8] { 80 | let ptr = self.data.as_ptr() as *const u8; 81 | let size = self.len() * mem::size_of::(); 82 | unsafe { slice::from_raw_parts(ptr, size) } 83 | } 84 | } 85 | 86 | impl AsMut<[u8]> for List 87 | where 88 | T: StatBuffer + Unmarshal + Marshal + Copy + Sync + 'static, 89 | { 90 | fn as_mut(&mut self) -> &mut [u8] { 91 | let ptr = self.data.as_mut_ptr() as *mut u8; 92 | let size = self.len() * mem::size_of::(); 93 | unsafe { slice::from_raw_parts_mut(ptr, size) } 94 | } 95 | } 96 | 97 | #[cfg(test)] 98 | mod test { 99 | use rand::random; 100 | 101 | use super::*; 102 | use std::f64::EPSILON; 103 | 104 | #[test] 105 | fn test_push_and_len() { 106 | let mut list = List::with_capacity(100_500); 107 | 108 | for i in 0..100_500 { 109 | assert_eq!(list.len(), i as usize); 110 | list.push(i); 111 | } 112 | } 113 | 114 | #[test] 115 | fn test_push_and_get() { 116 | let mut list = List::::new(); 117 | let mut vs = vec![0.0_f64; 100]; 118 | 119 | for (count, _) in (0..100).enumerate() { 120 | assert_eq!(list.len(), count); 121 | 122 | for (i, v) in vs.iter().take(count).enumerate() { 123 | assert!((list.at(i) - *v).abs() < EPSILON); 124 | } 125 | 126 | let k = random(); 127 | list.push(k); 128 | vs[count] = k; 129 | } 130 | } 131 | } 132 | -------------------------------------------------------------------------------- /src/types/column/nullable.rs: -------------------------------------------------------------------------------- 1 | use std::sync::Arc; 2 | 3 | use crate::{ 4 | binary::{Encoder, ReadEx}, 5 | errors::Result, 6 | types::{ 7 | column::{ 8 | column_data::{ArcColumnData, BoxColumnData}, 9 | ArcColumnWrapper, ColumnData, 10 | }, 11 | SqlType, Value, ValueRef, 12 | }, 13 | }; 14 | 15 | use chrono_tz::Tz; 16 | use either::Either; 17 | 18 | pub(crate) struct NullableColumnData { 19 | pub(crate) inner: ArcColumnData, 20 | pub(crate) nulls: Vec, 21 | } 22 | 23 | impl NullableColumnData { 24 | pub(crate) fn load( 25 | reader: &mut R, 26 | type_name: &str, 27 | size: usize, 28 | tz: Tz, 29 | ) -> Result { 30 | let mut nulls = vec![0; size]; 31 | reader.read_bytes(nulls.as_mut())?; 32 | let inner = 33 | ::load_data::(reader, type_name, size, tz)?; 34 | Ok(NullableColumnData { inner, nulls }) 35 | } 36 | } 37 | 38 | impl ColumnData for NullableColumnData { 39 | fn sql_type(&self) -> SqlType { 40 | let inner_type = self.inner.sql_type(); 41 | SqlType::Nullable(inner_type.into()) 42 | } 43 | 44 | fn save(&self, encoder: &mut Encoder, start: usize, end: usize) { 45 | let nulls: &[u8] = self.nulls.as_ref(); 46 | encoder.write_bytes(&nulls[start..end]); 47 | self.inner.save(encoder, start, end); 48 | } 49 | 50 | fn len(&self) -> usize { 51 | assert_eq!(self.nulls.len(), self.inner.len()); 52 | self.inner.len() 53 | } 54 | 55 | fn push(&mut self, value: Value) { 56 | let inner_column: &mut dyn ColumnData = Arc::get_mut(&mut self.inner).unwrap(); 57 | 58 | if let Value::Nullable(e) = value { 59 | match e { 60 | Either::Left(sql_type) => { 61 | let default_value = Value::default(sql_type.clone()); 62 | inner_column.push(default_value); 63 | self.nulls.push(true as u8); 64 | } 65 | Either::Right(inner) => { 66 | inner_column.push(*inner); 67 | self.nulls.push(false as u8); 68 | } 69 | } 70 | } else { 71 | inner_column.push(value); 72 | self.nulls.push(false as u8); 73 | } 74 | } 75 | 76 | fn at(&self, index: usize) -> ValueRef { 77 | if self.nulls[index] == 1 { 78 | let sql_type = self.inner.sql_type(); 79 | ValueRef::Nullable(Either::Left(sql_type.into())) 80 | } else { 81 | let inner_value = self.inner.at(index); 82 | ValueRef::Nullable(Either::Right(Box::new(inner_value))) 83 | } 84 | } 85 | 86 | fn clone_instance(&self) -> BoxColumnData { 87 | Box::new(Self { 88 | inner: self.inner.clone(), 89 | nulls: self.nulls.clone(), 90 | }) 91 | } 92 | 93 | unsafe fn get_internal( 94 | &self, 95 | pointers: &[*mut *const u8], 96 | level: u8, 97 | props: u32, 98 | ) -> Result<()> { 99 | if level == self.sql_type().level() { 100 | *pointers[0] = self.nulls.as_ptr(); 101 | *(pointers[1] as *mut usize) = self.len(); 102 | Ok(()) 103 | } else { 104 | self.inner.get_internal(pointers, level, props) 105 | } 106 | } 107 | 108 | fn cast_to(&self, _this: &ArcColumnData, target: &SqlType) -> Option { 109 | if let SqlType::Nullable(inner_target) = target { 110 | if let Some(inner) = self.inner.cast_to(&self.inner, inner_target) { 111 | return Some(Arc::new(NullableColumnData { 112 | inner, 113 | nulls: self.nulls.clone(), 114 | })); 115 | } 116 | } 117 | None 118 | } 119 | 120 | fn get_timezone(&self) -> Option { 121 | self.inner.get_timezone() 122 | } 123 | } 124 | -------------------------------------------------------------------------------- /src/types/column/numeric.rs: -------------------------------------------------------------------------------- 1 | use chrono_tz::Tz; 2 | use std::{mem, sync::Arc}; 3 | 4 | use crate::{ 5 | binary::{Encoder, ReadEx}, 6 | errors::Result, 7 | types::{ 8 | column::{ 9 | array::ArrayColumnData, nullable::NullableColumnData, ArcColumnWrapper, ColumnWrapper, 10 | }, 11 | HasSqlType, Marshal, SqlType, StatBuffer, Unmarshal, Value, ValueRef, 12 | }, 13 | }; 14 | 15 | use super::{ 16 | column_data::{BoxColumnData, ColumnData}, 17 | list::List, 18 | ColumnFrom, 19 | }; 20 | 21 | #[derive(Clone)] 22 | pub struct VectorColumnData 23 | where 24 | T: StatBuffer 25 | + Unmarshal 26 | + Marshal 27 | + Copy 28 | + Clone 29 | + Into 30 | + From 31 | + Sync 32 | + HasSqlType 33 | + 'static, 34 | { 35 | pub(crate) data: List, 36 | } 37 | 38 | impl ColumnFrom for Vec 39 | where 40 | T: StatBuffer 41 | + Unmarshal 42 | + Marshal 43 | + Copy 44 | + Into 45 | + From 46 | + Send 47 | + Sync 48 | + HasSqlType 49 | + 'static, 50 | { 51 | fn column_from(source: Self) -> W::Wrapper { 52 | let mut data = List::with_capacity(source.len()); 53 | for s in source { 54 | data.push(s); 55 | } 56 | W::wrap(VectorColumnData { data }) 57 | } 58 | } 59 | 60 | impl ColumnFrom for Vec> 61 | where 62 | Value: From, 63 | T: StatBuffer 64 | + Unmarshal 65 | + Marshal 66 | + Copy 67 | + Into 68 | + From 69 | + Send 70 | + Sync 71 | + HasSqlType 72 | + 'static, 73 | { 74 | fn column_from(source: Self) -> W::Wrapper { 75 | let fake: Vec = Vec::with_capacity(source.len()); 76 | let inner = Vec::column_from::(fake); 77 | 78 | let mut data = NullableColumnData { 79 | inner, 80 | nulls: Vec::with_capacity(source.len()), 81 | }; 82 | 83 | for value in source { 84 | data.push(value.into()); 85 | } 86 | 87 | W::wrap(data) 88 | } 89 | } 90 | 91 | impl ColumnFrom for Vec> 92 | where 93 | Value: From, 94 | T: StatBuffer 95 | + Unmarshal 96 | + Marshal 97 | + Copy 98 | + Into 99 | + From 100 | + Send 101 | + Sync 102 | + HasSqlType 103 | + 'static, 104 | { 105 | fn column_from(source: Self) -> W::Wrapper { 106 | let fake: Vec = Vec::with_capacity(source.len()); 107 | let inner = Vec::column_from::(fake); 108 | let sql_type = inner.sql_type(); 109 | 110 | let mut data = ArrayColumnData { 111 | inner, 112 | offsets: List::with_capacity(source.len()), 113 | }; 114 | 115 | for array in source { 116 | data.push(to_array(sql_type.clone(), array)); 117 | } 118 | 119 | W::wrap(data) 120 | } 121 | } 122 | 123 | fn to_array(sql_type: SqlType, vs: Vec) -> Value 124 | where 125 | Value: From, 126 | T: StatBuffer 127 | + Unmarshal 128 | + Marshal 129 | + Copy 130 | + Into 131 | + From 132 | + Send 133 | + Sync 134 | + HasSqlType 135 | + 'static, 136 | { 137 | let mut inner = Vec::with_capacity(vs.len()); 138 | for v in vs { 139 | let value: Value = v.into(); 140 | inner.push(value) 141 | } 142 | Value::Array(sql_type.into(), Arc::new(inner)) 143 | } 144 | 145 | impl VectorColumnData 146 | where 147 | T: StatBuffer 148 | + Unmarshal 149 | + Marshal 150 | + Copy 151 | + Into 152 | + From 153 | + Sync 154 | + HasSqlType 155 | + 'static, 156 | { 157 | pub(crate) fn with_capacity(capacity: usize) -> VectorColumnData { 158 | VectorColumnData { 159 | data: List::with_capacity(capacity), 160 | } 161 | } 162 | 163 | pub(crate) fn load(reader: &mut R, size: usize) -> Result> { 164 | let mut data = List::with_capacity(size); 165 | unsafe { 166 | data.set_len(size); 167 | } 168 | reader.read_bytes(data.as_mut())?; 169 | Ok(Self { data }) 170 | } 171 | 172 | pub(crate) fn get_by_index(&self, index: usize) -> T { 173 | self.data.at(index) 174 | } 175 | } 176 | 177 | impl ColumnData for VectorColumnData 178 | where 179 | T: StatBuffer 180 | + Unmarshal 181 | + Marshal 182 | + Copy 183 | + Into 184 | + From 185 | + Send 186 | + Sync 187 | + HasSqlType 188 | + 'static, 189 | { 190 | fn sql_type(&self) -> SqlType { 191 | T::sql_type() 192 | } 193 | 194 | fn save(&self, encoder: &mut Encoder, start: usize, end: usize) { 195 | save_data::(self.data.as_ref(), encoder, start, end); 196 | } 197 | 198 | fn len(&self) -> usize { 199 | self.data.len() 200 | } 201 | 202 | fn push(&mut self, value: Value) { 203 | self.data.push(T::from(value)); 204 | } 205 | 206 | fn at(&self, index: usize) -> ValueRef { 207 | let v: Value = self.data.at(index).into(); 208 | match v { 209 | Value::Bool(x) => ValueRef::Bool(x), 210 | Value::UInt8(x) => ValueRef::UInt8(x), 211 | Value::UInt16(x) => ValueRef::UInt16(x), 212 | Value::UInt32(x) => ValueRef::UInt32(x), 213 | Value::UInt64(x) => ValueRef::UInt64(x), 214 | Value::UInt128(x) => ValueRef::UInt128(x), 215 | 216 | Value::Int8(x) => ValueRef::Int8(x), 217 | Value::Int16(x) => ValueRef::Int16(x), 218 | Value::Int32(x) => ValueRef::Int32(x), 219 | Value::Int64(x) => ValueRef::Int64(x), 220 | Value::Int128(x) => ValueRef::Int128(x), 221 | 222 | Value::Float32(x) => ValueRef::Float32(x), 223 | Value::Float64(x) => ValueRef::Float64(x), 224 | 225 | _ => panic!("can't convert value to value_ref."), 226 | } 227 | } 228 | 229 | fn clone_instance(&self) -> BoxColumnData { 230 | Box::new(Self { 231 | data: self.data.clone(), 232 | }) 233 | } 234 | 235 | unsafe fn get_internal( 236 | &self, 237 | pointers: &[*mut *const u8], 238 | level: u8, 239 | _props: u32, 240 | ) -> Result<()> { 241 | assert_eq!(level, 0); 242 | *pointers[0] = self.data.as_ptr() as *const u8; 243 | *(pointers[1] as *mut usize) = self.len(); 244 | Ok(()) 245 | } 246 | 247 | fn get_timezone(&self) -> Option { 248 | None 249 | } 250 | } 251 | 252 | pub(crate) fn save_data(data: &[u8], encoder: &mut Encoder, start: usize, end: usize) { 253 | let start_index = start * mem::size_of::(); 254 | let end_index = end * mem::size_of::(); 255 | encoder.write_bytes(&data[start_index..end_index]); 256 | } 257 | -------------------------------------------------------------------------------- /src/types/column/simple_agg_func.rs: -------------------------------------------------------------------------------- 1 | use crate::{ 2 | binary::{Encoder, ReadEx}, 3 | errors::Result, 4 | types::{ 5 | column::{ 6 | column_data::{ArcColumnData, BoxColumnData}, 7 | ArcColumnWrapper, ColumnData, 8 | }, 9 | SimpleAggFunc, SqlType, Value, ValueRef, 10 | }, 11 | }; 12 | 13 | use chrono_tz::Tz; 14 | use std::sync::Arc; 15 | 16 | pub(crate) struct SimpleAggregateFunctionColumnData { 17 | pub(crate) inner: ArcColumnData, 18 | pub(crate) func: SimpleAggFunc, 19 | } 20 | 21 | impl SimpleAggregateFunctionColumnData { 22 | pub(crate) fn load( 23 | reader: &mut R, 24 | func: SimpleAggFunc, 25 | type_name: &str, 26 | size: usize, 27 | tz: Tz, 28 | ) -> Result { 29 | let inner = 30 | ::load_data::(reader, type_name, size, tz)?; 31 | Ok(SimpleAggregateFunctionColumnData { inner, func }) 32 | } 33 | } 34 | 35 | impl ColumnData for SimpleAggregateFunctionColumnData { 36 | fn sql_type(&self) -> SqlType { 37 | let inner_type = self.inner.sql_type(); 38 | SqlType::SimpleAggregateFunction(self.func, inner_type.into()) 39 | } 40 | 41 | fn save(&self, encoder: &mut Encoder, start: usize, end: usize) { 42 | self.inner.save(encoder, start, end); 43 | } 44 | 45 | fn len(&self) -> usize { 46 | self.inner.len() 47 | } 48 | 49 | fn push(&mut self, value: Value) { 50 | let inner_column: &mut dyn ColumnData = Arc::get_mut(&mut self.inner).unwrap(); 51 | inner_column.push(value); 52 | } 53 | 54 | fn at(&self, index: usize) -> ValueRef { 55 | self.inner.at(index) 56 | } 57 | 58 | fn clone_instance(&self) -> BoxColumnData { 59 | Box::new(Self { 60 | inner: self.inner.clone(), 61 | func: self.func, 62 | }) 63 | } 64 | 65 | unsafe fn get_internal( 66 | &self, 67 | pointers: &[*mut *const u8], 68 | level: u8, 69 | props: u32, 70 | ) -> Result<()> { 71 | self.inner.get_internal(pointers, level, props) 72 | } 73 | 74 | fn cast_to(&self, _this: &ArcColumnData, target: &SqlType) -> Option { 75 | if let SqlType::SimpleAggregateFunction(func, inner_target) = target { 76 | if let Some(inner) = self.inner.cast_to(&self.inner, inner_target) { 77 | return Some(Arc::new(SimpleAggregateFunctionColumnData { 78 | inner, 79 | func: *func, 80 | })); 81 | } 82 | } 83 | None 84 | } 85 | 86 | fn get_timezone(&self) -> Option { 87 | self.inner.get_timezone() 88 | } 89 | } 90 | -------------------------------------------------------------------------------- /src/types/column/string.rs: -------------------------------------------------------------------------------- 1 | use chrono_tz::Tz; 2 | use std::{io::Write, string::ToString, sync::Arc}; 3 | 4 | use either::Either; 5 | 6 | use crate::{ 7 | binary::{Encoder, ReadEx}, 8 | errors::Result, 9 | types::{ 10 | column::{ 11 | array::ArrayColumnData, column_data::LowCardinalityAccessor, list::List, 12 | nullable::NullableColumnData, ArcColumnWrapper, ColumnWrapper, StringPool, 13 | }, 14 | Column, ColumnType, FromSql, SqlType, Value, ValueRef, 15 | }, 16 | }; 17 | 18 | use super::{ 19 | column_data::{BoxColumnData, ColumnData}, 20 | ColumnFrom, 21 | }; 22 | 23 | pub(crate) struct StringColumnData { 24 | pool: StringPool, 25 | } 26 | 27 | pub(crate) struct StringAdapter { 28 | pub(crate) column: Column, 29 | } 30 | 31 | impl StringColumnData { 32 | pub(crate) fn with_capacity(capacity: usize) -> Self { 33 | Self { 34 | pool: StringPool::with_capacity(capacity), 35 | } 36 | } 37 | 38 | pub(crate) fn load(reader: &mut T, size: usize) -> Result { 39 | let mut data = Self::with_capacity(size); 40 | 41 | for _ in 0..size { 42 | reader.read_str_into_buffer(&mut data.pool)?; 43 | } 44 | 45 | Ok(data) 46 | } 47 | } 48 | 49 | impl ColumnFrom for Vec { 50 | fn column_from(data: Self) -> W::Wrapper { 51 | W::wrap(StringColumnData { pool: data.into() }) 52 | } 53 | } 54 | 55 | impl<'a> ColumnFrom for Vec<&'a str> { 56 | fn column_from(source: Self) -> W::Wrapper { 57 | let data: Vec<_> = source.iter().map(ToString::to_string).collect(); 58 | W::wrap(StringColumnData { pool: data.into() }) 59 | } 60 | } 61 | 62 | impl<'a> ColumnFrom for Vec<&'a [u8]> { 63 | fn column_from(data: Self) -> W::Wrapper { 64 | W::wrap(StringColumnData { pool: data.into() }) 65 | } 66 | } 67 | 68 | trait StringSource { 69 | fn into_value(self) -> Value; 70 | } 71 | 72 | impl StringSource for String { 73 | fn into_value(self) -> Value { 74 | self.into() 75 | } 76 | } 77 | 78 | impl StringSource for &str { 79 | fn into_value(self) -> Value { 80 | self.into() 81 | } 82 | } 83 | 84 | impl StringSource for Vec { 85 | fn into_value(self) -> Value { 86 | Value::String(Arc::new(self)) 87 | } 88 | } 89 | 90 | impl ColumnFrom for Vec> { 91 | fn column_from(source: Self) -> ::Wrapper { 92 | make_array_of_array::(source) 93 | } 94 | } 95 | 96 | impl ColumnFrom for Vec> { 97 | fn column_from(source: Self) -> ::Wrapper { 98 | make_array_of_array::(source) 99 | } 100 | } 101 | 102 | fn make_array_of_array( 103 | source: Vec>, 104 | ) -> ::Wrapper { 105 | let fake: Vec = Vec::with_capacity(source.len()); 106 | let inner = Vec::column_from::(fake); 107 | let sql_type = inner.sql_type(); 108 | 109 | let mut data = ArrayColumnData { 110 | inner, 111 | offsets: List::with_capacity(source.len()), 112 | }; 113 | 114 | for vs in source { 115 | let mut inner = Vec::with_capacity(vs.len()); 116 | for v in vs { 117 | let value: Value = v.into_value(); 118 | inner.push(value) 119 | } 120 | data.push(Value::Array(sql_type.clone().into(), Arc::new(inner))); 121 | } 122 | 123 | W::wrap(data) 124 | } 125 | 126 | impl ColumnFrom for Vec>> { 127 | fn column_from(source: Self) -> W::Wrapper { 128 | make_opt_column::>(source) 129 | } 130 | } 131 | 132 | impl ColumnFrom for Vec> { 133 | fn column_from(source: Self) -> W::Wrapper { 134 | make_opt_column::(source) 135 | } 136 | } 137 | 138 | impl ColumnFrom for Vec> { 139 | fn column_from(source: Self) -> W::Wrapper { 140 | make_opt_column::(source) 141 | } 142 | } 143 | 144 | fn make_opt_column(source: Vec>) -> W::Wrapper { 145 | let inner = Arc::new(StringColumnData::with_capacity(source.len())); 146 | 147 | let mut data = NullableColumnData { 148 | inner, 149 | nulls: Vec::with_capacity(source.len()), 150 | }; 151 | 152 | for value in source { 153 | let item = if let Some(v) = value { 154 | let inner = v.into_value(); 155 | Value::Nullable(Either::Right(Box::new(inner))) 156 | } else { 157 | Value::Nullable(Either::Left(SqlType::String.into())) 158 | }; 159 | data.push(item); 160 | } 161 | 162 | W::wrap(data) 163 | } 164 | 165 | impl LowCardinalityAccessor for StringColumnData { 166 | fn get_string(&self, index: usize) -> &[u8] { 167 | self.pool.get(index) 168 | } 169 | } 170 | 171 | impl ColumnData for StringColumnData { 172 | fn sql_type(&self) -> SqlType { 173 | SqlType::String 174 | } 175 | 176 | fn save(&self, encoder: &mut Encoder, start: usize, end: usize) { 177 | let strings = self.pool.strings(); 178 | for v in strings.skip(start).take(end - start) { 179 | encoder.byte_string(v); 180 | } 181 | } 182 | 183 | fn len(&self) -> usize { 184 | self.pool.len() 185 | } 186 | 187 | fn push(&mut self, value: Value) { 188 | let s: Vec = value.into(); 189 | let mut b = self.pool.allocate(s.len()); 190 | b.write_all(s.as_ref()).unwrap(); 191 | } 192 | 193 | fn at(&self, index: usize) -> ValueRef { 194 | let s = self.pool.get(index); 195 | ValueRef::from(s) 196 | } 197 | 198 | fn clone_instance(&self) -> BoxColumnData { 199 | Box::new(Self { 200 | pool: self.pool.clone(), 201 | }) 202 | } 203 | 204 | unsafe fn get_internal( 205 | &self, 206 | pointers: &[*mut *const u8], 207 | level: u8, 208 | _props: u32, 209 | ) -> Result<()> { 210 | assert_eq!(level, 0); 211 | *pointers[0] = &self.pool as *const StringPool as *const u8; 212 | *(pointers[1] as *mut usize) = self.len(); 213 | Ok(()) 214 | } 215 | 216 | fn get_timezone(&self) -> Option { 217 | None 218 | } 219 | 220 | fn get_low_cardinality_accessor(&self) -> Option<&dyn LowCardinalityAccessor> { 221 | Some(self) 222 | } 223 | } 224 | 225 | impl ColumnData for StringAdapter { 226 | fn sql_type(&self) -> SqlType { 227 | SqlType::String 228 | } 229 | 230 | fn save(&self, encoder: &mut Encoder, start: usize, end: usize) { 231 | for index in start..end { 232 | let buf: Vec = Vec::from_sql(self.column.at(index)).unwrap(); 233 | encoder.byte_string(buf); 234 | } 235 | } 236 | 237 | fn len(&self) -> usize { 238 | self.column.len() 239 | } 240 | 241 | fn push(&mut self, _value: Value) { 242 | unimplemented!() 243 | } 244 | 245 | fn at(&self, index: usize) -> ValueRef { 246 | self.column.at(index) 247 | } 248 | 249 | fn clone_instance(&self) -> BoxColumnData { 250 | unimplemented!() 251 | } 252 | 253 | fn get_timezone(&self) -> Option { 254 | None 255 | } 256 | } 257 | -------------------------------------------------------------------------------- /src/types/column/string_pool.rs: -------------------------------------------------------------------------------- 1 | use std::{io::Write, slice}; 2 | 3 | const AVG_STR_SIZE: usize = 80; 4 | 5 | #[derive(Copy, Clone)] 6 | struct StringPtr { 7 | chunk: usize, 8 | shift: usize, 9 | len: usize, 10 | } 11 | 12 | #[derive(Clone)] 13 | pub(crate) struct StringPool { 14 | chunks: Vec>, 15 | pointers: Vec, 16 | position: usize, 17 | capacity: usize, 18 | } 19 | 20 | pub(crate) struct StringIter<'a> { 21 | pool: &'a StringPool, 22 | index: usize, 23 | } 24 | 25 | impl<'a> Iterator for StringIter<'a> { 26 | type Item = &'a [u8]; 27 | 28 | fn next(&mut self) -> Option { 29 | if self.index < self.pool.len() { 30 | let result = self.pool.get(self.index); 31 | self.index += 1; 32 | return Some(result); 33 | } 34 | 35 | None 36 | } 37 | } 38 | 39 | impl From> for StringPool 40 | where 41 | T: AsRef<[u8]>, 42 | { 43 | fn from(source: Vec) -> Self { 44 | let mut pool = StringPool::with_capacity(source.len()); 45 | for s in source.iter() { 46 | let mut b = pool.allocate(s.as_ref().len()); 47 | b.write_all(s.as_ref()).unwrap(); 48 | } 49 | pool 50 | } 51 | } 52 | 53 | impl StringPool { 54 | pub(crate) fn with_capacity(capacity: usize) -> StringPool { 55 | StringPool { 56 | pointers: Vec::with_capacity(capacity), 57 | chunks: Vec::new(), 58 | position: 0, 59 | capacity, 60 | } 61 | } 62 | 63 | pub(crate) fn allocate(&mut self, size: usize) -> &mut [u8] { 64 | if self.free_space() < size || self.chunks.is_empty() { 65 | self.reserve(size); 66 | return self.allocate(size); 67 | } 68 | 69 | self.try_allocate(size).unwrap() 70 | } 71 | 72 | fn free_space(&self) -> usize { 73 | if let Some(buffer) = self.chunks.last() { 74 | return buffer.len() - self.position; 75 | } 76 | 77 | 0 78 | } 79 | 80 | fn try_allocate(&mut self, size: usize) -> Option<&mut [u8]> { 81 | if !self.chunks.is_empty() { 82 | let chunk = self.chunks.len() - 1; 83 | 84 | let position = self.position; 85 | self.position += size; 86 | self.pointers.push(StringPtr { 87 | len: size, 88 | shift: position, 89 | chunk, 90 | }); 91 | 92 | let buffer = &mut self.chunks[chunk]; 93 | return Some(&mut buffer[position..position + size]); 94 | } 95 | 96 | None 97 | } 98 | 99 | fn reserve(&mut self, size: usize) { 100 | use std::cmp::max; 101 | self.position = 0; 102 | self.chunks 103 | .push(vec![0_u8; max(self.capacity * AVG_STR_SIZE, size)]); 104 | } 105 | 106 | #[inline(always)] 107 | pub(crate) fn get(&self, index: usize) -> &[u8] { 108 | let pointer = &self.pointers[index]; 109 | unsafe { self.get_by_pointer(pointer) } 110 | } 111 | 112 | #[inline(always)] 113 | pub(crate) unsafe fn get_unchecked(&self, index: usize) -> &[u8] { 114 | let pointer = self.pointers.get_unchecked(index); 115 | self.get_by_pointer(pointer) 116 | } 117 | 118 | #[inline(always)] 119 | unsafe fn get_by_pointer(&self, pointer: &StringPtr) -> &[u8] { 120 | let chunk = &self.chunks.get_unchecked(pointer.chunk); 121 | 122 | let ptr = chunk.as_ptr().add(pointer.shift); 123 | slice::from_raw_parts(ptr, pointer.len) 124 | } 125 | 126 | #[inline(always)] 127 | pub(crate) fn len(&self) -> usize { 128 | self.pointers.len() 129 | } 130 | 131 | pub(crate) fn strings(&self) -> StringIter { 132 | StringIter { 133 | pool: self, 134 | index: 0, 135 | } 136 | } 137 | } 138 | 139 | #[cfg(test)] 140 | mod test { 141 | use super::*; 142 | use std::io::Write; 143 | 144 | #[test] 145 | fn test_allocate() { 146 | let mut pool = StringPool::with_capacity(10); 147 | for i in 1..1000 { 148 | let buffer = pool.allocate(i); 149 | assert_eq!(buffer.len(), i); 150 | assert_eq!(buffer[0], 0); 151 | buffer[0] = 1 152 | } 153 | } 154 | 155 | #[test] 156 | fn test_get() { 157 | let mut pool = StringPool::with_capacity(10); 158 | 159 | for i in 0..1000 { 160 | let s = format!("text-{i}"); 161 | let mut buffer = pool.allocate(s.len()); 162 | 163 | assert_eq!(buffer.len(), s.len()); 164 | buffer.write_all(s.as_bytes()).unwrap(); 165 | } 166 | 167 | for i in 0..1000 { 168 | let s = String::from_utf8(Vec::from(pool.get(i))).unwrap(); 169 | assert_eq!(s, format!("text-{i}")); 170 | } 171 | } 172 | } 173 | -------------------------------------------------------------------------------- /src/types/column/util.rs: -------------------------------------------------------------------------------- 1 | use crate::{ 2 | types::{column::ColumnData, FromSql}, 3 | Result, 4 | }; 5 | 6 | pub(crate) fn extract_nulls_and_values<'a, T, C>( 7 | column_data: &'a C, 8 | start: usize, 9 | end: usize, 10 | ) -> Result<(Vec, Vec>)> 11 | where 12 | C: ColumnData, 13 | T: FromSql<'a>, 14 | { 15 | let size = end - start; 16 | let mut nulls = vec![0_u8; size]; 17 | let mut values: Vec> = Vec::with_capacity(size); 18 | 19 | for (i, index) in (start..end).enumerate() { 20 | let value = Option::from_sql(column_data.at(index))?; 21 | if value.is_none() { 22 | nulls[i] = 1; 23 | } 24 | values.push(value); 25 | } 26 | 27 | Ok((nulls, values)) 28 | } 29 | -------------------------------------------------------------------------------- /src/types/date_converter.rs: -------------------------------------------------------------------------------- 1 | use chrono::prelude::*; 2 | use chrono_tz::Tz; 3 | 4 | use crate::types::{DateTimeType, SqlType, Value, ValueRef}; 5 | 6 | pub trait DateConverter { 7 | fn to_date(&self, tz: Tz) -> ValueRef<'static>; 8 | fn get_stamp(source: Value) -> Self; 9 | fn date_type() -> SqlType; 10 | 11 | fn get_days(date: NaiveDate) -> u16 { 12 | const UNIX_EPOCH_DAY: i64 = 719_163; 13 | let gregorian_day = i64::from(date.num_days_from_ce()); 14 | (gregorian_day - UNIX_EPOCH_DAY) as u16 15 | } 16 | } 17 | 18 | impl DateConverter for u16 { 19 | fn to_date(&self, _tz: Tz) -> ValueRef<'static> { 20 | ValueRef::Date(*self) 21 | } 22 | 23 | fn get_stamp(source: Value) -> Self { 24 | Self::get_days(NaiveDate::from(source)) 25 | } 26 | 27 | fn date_type() -> SqlType { 28 | SqlType::Date 29 | } 30 | } 31 | 32 | impl DateConverter for u32 { 33 | fn to_date(&self, tz: Tz) -> ValueRef<'static> { 34 | ValueRef::DateTime(*self, tz) 35 | } 36 | 37 | fn get_stamp(source: Value) -> Self { 38 | DateTime::::from(source).timestamp() as Self 39 | } 40 | 41 | fn date_type() -> SqlType { 42 | SqlType::DateTime(DateTimeType::DateTime32) 43 | } 44 | } 45 | -------------------------------------------------------------------------------- /src/types/decimal.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | cmp::Ordering, 3 | fmt, 4 | hash::{Hash, Hasher}, 5 | }; 6 | 7 | static FACTORS10: &[i64] = &[ 8 | 1, 9 | 10, 10 | 100, 11 | 1000, 12 | 10000, 13 | 100_000, 14 | 1_000_000, 15 | 10_000_000, 16 | 100_000_000, 17 | 1_000_000_000, 18 | 10_000_000_000, 19 | 100_000_000_000, 20 | 1_000_000_000_000, 21 | 10_000_000_000_000, 22 | 100_000_000_000_000, 23 | 1_000_000_000_000_000, 24 | 10_000_000_000_000_000, 25 | 100_000_000_000_000_000, 26 | 1_000_000_000_000_000_000, 27 | ]; 28 | 29 | pub trait Base { 30 | fn scale(self, scale: i64) -> i64; 31 | } 32 | 33 | pub trait InternalResult { 34 | fn get(underlying: i64) -> Self; 35 | } 36 | 37 | #[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)] 38 | pub(crate) enum NoBits { 39 | N32, 40 | N64, 41 | } 42 | 43 | /// Provides arbitrary-precision floating point decimal. 44 | #[derive(Clone)] 45 | pub struct Decimal { 46 | pub(crate) underlying: i64, 47 | pub(crate) nobits: NoBits, // its domain is {32, 64} 48 | pub(crate) precision: u8, 49 | pub(crate) scale: u8, 50 | } 51 | 52 | impl Default for Decimal { 53 | fn default() -> Self { 54 | Decimal { 55 | underlying: 0, 56 | precision: 9, 57 | scale: 4, 58 | nobits: NoBits::N32, 59 | } 60 | } 61 | } 62 | 63 | impl Hash for Decimal { 64 | fn hash(&self, state: &mut H) { 65 | self.underlying.hash(state); 66 | self.nobits.hash(state); 67 | self.precision.hash(state); 68 | self.scale.hash(state); 69 | } 70 | } 71 | 72 | macro_rules! base_for { 73 | ( $( $t:ty: $cast:expr ),* ) => { 74 | $( 75 | impl Base for $t { 76 | fn scale(self, scale: i64) -> i64 { 77 | $cast(self * (scale as $t)) as i64 78 | } 79 | } 80 | )* 81 | }; 82 | } 83 | 84 | base_for! { 85 | f32: std::convert::identity, 86 | f64: std::convert::identity, 87 | i8: i64::from, 88 | i16: i64::from, 89 | i32: i64::from, 90 | i64: std::convert::identity, 91 | u8: i64::from, 92 | u16: i64::from, 93 | u32: i64::from, 94 | u64 : std::convert::identity 95 | } 96 | 97 | impl InternalResult for i32 { 98 | #[inline(always)] 99 | fn get(underlying: i64) -> Self { 100 | underlying as Self 101 | } 102 | } 103 | 104 | impl InternalResult for i64 { 105 | #[inline(always)] 106 | fn get(underlying: i64) -> Self { 107 | underlying 108 | } 109 | } 110 | 111 | impl NoBits { 112 | pub(crate) fn from_precision(precision: u8) -> Option { 113 | if precision <= 9 { 114 | Some(NoBits::N32) 115 | } else if precision <= 18 { 116 | Some(NoBits::N64) 117 | } else { 118 | None 119 | } 120 | } 121 | } 122 | 123 | impl PartialEq for Decimal { 124 | fn eq(&self, other: &Self) -> bool { 125 | match self.scale.cmp(&other.scale) { 126 | Ordering::Less => { 127 | let delta = other.scale() - self.scale(); 128 | let underlying = self.underlying * FACTORS10[delta]; 129 | other.underlying == underlying 130 | } 131 | Ordering::Equal => self.underlying == other.underlying, 132 | Ordering::Greater => { 133 | let delta = self.scale() - other.scale(); 134 | let underlying = other.underlying * FACTORS10[delta]; 135 | self.underlying == underlying 136 | } 137 | } 138 | } 139 | } 140 | 141 | fn decimal2str(decimal: &Decimal) -> String { 142 | let mut r = format!("{}", decimal.underlying); 143 | while r.len() < decimal.scale() { 144 | r.insert(0, '0'); 145 | } 146 | let pos = r.len() - decimal.scale(); 147 | r.insert(pos, '.'); 148 | if r.starts_with('.') { 149 | r.insert(0, '0'); 150 | } 151 | r 152 | } 153 | 154 | impl fmt::Display for Decimal { 155 | fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> { 156 | write!(f, "{}", decimal2str(self)) 157 | } 158 | } 159 | 160 | impl fmt::Debug for Decimal { 161 | fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> { 162 | write!(f, "{}", decimal2str(self)) 163 | } 164 | } 165 | 166 | impl From for f32 { 167 | fn from(value: Decimal) -> Self { 168 | value.underlying as f32 / FACTORS10[value.scale()] as f32 169 | } 170 | } 171 | 172 | impl From for f64 { 173 | fn from(value: Decimal) -> Self { 174 | value.underlying as f64 / FACTORS10[value.scale()] as f64 175 | } 176 | } 177 | 178 | impl Decimal { 179 | /// Method of creating a Decimal. 180 | pub fn new(underlying: i64, scale: u8) -> Decimal { 181 | let precision = 18; 182 | if scale > precision { 183 | panic!("scale can't be greater than 18"); 184 | } 185 | 186 | Decimal { 187 | underlying, 188 | precision, 189 | scale, 190 | nobits: NoBits::N64, 191 | } 192 | } 193 | 194 | pub fn of(source: B, scale: u8) -> Decimal { 195 | let precision = 18; 196 | if scale > precision { 197 | panic!("scale can't be greater than 18"); 198 | } 199 | 200 | let underlying = source.scale(FACTORS10[scale as usize]); 201 | if underlying > FACTORS10[precision as usize] { 202 | panic!("{} > {}", underlying, FACTORS10[precision as usize]); 203 | } 204 | 205 | Decimal { 206 | underlying, 207 | precision, 208 | scale, 209 | nobits: NoBits::N64, 210 | } 211 | } 212 | 213 | /// Get the internal representation of decimal as [`i32`] or [`i64`]. 214 | /// 215 | /// example: 216 | /// ```rust 217 | /// # use std::env; 218 | /// # use clickhouse_rs::{Pool, types::Decimal, errors::Result}; 219 | /// # let mut rt = tokio::runtime::Runtime::new().unwrap(); 220 | /// # let ret: Result<()> = rt.block_on(async { 221 | /// # let database_url = env::var("DATABASE_URL") 222 | /// # .unwrap_or("tcp://localhost:9000?compression=lz4".into()); 223 | /// # let pool = Pool::new(database_url); 224 | /// let mut c = pool.get_handle().await?; 225 | /// let block = c.query("SELECT toDecimal32(2, 4) AS x").fetch_all().await?; 226 | /// 227 | /// let x: Decimal = block.get(0, "x")?; 228 | /// let actual: i32 = x.internal(); 229 | /// assert_eq!(20000, actual); 230 | /// # Ok(()) 231 | /// # }); 232 | /// # ret.unwrap() 233 | /// ``` 234 | #[inline(always)] 235 | pub fn internal(&self) -> I { 236 | InternalResult::get(self.underlying) 237 | } 238 | 239 | /// Determines how many decimal digits fraction can have. 240 | pub fn scale(&self) -> usize { 241 | self.scale as usize 242 | } 243 | 244 | pub(crate) fn set_scale(self, scale: u8) -> Self { 245 | let underlying = match scale.cmp(&self.scale) { 246 | Ordering::Less => { 247 | let delta = self.scale() - scale as usize; 248 | self.underlying / FACTORS10[delta] 249 | } 250 | Ordering::Equal => return self, 251 | Ordering::Greater => { 252 | let delta = scale as usize - self.scale(); 253 | self.underlying * FACTORS10[delta] 254 | } 255 | }; 256 | 257 | Decimal { 258 | underlying, 259 | precision: self.precision, 260 | scale, 261 | nobits: self.nobits, 262 | } 263 | } 264 | } 265 | 266 | #[cfg(test)] 267 | mod test { 268 | use super::*; 269 | 270 | #[test] 271 | fn test_new() { 272 | assert_eq!(Decimal::new(2, 1), Decimal::of(0.2_f64, 1)); 273 | assert_eq!(Decimal::new(2, 5), Decimal::of(0.00002_f64, 5)); 274 | } 275 | 276 | #[test] 277 | fn test_display() { 278 | assert_eq!(format!("{}", Decimal::of(2.1_f32, 4)), "2.1000"); 279 | assert_eq!(format!("{}", Decimal::of(0.2_f64, 4)), "0.2000"); 280 | assert_eq!(format!("{}", Decimal::of(2, 4)), "2.0000"); 281 | assert_eq!(format!("{:?}", Decimal::of(2, 4)), "2.0000"); 282 | } 283 | 284 | #[test] 285 | fn test_eq() { 286 | assert_eq!(Decimal::of(2.0_f64, 4), Decimal::of(2.0_f64, 4)); 287 | assert_ne!(Decimal::of(3.0_f64, 4), Decimal::of(2.0_f64, 4)); 288 | 289 | assert_eq!(Decimal::of(2.0_f64, 4), Decimal::of(2.0_f64, 2)); 290 | assert_ne!(Decimal::of(2.0_f64, 4), Decimal::of(3.0_f64, 2)); 291 | 292 | assert_eq!(Decimal::of(2.0_f64, 2), Decimal::of(2.0_f64, 4)); 293 | assert_ne!(Decimal::of(3.0_f64, 2), Decimal::of(2.0_f64, 4)); 294 | } 295 | 296 | #[test] 297 | fn test_internal32() { 298 | let internal: i32 = Decimal::of(2, 4).internal(); 299 | assert_eq!(internal, 20000_i32); 300 | } 301 | 302 | #[test] 303 | fn test_internal64() { 304 | let internal: i64 = Decimal::of(2, 4).internal(); 305 | assert_eq!(internal, 20000_i64); 306 | } 307 | 308 | #[test] 309 | fn test_scale() { 310 | assert_eq!(Decimal::of(2, 4).scale(), 4); 311 | } 312 | 313 | #[test] 314 | fn test_from_f32() { 315 | let value: f32 = Decimal::of(2, 4).into(); 316 | assert!((value - 2.0_f32).abs() <= std::f32::EPSILON); 317 | } 318 | 319 | #[test] 320 | fn test_from_f64() { 321 | let value: f64 = Decimal::of(2, 4).into(); 322 | assert!((value - 2.0_f64).abs() < std::f64::EPSILON); 323 | } 324 | 325 | #[test] 326 | fn set_scale1() { 327 | let a = Decimal::of(12, 3); 328 | let b = a.set_scale(2); 329 | 330 | assert_eq!(2, b.scale); 331 | assert_eq!(1200, b.underlying); 332 | } 333 | 334 | #[test] 335 | fn set_scale2() { 336 | let a = Decimal::of(12, 3); 337 | let b = a.set_scale(4); 338 | 339 | assert_eq!(4, b.scale); 340 | assert_eq!(120_000, b.underlying); 341 | } 342 | 343 | #[test] 344 | fn test_decimal2str() { 345 | let d = Decimal::of(0.00001, 5); 346 | let actual = decimal2str(&d); 347 | assert_eq!(actual, "0.00001".to_string()); 348 | } 349 | } 350 | -------------------------------------------------------------------------------- /src/types/enums.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | fmt, 3 | hash::{Hash, Hasher}, 4 | }; 5 | 6 | // TODO Using strings as a keys 7 | #[derive(Clone, Copy, Default)] 8 | pub struct Enum8(pub(crate) i8); 9 | 10 | #[derive(Clone, Copy, Default)] 11 | pub struct Enum16(pub(crate) i16); 12 | 13 | impl PartialEq for Enum16 { 14 | fn eq(&self, other: &Self) -> bool { 15 | self.0 == other.0 16 | } 17 | } 18 | 19 | impl fmt::Display for Enum16 { 20 | fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> { 21 | write!(f, "Enum({})", self.0) 22 | } 23 | } 24 | 25 | impl fmt::Debug for Enum16 { 26 | fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> { 27 | write!(f, "Enum({})", self.0) 28 | } 29 | } 30 | 31 | impl Enum16 { 32 | pub fn of(source: i16) -> Self { 33 | Self(source) 34 | } 35 | #[inline(always)] 36 | pub fn internal(self) -> i16 { 37 | self.0 38 | } 39 | } 40 | 41 | impl Hash for Enum16 { 42 | fn hash(&self, state: &mut H) { 43 | self.0.hash(state); 44 | } 45 | } 46 | 47 | impl PartialEq for Enum8 { 48 | fn eq(&self, other: &Self) -> bool { 49 | self.0 == other.0 50 | } 51 | } 52 | 53 | impl Hash for Enum8 { 54 | fn hash(&self, state: &mut H) { 55 | self.0.hash(state); 56 | } 57 | } 58 | 59 | impl fmt::Display for Enum8 { 60 | fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> { 61 | write!(f, "Enum8({})", self.0) 62 | } 63 | } 64 | 65 | impl fmt::Debug for Enum8 { 66 | fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> { 67 | write!(f, "Enum8({})", self.0) 68 | } 69 | } 70 | 71 | impl Enum8 { 72 | pub fn of(source: i8) -> Self { 73 | Self(source) 74 | } 75 | #[inline(always)] 76 | pub fn internal(self) -> i8 { 77 | self.0 78 | } 79 | } 80 | -------------------------------------------------------------------------------- /src/types/marshal.rs: -------------------------------------------------------------------------------- 1 | pub trait Marshal { 2 | fn marshal(&self, scratch: &mut [u8]); 3 | } 4 | 5 | macro_rules! int_marshals { 6 | ( $( $t:ident ),* ) => { 7 | $( 8 | impl Marshal for $t { 9 | fn marshal(&self, scratch: &mut [u8]) { 10 | scratch.clone_from_slice(&self.to_le_bytes()); 11 | } 12 | } 13 | )* 14 | }; 15 | } 16 | 17 | macro_rules! float_marshals { 18 | ( $( $t:ident ),* ) => { 19 | $( 20 | impl Marshal for $t { 21 | fn marshal(&self, scratch: &mut [u8]) { 22 | let bits = self.to_bits(); 23 | scratch.clone_from_slice(&bits.to_le_bytes()); 24 | } 25 | } 26 | )* 27 | }; 28 | } 29 | 30 | int_marshals! { u8, u16, u32, u64, u128, i8, i16, i32, i64, i128 } 31 | float_marshals! { f32, f64 } 32 | 33 | impl Marshal for bool { 34 | fn marshal(&self, scratch: &mut [u8]) { 35 | scratch[0] = *self as u8; 36 | } 37 | } 38 | 39 | #[cfg(test)] 40 | mod test { 41 | use std::fmt; 42 | 43 | use crate::types::{Marshal, StatBuffer, Unmarshal}; 44 | use rand::distributions::{Distribution, Standard}; 45 | use rand::random; 46 | 47 | fn test_some() 48 | where 49 | T: Copy + fmt::Debug + StatBuffer + Marshal + Unmarshal + PartialEq, 50 | Standard: Distribution, 51 | { 52 | for _ in 0..100 { 53 | let mut buffer = T::buffer(); 54 | let v = random::(); 55 | 56 | v.marshal(buffer.as_mut()); 57 | let u = T::unmarshal(buffer.as_ref()); 58 | 59 | assert_eq!(v, u); 60 | } 61 | } 62 | 63 | #[test] 64 | fn test_u8() { 65 | test_some::() 66 | } 67 | 68 | #[test] 69 | fn test_u16() { 70 | test_some::() 71 | } 72 | 73 | #[test] 74 | fn test_u32() { 75 | test_some::() 76 | } 77 | 78 | #[test] 79 | fn test_u64() { 80 | test_some::() 81 | } 82 | 83 | #[test] 84 | fn test_u128() { 85 | test_some::() 86 | } 87 | 88 | #[test] 89 | fn test_i8() { 90 | test_some::() 91 | } 92 | 93 | #[test] 94 | fn test_i16() { 95 | test_some::() 96 | } 97 | 98 | #[test] 99 | fn test_i32() { 100 | test_some::() 101 | } 102 | 103 | #[test] 104 | fn test_i64() { 105 | test_some::() 106 | } 107 | 108 | #[test] 109 | fn test_i128() { 110 | test_some::() 111 | } 112 | 113 | #[test] 114 | fn test_f32() { 115 | test_some::() 116 | } 117 | 118 | #[test] 119 | fn test_f64() { 120 | test_some::() 121 | } 122 | 123 | #[test] 124 | fn test_bool() { 125 | test_some::() 126 | } 127 | } 128 | -------------------------------------------------------------------------------- /src/types/query.rs: -------------------------------------------------------------------------------- 1 | #[derive(Clone, Debug)] 2 | pub struct Query { 3 | sql: String, 4 | id: String, 5 | } 6 | 7 | impl Query { 8 | pub fn new(sql: impl AsRef) -> Self { 9 | Self { 10 | sql: sql.as_ref().to_string(), 11 | id: "".to_string(), 12 | } 13 | } 14 | 15 | pub fn id(self, id: impl AsRef) -> Self { 16 | Self { 17 | id: id.as_ref().to_string(), 18 | ..self 19 | } 20 | } 21 | 22 | pub(crate) fn get_sql(&self) -> &str { 23 | &self.sql 24 | } 25 | 26 | pub(crate) fn get_id(&self) -> &str { 27 | &self.id 28 | } 29 | 30 | pub(crate) fn map_sql(self, f: F) -> Self 31 | where 32 | F: Fn(&str) -> String, 33 | { 34 | Self { 35 | sql: f(&self.sql), 36 | ..self 37 | } 38 | } 39 | } 40 | 41 | impl From for Query 42 | where 43 | T: AsRef, 44 | { 45 | fn from(source: T) -> Self { 46 | Self::new(source) 47 | } 48 | } 49 | -------------------------------------------------------------------------------- /src/types/query_result/mod.rs: -------------------------------------------------------------------------------- 1 | use futures_util::{ 2 | future, 3 | stream::{self, BoxStream, StreamExt}, 4 | TryStreamExt, 5 | }; 6 | use log::info; 7 | use std::{marker::PhantomData, sync::Arc}; 8 | 9 | use crate::{ 10 | errors::Result, 11 | try_opt, 12 | types::{ 13 | block::BlockRef, query_result::stream_blocks::BlockStream, Block, Cmd, Complex, Query, Row, 14 | Rows, Simple, 15 | }, 16 | with_timeout, ClientHandle, 17 | }; 18 | 19 | pub(crate) mod stream_blocks; 20 | 21 | /// Result of a query or statement execution. 22 | pub struct QueryResult<'a> { 23 | pub(crate) client: &'a mut ClientHandle, 24 | pub(crate) query: Query, 25 | } 26 | 27 | impl<'a> QueryResult<'a> { 28 | /// Fetch data from table. It returns a block that contains all rows. 29 | pub async fn fetch_all(self) -> Result> { 30 | let timeout = try_opt!(self.client.context.options.get()).query_timeout; 31 | 32 | with_timeout( 33 | async { 34 | let blocks = self 35 | .stream_blocks_(false) 36 | .try_fold(Vec::new(), |mut blocks, block| { 37 | if !block.is_empty() { 38 | blocks.push(block); 39 | } 40 | future::ready(Ok(blocks)) 41 | }) 42 | .await?; 43 | Ok(Block::concat(blocks.as_slice())) 44 | }, 45 | timeout, 46 | ) 47 | .await 48 | } 49 | 50 | /// Method that produces a stream of blocks containing rows 51 | /// 52 | /// example: 53 | /// 54 | /// ```rust 55 | /// # use std::env; 56 | /// # use clickhouse_rs::{Pool, errors::Result}; 57 | /// # use futures_util::{future, TryStreamExt}; 58 | /// # 59 | /// # let mut rt = tokio::runtime::Runtime::new().unwrap(); 60 | /// # let ret: Result<()> = rt.block_on(async { 61 | /// # 62 | /// # let database_url = env::var("DATABASE_URL") 63 | /// # .unwrap_or("tcp://localhost:9000?compression=lz4".into()); 64 | /// # 65 | /// # let sql_query = "SELECT number FROM system.numbers LIMIT 100000"; 66 | /// # let pool = Pool::new(database_url); 67 | /// # 68 | /// let mut c = pool.get_handle().await?; 69 | /// let mut result = c.query(sql_query) 70 | /// .stream_blocks() 71 | /// .try_for_each(|block| { 72 | /// println!("{:?}\nblock counts: {} rows", block, block.row_count()); 73 | /// future::ready(Ok(())) 74 | /// }).await?; 75 | /// # Ok(()) 76 | /// # }); 77 | /// # ret.unwrap() 78 | /// ``` 79 | pub fn stream_blocks(self) -> BoxStream<'a, Result> { 80 | self.stream_blocks_(true) 81 | } 82 | 83 | fn stream_blocks_(self, skip_first_block: bool) -> BoxStream<'a, Result> { 84 | let query = self.query.clone(); 85 | 86 | self.client 87 | .wrap_stream::<'a, _>(move |c: &'a mut ClientHandle| { 88 | info!("[send query] {}", query.get_sql()); 89 | c.pool.detach(); 90 | 91 | let context = c.context.clone(); 92 | 93 | let inner = c.get_inner()?.call(Cmd::SendQuery(query, context)); 94 | 95 | Ok(BlockStream::<'a>::new(c, inner, skip_first_block)) 96 | }) 97 | } 98 | 99 | /// Method that produces a stream of rows 100 | pub fn stream(self) -> BoxStream<'a, Result>> { 101 | Box::pin( 102 | self.stream_blocks() 103 | .map(|block_ret| { 104 | let result: BoxStream<'a, Result>> = match block_ret { 105 | Ok(block) => { 106 | let block = Arc::new(block); 107 | let block_ref = BlockRef::Owned(block); 108 | 109 | Box::pin( 110 | stream::iter(Rows { 111 | row: 0, 112 | block_ref, 113 | kind: PhantomData, 114 | }) 115 | .map(|row| -> Result> { Ok(row) }), 116 | ) 117 | } 118 | Err(err) => Box::pin(stream::once(future::err(err))), 119 | }; 120 | result 121 | }) 122 | .flatten(), 123 | ) 124 | } 125 | } 126 | -------------------------------------------------------------------------------- /src/types/query_result/stream_blocks.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | borrow::Cow, 3 | io::ErrorKind, 4 | pin::Pin, 5 | task::{self, Poll}, 6 | }; 7 | 8 | use futures_core::Stream; 9 | use futures_util::StreamExt; 10 | 11 | use crate::{ 12 | errors::{DriverError, Error, Result}, 13 | io::transport::PacketStream, 14 | types::{Block, Packet}, 15 | ClientHandle, 16 | }; 17 | 18 | pub(crate) struct BlockStream<'a> { 19 | client: &'a mut ClientHandle, 20 | inner: PacketStream, 21 | state: BlockStreamState, 22 | block_index: usize, 23 | skip_first_block: bool, 24 | } 25 | 26 | #[derive(Clone, Copy)] 27 | pub(crate) enum BlockStreamState { 28 | /// Currently reading from block packet stream; some further packets may be pending 29 | Reading, 30 | /// Completely finished reading from block packet stream; connection is now idle 31 | Finished, 32 | /// There was an error reading packet; transport is broken 33 | Error, 34 | } 35 | 36 | impl<'a> Drop for BlockStream<'a> { 37 | fn drop(&mut self) { 38 | match self.state { 39 | BlockStreamState::Reading => { 40 | if !self.client.pool.is_attached() { 41 | self.client.pool.attach(); 42 | } 43 | 44 | if let Some(mut transport) = self.inner.take_transport() { 45 | transport.inconsistent = true; 46 | self.client.inner = Some(transport); 47 | } 48 | } 49 | BlockStreamState::Finished => {} 50 | BlockStreamState::Error => { 51 | // drop broken transport; don't return it to pool to prevent pool poisoning 52 | } 53 | } 54 | } 55 | } 56 | 57 | impl<'a> BlockStream<'a> { 58 | pub(crate) fn new( 59 | client: &mut ClientHandle, 60 | inner: PacketStream, 61 | skip_first_block: bool, 62 | ) -> BlockStream { 63 | BlockStream { 64 | client, 65 | inner, 66 | state: BlockStreamState::Reading, 67 | block_index: 0, 68 | skip_first_block, 69 | } 70 | } 71 | } 72 | 73 | impl<'a> Stream for BlockStream<'a> { 74 | type Item = Result; 75 | 76 | fn poll_next(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll> { 77 | loop { 78 | match self.state { 79 | BlockStreamState::Reading => {} 80 | BlockStreamState::Finished => return Poll::Ready(None), 81 | BlockStreamState::Error => { 82 | return Poll::Ready(Some(Err(Error::Other(Cow::Borrowed( 83 | "Attempt to read from broken transport", 84 | ))))) 85 | } 86 | }; 87 | 88 | let packet = match self.inner.poll_next_unpin(cx) { 89 | Poll::Ready(Some(Err(err))) => return Poll::Ready(Some(Err(err.into()))), 90 | Poll::Pending => return Poll::Pending, 91 | Poll::Ready(None) => { 92 | self.state = BlockStreamState::Error; 93 | return Poll::Ready(Some(Err(Error::Io(std::io::Error::from( 94 | ErrorKind::UnexpectedEof, 95 | ))))); 96 | } 97 | Poll::Ready(Some(Ok(packet))) => packet, 98 | }; 99 | 100 | match packet { 101 | Packet::Eof(inner) => { 102 | self.client.inner = Some(inner); 103 | if !self.client.pool.is_attached() { 104 | self.client.pool.attach(); 105 | } 106 | self.state = BlockStreamState::Finished; 107 | } 108 | Packet::ProfileInfo(_) | Packet::Progress(_) => {} 109 | Packet::Exception(exception) => { 110 | self.state = BlockStreamState::Finished; 111 | return Poll::Ready(Some(Err(Error::Server(exception)))); 112 | } 113 | Packet::Block(block) => { 114 | self.block_index += 1; 115 | if (self.block_index > 1 || !self.skip_first_block) && !block.is_empty() { 116 | return Poll::Ready(Some(Ok(block))); 117 | } 118 | } 119 | _ => return Poll::Ready(Some(Err(Error::Driver(DriverError::UnexpectedPacket)))), 120 | } 121 | } 122 | } 123 | } 124 | -------------------------------------------------------------------------------- /src/types/stat_buffer.rs: -------------------------------------------------------------------------------- 1 | use crate::types::SqlType; 2 | 3 | pub trait StatBuffer { 4 | type Buffer: AsMut<[u8]> + AsRef<[u8]> + Copy + Sync; 5 | fn buffer() -> Self::Buffer; 6 | fn sql_type() -> SqlType; 7 | } 8 | 9 | impl StatBuffer for u8 { 10 | type Buffer = [Self; 1]; 11 | 12 | fn buffer() -> Self::Buffer { 13 | [0; 1] 14 | } 15 | 16 | fn sql_type() -> SqlType { 17 | SqlType::UInt8 18 | } 19 | } 20 | 21 | impl StatBuffer for u16 { 22 | type Buffer = [u8; 2]; 23 | 24 | fn buffer() -> Self::Buffer { 25 | [0; 2] 26 | } 27 | 28 | fn sql_type() -> SqlType { 29 | SqlType::UInt16 30 | } 31 | } 32 | 33 | impl StatBuffer for u32 { 34 | type Buffer = [u8; 4]; 35 | 36 | fn buffer() -> Self::Buffer { 37 | [0; 4] 38 | } 39 | 40 | fn sql_type() -> SqlType { 41 | SqlType::UInt32 42 | } 43 | } 44 | 45 | impl StatBuffer for u64 { 46 | type Buffer = [u8; 8]; 47 | 48 | fn buffer() -> Self::Buffer { 49 | [0; 8] 50 | } 51 | 52 | fn sql_type() -> SqlType { 53 | SqlType::UInt64 54 | } 55 | } 56 | 57 | impl StatBuffer for u128 { 58 | type Buffer = [u8; 16]; 59 | 60 | fn buffer() -> Self::Buffer { 61 | [0; 16] 62 | } 63 | 64 | fn sql_type() -> SqlType { 65 | SqlType::UInt128 66 | } 67 | } 68 | 69 | impl StatBuffer for i8 { 70 | type Buffer = [u8; 1]; 71 | 72 | fn buffer() -> Self::Buffer { 73 | [0; 1] 74 | } 75 | 76 | fn sql_type() -> SqlType { 77 | SqlType::Int8 78 | } 79 | } 80 | 81 | impl StatBuffer for i16 { 82 | type Buffer = [u8; 2]; 83 | 84 | fn buffer() -> Self::Buffer { 85 | [0; 2] 86 | } 87 | 88 | fn sql_type() -> SqlType { 89 | SqlType::Int16 90 | } 91 | } 92 | 93 | impl StatBuffer for i32 { 94 | type Buffer = [u8; 4]; 95 | 96 | fn buffer() -> Self::Buffer { 97 | [0; 4] 98 | } 99 | 100 | fn sql_type() -> SqlType { 101 | SqlType::Int32 102 | } 103 | } 104 | 105 | impl StatBuffer for i64 { 106 | type Buffer = [u8; 8]; 107 | 108 | fn buffer() -> Self::Buffer { 109 | [0; 8] 110 | } 111 | 112 | fn sql_type() -> SqlType { 113 | SqlType::Int64 114 | } 115 | } 116 | 117 | impl StatBuffer for i128 { 118 | type Buffer = [u8; 16]; 119 | 120 | fn buffer() -> Self::Buffer { 121 | [0; 16] 122 | } 123 | 124 | fn sql_type() -> SqlType { 125 | SqlType::Int128 126 | } 127 | } 128 | 129 | impl StatBuffer for f32 { 130 | type Buffer = [u8; 4]; 131 | 132 | fn buffer() -> Self::Buffer { 133 | [0; 4] 134 | } 135 | 136 | fn sql_type() -> SqlType { 137 | SqlType::Float32 138 | } 139 | } 140 | 141 | impl StatBuffer for f64 { 142 | type Buffer = [u8; 8]; 143 | 144 | fn buffer() -> Self::Buffer { 145 | [0; 8] 146 | } 147 | 148 | fn sql_type() -> SqlType { 149 | SqlType::Float64 150 | } 151 | } 152 | 153 | impl StatBuffer for bool { 154 | type Buffer = [u8; 1]; 155 | 156 | fn buffer() -> Self::Buffer { 157 | [0; 1] 158 | } 159 | 160 | fn sql_type() -> SqlType { 161 | SqlType::Bool 162 | } 163 | } 164 | -------------------------------------------------------------------------------- /src/types/unmarshal.rs: -------------------------------------------------------------------------------- 1 | pub trait Unmarshal { 2 | fn unmarshal(scratch: &[u8]) -> T; 3 | } 4 | 5 | macro_rules! int_unmarshals { 6 | ( $( $t:ident ),* ) => { 7 | $( 8 | impl Unmarshal<$t> for $t { 9 | fn unmarshal(scratch: &[u8]) -> Self { 10 | let mut buffer = [0_u8; std::mem::size_of::()]; 11 | buffer.clone_from_slice(scratch); 12 | Self::from_le_bytes(buffer) 13 | } 14 | } 15 | )* 16 | }; 17 | } 18 | 19 | macro_rules! float_unmarshals { 20 | ( $( $t:ident: $b:ident ),* ) => { 21 | $( 22 | impl Unmarshal<$t> for $t { 23 | fn unmarshal(scratch: &[u8]) -> Self { 24 | let mut buffer = [0_u8; std::mem::size_of::()]; 25 | buffer.clone_from_slice(scratch); 26 | let bits = $b::from_le_bytes(buffer); 27 | Self::from_bits(bits) 28 | } 29 | } 30 | )* 31 | }; 32 | } 33 | 34 | int_unmarshals! { u8, u16, u32, u64, u128, i8, i16, i32, i64, i128 } 35 | float_unmarshals! { f32: u32, f64: u64 } 36 | 37 | impl Unmarshal for bool { 38 | fn unmarshal(scratch: &[u8]) -> Self { 39 | scratch[0] != 0 40 | } 41 | } 42 | --------------------------------------------------------------------------------