Files
sea-orm-generator/src/generator/modules/discovery/mod.rs
2025-04-08 08:13:41 +04:00

138 lines
4.0 KiB
Rust

pub mod column;
pub mod db;
pub mod table;
use crate::generator::DatabaseUrl;
use super::{Module, ModulesContext};
use color_eyre::Result;
use db::DbType;
use serde::Deserialize;
use serde_inline_default::serde_inline_default;
use table::Table;
#[derive(Debug, Clone, Deserialize)]
#[serde(default)]
pub struct DiscoveryConfig {
pub enable: bool,
pub database_schema: Option<String>,
pub max_connections: u32,
pub acquire_timeout: u64,
pub filter: DiscoveryFilterConfig,
}
impl Default for DiscoveryConfig {
fn default() -> Self {
Self {
enable: false,
database_schema: None,
max_connections: 10,
acquire_timeout: 30,
filter: DiscoveryFilterConfig::default(),
}
}
}
#[derive(Debug, Clone, Deserialize)]
#[serde(default)]
pub struct DiscoveryFilterConfig {
pub include_hidden: bool,
pub skip_seaql_migrations: bool,
#[serde(flatten)]
pub table: Option<TableConfig>,
}
impl Default for DiscoveryFilterConfig {
fn default() -> Self {
Self {
include_hidden: false,
skip_seaql_migrations: true,
table: None,
}
}
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "snake_case")]
#[serde(untagged)]
pub enum TableConfig {
Specific { only: Vec<String> },
Exclude { exclude: Vec<String> },
}
impl DiscoveryFilterConfig {
pub fn get_filter(&self) -> Box<dyn Fn(&String) -> bool + Send> {
let include_hidden = self.include_hidden;
if let Some(table) = &self.table {
match table {
TableConfig::Specific { only } => {
let specific = only.clone();
Box::new(move |table: &String| {
(include_hidden || !table.starts_with('_')) && specific.contains(table)
})
}
TableConfig::Exclude { exclude } => {
let exclude = exclude.clone();
Box::new(move |table: &String| {
(include_hidden || !table.starts_with('_')) && !exclude.contains(table)
})
}
}
} else if self.skip_seaql_migrations {
Box::new(move |table: &String| {
(include_hidden || !table.starts_with('_'))
&& !table.starts_with("seaql_migrations")
})
} else {
Box::new(move |table: &String| (include_hidden || !table.starts_with('_')))
}
}
}
#[derive(Debug, Clone)]
pub struct DiscoveredSchema {
pub tables: Vec<Table>,
pub database_type: DbType,
}
#[derive(Debug)]
pub struct DiscoveryModule;
#[async_trait::async_trait]
impl Module for DiscoveryModule {
fn init(&mut self, ctx: &mut ModulesContext) -> Result<()> {
ctx.get_config_auto::<DiscoveryConfig>("modules.discovery")?;
Ok(())
}
async fn validate(&mut self, ctx: &mut ModulesContext) -> Result<bool> {
let map = ctx.get_anymap();
if let (Some(config), Some(_)) = (map.get::<DiscoveryConfig>(), map.get::<DatabaseUrl>()) {
Ok(config.enable)
} else {
// One or both keys are missing
Ok(false)
}
}
async fn execute(&mut self, ctx: &mut ModulesContext) -> Result<()> {
if let (Some(config), Some(url)) = (
ctx.get_anymap().get::<DiscoveryConfig>(),
ctx.get_anymap().get::<DatabaseUrl>(),
) {
let url = url.clone();
let (stmts, db_type) = db::get_tables(url.0, config).await?;
let tables = stmts
.into_iter()
.map(Table::new)
.collect::<Result<Vec<Table>>>()?;
tracing::info!(?tables, ?db_type);
let discovered = DiscoveredSchema {
tables,
database_type: db_type,
};
ctx.get_anymap_mut().insert(discovered);
// db::generate(ctx).await?;
}
Ok(())
}
}