use super::partition::{
HistoricalAttribution, ModulePartition, UNMODULED_DISPLAY, detect_partition_divergence,
display_module, validate_module_references,
};
use crate::catalog::Catalog;
use crate::catalog::file_dependencies::FileToObjectMapping;
use crate::catalog::id::DbObjectId;
use crate::config::Config;
use anyhow::Result;
use std::collections::BTreeMap;
#[derive(Debug, Clone)]
pub struct StepSection {
pub name: String,
pub module: Option<String>,
pub remaps: Option<String>,
pub comment: Option<String>,
pub steps: Vec<crate::diff::operations::MigrationStep>,
}
pub(crate) fn provenance_cut_baseline_sections(
sections: Vec<StepSection>,
historical: &HistoricalAttribution,
) -> Vec<StepSection> {
use crate::diff::operations::MigrationStep;
use crate::diff::operations::SqlRenderer;
struct Run {
module: Option<String>,
remap: Option<String>,
steps: Vec<MigrationStep>,
}
let mut runs: Vec<Run> = Vec::new();
for section in sections {
let own = section.module.clone();
for step in section.steps {
let id = step.db_object_id();
let remap = match historical.object_modules.get(&id) {
Some(prior) if prior.as_deref() != own.as_deref() => {
Some(display_module(prior.as_deref()).to_string())
}
_ => None,
};
match runs.last_mut() {
Some(last) if last.module == own && last.remap == remap => last.steps.push(step),
_ => runs.push(Run {
module: own.clone(),
remap,
steps: vec![step],
}),
}
}
}
let mut out: Vec<StepSection> = runs
.into_iter()
.map(|run| StepSection {
name: String::new(),
module: run.module,
remaps: run.remap,
comment: None,
steps: run.steps,
})
.collect();
assign_section_names(&mut out);
out
}
pub(crate) fn assign_section_names(sections: &mut [StepSection]) {
let mut name_counts: BTreeMap<String, usize> = BTreeMap::new();
for section in sections.iter_mut() {
let base_name = section.module.as_deref().unwrap_or("default").to_string();
let count = name_counts.entry(base_name.clone()).or_insert(0);
*count += 1;
section.name = if *count == 1 {
base_name
} else {
format!("{}_{}", base_name, count)
};
}
}
pub fn sectionize_steps(
old_catalog: &Catalog,
new_catalog: &Catalog,
partition: &ModulePartition,
desired_mapping: &FileToObjectMapping,
historical: &HistoricalAttribution,
) -> Result<Vec<StepSection>> {
use crate::diff::operations::SqlRenderer;
let mut grant_targets: BTreeMap<DbObjectId, DbObjectId> = BTreeMap::new();
for grant in old_catalog.grants.iter().chain(new_catalog.grants.iter()) {
grant_targets.insert(
DbObjectId::Grant { id: grant.id() },
grant.target.object.clone(),
);
}
fn resolve_module(
id: &DbObjectId,
partition: &ModulePartition,
desired_mapping: &FileToObjectMapping,
historical: &HistoricalAttribution,
grant_targets: &BTreeMap<DbObjectId, DbObjectId>,
) -> Result<Option<String>> {
if let DbObjectId::Comment { object_id } = id {
return resolve_module(
object_id,
partition,
desired_mapping,
historical,
grant_targets,
);
}
if let DbObjectId::Grant { .. } = id
&& !desired_mapping.object_files.contains_key(id)
&& let Some(target) = grant_targets.get(id)
{
return resolve_module(
target,
partition,
desired_mapping,
historical,
grant_targets,
);
}
if desired_mapping.object_files.contains_key(id) {
return Ok(partition
.module_for_object(id, desired_mapping)?
.map(str::to_string));
}
if historical.object_modules.contains_key(id) {
return Ok(historical.module_of(id).map(str::to_string));
}
if let DbObjectId::Policy { schema, table, .. } = id {
return resolve_module(
&DbObjectId::Table {
schema: schema.clone(),
name: table.clone(),
},
partition,
desired_mapping,
historical,
grant_targets,
);
}
Ok(None)
}
let mut module_of = |step: &crate::diff::operations::MigrationStep| {
resolve_module(
&step.db_object_id(),
partition,
desired_mapping,
historical,
&grant_targets,
)
};
let ordered = crate::diff::plan_annotated(old_catalog, new_catalog, &mut module_of)?;
let mut sections: Vec<StepSection> = Vec::new();
for node in ordered {
match sections.last_mut() {
Some(last) if last.module == node.module => last.steps.push(node.step),
_ => sections.push(StepSection {
name: String::new(),
module: node.module,
remaps: None,
comment: None,
steps: vec![node.step],
}),
}
}
assign_section_names(&mut sections);
Ok(sections)
}
pub fn render_sectioned_migration(sections: &[StepSection]) -> String {
use crate::diff::operations::SqlRenderer;
let mut parts: Vec<String> = Vec::new();
for section in sections {
let mut header = format!("-- pgmt:section name=\"{}\"", section.name);
if let Some(module) = §ion.module {
header.push_str(&format!(" module=\"{}\"", module));
}
if let Some(source) = §ion.remaps {
header.push_str(&format!(" remaps=\"{}\"", source));
}
if let Some(comment) = §ion.comment {
header = format!("{}\n{}", comment, header);
}
let step_sql: Vec<String> = section
.steps
.iter()
.flat_map(|step| step.to_sql().into_iter().map(|r| r.sql))
.collect();
parts.push(format!("{}\n{}", header, step_sql.join("\n\n")));
}
parts.join("\n\n")
}
pub struct ModuleGeneration {
pub partition: ModulePartition,
pub diverged: bool,
}
pub fn evaluate_module_generation(
config: &Config,
old_catalog: &Catalog,
new_catalog: &Catalog,
file_mapping: &FileToObjectMapping,
historical: &HistoricalAttribution,
re_anchor_allowed: bool,
re_anchor_guidance: &str,
) -> Result<Option<ModuleGeneration>> {
if !config.modules.is_enabled() {
return Ok(None);
}
let partition = ModulePartition::from_config(config)?;
let report = validate_module_references(new_catalog, file_mapping, &partition, config)?;
if !report.is_clean() {
for warning in &report.warnings {
eprintln!("Warning: {}", warning);
}
if !report.errors.is_empty() {
anyhow::bail!(
"module reference validation failed:\n - {}",
report.errors.join("\n - ")
);
}
}
let divergence =
detect_partition_divergence(old_catalog, &partition, file_mapping, historical)?;
if !divergence.is_empty() {
if !re_anchor_allowed {
anyhow::bail!(
"partition re-anchor required:\n - {}\n\n\
Replaying module history would reproduce the old ownership; \
{}",
divergence.reasons.join("\n - "),
re_anchor_guidance
);
}
println!("Re-anchoring the module partition:");
for reason in &divergence.reasons {
println!(" - {}", reason);
}
}
Ok(Some(ModuleGeneration {
partition,
diverged: !divergence.is_empty(),
}))
}
pub(crate) fn write_sectioned_baseline(
path: &std::path::Path,
new_catalog: &Catalog,
partition: &ModulePartition,
file_mapping: &FileToObjectMapping,
historical: &HistoricalAttribution,
) -> Result<Vec<StepSection>> {
let sections = sectionize_steps(
&Catalog::empty(),
new_catalog,
partition,
file_mapping,
historical,
)?;
let sections = provenance_cut_baseline_sections(sections, historical);
std::fs::write(path, render_sectioned_migration(§ions))?;
Ok(sections)
}
pub(crate) fn acquisition_sections(provenance_cut_sections: &[StepSection]) -> Vec<StepSection> {
provenance_cut_sections
.iter()
.filter(|s| {
s.remaps
.as_deref()
.is_some_and(|source| source != UNMODULED_DISPLAY)
})
.map(|s| {
let source = s.remaps.as_ref().expect("filtered on non-empty remaps");
StepSection {
name: String::new(), module: s.module.clone(),
remaps: s.remaps.clone(),
comment: Some(
[
format!(
"-- objects moved from module '{source}'; runs only on targets without"
),
format!(
"-- it — targets holding '{source}' already have them (satisfied)."
),
]
.join("\n"),
),
steps: s.steps.clone(),
}
})
.collect()
}
fn acquisition_sections_from_starting_catalog(
old_catalog: &Catalog,
partition: &ModulePartition,
file_mapping: &FileToObjectMapping,
historical: &HistoricalAttribution,
) -> Result<Vec<StepSection>> {
let sections = sectionize_steps(
&Catalog::empty(),
old_catalog,
partition,
file_mapping,
historical,
)?;
let sections = provenance_cut_baseline_sections(sections, historical);
Ok(acquisition_sections(§ions))
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn render_migration_with_acquisitions(
has_changes: bool,
re_anchored: bool,
old_catalog: &Catalog,
new_catalog: &Catalog,
partition: &ModulePartition,
file_mapping: &FileToObjectMapping,
historical: &HistoricalAttribution,
) -> Result<Option<String>> {
let ordinary = if has_changes {
sectionize_steps(
old_catalog,
new_catalog,
partition,
file_mapping,
historical,
)?
} else {
Vec::new()
};
let mut sections = if re_anchored {
acquisition_sections_from_starting_catalog(
old_catalog,
partition,
file_mapping,
historical,
)?
} else {
Vec::new()
};
sections.extend(ordinary);
if sections.is_empty() {
return Ok(None);
}
assign_section_names(&mut sections);
Ok(Some(render_sectioned_migration(§ions)))
}
pub fn section_baseline_if_moduled(
module_gen: Option<&ModuleGeneration>,
baseline_path: &std::path::Path,
new_catalog: &Catalog,
file_mapping: &FileToObjectMapping,
historical: &HistoricalAttribution,
) -> Result<()> {
if let Some(module_gen) = module_gen {
write_sectioned_baseline(
baseline_path,
new_catalog,
&module_gen.partition,
file_mapping,
historical,
)?;
}
Ok(())
}
pub fn render_generated_migration(
module_gen: Option<&ModuleGeneration>,
has_changes: bool,
plain_migration_sql: &str,
old_catalog: &Catalog,
new_catalog: &Catalog,
file_mapping: &FileToObjectMapping,
historical: &HistoricalAttribution,
) -> Result<Option<String>> {
match module_gen {
Some(module_gen) => render_migration_with_acquisitions(
has_changes,
module_gen.diverged,
old_catalog,
new_catalog,
&module_gen.partition,
file_mapping,
historical,
),
None => Ok(has_changes.then(|| plain_migration_sql.to_string())),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::catalog::id::DbObjectId;
fn step_section(module: Option<&str>, objects: &[&str]) -> StepSection {
use crate::diff::operations::{MigrationStep, SchemaOperation};
StepSection {
name: module.unwrap_or("default").to_string(),
module: module.map(str::to_string),
remaps: None,
comment: None,
steps: objects
.iter()
.map(|name| {
MigrationStep::Schema(SchemaOperation::Create {
name: name.to_string(),
})
})
.collect(),
}
}
fn historical(entries: &[(&str, Option<&str>)]) -> HistoricalAttribution {
let mut h = HistoricalAttribution::default();
for (name, module) in entries {
h.object_modules.insert(
DbObjectId::Schema {
name: name.to_string(),
},
module.map(str::to_string),
);
}
h
}
fn shape(section: &StepSection) -> (String, Option<String>, Option<String>) {
(
section.name.clone(),
section.module.clone(),
section.remaps.clone(),
)
}
#[test]
fn test_provenance_cut_move_into_brand_new_module() {
let sections = vec![step_section(Some("b"), &["x", "y"])];
let hist = historical(&[("x", Some("a")), ("y", Some("a"))]);
let out = provenance_cut_baseline_sections(sections, &hist);
assert_eq!(out.len(), 1);
assert_eq!(
shape(&out[0]),
(
"b".to_string(),
Some("b".to_string()),
Some("a".to_string())
)
);
}
#[test]
fn test_provenance_cut_move_into_pre_existing_module() {
let sections = vec![step_section(Some("b"), &["own", "acquired"])];
let hist = historical(&[("own", Some("b")), ("acquired", Some("a"))]);
let out = provenance_cut_baseline_sections(sections, &hist);
assert_eq!(
out.len(),
2,
"{:?}",
out.iter().map(shape).collect::<Vec<_>>()
);
assert_eq!(
shape(&out[0]),
("b".to_string(), Some("b".to_string()), None)
);
assert_eq!(
shape(&out[1]),
(
"b_2".to_string(),
Some("b".to_string()),
Some("a".to_string())
)
);
}
#[test]
fn test_provenance_cut_untouched_module_stays_plain() {
let sections = vec![step_section(Some("b"), &["p", "q"])];
let hist = historical(&[("p", Some("b")), ("q", Some("b"))]);
let out = provenance_cut_baseline_sections(sections, &hist);
assert_eq!(out.len(), 1);
assert_eq!(
shape(&out[0]),
("b".to_string(), Some("b".to_string()), None)
);
}
#[test]
fn test_provenance_cut_untouched_with_brand_new_objects_stays_plain() {
let sections = vec![step_section(Some("b"), &["kept", "fresh"])];
let hist = historical(&[("kept", Some("b"))]); let out = provenance_cut_baseline_sections(sections, &hist);
assert_eq!(out.len(), 1);
assert!(out[0].remaps.is_none(), "{:?}", out[0].remaps);
}
#[test]
fn test_provenance_cut_demotion_into_base() {
let sections = vec![step_section(None, &["x"])];
let hist = historical(&[("x", Some("a"))]);
let out = provenance_cut_baseline_sections(sections, &hist);
assert_eq!(out.len(), 1);
assert_eq!(
shape(&out[0]),
("default".to_string(), None, Some("a".to_string()))
);
}
#[test]
fn test_provenance_cut_modularization_from_base() {
let sections = vec![step_section(Some("app"), &["users"])];
let hist = historical(&[("users", None)]);
let out = provenance_cut_baseline_sections(sections, &hist);
assert_eq!(out.len(), 1);
assert_eq!(
shape(&out[0]),
(
"app".to_string(),
Some("app".to_string()),
Some(UNMODULED_DISPLAY.to_string())
)
);
}
#[test]
fn test_provenance_cut_acquisition_from_two_sources_splits() {
let sections = vec![step_section(Some("c"), &["from_a", "from_b"])];
let hist = historical(&[("from_a", Some("a")), ("from_b", Some("b"))]);
let out = provenance_cut_baseline_sections(sections, &hist);
assert_eq!(
out.len(),
2,
"{:?}",
out.iter().map(shape).collect::<Vec<_>>()
);
assert_eq!(
shape(&out[0]),
(
"c".to_string(),
Some("c".to_string()),
Some("a".to_string())
)
);
assert_eq!(
shape(&out[1]),
(
"c_2".to_string(),
Some("c".to_string()),
Some("b".to_string())
)
);
}
}