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, COALESCE(partition_expression, '') AS partition_expression, \
25COALESCE(partition_description, '') AS partition_description, greptime_partition_id, \
26partition_ordinal_position FROM information_schema.partitions WHERE table_name = ? \
27ORDER BY partition_ordinal_position;";
28
29#[derive(Debug, Clone, sqlx::FromRow)]
30pub struct PartitionInfo {
31    pub table_catalog: String,
32    pub table_schema: String,
33    pub table_name: String,
34    pub partition_name: String,
35    pub partition_expression: String,
36    pub partition_description: String,
37    pub greptime_partition_id: u64,
38    pub partition_ordinal_position: i64,
39}
40
41/// Fetches the partitions info from the information_schema.partitions table.
42pub async fn fetch_partitions_info_schema(
43    db: &MySqlPool,
44    _schema_name: Ident,
45    table: &Ident,
46) -> Result<Vec<PartitionInfo>> {
47    sqlx::query_as::<_, PartitionInfo>(PARTITIONS_INFO_SCHEMA_SQL)
48        .bind(&table.value)
49        .fetch_all(db)
50        .await
51        .context(error::ExecuteQuerySnafu {
52            sql: PARTITIONS_INFO_SCHEMA_SQL,
53        })
54}
55
56fn normalize(s: &str) -> String {
57    s.replace("\\\"", "\"").replace("\\\\", "\\")
58}
59
60/// Asserts the partitions are equal to the expected partitions.
61pub fn assert_partitions(expected: &PartitionDef, actual: &[PartitionInfo]) -> Result<()> {
62    ensure!(
63        expected.exprs.len() == actual.len(),
64        error::AssertSnafu {
65            reason: format!(
66                "Expected partitions length: {}, got: {}",
67                expected.exprs.len(),
68                actual.len()
69            ),
70        }
71    );
72
73    let expected_exprs = expected.exprs.iter().map(|expr| expr.to_string());
74    for expr in expected_exprs {
75        let actual_expr = actual
76            .iter()
77            .find(|info| normalize(&info.partition_description) == normalize(&expr));
78        ensure!(
79            actual_expr.is_some(),
80            error::AssertSnafu {
81                reason: format!(
82                    "Expected partition expression: '{expr:?}' not found, actual: {:?}",
83                    actual
84                        .iter()
85                        .map(|info| format!("'{:?}'", info.partition_description.clone()))
86                        .collect::<Vec<_>>()
87                        .join("; ")
88                ),
89            }
90        );
91    }
92
93    Ok(())
94}
95
96/// Asserts that the table has no partition metadata in information schema.
97pub fn assert_unpartitioned(actual: &[PartitionInfo]) -> Result<()> {
98    let has_no_partition_metadata = actual.is_empty()
99        || (actual.len() == 1
100            && actual[0].partition_expression.is_empty()
101            && actual[0].partition_description.is_empty());
102
103    ensure!(
104        has_no_partition_metadata,
105        error::AssertSnafu {
106            reason: format!("Expected unpartitioned table, got partitions: {actual:?}"),
107        }
108    );
109
110    Ok(())
111}
tests_fuzz/validator/partition.rs

tests_fuzz/validator/
partition.rs
