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;
18/// CSV dump writer utilities for fuzz tests.
19pub mod csv_dump_writer;
20pub mod health;
21pub mod kafka_wal_http;
22pub mod migration;
23pub mod network_chaos;
24pub mod partition;
25pub mod pod_failure;
26pub mod procedure;
27#[cfg(feature = "unstable")]
28pub mod process;
29pub mod retry;
30/// SQL dump writer utilities for fuzz tests.
31pub mod sql_dump_writer;
32pub mod wait;
33
34use std::env;
35use std::str::FromStr;
36
37use common_base::readable_size::ReadableSize;
38use common_telemetry::info;
39use common_telemetry::tracing::log::LevelFilter;
40use paste::paste;
41use snafu::ResultExt;
42use sqlx::mysql::{MySqlConnectOptions, MySqlPoolOptions};
43use sqlx::{ConnectOptions, MySql, Pool};
44
45use crate::error::{self, Result};
46use crate::ir::Ident;
47
48/// Database connections
49pub struct Connections {
50    pub mysql: Option<Pool<MySql>>,
51}
52
53const GT_MYSQL_ADDR: &str = "GT_MYSQL_ADDR";
54
55/// Connects to GreptimeDB via env variables.
56pub async fn init_greptime_connections_via_env() -> Connections {
57    crate::install_rustls_crypto_provider();
58
59    let _ = dotenv::dotenv();
60    let mysql = if let Ok(addr) = env::var(GT_MYSQL_ADDR) {
61        Some(addr)
62    } else {
63        info!("GT_MYSQL_ADDR is empty, ignores test");
64        None
65    };
66
67    init_greptime_connections(mysql).await
68}
69
70/// Connects to GreptimeDB.
71pub async fn init_greptime_connections(mysql: Option<String>) -> Connections {
72    let mysql = if let Some(addr) = mysql {
73        let opts = format!("mysql://{addr}/public")
74            .parse::<MySqlConnectOptions>()
75            .unwrap()
76            .log_statements(LevelFilter::Off);
77
78        Some(MySqlPoolOptions::new().connect_with(opts).await.unwrap())
79    } else {
80        None
81    };
82
83    Connections { mysql }
84}
85
86const GT_FUZZ_BINARY_PATH: &str = "GT_FUZZ_BINARY_PATH";
87const GT_FUZZ_INSTANCE_ROOT_DIR: &str = "GT_FUZZ_INSTANCE_ROOT_DIR";
88
89/// The variables for unstable test
90pub struct UnstableTestVariables {
91    pub binary_path: String,
92    pub root_dir: Option<String>,
93}
94
95/// Loads env variables for unstable test
96pub fn load_unstable_test_env_variables() -> UnstableTestVariables {
97    let _ = dotenv::dotenv();
98    let binary_path = env::var(GT_FUZZ_BINARY_PATH).expect("GT_FUZZ_BINARY_PATH not found");
99    let root_dir = env::var(GT_FUZZ_INSTANCE_ROOT_DIR).ok();
100
101    UnstableTestVariables {
102        binary_path,
103        root_dir,
104    }
105}
106
107pub const GT_FUZZ_CLUSTER_NAMESPACE: &str = "GT_FUZZ_CLUSTER_NAMESPACE";
108pub const GT_FUZZ_CLUSTER_NAME: &str = "GT_FUZZ_CLUSTER_NAME";
109
110/// Flushes memtable to SST file.
111pub async fn flush_memtable(e: &Pool<MySql>, table_name: &Ident) -> Result<()> {
112    let sql = format!("admin flush_table(\"{}\")", table_name);
113    let result = sqlx::query(&sql)
114        .execute(e)
115        .await
116        .context(error::ExecuteQuerySnafu { sql })?;
117    info!("Flush table: {}\n\nResult: {result:?}\n\n", table_name);
118
119    Ok(())
120}
121
122/// Triggers a compaction for table
123pub async fn compact_table(e: &Pool<MySql>, table_name: &Ident) -> Result<()> {
124    let sql = format!("admin compact_table(\"{}\")", table_name);
125    let result = sqlx::query(&sql)
126        .execute(e)
127        .await
128        .context(error::ExecuteQuerySnafu { sql })?;
129    info!("Compact table: {}\n\nResult: {result:?}\n\n", table_name);
130
131    Ok(())
132}
133
134pub const GT_FUZZ_INPUT_MAX_ROWS: &str = "GT_FUZZ_INPUT_MAX_ROWS";
135pub const GT_FUZZ_INPUT_MAX_TABLES: &str = "GT_FUZZ_INPUT_MAX_TABLES";
136pub const GT_FUZZ_INPUT_MAX_COLUMNS: &str = "GT_FUZZ_INPUT_MAX_COLUMNS";
137pub const GT_FUZZ_INPUT_MAX_ALTER_ACTIONS: &str = "GT_FUZZ_INPUT_MAX_ALTER_ACTIONS";
138pub const GT_FUZZ_INPUT_MAX_INSERT_ACTIONS: &str = "GT_FUZZ_INPUT_MAX_INSERT_ACTIONS";
139pub const FUZZ_OVERRIDE_PREFIX: &str = "GT_FUZZ_OVERRIDE_";
140/// Enables CSV dump generation for fuzz runs.
141pub const GT_FUZZ_DUMP_TABLE_CSV: &str = "GT_FUZZ_DUMP_TABLE_CSV";
142/// Base directory for CSV dump sessions.
143pub const GT_FUZZ_DUMP_DIR: &str = "GT_FUZZ_DUMP_DIR";
144/// Directory suffix used by one CSV dump session.
145pub const GT_FUZZ_DUMP_SUFFIX: &str = "GT_FUZZ_DUMP_SUFFIX";
146/// Max in-memory CSV buffer size before auto flush.
147pub const GT_FUZZ_DUMP_BUFFER_MAX_BYTES: &str = "GT_FUZZ_DUMP_BUFFER_MAX_BYTES";
148
149/// Reads an override value for a fuzz parameter from env `GT_FUZZ_OVERRIDE_<NAME>`.
150pub fn get_fuzz_override<T>(name: &str) -> Option<T>
151where
152    T: std::str::FromStr,
153{
154    let _ = dotenv::dotenv();
155    let key = format!("{}{}", FUZZ_OVERRIDE_PREFIX, name.to_ascii_uppercase());
156    env::var(&key).ok().and_then(|v| v.parse().ok())
157}
158
159/// Returns CSV dump base directory.
160pub fn get_gt_fuzz_dump_dir() -> String {
161    let _ = dotenv::dotenv();
162    env::var(GT_FUZZ_DUMP_DIR).unwrap_or_else(|_| "/tmp/greptime-fuzz-dumps".to_string())
163}
164
165/// Returns CSV dump directory suffix.
166pub fn get_gt_fuzz_dump_suffix() -> String {
167    let _ = dotenv::dotenv();
168    env::var(GT_FUZZ_DUMP_SUFFIX).unwrap_or_else(|_| ".repartition-metric-csv".to_string())
169}
170
171/// Returns max CSV in-memory buffer size.
172pub fn get_gt_fuzz_dump_buffer_max_bytes() -> usize {
173    let _ = dotenv::dotenv();
174    env::var(GT_FUZZ_DUMP_BUFFER_MAX_BYTES)
175        .ok()
176        .and_then(|value| {
177            value.parse::<usize>().ok().or_else(|| {
178                ReadableSize::from_str(&value)
179                    .ok()
180                    .map(|size| size.as_bytes() as usize)
181            })
182        })
183        .unwrap_or(8 * 1024 * 1024)
184}
185
186macro_rules! make_get_from_env_helper {
187    ($key:expr, $default: expr) => {
188        paste! {
189            #[doc = "Retrieves `" $key "` environment variable \
190                     or returns a default value (`" $default "`) if the environment variable is not set.
191            "]
192            pub fn [<get_ $key:lower>]() -> usize {
193                get_from_env_or_default_value($key, $default)
194            }
195        }
196    };
197}
198
199make_get_from_env_helper!(GT_FUZZ_INPUT_MAX_ALTER_ACTIONS, 256);
200make_get_from_env_helper!(GT_FUZZ_INPUT_MAX_INSERT_ACTIONS, 4);
201make_get_from_env_helper!(GT_FUZZ_INPUT_MAX_ROWS, 512);
202make_get_from_env_helper!(GT_FUZZ_INPUT_MAX_TABLES, 32);
203make_get_from_env_helper!(GT_FUZZ_INPUT_MAX_COLUMNS, 16);
204
205/// Retrieves a value from the environment variables
206/// or returns a default value if the environment variable is not set.
207fn get_from_env_or_default_value(key: &str, default_value: usize) -> usize {
208    let _ = dotenv::dotenv();
209    if let Ok(value) = env::var(key) {
210        value.parse().unwrap()
211    } else {
212        default_value
213    }
214}
tests_fuzz/utils.rs

tests_fuzz/
utils.rs
