tests_fuzz/validator/
partition.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 snafu::{ResultExt, ensure};
16use sqlx::MySqlPool;
17
18use crate::error;
19use crate::error::Result;
20use crate::ir::Ident;
21use crate::ir::create_expr::PartitionDef;
22
23const PARTITIONS_INFO_SCHEMA_SQL: &str = "SELECT table_catalog, table_schema, table_name, \
24partition_name, partition_expression, partition_description, greptime_partition_id, \
25partition_ordinal_position FROM information_schema.partitions WHERE table_name = ? \
26ORDER BY partition_ordinal_position;";
27
28#[derive(Debug, Clone, sqlx::FromRow)]
29pub struct PartitionInfo {
30    pub table_catalog: String,
31    pub table_schema: String,
32    pub table_name: String,
33    pub partition_name: String,
34    pub partition_expression: String,
35    pub partition_description: String,
36    pub greptime_partition_id: u64,
37    pub partition_ordinal_position: i64,
38}
39
40/// Fetches the partitions info from the information_schema.partitions table.
41pub async fn fetch_partitions_info_schema(
42    db: &MySqlPool,
43    _schema_name: Ident,
44    table: &Ident,
45) -> Result<Vec<PartitionInfo>> {
46    sqlx::query_as::<_, PartitionInfo>(PARTITIONS_INFO_SCHEMA_SQL)
47        .bind(&table.value)
48        .fetch_all(db)
49        .await
50        .context(error::ExecuteQuerySnafu {
51            sql: PARTITIONS_INFO_SCHEMA_SQL,
52        })
53}
54
55/// Asserts the partitions are equal to the expected partitions.
56pub fn assert_partitions(expected: &PartitionDef, actual: &[PartitionInfo]) -> Result<()> {
57    ensure!(
58        expected.exprs.len() == actual.len(),
59        error::AssertSnafu {
60            reason: format!(
61                "Expected partitions length: {}, got: {}",
62                expected.exprs.len(),
63                actual.len()
64            ),
65        }
66    );
67
68    let expected_exprs = expected.exprs.iter().map(|expr| expr.to_string());
69    for expr in expected_exprs {
70        let actual_expr = actual
71            .iter()
72            .find(|info| info.partition_description == expr);
73        ensure!(
74            actual_expr.is_some(),
75            error::AssertSnafu {
76                reason: format!(
77                    "Expected partition expression: {expr} not found, actual: {:?}",
78                    actual
79                        .iter()
80                        .map(|info| info.partition_description.clone())
81                        .collect::<Vec<_>>()
82                        .join(", ")
83                ),
84            }
85        );
86    }
87
88    Ok(())
89}
tests_fuzz/validator/partition.rs

tests_fuzz/validator/
partition.rs
