pub mod codegen;
pub mod config;
pub mod error;
pub mod parser;
use std::collections::HashSet;
use std::path::{Path, PathBuf};
use tracing::{debug, info};
pub use config::CodegenConfig;
pub use error::{CodegenError, Result};
pub fn generate(config: &CodegenConfig) -> Result<()> {
info!("Parsing schema: {:?}", config.schema_file);
let schema_sql = std::fs::read_to_string(&config.schema_file)?;
let tables = parser::parse_schema(&schema_sql)?;
info!("Found {} tables", tables.len());
let tables = filter_tables(tables, &config.include_tables, &config.exclude_tables);
debug!(
"After filtering: {} tables (include={}, exclude={})",
tables.len(),
config.include_tables,
config.exclude_tables
);
if config.generate_structs {
info!("Generating structs in {:?}", config.output_structs_dir);
codegen::generate_structs(&tables, config)?;
}
if config.generate_dao {
info!("Generating DAOs in {:?}", config.output_dao_dir);
codegen::generate_daos(&tables, config)?;
}
if config.generate_structs && config.generate_dao {
if let (Some(structs_parent), Some(dao_parent)) = (
config.output_structs_dir.parent(),
config.output_dao_dir.parent(),
) {
if structs_parent == dao_parent {
let structs_dir_name = config
.output_structs_dir
.file_name()
.and_then(|n| n.to_str())
.unwrap_or("models");
let dao_dir_name = config
.output_dao_dir
.file_name()
.and_then(|n| n.to_str())
.unwrap_or("dao");
let parent_mod_path = structs_parent.join("mod.rs");
let parent_mod_content = format!(
"// Generated by rdbi-codegen - do not edit manually\n#![allow(dead_code)]\n#![allow(clippy::all)]\n\npub mod {};\npub mod {};\n",
dao_dir_name, structs_dir_name
);
if !config.dry_run {
std::fs::write(&parent_mod_path, parent_mod_content)?;
info!("Generated parent mod.rs at {:?}", parent_mod_path);
}
}
}
}
info!("Code generation complete");
Ok(())
}
fn filter_tables(
tables: Vec<parser::TableMetadata>,
include: &str,
exclude: &str,
) -> Vec<parser::TableMetadata> {
let include_all = include.trim() == "*" || include.trim().is_empty();
let include_set: HashSet<String> = if include_all {
HashSet::new()
} else {
include.split(',').map(|s| s.trim().to_string()).collect()
};
let exclude_set: HashSet<String> = exclude
.split(',')
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect();
tables
.into_iter()
.filter(|t| {
let name = &t.name;
let included = include_all || include_set.contains(name);
let excluded = exclude_set.contains(name);
included && !excluded
})
.collect()
}
pub struct CodegenBuilder {
config: CodegenConfig,
}
impl CodegenBuilder {
pub fn new(schema_file: impl AsRef<Path>) -> Self {
Self {
config: CodegenConfig::default_with_schema(schema_file.as_ref().to_path_buf()),
}
}
pub fn output_dir(mut self, dir: impl AsRef<Path>) -> Self {
let dir = dir.as_ref();
self.config.output_structs_dir = dir.join("models");
self.config.output_dao_dir = dir.join("dao");
self
}
pub fn output_structs_dir(mut self, dir: impl AsRef<Path>) -> Self {
self.config.output_structs_dir = dir.as_ref().to_path_buf();
self
}
pub fn output_dao_dir(mut self, dir: impl AsRef<Path>) -> Self {
self.config.output_dao_dir = dir.as_ref().to_path_buf();
self
}
pub fn include_tables(mut self, tables: &[&str]) -> Self {
self.config.include_tables = tables.join(",");
self
}
pub fn exclude_tables(mut self, tables: &[&str]) -> Self {
self.config.exclude_tables = tables.join(",");
self
}
pub fn structs_only(mut self) -> Self {
self.config.generate_dao = false;
self
}
pub fn dao_only(mut self) -> Self {
self.config.generate_structs = false;
self
}
pub fn models_module(mut self, name: &str) -> Self {
self.config.models_module = name.to_string();
self
}
pub fn dao_module(mut self, name: &str) -> Self {
self.config.dao_module = name.to_string();
self
}
pub fn dry_run(mut self) -> Self {
self.config.dry_run = true;
self
}
pub fn generate(self) -> Result<()> {
generate(&self.config)
}
}
#[derive(Debug, Clone, Default, serde::Deserialize)]
struct CargoMetadataConfig {
schema_file: Option<String>,
#[serde(default)]
include_tables: Vec<String>,
#[serde(default)]
exclude_tables: Vec<String>,
generate_structs: Option<bool>,
generate_dao: Option<bool>,
output_structs_dir: Option<String>,
output_dao_dir: Option<String>,
}
#[derive(Debug, serde::Deserialize)]
struct CargoToml {
package: Option<CargoPackage>,
}
#[derive(Debug, serde::Deserialize)]
struct CargoPackage {
metadata: Option<CargoPackageMetadata>,
}
#[derive(Debug, serde::Deserialize)]
struct CargoPackageMetadata {
#[serde(rename = "rdbi-codegen")]
rdbi_codegen: Option<CargoMetadataConfig>,
}
pub fn generate_from_cargo_metadata() -> Result<()> {
let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").map_err(|_| {
CodegenError::ConfigError(
"CARGO_MANIFEST_DIR not set - are you running from build.rs?".into(),
)
})?;
let cargo_toml_path = PathBuf::from(&manifest_dir).join("Cargo.toml");
let cargo_toml_content = std::fs::read_to_string(&cargo_toml_path)?;
let cargo_toml: CargoToml = toml::from_str(&cargo_toml_content).map_err(|e| {
CodegenError::ConfigError(format!(
"Failed to parse {}: {}",
cargo_toml_path.display(),
e
))
})?;
let metadata_config = cargo_toml
.package
.and_then(|p| p.metadata)
.and_then(|m| m.rdbi_codegen)
.ok_or_else(|| {
CodegenError::ConfigError(
"Missing [package.metadata.rdbi-codegen] section in Cargo.toml".into(),
)
})?;
let schema_file = metadata_config.schema_file.ok_or_else(|| {
CodegenError::ConfigError(
"schema_file is required in [package.metadata.rdbi-codegen]".into(),
)
})?;
let schema_path = PathBuf::from(&manifest_dir).join(&schema_file);
let out_dir = std::env::var("OUT_DIR").map(PathBuf::from).map_err(|_| {
CodegenError::ConfigError("OUT_DIR not set - are you running from build.rs?".into())
})?;
let mut builder = CodegenBuilder::new(&schema_path);
if let Some(ref structs_dir) = metadata_config.output_structs_dir {
builder = builder.output_structs_dir(PathBuf::from(&manifest_dir).join(structs_dir));
let module_path = structs_dir
.strip_prefix("src/")
.unwrap_or(structs_dir)
.replace('/', "::");
builder = builder.models_module(&module_path);
} else {
builder = builder.output_structs_dir(out_dir.join("models"));
}
if let Some(ref dao_dir) = metadata_config.output_dao_dir {
builder = builder.output_dao_dir(PathBuf::from(&manifest_dir).join(dao_dir));
let module_path = dao_dir
.strip_prefix("src/")
.unwrap_or(dao_dir)
.replace('/', "::");
builder = builder.dao_module(&module_path);
} else {
builder = builder.output_dao_dir(out_dir.join("dao"));
}
if !metadata_config.include_tables.is_empty() {
let tables: Vec<&str> = metadata_config
.include_tables
.iter()
.map(|s| s.as_str())
.collect();
builder = builder.include_tables(&tables);
}
if !metadata_config.exclude_tables.is_empty() {
let tables: Vec<&str> = metadata_config
.exclude_tables
.iter()
.map(|s| s.as_str())
.collect();
builder = builder.exclude_tables(&tables);
}
if let Some(false) = metadata_config.generate_structs {
builder = builder.dao_only();
}
if let Some(false) = metadata_config.generate_dao {
builder = builder.structs_only();
}
println!("cargo:rerun-if-changed={}", schema_path.display());
println!("cargo:rerun-if-changed={}", cargo_toml_path.display());
builder.generate()
}