hayson-engine 0.27.0

use std::collections::HashMap;
use std::fs;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;

use serde_json::Value;

use crate::schema::{assert_partition_value, is_partitioned, TableSchema};

use super::log_file::{LineCodec, LogFile, Slot};
use super::log_set::LOG_SUFFIX;
use super::record::LogOp;

const STAGING_SUFFIX: &str = ".repartition";

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RepartitionResult {
    pub records: usize,
    pub partitions: Vec<String>,
}

fn staging_of(partition_dir: &Path) -> PathBuf {
    let mut staging = partition_dir.as_os_str().to_os_string();
    staging.push(STAGING_SUFFIX);
    PathBuf::from(staging)
}

/// Brings a table that declares `partitionBy` into the layout its schema asks for, moving the live
/// records of an existing flat log into one file per partition value, and finishes a move an
/// earlier crash left half-done.
///
/// Idempotent by design: the staging directory is thrown away on sight, a flat log left next to a
/// finished partition directory is the copy the move did not get to delete, and a table with
/// neither is simply new. Returns `None` when there was nothing to move.
pub fn ensure_partitioned(
    flat: &Path,
    partition_dir: &Path,
    schema: &TableSchema,
) -> io::Result<Option<RepartitionResult>> {
    ensure_partitioned_with(flat, partition_dir, schema, None)
}

/// [`ensure_partitioned`] for a table whose log lines are sealed by `codec`: the flat log is read
/// through it and the partition files are written through it.
pub fn ensure_partitioned_with(
    flat: &Path,
    partition_dir: &Path,
    schema: &TableSchema,
    codec: Option<Arc<dyn LineCodec>>,
) -> io::Result<Option<RepartitionResult>> {
    if !is_partitioned(schema) {
        return Ok(None);
    }
    let staging = staging_of(partition_dir);
    match fs::remove_dir_all(&staging) {
        Ok(()) => {}
        Err(error) if error.kind() == io::ErrorKind::NotFound => {}
        Err(error) => return Err(error),
    }

    if partition_dir.exists() {
        // NOTE: a flat log beside a finished partition directory is the copy the move did not delete.
        remove_if_present(flat)?;
        return Ok(None);
    }
    if !flat.exists() {
        return Ok(None);
    }

    let result = move_into_partitions(flat, &staging, schema, codec)?;
    fs::rename(&staging, partition_dir)?;
    remove_if_present(flat)?;
    Ok(Some(result))
}

fn remove_if_present(path: &Path) -> io::Result<()> {
    match fs::remove_file(path) {
        Ok(()) => Ok(()),
        Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()),
        Err(error) => Err(error),
    }
}

