use std::collections::BTreeSet;
use std::path::{Path, PathBuf};
use crate::application::{CreateDocumentInput, EntityUseCases};
use crate::catalog::CollectionModel;
use crate::presentation::entity_json::storage_json_bytes_to_json;
use crate::storage::schema::Value;
use crate::storage::EntityData;
use crate::{RedDBError, RedDBOptions, RedDBResult, RedDBRuntime};
const DATA_FILE: &str = "db.rdb";
const MIGRATING_SUFFIX: &str = "migrating";
const BACKUP_SUFFIX: &str = "pre-migration";
#[derive(Debug, Clone)]
pub struct CollectionMigration {
pub name: String,
pub source_documents: usize,
pub migrated_documents: usize,
pub auto_indexed_fields: Vec<String>,
}
#[derive(Debug, Clone)]
pub struct MigrationReport {
pub collections: Vec<CollectionMigration>,
pub backup_dir: PathBuf,
pub total_documents: usize,
}
struct SourceCollection {
name: String,
bodies: Vec<crate::json::Value>,
legacy_fields: Vec<String>,
}
pub fn migrate_store_to_binary_body(store_dir: &Path) -> RedDBResult<MigrationReport> {
if !store_dir.is_dir() {
return Err(RedDBError::InvalidOperation(format!(
"migration source is not a directory: {}",
store_dir.display()
)));
}
let collections = read_source_collections(store_dir)?;
let migrating_dir = sibling_dir(store_dir, MIGRATING_SUFFIX)?;
let outcome = build_migrated_store(&migrating_dir, &collections);
let report_collections = match outcome {
Ok(report) => report,
Err(err) => {
let _ = std::fs::remove_dir_all(&migrating_dir);
return Err(err);
}
};
let backup_dir = swap_in_migrated_store(store_dir, &migrating_dir)?;
let total_documents = report_collections
.iter()
.map(|c| c.migrated_documents)
.sum();
Ok(MigrationReport {
collections: report_collections,
backup_dir,
total_documents,
})
}
fn read_source_collections(store_dir: &Path) -> RedDBResult<Vec<SourceCollection>> {
let source = RedDBRuntime::with_options(RedDBOptions::persistent(store_dir.join(DATA_FILE)))?;
let mut document_collections: Vec<String> = source
.catalog()
.collections
.into_iter()
.filter(|descriptor| descriptor.model == CollectionModel::Document)
.map(|descriptor| descriptor.name)
.collect();
document_collections.sort();
let mut out = Vec::with_capacity(document_collections.len());
for name in document_collections {
let mut bodies = Vec::new();
let mut legacy_fields: BTreeSet<String> = BTreeSet::new();
let mut cursor = None;
loop {
let page = source.scan_collection(&name, cursor, 1_000)?;
for entity in &page.items {
let EntityData::Row(row) = &entity.data else {
continue;
};
for (field, _value) in row.iter_fields() {
if field != "body"
&& !crate::reserved_fields::is_reserved_public_item_field(field)
{
legacy_fields.insert(field.to_string());
}
}
bodies.push(decode_body(row)?);
}
match page.next {
Some(next) => cursor = Some(next),
None => break,
}
}
out.push(SourceCollection {
name,
bodies,
legacy_fields: legacy_fields.into_iter().collect(),
});
}
Ok(out)
}
fn decode_body(row: &crate::storage::RowData) -> RedDBResult<crate::json::Value> {
match row.get_field("body") {
Some(Value::Json(bytes)) => Ok(storage_json_bytes_to_json(bytes)),
Some(other) => Ok(crate::presentation::entity_json::storage_value_to_json(
other,
)),
None => Err(RedDBError::InvalidOperation(
"document row is missing its `body` field".to_string(),
)),
}
}
fn build_migrated_store(
migrating_dir: &Path,
collections: &[SourceCollection],
) -> RedDBResult<Vec<CollectionMigration>> {
std::fs::create_dir_all(migrating_dir).map_err(RedDBError::Io)?;
let target =
RedDBRuntime::with_options(RedDBOptions::persistent(migrating_dir.join(DATA_FILE)))?;
target.execute_query("SET CONFIG storage.binary_document_body = true")?;
let entities = EntityUseCases::new(&target);
let mut report = Vec::with_capacity(collections.len());
for collection in collections {
target.execute_query(&format!("CREATE DOCUMENT {}", collection.name))?;
for body in &collection.bodies {
entities.create_document(CreateDocumentInput {
collection: collection.name.clone(),
body: body.clone(),
metadata: Vec::new(),
node_links: Vec::new(),
vector_links: Vec::new(),
})?;
}
for field in &collection.legacy_fields {
target.execute_query(&format!(
"CREATE INDEX {index} ON {collection} ({field}) USING BTREE",
index = auto_index_name(&collection.name, field),
collection = collection.name,
field = field,
))?;
}
let migrated = target.scan_collection(&collection.name, None, 1)?.total;
let source = collection.bodies.len();
ensure_counts_match(&collection.name, source, migrated)?;
report.push(CollectionMigration {
name: collection.name.clone(),
source_documents: source,
migrated_documents: migrated,
auto_indexed_fields: collection.legacy_fields.clone(),
});
}
target.flush()?;
target.checkpoint()?;
drop(target);
Ok(report)
}
fn swap_in_migrated_store(store_dir: &Path, migrating_dir: &Path) -> RedDBResult<PathBuf> {
let backup_dir = sibling_dir(store_dir, BACKUP_SUFFIX)?;
std::fs::rename(store_dir, &backup_dir).map_err(RedDBError::Io)?;
if let Err(err) = std::fs::rename(migrating_dir, store_dir) {
let _ = std::fs::rename(&backup_dir, store_dir);
return Err(RedDBError::Io(err));
}
Ok(backup_dir)
}
fn sibling_dir(store_dir: &Path, suffix: &str) -> RedDBResult<PathBuf> {
let file_name = store_dir
.file_name()
.and_then(|name| name.to_str())
.ok_or_else(|| {
RedDBError::InvalidOperation(format!(
"store directory has no usable name: {}",
store_dir.display()
))
})?;
let parent = store_dir.parent().unwrap_or_else(|| Path::new("."));
let candidate = parent.join(format!("{file_name}.{suffix}"));
if candidate.exists() {
return Err(RedDBError::InvalidOperation(format!(
"migration sibling directory already exists: {}",
candidate.display()
)));
}
Ok(candidate)
}
fn ensure_counts_match(collection: &str, source: usize, migrated: usize) -> RedDBResult<()> {
if migrated != source {
return Err(RedDBError::InvalidOperation(format!(
"document count mismatch for collection '{collection}': source {source} != \
migrated {migrated}; aborting before swap"
)));
}
Ok(())
}
fn auto_index_name(collection: &str, field: &str) -> String {
fn sanitize(input: &str) -> String {
input
.chars()
.map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
.collect()
}
format!("mig_{}_{}", sanitize(collection), sanitize(field))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn matching_counts_pass_mismatch_aborts() {
assert!(ensure_counts_match("docs", 4, 4).is_ok());
let err = ensure_counts_match("docs", 4, 3).expect_err("mismatch aborts");
assert!(matches!(err, RedDBError::InvalidOperation(_)));
}
#[test]
fn auto_index_name_is_identifier_safe() {
assert_eq!(auto_index_name("docs", "score"), "mig_docs_score");
assert_eq!(
auto_index_name("my-docs", "user.name"),
"mig_my_docs_user_name"
);
}
#[test]
fn sibling_dir_rejects_existing() {
let base = std::env::temp_dir().join(format!(
"reddb-mig-sibling-{}-{}",
std::process::id(),
"store"
));
let _ = std::fs::remove_dir_all(&base);
std::fs::create_dir_all(&base).unwrap();
let sib = sibling_dir(&base, "migrating").unwrap();
assert!(sib
.file_name()
.and_then(|name| name.to_str())
.is_some_and(|name| name.ends_with(".migrating")));
std::fs::create_dir_all(&sib).unwrap();
assert!(sibling_dir(&base, "migrating").is_err());
let _ = std::fs::remove_dir_all(&base);
let _ = std::fs::remove_dir_all(&sib);
}
}