use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use anyhow::{Context, bail};
use heck::ToSnakeCase;
use crate::config::{Config, Derive};
use crate::output::is_generated_file;
use crate::queries::AnalyzedQuery;
use crate::report::{Error, QueryFailure, RefusedFile, Report, Result, TableFailure, display_path};
use crate::schema::{BaseType, EnumType, Shape, Table};
use crate::{generator, output, parser, queries};
pub(crate) fn generate(config: &Config) -> Result<Report> {
let mut report = Report::default();
match run(config, &mut report) {
Ok(()) => Ok(report),
Err(err) => Err(Error::new(err, report)),
}
}
fn run(config: &Config, report: &mut Report) -> anyhow::Result<()> {
let output_dir = &config.output;
let resource = config.resource.as_deref();
let (dry_run, force, lenient) = (config.dry_run, config.force, config.lenient);
check_options(config)?;
for (key, module) in &config.modules {
if parser::naming::rust_identifier(module).as_deref() != Ok(module.as_str()) {
bail!(
"{} in the configuration renames the module to `{module}`, which is not a snake_case Rust identifier",
module_location(key)
);
}
}
let sources = read_schema(config)?;
let mut parsed = parser::parse_sources_with_modules(
&sources
.iter()
.map(|(path, sql)| parser::SchemaSource { path, sql })
.collect::<Vec<_>>(),
&config.modules,
)
.context("failed to parse SQL schema")?;
if let Some(key) = parsed.unmatched_modules.first() {
bail!(
"{} in the configuration names no table of the schema",
module_location(key)
);
}
parsed.fail_tables(|table| generator::table_failures(table, config));
let sql = schema_snapshot(&sources);
let matches_resource =
|table: &str| resource.is_none_or(|name| table == name || table.to_snake_case() == name);
report.table_failures = parsed
.failures
.iter()
.filter(|failure| matches_resource(&failure.table))
.map(|failure| TableFailure {
table: failure.table.clone(),
reason: failure.reason.clone(),
path: failure.path.clone(),
line: failure.line,
})
.collect();
if !report.table_failures.is_empty() && !lenient {
bail!(
"{} unsupported table(s) found; pass --lenient to skip them",
report.table_failures.len()
);
}
let mut include: HashSet<String> = parsed
.schema
.tables
.iter()
.filter(|table| {
resource.is_none_or(|name| table.sql_name == name || table.module_name == name)
})
.map(|table| table.module_name.clone())
.collect();
if let Ok(entries) = std::fs::read_dir(output_dir) {
for entry in entries.flatten() {
let module = entry.file_name().to_string_lossy().into_owned();
let already_generated = entry.path().join("models.rs").is_file();
if already_generated
&& parsed
.schema
.tables
.iter()
.any(|table| table.module_name == module)
{
include.insert(module);
}
}
}
let selected: Vec<_> = parsed
.schema
.tables
.iter()
.filter(|table| include.contains(&table.module_name))
.cloned()
.collect();
if selected.is_empty() && report.table_failures.is_empty() {
bail!("no matching resource found for {resource:?}");
}
if selected.is_empty() {
if resource.is_some() {
bail!("no generatable table matches {resource:?}");
}
bail!("all tables are unsupported; see warnings above");
}
generator::conventions::check(&selected, &config.conventions).map_err(anyhow::Error::msg)?;
let kept = kept_resources(
output_dir,
&parsed.failures,
&selected,
&parsed.schema.enums,
);
let schema_tables: Vec<Table> = parsed
.schema
.tables
.iter()
.chain(&kept.tables)
.cloned()
.collect();
generator::relationships::check_changes(
&schema_tables,
&config.conventions,
&config.relationships,
)
.map_err(anyhow::Error::msg)?;
let query_inputs = query_inputs(&config.queries)?;
if !query_inputs.is_empty() && selected.iter().any(|table| table.module_name == "queries") {
bail!(
"table module `queries` collides with the generated `queries` module; rename the table or omit --queries"
);
}
let mut analyzed = Vec::new();
let mut failed_files = Vec::new();
for input in &query_inputs {
let (queries, failures) = analyze_query_file(&input.path, &parsed, &selected)?;
if !failures.is_empty() {
failed_files.push(format!("`{}`", display_path(&input.path)));
}
report.query_failures.extend(failures);
analyzed.push(generator::QueryModule {
name: input.module.clone().unwrap_or_default(),
queries,
});
}
if !report.query_failures.is_empty() && !lenient {
let count = report.query_failures.len();
bail!(
"{count} unsupported quer{} in {}; pass --lenient to skip them",
if count == 1 { "y" } else { "ies" },
failed_files.join(", ")
);
}
let queries_module = match analyzed.as_slice() {
[single] if query_inputs[0].module.is_none() => {
generator::QueriesModule::Generate(&single.queries)
}
[] if is_generated_file(&output_dir.join("queries.rs"))
|| is_generated_file(&output_dir.join("queries").join("mod.rs")) =>
{
generator::QueriesModule::Keep
}
[] => generator::QueriesModule::None,
modules => generator::QueriesModule::GenerateModules(modules),
};
let files = generator::plan_files_with_queries(
config,
&sql,
&selected,
&parsed.schema.enums,
queries_module,
&kept,
)
.map_err(anyhow::Error::msg)?;
let keep: HashSet<PathBuf> = files.iter().map(|(path, _)| path.clone()).collect();
let mut stale = output::find_stale(output_dir, &keep)?;
if matches!(queries_module, generator::QueriesModule::Keep) {
let (single, dir) = (output_dir.join("queries.rs"), output_dir.join("queries"));
stale.retain(|path| *path != single && !path.starts_with(&dir));
}
for table in &kept.tables {
let dir = output_dir.join(&table.module_name);
stale.retain(|path| !path.starts_with(&dir));
}
if config.check {
(report.missing, report.changed) = output::compare_files(&files)?;
report.stale = stale;
if report.has_drift() {
let count = report.missing.len() + report.changed.len() + report.stale.len();
bail!(
"generated code is out of date: {count} file(s) drifted; run `rok-sqlc generate` to update them"
);
}
report.unchanged = files.into_iter().map(|(path, _)| path).collect();
return Ok(());
}
let applied = output::apply_files(&files, dry_run, force)?;
report.unchanged = applied.unchanged;
report.refused = applied
.refused
.into_iter()
.map(|(path, reason)| RefusedFile { path, reason })
.collect();
if !report.refused.is_empty() && !dry_run {
bail!(
"nothing was written; {} file(s) refused (use --force to overwrite)",
report.refused.len()
);
}
report.written = applied.written;
if !dry_run {
output::remove_stale(output_dir, &stale)?;
}
report.stale = stale;
Ok(())
}
fn check_options(config: &Config) -> anyhow::Result<()> {
let validate = config.derive.contains(&Derive::Validate);
if validate && !config.derive.contains(&Derive::Serde) {
bail!(
"`validate` in `derive` needs `serde`, since validated inputs and patches are deserialized: add `serde` to `derive`"
);
}
if !config.extractors.is_empty() && !validate {
bail!(
"`extractors` needs `validate` in `derive`, since the extractors validate the bodies they take: add `serde` and `validate` to `derive`"
);
}
Ok(())
}
fn kept_resources(
output_dir: &Path,
failures: &[parser::TableFailure],
selected: &[Table],
enums: &[EnumType],
) -> generator::KeptResources {
let mut kept = generator::KeptResources::default();
for last_good in failures
.iter()
.filter_map(|failure| failure.last_good.as_ref())
{
let table = &last_good.table;
let module = &table.module_name;
let taken = selected
.iter()
.chain(&kept.tables)
.any(|other| other.module_name == *module);
if taken || !is_generated_file(&output_dir.join(module).join("models.rs")) {
continue;
}
let defined = |path: &str| {
enums
.iter()
.chain(&last_good.enums)
.any(|enum_type| enum_type.rust_path() == path)
};
let enums_generate = table.columns.iter().all(|column| {
let value = match Shape::of(&column.sqlx_type) {
Shape::List(item) => item,
Shape::Optional(_) | Shape::Value(_) => column.sqlx_type.as_str(),
};
BaseType::of(value) != BaseType::Enum || defined(value)
});
if enums_generate {
kept.tables.push(table.clone());
kept.enums.extend(last_good.enums.iter().cloned());
}
}
kept
}
fn read_schema(config: &Config) -> anyhow::Result<Vec<(PathBuf, String)>> {
if config.schema.is_empty() {
bail!("no schema input given");
}
let mut files = Vec::new();
for input in &config.schema {
if input.is_dir() {
let migrations = migration_files(input)?;
if migrations.is_empty() {
bail!(
"schema directory `{}` contains no `*.sql` migrations",
display_path(input)
);
}
files.extend(migrations);
} else {
files.push(input.clone());
}
}
files
.into_iter()
.map(|path| {
let sql = read_text(&path)
.with_context(|| format!("failed to read schema `{}`", display_path(&path)))?;
Ok((path, sql))
})
.collect()
}
fn read_text(path: &Path) -> std::io::Result<String> {
let text = std::fs::read_to_string(path)?;
Ok(if text.contains("\r\n") {
text.replace("\r\n", "\n")
} else {
text
})
}
fn migration_files(dir: &Path) -> anyhow::Result<Vec<PathBuf>> {
let read_error = || format!("failed to read schema directory `{}`", display_path(dir));
let mut files = Vec::new();
for entry in std::fs::read_dir(dir).with_context(read_error)? {
let entry = entry.with_context(read_error)?;
let name = entry.file_name().to_string_lossy().to_lowercase();
if name.ends_with(".sql") && !name.ends_with(".down.sql") && entry.path().is_file() {
files.push(entry.path());
}
}
files.sort_by(|a, b| a.file_name().cmp(&b.file_name()));
Ok(files)
}
fn schema_snapshot(sources: &[(PathBuf, String)]) -> String {
let mut sql = String::new();
for (_, text) in sources {
if !sql.is_empty() && !sql.ends_with('\n') {
sql.push('\n');
}
sql.push_str(text);
}
sql
}
struct QueryInput {
path: PathBuf,
module: Option<String>,
}
fn query_inputs(inputs: &[PathBuf]) -> anyhow::Result<Vec<QueryInput>> {
if let [path] = inputs
&& !path.is_dir()
{
return Ok(vec![QueryInput {
path: path.clone(),
module: None,
}]);
}
let mut files = Vec::new();
for input in inputs {
if !input.is_dir() {
files.push(input.clone());
continue;
}
let read_error = || format!("failed to read query directory `{}`", display_path(input));
let mut found = std::fs::read_dir(input)
.with_context(read_error)?
.map(|entry| entry.map(|entry| entry.path()))
.collect::<std::io::Result<Vec<_>>>()
.with_context(read_error)?;
found.retain(|path| path.is_file() && path.extension().is_some_and(|ext| ext == "sql"));
if found.is_empty() {
bail!(
"query directory `{}` contains no `.sql` files",
display_path(input)
);
}
found.sort();
files.extend(found);
}
let mut seen: HashMap<String, PathBuf> = HashMap::new();
let mut modules = Vec::with_capacity(files.len());
for path in files {
let module = query_module_name(&path)?;
if let Some(first) = seen.insert(module.clone(), path.clone()) {
bail!(
"query files `{}` and `{}` would both generate module `{module}`; rename one of them",
display_path(&first),
display_path(&path)
);
}
modules.push(QueryInput {
path,
module: Some(module),
});
}
Ok(modules)
}
fn query_module_name(path: &Path) -> anyhow::Result<String> {
let stem = path
.file_stem()
.map(|stem| stem.to_string_lossy().into_owned())
.unwrap_or_default();
match parser::naming::rust_identifier(&stem) {
Ok(ident) if ident == stem => Ok(stem),
suggestion => {
let hint = match suggestion {
Ok(ident) => format!("; rename it to `{ident}.sql`"),
Err(_) => String::new(),
};
bail!(
"query file `{}`: `{stem}` is not a valid Rust module name (use lowercase letters, digits and underscores, not starting with a digit or naming a keyword){hint}",
display_path(path)
)
}
}
}
fn analyze_query_file(
path: &Path,
parsed: &parser::ParsedSchema,
selected: &[Table],
) -> anyhow::Result<(Vec<AnalyzedQuery>, Vec<QueryFailure>)> {
let text = read_text(path)
.with_context(|| format!("failed to read query file `{}`", display_path(path)))?;
let (defs, mut failures) = queries::parse_query_file(&text);
if defs.is_empty() && failures.is_empty() {
bail!(
"query file `{}` contains no `-- name:` annotated queries",
display_path(path)
);
}
let generated_modules: HashSet<String> = selected
.iter()
.map(|table| table.module_name.clone())
.collect();
let mut tables = generator::keys::resolve(&parsed.schema.tables);
for column in tables.iter_mut().flat_map(|table| &mut table.columns) {
if column
.typed_id
.as_ref()
.is_some_and(|id| !generated_modules.contains(&id.module))
{
column.typed_id = None;
}
}
let (analyzed, analysis_failures) = queries::analyze_queries(
&defs,
&tables,
&parsed.schema.enums,
&generated_modules,
&parsed.failures,
);
failures.extend(analysis_failures);
let failures = failures
.into_iter()
.map(|failure| QueryFailure {
query: failure.query,
reason: failure.reason,
path: failure.line.map(|_| path.to_path_buf()),
line: failure.line,
})
.collect();
Ok((analyzed, failures))
}
fn module_location(key: &str) -> String {
if key
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
{
format!("`modules.{key}`")
} else {
format!("`modules.\"{key}\"`")
}
}