fn move_into_partitions(
    flat: &Path,
    staging: &Path,
    schema: &TableSchema,
    codec: Option<Arc<dyn LineCodec>>,
) -> io::Result<RepartitionResult> {
    fs::create_dir_all(staging)?;
    let mut log = LogFile::open(flat, true)?;
    log.set_codec(codec);
    let scan = log.scan(0)?;

    // NOTE: keeps first-write order, so the moved file reads in the order the flat log did.
    let mut position_of: HashMap<String, usize> = HashMap::new();
    let mut live: Vec<Option<Slot>> = Vec::new();
    for scanned in &scan.entries {
        let key = &scanned.entry.meta.key;
        match (scanned.entry.meta.op, &scanned.entry.data) {
            (LogOp::Put, Some(_)) => match position_of.get(key) {
                Some(&position) => live[position] = Some(scanned.slot),
                None => {
                    position_of.insert(key.clone(), live.len());
                    live.push(Some(scanned.slot));
                }
            },
            _ => {
                if let Some(position) = position_of.remove(key) {
                    live[position] = None;
                }
            }
        }
    }

    let mut records = 0usize;
    let mut buffers: Vec<(String, Vec<String>)> = Vec::new();
    for slot in live {
        let Some(slot) = slot else { continue };
        let line = log.read_slice(&slot)?;
        let parsed: Value =
            serde_json::from_str(&line).map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
        let part = assert_partition_value(schema, parsed.get(schema.partition_by.as_deref().unwrap_or("")))
            .map_err(io::Error::other)?;
        records += 1;
        match buffers.iter_mut().find(|(name, _)| name == &part) {
            Some((_, lines)) => lines.push(line),
            None => buffers.push((part, vec![line])),
        }
    }

    let mut partitions: Vec<String> = Vec::with_capacity(buffers.len());
    for (part, lines) in buffers {
        let mut file = fs::File::create(staging.join(format!("{part}{LOG_SUFFIX}")))?;
        file.write_all(format!("{}\n", log.encode_lines(&lines).join("\n")).as_bytes())?;
        file.sync_all()?;
        partitions.push(part);
    }
    partitions.sort_unstable();

    Ok(RepartitionResult { records, partitions })
}

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::super::record::{serialize_entry, LogMeta};
    use super::*;
    use crate::schema::{parse_schema, HaysonRecord};

    fn notes_schema() -> TableSchema {
        parse_schema(&json!({
            "schemaVersion": 1,
            "table": "notes",
            "primaryKey": "id",
            "partitionBy": "kind",
            "fields": {
                "id": { "type": "string" },
                "kind": { "type": "string", "required": true },
                "title": { "type": "string" },
            },
        }))
        .unwrap()
    }

    fn put(key: &str, seq: u64, kind: &str) -> String {
        let mut data = HaysonRecord::new();
        data.insert("id".into(), json!(key));
        data.insert("kind".into(), json!(kind));
        serialize_entry(
            &LogMeta {
                key: key.into(),
                op: LogOp::Put,
                seq,
                ts: 1,
            },
            Some(&data),
        )
    }

    fn del(key: &str, seq: u64) -> String {
        serialize_entry(
            &LogMeta {
                key: key.into(),
                op: LogOp::Del,
                seq,
                ts: 1,
            },
            None,
        )
    }

    fn seeded(dir: &Path) -> PathBuf {
        let flat = dir.join("notes.jsonl");
        let mut log = LogFile::open(&flat, false).unwrap();
        log.append(&[
            put("n1", 1, "note"),
            put("n2", 2, "note"),
            put("a1", 3, "adr"),
            del("n2", 4),
        ])
        .unwrap();
        log.sync().unwrap();
        flat
    }

    fn lines_in(path: &Path) -> Vec<String> {
        fs::read_to_string(path)
            .unwrap()
            .trim()
            .lines()
            .map(str::to_string)
            .collect()
    }

    #[test]
    fn moves_the_live_records_into_one_file_per_partition_value() {
        let dir = tempfile::tempdir().unwrap();
        let flat = seeded(dir.path());
        let partition_dir = dir.path().join("notes");

        let result = ensure_partitioned(&flat, &partition_dir, &notes_schema())
            .unwrap()
            .unwrap();

        assert_eq!(result.records, 2);
        assert_eq!(result.partitions, vec!["adr".to_string(), "note".to_string()]);
        assert!(!flat.exists());
        assert_eq!(lines_in(&partition_dir.join("note.jsonl")).len(), 1);
        assert_eq!(lines_in(&partition_dir.join("adr.jsonl")).len(), 1);
        assert!(lines_in(&partition_dir.join("note.jsonl"))[0].contains("\"_k\":\"n1\""));
    }

    #[test]
    fn a_table_with_no_flat_log_and_no_partitions_has_nothing_to_move() {
        let dir = tempfile::tempdir().unwrap();
        let outcome = ensure_partitioned(
            &dir.path().join("notes.jsonl"),
            &dir.path().join("notes"),
            &notes_schema(),
        )
        .unwrap();
        assert_eq!(outcome, None);
        assert!(!dir.path().join("notes").exists());
    }

    #[test]
    fn an_unpartitioned_schema_is_left_exactly_as_it_is() {
        let dir = tempfile::tempdir().unwrap();
        let flat = seeded(dir.path());
        let before = fs::read_to_string(&flat).unwrap();
        let mut schema = notes_schema();
        schema.partition_by = None;

        assert_eq!(
            ensure_partitioned(&flat, &dir.path().join("notes"), &schema).unwrap(),
            None
        );
        assert_eq!(fs::read_to_string(&flat).unwrap(), before);
    }

    #[test]
    fn throws_away_a_staging_directory_an_interrupted_move_left_behind() {
        let dir = tempfile::tempdir().unwrap();
        let flat = seeded(dir.path());
        let partition_dir = dir.path().join("notes");
        let staging = staging_of(&partition_dir);
        fs::create_dir_all(&staging).unwrap();
        fs::write(staging.join("note.jsonl"), "{\"broken\"\n").unwrap();

        let result = ensure_partitioned(&flat, &partition_dir, &notes_schema())
            .unwrap()
            .unwrap();

        assert_eq!(result.records, 2);
        assert!(!staging.exists());
        assert_eq!(lines_in(&partition_dir.join("note.jsonl")).len(), 1);
    }

    #[test]
    fn drops_the_flat_log_an_interrupted_move_left_beside_the_partitions() {
        let dir = tempfile::tempdir().unwrap();
        let flat = seeded(dir.path());
        let partition_dir = dir.path().join("notes");
        ensure_partitioned(&flat, &partition_dir, &notes_schema()).unwrap();
        fs::write(&flat, "{\"stale\":true}\n").unwrap();

        assert_eq!(
            ensure_partitioned(&flat, &partition_dir, &notes_schema()).unwrap(),
            None
        );
        assert!(!flat.exists());
        assert_eq!(lines_in(&partition_dir.join("note.jsonl")).len(), 1);
    }

    #[test]
    fn refuses_a_flat_log_holding_a_value_that_cannot_name_a_file() {
        let dir = tempfile::tempdir().unwrap();
        let flat = dir.path().join("notes.jsonl");
        let mut log = LogFile::open(&flat, false).unwrap();
        log.append(&[put("n1", 1, "../escape")]).unwrap();

        let error = ensure_partitioned(&flat, &dir.path().join("notes"), &notes_schema()).unwrap_err();

        assert!(error.to_string().contains("cannot name a partition file"));
        assert!(flat.exists(), "the flat log must survive a refused move");
        assert!(!dir.path().join("notes").exists());
    }
}