tests_fuzz/
context.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use common_query::AddColumnLocation;
18use datatypes::types::cast;
19use rand::Rng;
20use snafu::{OptionExt, ensure};
21
22use crate::error::{self, Result};
23use crate::generator::Random;
24use crate::ir::alter_expr::{AlterTableOperation, AlterTableOption};
25use crate::ir::create_expr::{ColumnOption, PartitionDef};
26use crate::ir::partition_expr::SimplePartitions;
27use crate::ir::repartition_expr::RepartitionExpr;
28use crate::ir::{AlterTableExpr, Column, CreateTableExpr, Ident};
29
30pub type TableContextRef = Arc<TableContext>;
31
32/// TableContext stores table info.
33#[derive(Debug, Clone)]
34pub struct TableContext {
35    pub name: Ident,
36    pub columns: Vec<Column>,
37
38    // GreptimeDB specific options
39    pub partition: Option<PartitionDef>,
40    pub primary_keys: Vec<usize>,
41    pub table_options: Vec<AlterTableOption>,
42}
43
44impl From<&CreateTableExpr> for TableContext {
45    fn from(
46        CreateTableExpr {
47            table_name: name,
48            columns,
49            partition,
50            primary_keys,
51            ..
52        }: &CreateTableExpr,
53    ) -> Self {
54        Self {
55            name: name.clone(),
56            columns: columns.clone(),
57            partition: partition.clone(),
58            primary_keys: primary_keys.clone(),
59            table_options: vec![],
60        }
61    }
62}
63
64impl TableContext {
65    /// Returns the timestamp column
66    pub fn timestamp_column(&self) -> Option<Column> {
67        self.columns.iter().find(|c| c.is_time_index()).cloned()
68    }
69
70    /// Applies the [AlterTableExpr].
71    pub fn alter(mut self, expr: AlterTableExpr) -> Result<TableContext> {
72        match expr.alter_kinds {
73            AlterTableOperation::AddColumn { column, location } => {
74                ensure!(
75                    !self.columns.iter().any(|col| col.name == column.name),
76                    error::UnexpectedSnafu {
77                        violated: format!("Column {} exists", column.name),
78                    }
79                );
80                match location {
81                    Some(AddColumnLocation::First) => {
82                        let mut columns = Vec::with_capacity(self.columns.len() + 1);
83                        columns.push(column);
84                        columns.extend(self.columns);
85                        self.columns = columns;
86                    }
87                    Some(AddColumnLocation::After { column_name }) => {
88                        let index = self
89                            .columns
90                            .iter()
91                            // TODO(weny): find a better way?
92                            .position(|col| col.name.to_string() == column_name)
93                            .context(error::UnexpectedSnafu {
94                                violated: format!("Column: {column_name} not found"),
95                            })?;
96                        self.columns.insert(index + 1, column);
97                    }
98                    None => self.columns.push(column),
99                }
100                // Re-generates the primary_keys
101                self.primary_keys = self
102                    .columns
103                    .iter()
104                    .enumerate()
105                    .flat_map(|(idx, col)| {
106                        if col.is_primary_key() {
107                            Some(idx)
108                        } else {
109                            None
110                        }
111                    })
112                    .collect();
113                Ok(self)
114            }
115            AlterTableOperation::DropColumn { name } => {
116                self.columns.retain(|col| col.name != name);
117                // Re-generates the primary_keys
118                self.primary_keys = self
119                    .columns
120                    .iter()
121                    .enumerate()
122                    .flat_map(|(idx, col)| {
123                        if col.is_primary_key() {
124                            Some(idx)
125                        } else {
126                            None
127                        }
128                    })
129                    .collect();
130                Ok(self)
131            }
132            AlterTableOperation::RenameTable { new_table_name } => {
133                ensure!(
134                    new_table_name != self.name,
135                    error::UnexpectedSnafu {
136                        violated: "The new table name is equal the current name",
137                    }
138                );
139                self.name = new_table_name;
140                Ok(self)
141            }
142            AlterTableOperation::ModifyDataType { column } => {
143                if let Some(idx) = self.columns.iter().position(|col| col.name == column.name) {
144                    self.columns[idx].column_type = column.column_type.clone();
145                    for opt in self.columns[idx].options.iter_mut() {
146                        if let ColumnOption::DefaultValue(value) = opt {
147                            *value = cast(value.clone(), &column.column_type).unwrap();
148                        }
149                    }
150                }
151                Ok(self)
152            }
153            AlterTableOperation::SetTableOptions { options } => {
154                for option in options {
155                    if let Some(idx) = self
156                        .table_options
157                        .iter()
158                        .position(|opt| opt.key() == option.key())
159                    {
160                        self.table_options[idx] = option;
161                    } else {
162                        self.table_options.push(option);
163                    }
164                }
165                Ok(self)
166            }
167            AlterTableOperation::UnsetTableOptions { keys } => {
168                self.table_options
169                    .retain(|opt| !keys.contains(&opt.key().to_string()));
170                Ok(self)
171            }
172        }
173    }
174
175    pub fn repartition(mut self, expr: RepartitionExpr) -> Result<TableContext> {
176        match expr {
177            RepartitionExpr::Split(split) => {
178                let partition_def = self.partition.as_mut().expect("expected partition def");
179                let insert_pos = partition_def
180                    .exprs
181                    .iter()
182                    .position(|expr| expr == &split.target)
183                    .unwrap();
184                partition_def.exprs[insert_pos] = split.into[0].clone();
185                partition_def
186                    .exprs
187                    .insert(insert_pos + 1, split.into[1].clone());
188            }
189            RepartitionExpr::Merge(merge) => {
190                let partition_def = self.partition.as_mut().expect("expected partition def");
191                let removed_idx = partition_def
192                    .exprs
193                    .iter()
194                    .position(|expr| expr == &merge.targets[0])
195                    .unwrap();
196                let mut partitions = SimplePartitions::from_exprs(
197                    partition_def.columns[0].clone(),
198                    &partition_def.exprs,
199                )?;
200                partitions.remove_bound(removed_idx)?;
201                partition_def.exprs = partitions.generate()?;
202            }
203            RepartitionExpr::AlterPartitions(partition) => {
204                ensure!(
205                    self.partition.is_none(),
206                    error::UnexpectedSnafu {
207                        violated: format!("Table {} already has partition", self.name),
208                    }
209                );
210                self.partition = Some(partition.partition);
211            }
212        }
213
214        Ok(self)
215    }
216
217    pub fn generate_unique_column_name<R: Rng>(
218        &self,
219        rng: &mut R,
220        generator: &dyn Random<Ident, R>,
221    ) -> Ident {
222        let mut name = generator.generate(rng);
223        while self.columns.iter().any(|col| col.name.value == name.value) {
224            name = generator.generate(rng);
225        }
226        name
227    }
228
229    pub fn generate_unique_table_name<R: Rng>(
230        &self,
231        rng: &mut R,
232        generator: &dyn Random<Ident, R>,
233    ) -> Ident {
234        let mut name = generator.generate(rng);
235        while self.name.value == name.value {
236            name = generator.generate(rng);
237        }
238        name
239    }
240}
241
242#[cfg(test)]
243mod tests {
244    use common_query::AddColumnLocation;
245    use common_time::Duration;
246    use datatypes::data_type::ConcreteDataType;
247    use datatypes::value::Value;
248    use rand::SeedableRng;
249
250    use super::TableContext;
251    use crate::generator::Generator;
252    use crate::generator::create_expr::CreateTableExprGeneratorBuilder;
253    use crate::ir::alter_expr::{AlterTableOperation, AlterTableOption, Ttl};
254    use crate::ir::create_expr::ColumnOption;
255    use crate::ir::partition_expr::SimplePartitions;
256    use crate::ir::repartition_expr::{MergePartitionExpr, RepartitionExpr, SplitPartitionExpr};
257    use crate::ir::{AlterTableExpr, Column, Ident};
258
259    #[test]
260    fn test_table_context_alter() {
261        let table_ctx = TableContext {
262            name: "foo".into(),
263            columns: vec![],
264            partition: None,
265            primary_keys: vec![],
266            table_options: vec![],
267        };
268        // Add a column
269        let expr = AlterTableExpr {
270            table_name: "foo".into(),
271            alter_kinds: AlterTableOperation::AddColumn {
272                column: Column {
273                    name: "a".into(),
274                    column_type: ConcreteDataType::timestamp_microsecond_datatype(),
275                    options: vec![ColumnOption::PrimaryKey],
276                },
277                location: None,
278            },
279        };
280        let table_ctx = table_ctx.alter(expr).unwrap();
281        assert_eq!(table_ctx.columns[0].name, Ident::new("a"));
282        assert_eq!(table_ctx.primary_keys, vec![0]);
283
284        // Add a column at first
285        let expr = AlterTableExpr {
286            table_name: "foo".into(),
287            alter_kinds: AlterTableOperation::AddColumn {
288                column: Column {
289                    name: "b".into(),
290                    column_type: ConcreteDataType::timestamp_microsecond_datatype(),
291                    options: vec![ColumnOption::PrimaryKey],
292                },
293                location: Some(AddColumnLocation::First),
294            },
295        };
296        let table_ctx = table_ctx.alter(expr).unwrap();
297        assert_eq!(table_ctx.columns[0].name, Ident::new("b"));
298        assert_eq!(table_ctx.primary_keys, vec![0, 1]);
299
300        // Add a column after "b"
301        let expr = AlterTableExpr {
302            table_name: "foo".into(),
303            alter_kinds: AlterTableOperation::AddColumn {
304                column: Column {
305                    name: "c".into(),
306                    column_type: ConcreteDataType::timestamp_microsecond_datatype(),
307                    options: vec![ColumnOption::PrimaryKey],
308                },
309                location: Some(AddColumnLocation::After {
310                    column_name: "b".into(),
311                }),
312            },
313        };
314        let table_ctx = table_ctx.alter(expr).unwrap();
315        assert_eq!(table_ctx.columns[1].name, Ident::new("c"));
316        assert_eq!(table_ctx.primary_keys, vec![0, 1, 2]);
317
318        // Drop the column "b"
319        let expr = AlterTableExpr {
320            table_name: "foo".into(),
321            alter_kinds: AlterTableOperation::DropColumn { name: "b".into() },
322        };
323        let table_ctx = table_ctx.alter(expr).unwrap();
324        assert_eq!(table_ctx.columns[1].name, Ident::new("a"));
325        assert_eq!(table_ctx.primary_keys, vec![0, 1]);
326
327        // Set table options
328        let ttl_option = AlterTableOption::Ttl(Ttl::Duration(Duration::new_second(60)));
329        let expr = AlterTableExpr {
330            table_name: "foo".into(),
331            alter_kinds: AlterTableOperation::SetTableOptions {
332                options: vec![ttl_option.clone()],
333            },
334        };
335        let table_ctx = table_ctx.alter(expr).unwrap();
336        assert_eq!(table_ctx.table_options.len(), 1);
337        assert_eq!(table_ctx.table_options[0], ttl_option);
338
339        // Unset table options
340        let expr = AlterTableExpr {
341            table_name: "foo".into(),
342            alter_kinds: AlterTableOperation::UnsetTableOptions {
343                keys: vec![ttl_option.key().to_string()],
344            },
345        };
346        let table_ctx = table_ctx.alter(expr).unwrap();
347        assert_eq!(table_ctx.table_options.len(), 0);
348    }
349
350    #[test]
351    fn test_apply_split_partition_expr() {
352        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0);
353        let expr = CreateTableExprGeneratorBuilder::default()
354            .columns(10)
355            .partition(10)
356            .if_not_exists(true)
357            .engine("mito2")
358            .build()
359            .unwrap()
360            .generate(&mut rng)
361            .unwrap();
362        let mut table_ctx = TableContext::from(&expr);
363        // "age < 10"
364        // "age >= 10 AND age < 20"
365        // "age >= 20" (SPLIT) INTO (age >= 20 AND age < 30, age >= 30)
366        let partitions = SimplePartitions::new(
367            table_ctx.partition.as_ref().unwrap().columns[0].clone(),
368            vec![Value::from(10), Value::from(20)],
369        )
370        .generate()
371        .unwrap();
372        // "age < 10"
373        // "age >= 10 AND age < 20"
374        // "age >= 20" AND age < 30"
375        // "age >= 30"
376        let expected_exprs = SimplePartitions::new(
377            table_ctx.partition.as_ref().unwrap().columns[0].clone(),
378            vec![Value::from(10), Value::from(20), Value::from(30)],
379        )
380        .generate()
381        .unwrap();
382        table_ctx.partition.as_mut().unwrap().exprs = partitions.clone();
383        let table_ctx = table_ctx
384            .repartition(RepartitionExpr::Split(SplitPartitionExpr {
385                table_name: expr.table_name.clone(),
386                target: partitions.last().unwrap().clone(),
387                into: vec![expected_exprs[2].clone(), expected_exprs[3].clone()],
388                wait: true,
389            }))
390            .unwrap();
391        let partition_def = table_ctx.partition.as_ref().unwrap();
392        assert_eq!(partition_def.exprs, expected_exprs);
393    }
394
395    #[test]
396    fn test_apply_merge_partition_expr() {
397        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0);
398        let expr = CreateTableExprGeneratorBuilder::default()
399            .columns(10)
400            .partition(10)
401            .if_not_exists(true)
402            .engine("mito2")
403            .build()
404            .unwrap()
405            .generate(&mut rng)
406            .unwrap();
407        let mut table_ctx = TableContext::from(&expr);
408        // "age < 10"
409        // "age >= 10 AND age < 20" (MERGE)
410        // "age >= 20" (MERGE)
411        let partitions = SimplePartitions::new(
412            table_ctx.partition.as_ref().unwrap().columns[0].clone(),
413            vec![Value::from(10), Value::from(20)],
414        )
415        .generate()
416        .unwrap();
417        // "age < 10"
418        // "age >= 10
419        let expected_exprs = SimplePartitions::new(
420            table_ctx.partition.as_ref().unwrap().columns[0].clone(),
421            vec![Value::from(10)],
422        )
423        .generate()
424        .unwrap();
425        table_ctx.partition.as_mut().unwrap().exprs = partitions.clone();
426        let table_ctx = table_ctx
427            .repartition(RepartitionExpr::Merge(MergePartitionExpr {
428                table_name: expr.table_name.clone(),
429                targets: vec![partitions[1].clone(), partitions[2].clone()],
430                wait: true,
431            }))
432            .unwrap();
433        let partition_def = table_ctx.partition.as_ref().unwrap();
434        assert_eq!(partition_def.exprs, expected_exprs);
435    }
436}
tests_fuzz/context.rs

tests_fuzz/
context.rs
