clickhouse_testing/
lib.rs

1use clickhouse::Row;
2use dotenvy::dotenv;
3use serde::Deserialize;
4use std::fs::read_dir;
5use std::io::ErrorKind;
6use std::path::PathBuf;
7use std::{env, fs, io};
8
9pub type Client = clickhouse::Client;
10pub use clickhouse_testing_macros::test;
11
12#[derive(Debug)]
13pub enum Error {
14    Io(std::io::Error),
15    Env(std::env::VarError),
16    Clickhouse(clickhouse::error::Error),
17    Migration(String),
18}
19
20impl From<std::io::Error> for Error {
21    fn from(e: std::io::Error) -> Self {
22        Error::Io(e)
23    }
24}
25
26impl From<std::env::VarError> for Error {
27    fn from(e: std::env::VarError) -> Self {
28        Error::Env(e)
29    }
30}
31
32impl From<clickhouse::error::Error> for Error {
33    fn from(e: clickhouse::error::Error) -> Self {
34        Error::Clickhouse(e)
35    }
36}
37
38pub async fn init_test(module_path: &str, test_name: &str) -> Result<Client, Error> {
39    let (config, migrations_path) = if let Some(path) = env::var_os("CONFIG_PATH") {
40        read_file_config(PathBuf::from(path))?
41    } else {
42        _ = dotenv();
43        (read_clickhouse_config(), env::var("MIGRATIONS_DIR")?)
44    };
45    init_test_with_config(module_path, test_name, &config, &migrations_path).await
46}
47
48async fn init_test_with_config(
49    module_path: &str,
50    test_name: &str,
51    config: &ClickhouseConfig,
52    migrations_path: &str,
53) -> Result<Client, Error> {
54    let client = create_client(config);
55    let databases = get_dbs_list(&client).await?;
56    let db_name = next_db_version(&databases, module_path, test_name);
57
58    create_database(&client, &db_name).await?;
59
60    let test_client = client.with_database(db_name);
61    apply_migrations(&test_client, migrations_path).await?;
62
63    Ok(test_client)
64}
65
66pub async fn cleanup_test(client: &Client) -> Result<(), Error> {
67    let current_db = get_current_db(client).await?;
68    drop_db(client, &current_db).await?;
69
70    Ok(())
71}
72
73async fn apply_migrations(client: &Client, migrations_path: &str) -> Result<(), Error> {
74    let project_root = get_project_root()?;
75
76    let mut sql_files: Vec<_> = read_dir(project_root.join(migrations_path))?
77        .filter_map(|entry| entry.ok())
78        .map(|entry| entry.path())
79        .filter(|path| path.extension().and_then(|ext| ext.to_str()) == Some("sql"))
80        .collect();
81
82    sql_files.sort();
83
84    for file in sql_files {
85        let script = fs::read_to_string(&file)?;
86
87        // For avoiding "Multi-statements are not allowed" error
88        let script_parts: Vec<&str> = script.split(';').filter(|s| !s.trim().is_empty()).collect();
89
90        for script_part in script_parts {
91            client.query(script_part).execute().await?;
92        }
93    }
94
95    Ok(())
96}
97
98fn get_project_root() -> Result<PathBuf, Error> {
99    let path = env::current_dir()?;
100
101    for ancestor_path in path.ancestors() {
102        let has_cargo = read_dir(ancestor_path)?.any(|p| p.unwrap().file_name() == "Cargo.lock");
103        if has_cargo {
104            return Ok(PathBuf::from(ancestor_path));
105        }
106    }
107
108    Err(io::Error::new(ErrorKind::NotFound, "Cargo.lock not found").into())
109}
110
111fn create_client(config: &ClickhouseConfig) -> Client {
112    Client::default()
113        .with_url(&config.url)
114        .with_database(&config.db)
115        .with_user(&config.user)
116        .with_password(&config.password)
117}
118
119fn read_clickhouse_config() -> ClickhouseConfig {
120    ClickhouseConfig {
121        url: env::var("CLICKHOUSE_URL").unwrap_or("http://localhost:8123".into()),
122        db: env::var("CLICKHOUSE_DB").unwrap_or("default".into()),
123        user: env::var("CLICKHOUSE_USER").unwrap_or("default".into()),
124        password: env::var("CLICKHOUSE_PASSWORD").unwrap_or("".into()),
125    }
126}
127
128fn read_file_config(path: PathBuf) -> Result<(ClickhouseConfig, String), Error> {
129    let contents = fs::read(&path)?;
130    let file: FileConfig = serde_json::from_slice(&contents).map_err(|error| {
131        io::Error::new(
132            ErrorKind::InvalidData,
133            format!("invalid test configuration {}: {error}", path.display()),
134        )
135    })?;
136    Ok(file.with_overrides(|name| env::var(name).ok()))
137}
138
139#[derive(Deserialize)]
140struct FileConfig {
141    db: ClickhouseConfig,
142    migration_source: String,
143}
144
145impl FileConfig {
146    fn with_overrides(self, get_var: impl Fn(&str) -> Option<String>) -> (ClickhouseConfig, String) {
147        (
148            ClickhouseConfig {
149                url: get_var("CLICKHOUSE_URL").unwrap_or(self.db.url),
150                db: get_var("CLICKHOUSE_DB").unwrap_or(self.db.db),
151                user: get_var("CLICKHOUSE_USER").unwrap_or(self.db.user),
152                password: get_var("CLICKHOUSE_PASSWORD").unwrap_or(self.db.password),
153            },
154            get_var("MIGRATIONS_DIR").unwrap_or(self.migration_source),
155        )
156    }
157}
158
159#[cfg(test)]
160mod config_tests {
161    use super::FileConfig;
162
163    #[test]
164    fn json_config_uses_environment_overrides() {
165        let file: FileConfig = serde_json::from_str(
166            r#"{"db":{"url":"http://file:8123","db":"file_db","user":"file_user","password":"file_password"},"migration_source":"file_migrations","mongo":{"db":"ignored"}}"#,
167        )
168        .unwrap();
169
170        let (db, migrations) = file.with_overrides(|name| match name {
171            "CLICKHOUSE_URL" => Some("http://env:8123".into()),
172            "MIGRATIONS_DIR" => Some("env_migrations".into()),
173            _ => None,
174        });
175
176        assert_eq!(db.url, "http://env:8123");
177        assert_eq!(db.db, "file_db");
178        assert_eq!(db.user, "file_user");
179        assert_eq!(db.password, "file_password");
180        assert_eq!(migrations, "env_migrations");
181    }
182}
183
184async fn get_dbs_list(client: &Client) -> Result<Vec<Database>, Error> {
185    let databases = client
186        .query("SELECT name FROM system.databases")
187        .fetch_all::<Database>()
188        .await?;
189
190    Ok(databases)
191}
192
193async fn create_database(client: &Client, db_name: &str) -> Result<(), Error> {
194    let query = format!("CREATE DATABASE IF NOT EXISTS {}", db_name);
195    client.query(&query).execute().await?;
196
197    Ok(())
198}
199
200async fn get_current_db(client: &Client) -> Result<Database, Error> {
201    let database = client
202        .query("SELECT currentDatabase() AS name")
203        .fetch_one::<Database>()
204        .await?;
205
206    Ok(database)
207}
208
209async fn drop_db(client: &Client, database: &Database) -> Result<(), Error> {
210    client
211        .query(&format!("DROP DATABASE {}", database.name))
212        .execute()
213        .await?;
214
215    Ok(())
216}
217
218fn sanitize_name(name: &str) -> String {
219    name.replace("::", "_")
220        .replace("-", "_")
221        .replace(":", "_")
222        .trim_matches('_')
223        .to_lowercase()
224}
225
226fn next_db_version(tests_dbs: &[Database], module_path: &str, test_name: &str) -> String {
227    let sanitized_module = sanitize_name(module_path);
228    let sanitized_test = sanitize_name(test_name);
229    let current_test_db = format!("test_db_{sanitized_module}_{sanitized_test}_");
230
231    let db_version = tests_dbs
232        .iter()
233        .filter_map(|db| {
234            db.name
235                .strip_prefix(&current_test_db)?
236                .parse::<usize>()
237                .ok()
238        })
239        .max()
240        .map(|v| v + 1)
241        .unwrap_or(1);
242
243    format!("{current_test_db}{db_version}")
244}
245#[derive(Debug, Deserialize, Row)]
246struct Database {
247    name: String,
248}
249
250#[derive(Debug, Deserialize)]
251struct ClickhouseConfig {
252    url: String,
253    db: String,
254    user: String,
255    password: String,
256}
clickhouse_testing/lib.rs

clickhouse_testing/
lib.rs
