use core::time; use color_eyre::eyre::{eyre, ContextCompat, Report, Result}; use sea_schema::sea_query::TableCreateStatement; use url::Url; use crate::config::db::DbConfig; #[derive(Debug, Clone)] pub enum DbType { MySql, Postgres, Sqlite, } pub async fn get_tables( database_url: String, filter: Box bool>, database_config: &DbConfig, ) -> Result<(Vec, DbType)> { let url = Url::parse(&database_url)?; tracing::trace!(?url); let is_sqlite = url.scheme() == "sqlite"; let database_name: &str = (if !is_sqlite { let database_name = url .path_segments() .context("No database name as part of path")? .next() .context("No database name as part of path")?; if database_name.is_empty() { return Err(eyre!("Database path name is empty")); } Ok::<&str, Report>(database_name) } else { Ok(Default::default()) })?; let (table_stmts, db_type) = match url.scheme() { "mysql" => { use sea_schema::mysql::discovery::SchemaDiscovery; use sqlx::MySql; tracing::info!("Connecting to MySQL"); let connection = sqlx_connect::( database_config.max_connections, database_config.acquire_timeout, url.as_str(), None, ) .await?; tracing::info!("Discovering schema"); let schema_discovery = SchemaDiscovery::new(connection, database_name); let schema = schema_discovery.discover().await?; let table_stmts = schema .tables .into_iter() .filter(|schema| filter(&schema.info.name)) .map(|schema| schema.write()) .collect(); (table_stmts, DbType::MySql) } "sqlite" => { use sea_schema::sqlite::discovery::SchemaDiscovery; use sqlx::Sqlite; tracing::info!("Connecting to SQLite"); let connection = sqlx_connect::( database_config.max_connections, database_config.acquire_timeout, url.as_str(), None, ) .await?; tracing::info!("Discovering schema"); let schema_discovery = SchemaDiscovery::new(connection); let schema = schema_discovery .discover() .await? .merge_indexes_into_table(); let table_stmts = schema .tables .into_iter() .filter(|schema| filter(&schema.name)) .map(|schema| schema.write()) .collect(); (table_stmts, DbType::Sqlite) } "postgres" | "potgresql" => { use sea_schema::postgres::discovery::SchemaDiscovery; use sqlx::Postgres; tracing::info!("Connecting to Postgres"); let schema = &database_config .database_schema .as_deref() .unwrap_or("public"); let connection = sqlx_connect::( database_config.max_connections, database_config.acquire_timeout, url.as_str(), Some(schema), ) .await?; tracing::info!("Discovering schema"); let schema_discovery = SchemaDiscovery::new(connection, schema); let schema = schema_discovery.discover().await?; tracing::info!(?schema); let table_stmts = schema .tables .into_iter() .filter(|schema| filter(&schema.info.name)) .map(|schema| schema.write()) .collect(); (table_stmts, DbType::Postgres) } _ => unimplemented!("{} is not supported", url.scheme()), }; tracing::info!("Schema discovered"); Ok((table_stmts, db_type)) } async fn sqlx_connect( max_connections: u32, acquire_timeout: u64, url: &str, schema: Option<&str>, ) -> Result> where DB: sqlx::Database, for<'a> &'a mut ::Connection: sqlx::Executor<'a>, { let mut pool_options = sqlx::pool::PoolOptions::::new() .max_connections(max_connections) .acquire_timeout(time::Duration::from_secs(acquire_timeout)); // Set search_path for Postgres, E.g. Some("public") by default // MySQL & SQLite connection initialize with schema `None` if let Some(schema) = schema { let sql = format!("SET search_path = '{schema}'"); pool_options = pool_options.after_connect(move |conn, _| { let sql = sql.clone(); Box::pin(async move { sqlx::Executor::execute(conn, sql.as_str()) .await .map(|_| ()) }) }); } pool_options.connect(url).await.map_err(Into::into) }