Skip to main content

ydb/
table_requests.rs

1//! Request builders for Table service DDL and read-table operations.
2//!
3//! API shape follows [ydb-go-sdk](https://github.com/ydb-platform/ydb-go-sdk) `table/options`.
4
5use std::collections::HashMap;
6
7use crate::errors::{YdbError, YdbResult};
8use crate::grpc_wrapper::raw_table_service::create_table::RawCreateTableColumn;
9use crate::types::Value;
10
11/// Column specification for [`CreateTableRequest`] and [`AlterTableRequest`].
12#[derive(Clone, Debug)]
13pub struct TableColumn {
14    pub name: String,
15    pub type_example: Value,
16    pub not_null: bool,
17    pub family: String,
18}
19
20impl TableColumn {
21    pub fn new(name: impl Into<String>, type_example: Value) -> Self {
22        Self {
23            name: name.into(),
24            type_example,
25            not_null: true,
26            family: String::new(),
27        }
28    }
29
30    pub fn with_not_null(mut self, not_null: bool) -> Self {
31        self.not_null = not_null;
32        self
33    }
34
35    pub fn with_family(mut self, family: impl Into<String>) -> Self {
36        self.family = family.into();
37        self
38    }
39
40    pub(crate) fn into_raw(self) -> YdbResult<RawCreateTableColumn> {
41        let typed: crate::grpc_wrapper::raw_table_service::value::RawTypedValue =
42            self.type_example.try_into().map_err(YdbError::from)?;
43        Ok(RawCreateTableColumn {
44            name: self.name,
45            column_type: typed.r#type,
46            not_null: self.not_null,
47            family: self.family,
48        })
49    }
50}
51
52/// CreateTable RPC request (go-sdk: `Session.CreateTable`).
53#[derive(Clone, Debug, Default)]
54pub struct CreateTableRequest {
55    pub path: String,
56    pub columns: Vec<TableColumn>,
57    pub primary_key: Vec<String>,
58    pub attributes: HashMap<String, String>,
59}
60
61impl CreateTableRequest {
62    pub fn new(path: impl Into<String>) -> Self {
63        Self {
64            path: path.into(),
65            ..Default::default()
66        }
67    }
68
69    pub(crate) fn into_raw(
70        self,
71        session_id: String,
72        operation_params: crate::grpc_wrapper::raw_ydb_operation::RawOperationParams,
73    ) -> YdbResult<crate::grpc_wrapper::raw_table_service::create_table::RawCreateTableRequest>
74    {
75        let columns = self
76            .columns
77            .into_iter()
78            .map(|column| column.into_raw())
79            .collect::<YdbResult<Vec<_>>>()?;
80        Ok(
81            crate::grpc_wrapper::raw_table_service::create_table::RawCreateTableRequest {
82                session_id,
83                path: self.path,
84                columns,
85                primary_key: self.primary_key,
86                attributes: self.attributes,
87                operation_params,
88            },
89        )
90    }
91
92    pub fn with_column(mut self, column: TableColumn) -> Self {
93        self.columns.push(column);
94        self
95    }
96
97    pub fn with_primary_key(
98        mut self,
99        columns: impl IntoIterator<Item = impl Into<String>>,
100    ) -> Self {
101        self.primary_key = columns.into_iter().map(Into::into).collect();
102        self
103    }
104
105    pub fn with_attribute(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
106        self.attributes.insert(key.into(), value.into());
107        self
108    }
109}
110
111/// DropTable RPC request.
112#[derive(Clone, Debug)]
113pub struct DropTableRequest {
114    pub path: String,
115}
116
117impl DropTableRequest {
118    pub fn new(path: impl Into<String>) -> Self {
119        Self { path: path.into() }
120    }
121}
122
123/// ReadRows RPC request (go-sdk: `table.Client.ReadRows` + `options.ReadRowsOption`).
124#[derive(Clone, Debug, Default)]
125pub struct ReadRowsRequest {
126    pub path: String,
127    pub keys: Vec<Value>,
128    pub columns: Vec<String>,
129}
130
131impl ReadRowsRequest {
132    pub fn new(path: impl Into<String>) -> Self {
133        Self {
134            path: path.into(),
135            ..Default::default()
136        }
137    }
138
139    pub fn with_keys(mut self, keys: Vec<Value>) -> Self {
140        self.keys = keys;
141        self
142    }
143
144    pub fn with_column(mut self, name: impl Into<String>) -> Self {
145        self.columns.push(name.into());
146        self
147    }
148
149    pub(crate) fn into_raw(
150        self,
151        session_id: String,
152    ) -> YdbResult<crate::grpc_wrapper::raw_table_service::read_rows::RawReadRowsRequest> {
153        let keys = crate::types_converters::try_vec_to_list_of_structs(self.keys)?
154            .ok_or_else(|| YdbError::Custom("read rows keys must be a list of structs".into()))?;
155        Ok(
156            crate::grpc_wrapper::raw_table_service::read_rows::RawReadRowsRequest {
157                session_id,
158                path: self.path,
159                keys: keys.try_into().map_err(YdbError::from)?,
160                columns: self.columns,
161            },
162        )
163    }
164}
165
166/// AlterTable RPC request.
167#[derive(Clone, Debug, Default)]
168pub struct AlterTableRequest {
169    pub path: String,
170    pub add_columns: Vec<TableColumn>,
171    pub drop_columns: Vec<String>,
172    pub alter_columns: Vec<TableColumn>,
173    pub alter_attributes: HashMap<String, String>,
174}
175
176impl AlterTableRequest {
177    pub fn new(path: impl Into<String>) -> Self {
178        Self {
179            path: path.into(),
180            ..Default::default()
181        }
182    }
183
184    pub(crate) fn into_raw(
185        self,
186        session_id: String,
187        operation_params: crate::grpc_wrapper::raw_ydb_operation::RawOperationParams,
188    ) -> YdbResult<crate::grpc_wrapper::raw_table_service::alter_table::RawAlterTableRequest> {
189        let add_columns = self
190            .add_columns
191            .into_iter()
192            .map(|column| column.into_raw())
193            .collect::<YdbResult<Vec<_>>>()?;
194        let alter_columns = self
195            .alter_columns
196            .into_iter()
197            .map(|column| column.into_raw())
198            .collect::<YdbResult<Vec<_>>>()?;
199        Ok(
200            crate::grpc_wrapper::raw_table_service::alter_table::RawAlterTableRequest {
201                session_id,
202                path: self.path,
203                add_columns,
204                drop_columns: self.drop_columns,
205                alter_columns,
206                alter_attributes: self.alter_attributes,
207                operation_params,
208            },
209        )
210    }
211
212    pub fn add_column(mut self, column: TableColumn) -> Self {
213        self.add_columns.push(column);
214        self
215    }
216
217    pub fn drop_column(mut self, name: impl Into<String>) -> Self {
218        self.drop_columns.push(name.into());
219        self
220    }
221
222    pub fn alter_column(mut self, column: TableColumn) -> Self {
223        self.alter_columns.push(column);
224        self
225    }
226
227    /// Set or update a table attribute (go-sdk: `options.WithAlterAttribute`).
228    ///
229    /// To remove an attribute, use [`Self::drop_attribute`] or pass an empty `value`
230    /// (server drops keys with blank values in `alter_attributes`).
231    pub fn alter_attribute(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
232        self.alter_attributes.insert(key.into(), value.into());
233        self
234    }
235
236    /// Add a table attribute (go-sdk: `options.WithAddAttribute`).
237    pub fn add_attribute(self, key: impl Into<String>, value: impl Into<String>) -> Self {
238        self.alter_attribute(key, value)
239    }
240
241    /// Drop a table attribute (go-sdk: `options.WithDropAttribute`).
242    pub fn drop_attribute(mut self, key: impl Into<String>) -> Self {
243        self.alter_attributes.insert(key.into(), String::new());
244        self
245    }
246}
247
248/// Named policy preset from [`TableClient::describe_table_options`].
249#[derive(Clone, Debug)]
250pub struct NamedPolicyDescription {
251    pub name: String,
252    pub labels: HashMap<String, String>,
253}
254
255/// Cluster-wide table option presets (go-sdk: `options.TableOptionsDescription`).
256#[derive(Clone, Debug, Default)]
257pub struct TableOptionsDescription {
258    pub table_profile_presets: Vec<NamedPolicyDescription>,
259    pub storage_policy_presets: Vec<NamedPolicyDescription>,
260    pub compaction_policy_presets: Vec<NamedPolicyDescription>,
261    pub partitioning_policy_presets: Vec<NamedPolicyDescription>,
262    pub execution_policy_presets: Vec<NamedPolicyDescription>,
263    pub replication_policy_presets: Vec<NamedPolicyDescription>,
264    pub caching_policy_presets: Vec<NamedPolicyDescription>,
265}
266
267impl From<crate::grpc_wrapper::raw_table_service::describe_table_options::RawNamedPolicyDescription>
268    for NamedPolicyDescription
269{
270    fn from(
271        value: crate::grpc_wrapper::raw_table_service::describe_table_options::RawNamedPolicyDescription,
272    ) -> Self {
273        Self {
274            name: value.name,
275            labels: value.labels,
276        }
277    }
278}
279
280impl From<crate::grpc_wrapper::raw_table_service::describe_table_options::RawDescribeTableOptionsResult>
281    for TableOptionsDescription
282{
283    fn from(
284        value: crate::grpc_wrapper::raw_table_service::describe_table_options::RawDescribeTableOptionsResult,
285    ) -> Self {
286        Self {
287            table_profile_presets: value.table_profile_presets.into_iter().map_into().collect(),
288            storage_policy_presets: value.storage_policy_presets.into_iter().map_into().collect(),
289            compaction_policy_presets: value
290                .compaction_policy_presets
291                .into_iter()
292                .map_into()
293                .collect(),
294            partitioning_policy_presets: value
295                .partitioning_policy_presets
296                .into_iter()
297                .map_into()
298                .collect(),
299            execution_policy_presets: value
300                .execution_policy_presets
301                .into_iter()
302                .map_into()
303                .collect(),
304            replication_policy_presets: value
305                .replication_policy_presets
306                .into_iter()
307                .map_into()
308                .collect(),
309            caching_policy_presets: value.caching_policy_presets.into_iter().map_into().collect(),
310        }
311    }
312}
313
314use itertools::Itertools;
315
316#[cfg(test)]
317mod tests {
318    use super::AlterTableRequest;
319
320    #[test]
321    fn drop_attribute_sets_empty_value_for_server() {
322        let req = AlterTableRequest::new("t").drop_attribute("baz");
323        assert_eq!(req.alter_attributes.get("baz"), Some(&String::new()));
324    }
325
326    #[test]
327    fn add_attribute_same_as_alter_attribute() {
328        let add = AlterTableRequest::new("t").add_attribute("foo", "bar");
329        let alter = AlterTableRequest::new("t").alter_attribute("foo", "bar");
330        assert_eq!(add.alter_attributes, alter.alter_attributes);
331    }
332}