-
Notifications
You must be signed in to change notification settings - Fork 2
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
3 changed files
with
82 additions
and
80 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -9,57 +9,60 @@ | |
| #![cfg(feature = "backend")] | ||
| use std::collections::HashSet; | ||
| use std::env; | ||
| use std::sync::{Mutex, Once}; | ||
| use std::sync::{Mutex, OnceLock}; | ||
|
|
||
| use diesel::dsl::sql; | ||
| use diesel::mysql::MysqlConnection; | ||
| use diesel::prelude::*; | ||
| use diesel::sql_types::Untyped; | ||
| use diesel::Connection; | ||
| use lazy_static::lazy_static; | ||
| use mysql::prelude::*; | ||
| use mysql::{Pool, PooledConn}; | ||
|
|
||
| const URI_ENV: &str = "UDF_TEST_BACKEND_URI"; | ||
| const DEFAULT_DATABASE_URI: &str = "mysql://root:[email protected]:12305/udf_tests"; | ||
|
|
||
| static INIT: Once = Once::new(); | ||
| static POOL: OnceLock<Pool> = OnceLock::new(); | ||
|
|
||
| lazy_static! { | ||
| static ref SETUP_STATE: Mutex<HashSet<String>> = Mutex::new(HashSet::new()); | ||
| } | ||
|
|
||
| /// Allow overriding the default URI | ||
| fn get_database_uri() -> String { | ||
| match env::var(URI_ENV) { | ||
| Ok(s) => s, | ||
| Err(_) => DEFAULT_DATABASE_URI.to_owned(), | ||
| } | ||
| } | ||
|
|
||
| /// Ensures that init items have been run | ||
| pub fn get_db_connection(init: &[&str]) -> MysqlConnection { | ||
| fn build_pool() -> Pool { | ||
| let db_url = get_database_uri(); | ||
| let (url, db) = db_url.rsplit_once('/').unwrap(); | ||
| let pool = Pool::new(url).expect("pool failed"); | ||
| let mut conn = pool.get_conn().expect("initial connection failed"); | ||
|
|
||
| INIT.call_once(|| { | ||
| let mut conn = MysqlConnection::establish(db_url.rsplit_once('/').unwrap().0) | ||
| .expect("initial connection failed"); | ||
| // Create default database | ||
| conn.query_drop(format!("CREATE OR REPLACE DATABASE {db}")) | ||
| .unwrap(); | ||
|
|
||
| sql::<(Untyped,)>("create or replace database udf_tests") | ||
| .execute(&mut conn) | ||
| .expect("could not create databases"); | ||
| }); | ||
| pool | ||
| } | ||
|
|
||
| let hset = &mut *SETUP_STATE.lock().unwrap(); | ||
| let mut conn = MysqlConnection::establish(&db_url).expect("initial connection failed"); | ||
| /// Ensures that init items have been run | ||
| pub fn get_db_connection(init: &[&str]) -> PooledConn { | ||
| let mut conn = POOL | ||
| .get_or_init(build_pool) | ||
| .get_conn() | ||
| .expect("failed to get conn"); | ||
|
|
||
| let ran_stmts = &mut *SETUP_STATE.lock().unwrap(); | ||
|
|
||
| // Store a list of our init calls so we don't repeat them | ||
| for stmt in init { | ||
| if hset.contains(*stmt) { | ||
| if ran_stmts.contains(*stmt) { | ||
| continue; | ||
| } | ||
| sql::<(Untyped,)>(stmt) | ||
| .execute(&mut conn) | ||
| .expect("could not run setup"); | ||
|
|
||
| hset.insert((*stmt).to_owned()); | ||
| conn.query_drop(stmt).expect("could not run setup"); | ||
|
|
||
| ran_stmts.insert((*stmt).to_owned()); | ||
| } | ||
|
|
||
| conn | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters