1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
//! Support for executing an `INSERT` statement using ClickHouse's [Native columnar format].
//!
//! [Native columnar format]: https://clickhouse.com/docs/reference/formats/Native
use crate::insert_formatted::InsertFormatted;
use crate::native::Block;
use crate::native::writer::BlockWriter;
use crate::{Client, Compression, sql};
use std::time::Duration;
/// Executes an `INSERT ... FORMAT Native` statement.
///
/// The [`InsertNative::end`] must be called to finalize the `INSERT`.
/// Otherwise, the whole `INSERT` will be aborted.
#[must_use]
pub struct InsertNative {
writer: BlockWriter,
}
impl InsertNative {
pub(crate) fn new(client: &Client, table_name: &str, escape: bool) -> Self {
let mut sql = "INSERT INTO ".to_string();
if escape {
sql::escape::identifier(table_name, &mut sql).expect("error escaping table name");
} else {
sql.push_str(table_name);
}
sql.push_str(" FORMAT Native");
Self {
writer: BlockWriter::new(InsertFormatted::new(
// FIXME: use HTTP body compression instead of block-level compression
&client.clone().with_compression(Compression::None),
sql,
Some(table_name),
)),
}
}
/// Sets timeouts for different operations.
///
/// `send_timeout` restricts time on sending a data chunk to a socket.
/// `None` disables the timeout, it's a default.
/// It's roughly equivalent to `tokio::time::timeout(insert.write(...))`.
///
/// `end_timeout` restricts time on waiting for a response from the CH
/// server. Thus, it includes all work needed to handle `INSERT` by the
/// CH server, e.g. handling all materialized views and so on.
/// `None` disables the timeout, it's a default.
/// It's roughly equivalent to `tokio::time::timeout(insert.end(...))`.
///
/// These timeouts are much more performant (~x10) than wrapping `write()`
/// and `end()` calls into `tokio::time::timeout()`.
pub fn with_timeouts(
mut self,
send_timeout: Option<Duration>,
end_timeout: Option<Duration>,
) -> Self {
self.writer
.insert_mut()
.set_timeouts(send_timeout, end_timeout);
self
}
/// Configure the [roles] to use when executing `INSERT` statements.
///
/// Overrides any roles previously set by this method, [`InsertFormatted::with_setting`],
/// [`Client::with_roles`] or [`Client::with_setting`].
///
/// An empty iterator may be passed to clear the set roles.
///
/// [roles]: https://clickhouse.com/docs/operations/access-rights#role-management
///
/// # Panics
/// If called after the request is started, i.e., after [`InsertNative::write`].
pub fn with_roles(mut self, roles: impl IntoIterator<Item = impl Into<String>>) -> Self {
self.writer.expect_client_mut().set_roles(roles);
self
}
/// Clear any explicit [roles] previously set on this `Insert` or inherited from [`Client`].
///
/// Overrides any roles previously set by [`InsertFormatted::with_roles`], [`InsertFormatted::with_setting`],
/// [`Client::with_roles`] or [`Client::with_setting`].
///
/// [roles]: https://clickhouse.com/docs/operations/access-rights#role-management
///
/// # Panics
/// If called after the request is started, i.e., after [`InsertNative::write`].
pub fn with_default_roles(mut self) -> Self {
self.writer.expect_client_mut().clear_roles();
self
}
/// Similar to [`Client::with_setting`], but for this particular INSERT
/// statement only.
///
/// # Panics
/// If called after the request is started, i.e., after [`InsertNative::write`].
#[track_caller]
pub fn with_setting(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
self.writer.expect_client_mut().set_setting(name, value);
self
}
/// Similar to [`Client::with_product_info()`], but for this `INSERT` statement only.
///
/// # Panics
/// If called after the request is started, i.e., after [`InsertNative::write`].
pub fn with_product_info(
mut self,
product_name: impl Into<String>,
product_version: impl Into<String>,
) -> Self {
self.writer
.expect_client_mut()
.add_product_info(product_name.into(), product_version.into());
self
}
/// Send a block of data.
///
/// # NOT Cancel Safe
/// If this `async` method is canceled (i.e. by dropping the resulting `Future`),
/// the insert is automatically aborted.
///
/// This is because the block data is not sent in a single write, since that would require
/// copying it into a separate buffer. There is no way to resynchronize the stream
/// once a block has been partially sent. Resuming a write would corrupt the stream.
pub async fn write(&mut self, block: &Block) -> crate::Result<()> {
self.writer.write(block).await
}
/// Finish the current `INSERT` request.
pub async fn end(self) -> crate::Result<()> {
self.writer.end().await
}
}
#[cfg(test)]
mod tests {
use crate::{Client, ProductInfo};
// These various setters are already covered functionally in `tests/it/insert_formatted.rs`;
// we just need to check that they're forwarded correctly for `InsertNative`.
#[test]
fn test_with_roles() {
let client = Client::default();
let insert = client.insert_native("foo");
assert!(insert.writer.expect_client().roles.is_empty());
let insert = insert.with_roles(["bar", "baz"]);
let roles = &insert.writer.expect_client().roles;
assert_eq!(roles.len(), 2, "unexpected roles: {roles:?}");
assert!(roles.contains("bar"), "missing role `bar`: {roles:?}");
assert!(roles.contains("baz"), "missing role `baz`: {roles:?}");
let insert = insert.with_default_roles();
assert!(insert.writer.expect_client().roles.is_empty());
}
#[test]
fn test_with_setting() {
let client = Client::default();
let insert = client.insert_native("foo");
let settings = &insert.writer.expect_client().settings;
assert!(settings.is_empty(), "unexpected settings: {settings:?}");
let insert = insert.with_setting("foo", "bar").with_setting("bar", "baz");
let settings = &insert.writer.expect_client().settings;
assert_eq!(settings.len(), 2, "unexpected settings: {settings:?}");
assert_eq!(settings.get("foo"), Some(&"bar".to_string()));
assert_eq!(settings.get("bar"), Some(&"baz".to_string()));
}
#[test]
fn test_with_product_info() {
let client = Client::default();
let insert = client.insert_native("foo");
let product_info = &insert.writer.expect_client().products_info;
assert!(
product_info.is_empty(),
"unexpected product_info: {product_info:?}"
);
let insert = insert
.with_product_info("foo", "1.0.0")
.with_product_info("bar", "0.1.0-alpha.1");
let product_info = &insert.writer.expect_client().products_info;
assert_eq!(
*product_info,
[
ProductInfo {
name: "foo".to_string(),
version: "1.0.0".to_string(),
},
ProductInfo {
name: "bar".to_string(),
version: "0.1.0-alpha.1".to_string()
}
]
);
}
}