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
55fn normalize(s: &str) -> String {
56    s.replace("\\\"", "\"").replace("\\\\", "\\")
57}
58
59/// Asserts the partitions are equal to the expected partitions.
60pub fn assert_partitions(expected: &PartitionDef, actual: &[PartitionInfo]) -> Result<()> {
61    ensure!(
62        expected.exprs.len() == actual.len(),
63        error::AssertSnafu {
64            reason: format!(
65                "Expected partitions length: {}, got: {}",
66                expected.exprs.len(),
67                actual.len()
68            ),
69        }
70    );
71
72    let expected_exprs = expected.exprs.iter().map(|expr| expr.to_string());
73    for expr in expected_exprs {
74        let actual_expr = actual
75            .iter()
76            .find(|info| normalize(&info.partition_description) == normalize(&expr));
77        ensure!(
78            actual_expr.is_some(),
79            error::AssertSnafu {
80                reason: format!(
81                    "Expected partition expression: '{expr:?}' not found, actual: {:?}",
82                    actual
83                        .iter()
84                        .map(|info| format!("'{:?}'", info.partition_description.clone()))
85                        .collect::<Vec<_>>()
86                        .join("; ")
87                ),
88            }
89        );
90    }
91
92    Ok(())
93}