tests_fuzz/
utils.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
15pub mod cluster_info;
16pub mod config;
17pub mod crd;
18pub mod health;
19pub mod migration;
20pub mod partition;
21pub mod pod_failure;
22pub mod procedure;
23#[cfg(feature = "unstable")]
24pub mod process;
25pub mod wait;
26
27use std::env;
28
29use common_telemetry::info;
30use common_telemetry::tracing::log::LevelFilter;
31use paste::paste;
32use snafu::ResultExt;
33use sqlx::mysql::{MySqlConnectOptions, MySqlPoolOptions};
34use sqlx::{ConnectOptions, MySql, Pool};
35
36use crate::error::{self, Result};
37use crate::ir::Ident;
38
39/// Database connections
40pub struct Connections {
41    pub mysql: Option<Pool<MySql>>,
42}
43
44const GT_MYSQL_ADDR: &str = "GT_MYSQL_ADDR";
45
46/// Connects to GreptimeDB via env variables.
47pub async fn init_greptime_connections_via_env() -> Connections {
48    let _ = dotenv::dotenv();
49    let mysql = if let Ok(addr) = env::var(GT_MYSQL_ADDR) {
50        Some(addr)
51    } else {
52        info!("GT_MYSQL_ADDR is empty, ignores test");
53        None
54    };
55
56    init_greptime_connections(mysql).await
57}
58
59/// Connects to GreptimeDB.
60pub async fn init_greptime_connections(mysql: Option<String>) -> Connections {
61    let mysql = if let Some(addr) = mysql {
62        let opts = format!("mysql://{addr}/public")
63            .parse::<MySqlConnectOptions>()
64            .unwrap()
65            .log_statements(LevelFilter::Off);
66
67        Some(MySqlPoolOptions::new().connect_with(opts).await.unwrap())
68    } else {
69        None
70    };
71
72    Connections { mysql }
73}
74
75const GT_FUZZ_BINARY_PATH: &str = "GT_FUZZ_BINARY_PATH";
76const GT_FUZZ_INSTANCE_ROOT_DIR: &str = "GT_FUZZ_INSTANCE_ROOT_DIR";
77
78/// The variables for unstable test
79pub struct UnstableTestVariables {
80    pub binary_path: String,
81    pub root_dir: Option<String>,
82}
83
84/// Loads env variables for unstable test
85pub fn load_unstable_test_env_variables() -> UnstableTestVariables {
86    let _ = dotenv::dotenv();
87    let binary_path = env::var(GT_FUZZ_BINARY_PATH).expect("GT_FUZZ_BINARY_PATH not found");
88    let root_dir = env::var(GT_FUZZ_INSTANCE_ROOT_DIR).ok();
89
90    UnstableTestVariables {
91        binary_path,
92        root_dir,
93    }
94}
95
96pub const GT_FUZZ_CLUSTER_NAMESPACE: &str = "GT_FUZZ_CLUSTER_NAMESPACE";
97pub const GT_FUZZ_CLUSTER_NAME: &str = "GT_FUZZ_CLUSTER_NAME";
98
99/// Flushes memtable to SST file.
100pub async fn flush_memtable(e: &Pool<MySql>, table_name: &Ident) -> Result<()> {
101    let sql = format!("admin flush_table(\"{}\")", table_name);
102    let result = sqlx::query(&sql)
103        .execute(e)
104        .await
105        .context(error::ExecuteQuerySnafu { sql })?;
106    info!("Flush table: {}\n\nResult: {result:?}\n\n", table_name);
107
108    Ok(())
109}
110
111/// Triggers a compaction for table
112pub async fn compact_table(e: &Pool<MySql>, table_name: &Ident) -> Result<()> {
113    let sql = format!("admin compact_table(\"{}\")", table_name);
114    let result = sqlx::query(&sql)
115        .execute(e)
116        .await
117        .context(error::ExecuteQuerySnafu { sql })?;
118    info!("Compact table: {}\n\nResult: {result:?}\n\n", table_name);
119
120    Ok(())
121}
122
123pub const GT_FUZZ_INPUT_MAX_ROWS: &str = "GT_FUZZ_INPUT_MAX_ROWS";
124pub const GT_FUZZ_INPUT_MAX_TABLES: &str = "GT_FUZZ_INPUT_MAX_TABLES";
125pub const GT_FUZZ_INPUT_MAX_COLUMNS: &str = "GT_FUZZ_INPUT_MAX_COLUMNS";
126pub const GT_FUZZ_INPUT_MAX_ALTER_ACTIONS: &str = "GT_FUZZ_INPUT_MAX_ALTER_ACTIONS";
127pub const GT_FUZZ_INPUT_MAX_INSERT_ACTIONS: &str = "GT_FUZZ_INPUT_MAX_INSERT_ACTIONS";
128
129macro_rules! make_get_from_env_helper {
130    ($key:expr, $default: expr) => {
131        paste! {
132            #[doc = "Retrieves `" $key "` environment variable \
133                     or returns a default value (`" $default "`) if the environment variable is not set.
134            "]
135            pub fn [<get_ $key:lower>]() -> usize {
136                get_from_env_or_default_value($key, $default)
137            }
138        }
139    };
140}
141
142make_get_from_env_helper!(GT_FUZZ_INPUT_MAX_ALTER_ACTIONS, 256);
143make_get_from_env_helper!(GT_FUZZ_INPUT_MAX_INSERT_ACTIONS, 4);
144make_get_from_env_helper!(GT_FUZZ_INPUT_MAX_ROWS, 512);
145make_get_from_env_helper!(GT_FUZZ_INPUT_MAX_TABLES, 32);
146make_get_from_env_helper!(GT_FUZZ_INPUT_MAX_COLUMNS, 16);
147
148/// Retrieves a value from the environment variables
149/// or returns a default value if the environment variable is not set.
150fn get_from_env_or_default_value(key: &str, default_value: usize) -> usize {
151    let _ = dotenv::dotenv();
152    if let Ok(value) = env::var(key) {
153        value.parse().unwrap()
154    } else {
155        default_value
156    }
157}
tests_fuzz/utils.rs

tests_fuzz/
utils.rs
