notmongo 0.1.7

In-process, bad database with an API similar to MongoDB
Documentation
use bson;
use bson::{Bson, Document};
use std::ffi::OsStr;
use std::fs;
use std::path::{Path, PathBuf};

use crate::Compression;

pub(crate) fn dump_to_json_string(
    bson_object: Bson,
    relaxed: bool,
    format: bool,
) -> Result<String, String> {
    let json_object: serde_json::Value = match relaxed {
        true => bson_object.into_relaxed_extjson(),
        false => bson_object.into_canonical_extjson(),
    };

    // TODO: Apparently serde_json's serializations can fail? This seems
    // unlikely given that the bson values are definitely valid. Investigate!
    if format {
        return Ok(serde_json::to_string_pretty(&json_object).unwrap());
    } else {
        return Ok(serde_json::to_string(&json_object).unwrap());
    }
}

fn backup_name(dir_path: &Path, stem: &OsStr, ext: &OsStr, backup_index: u32) -> PathBuf {
    dir_path.join(format!(
        "{}_backup{}.{}",
        // Is this conversion really an issue on any platform ever? Why don't we
        // add support for 6.3 bit bytes while we're at it?
        stem.to_string_lossy(),
        backup_index,
        ext.to_string_lossy()
    ))
}

fn safely_write_file_with_backups(
    content: &[u8],
    path: &Path,
    n_backups: u32,
    compression: Compression,
    subfile_name: &str,
) -> Result<(), String> {
    // Extract some path information, used down below
    let dir = path.parent().unwrap_or(Path::new("."));
    let extension = path.extension().unwrap_or(OsStr::new(""));
    let stem = match path.file_stem() {
        Some(stem) => stem,
        None => return Err(format!("The file path is missing a filename.")),
    };

    // Apply the compression
    let mut compression_buffer = Vec::new();

    let content = match compression {
        Compression::None => content,
        Compression::Zip => {
            {
                let buff = std::io::Cursor::new(&mut compression_buffer);
                let mut zw = zip::ZipWriter::new(buff);

                zw.start_file(
                    subfile_name,
                    zip::write::FileOptions::default()
                        .compression_method(zip::CompressionMethod::Deflated),
                )
                .map_err(|e| format!("Could not start the zip file: {}", e))?;

                std::io::Write::write_all(&mut zw, content)
                    .map_err(|e| format!("Could not write to the zip file: {}", e))?;

                zw.finish()
                    .map_err(|e| format!("Could not finish the zip file: {}", e))?;
            }

            &compression_buffer[..]
        }
    };

    // Write the file safely: First write to a temporary file, then rename
    // it to the final file.
    let partial_path: std::path::PathBuf = path.to_path_buf().with_file_name(format!(
        "{}_partial.{}",
        // Is this conversion really an issue on any platform ever? Why don't we
        // add support for 6.3 bit bytes while we're at it?
        stem.to_string_lossy(),
        extension.to_string_lossy()
    ));
    let partial_path = partial_path.as_path();
    match fs::write(partial_path, content) {
        std::io::Result::Ok(_) => {}
        std::io::Result::Err(val) => return Err(format!("Could not write the file: {}", val)),
    };

    // Back up old files if requested
    if n_backups > 0 {
        // How many existing backups are there?
        let mut move_paths = Vec::new();

        for ii in 1..n_backups + 1 {
            let backup_path = backup_name(dir, stem, extension, ii);

            if backup_path.exists() {
                move_paths.push(backup_path);
            } else {
                break;
            }
        }

        // Move all existing files back by one. Old backups are moved first,
        // so they aren't overwritten by the new ones.
        if move_paths.len() < n_backups as usize {
            move_paths.push(backup_name(
                dir,
                stem,
                extension,
                move_paths.len() as u32 + 1,
            ));
        }

        for ii in (1..move_paths.len()).rev() {
            let source_path = &move_paths[ii - 1];
            let dest_path = &move_paths[ii];

            // These are just backups. If writing fails, that's not worth
            // interrupting saving the main file.
            fs::rename(source_path, dest_path).ok();
        }

        // Create the new backup
        assert!(move_paths.len() > 0);
        if path.exists() {
            fs::rename(path, &move_paths[0]).ok();
        }
    }

    // Move the new file to the final location
    match fs::rename(partial_path, path) {
        std::io::Result::Ok(_) => Ok(()),
        std::io::Result::Err(val) => {
            // Try to delete the partial file again
            fs::remove_file(partial_path).ok(); // An error is about to be reported, ignore any additional ones here

            // TODO: This could be slightly improved: If moving the final file fails,
            //       restore the old file from a backup.

            Err(format!("Could not create the file: {}", val))
        }
    }
}

pub(crate) fn dump_to_json_file<P: AsRef<Path>>(
    bson_object: Bson,
    relaxed: bool,
    format: bool,
    path: P,
    n_backups: u32,
    compression: Compression,
) -> Result<(), String> {
    // First convert this to a string. As this can fail, it needs to be done
    // before messing with any files.
    let json_string = dump_to_json_string(bson_object, relaxed, format)?;

    // Write the file
    safely_write_file_with_backups(
        json_string.as_bytes(),
        path.as_ref(),
        n_backups,
        compression,
        "database.json",
    )
}

pub(crate) fn dump_to_bson_file<P: AsRef<Path>>(
    bson_document: &bson::Document,
    path: P,
    n_backups: u32,
    compression: Compression,
) -> Result<(), String> {
    // Serialize to a buffer
    let buffer = bson::ser::to_vec(bson_document).unwrap(); // TODO: Why should this be able to fail?

    // Write the new file
    safely_write_file_with_backups(
        &buffer,
        path.as_ref(),
        n_backups,
        compression,
        "database.bson",
    )
}

pub(crate) fn get_bson_field_honor_dot_notation<'a>(
    bson_object: &'a Bson,
    key: &str,
) -> Option<&'a Bson> {
    // Attempts to get the object referenced by the key. The difference to
    // simply using `get` on a BSON document, is that keys can consist of
    // multiple fields ("a.b.c"), and indices into arrays can also be numbers
    // (a.2.c)

    let mut cur_object = bson_object;

    for key_part in key.split('.') {
        cur_object = match bson_object {
            Bson::Document(doc) => match doc.get(key_part) {
                Some(val) => val,
                None => return None,
            },
            Bson::Array(arr) => {
                // Is the key a valid integer
                let index = match key_part.parse::<usize>() {
                    Ok(index) => index,
                    Err(_) => return None,
                };

                // Is the index in range
                if index >= arr.len() {
                    return None;
                }

                // Apply it
                &arr[index]
            }
            _ => return None,
        }
    }

    Some(cur_object)
}

pub(crate) fn get_document_field_honor_dot_notation<'a>(
    document: &'a Document,
    key: &str,
) -> Option<&'a Bson> {
    // The first access needs to be special cased, so the document doesn't have
    // to be copied. After that, just pass it on to the function for general
    // bson values.

    let (first_key, remainder) = match key.split_once(".") {
        Some((first_key, remainder)) => (first_key, Some(remainder)),
        None => (key, None),
    };

    match document.get(first_key) {
        Some(val) => match remainder {
            Some(remainder) => get_bson_field_honor_dot_notation(val, remainder),
            None => Some(val),
        },
        None => None,
    }
}

fn bson_type_sort_priority(object: &Bson) -> u8 {
    match object {
        // These aren't specified in the documentation
        Bson::JavaScriptCode(_) => 0,
        Bson::JavaScriptCodeWithScope(_) => 0,
        Bson::DbPointer(_) => 0,
        Bson::Undefined => 0,

        // These are taken straight from the documentation.
        Bson::MinKey => 1,
        Bson::Null => 2,
        Bson::Int32(_) | Bson::Int64(_) | Bson::Double(_) | Bson::Decimal128(_) => 3,
        Bson::Symbol(_) | Bson::String(_) => 3,
        Bson::Document(_) => 4,
        Bson::Array(_) => 5,
        Bson::Binary(_) => 6,
        Bson::ObjectId(_) => 7,
        Bson::Boolean(_) => 8,
        Bson::DateTime(_) => 9,
        Bson::Timestamp(_) => 10,
        Bson::RegularExpression(_) => 11,
        Bson::MaxKey => 12,
    }
}

fn bson_as_i64(value: &Bson) -> Option<i64> {
    match value {
        // TODO: What happens to other types? Test with actual mongodb.
        Bson::Boolean(val) => Some(*val as i64),
        Bson::Int32(val) => Some(*val as i64),
        Bson::Int64(val) => Some(*val),
        _ => None,
    }
}

fn bson_as_f64(value: &Bson) -> Option<f64> {
    match bson_as_i64(value) {
        Some(value) => Some(value as f64),
        None => match value {
            Bson::Double(val) => Some(*val),
            _ => None,
        },
    }
}

pub(crate) fn compare_bson(lhs: &Bson, rhs: &Bson) -> std::cmp::Ordering {
    // FIXME: This function is really mostly guesswork and common sense, rather
    // than following any actual specification. There is some documentation
    // available, but it's not detailed at all:
    // https://docs.mongodb.com/manual/reference/bson-type-comparison-order/

    // Identical according to BSON?
    if lhs == rhs {
        return std::cmp::Ordering::Equal;
    }

    // Both convertible to i64?
    let lhs_i64 = bson_as_i64(lhs);
    let rhs_i64 = bson_as_i64(rhs);

    if let (Some(lhs), Some(rhs)) = (lhs_i64, rhs_i64) {
        return lhs.cmp(&rhs);
    }

    // Both convertible to f64?
    let lhs_f64 = bson_as_f64(lhs);
    let rhs_f64 = bson_as_f64(rhs);

    if let (Some(lhs), Some(rhs)) = (lhs_f64, rhs_f64) {
        if lhs < rhs {
            return std::cmp::Ordering::Less;
        } else if lhs > rhs {
            return std::cmp::Ordering::Greater;
        } else {
            return std::cmp::Ordering::Equal;
        }
    }

    // Some other cases
    match (lhs, rhs) {
        // TODO: Binary and others?
        (Bson::ObjectId(lhs), Bson::ObjectId(rhs)) => return lhs.cmp(rhs),
        (Bson::String(lhs), Bson::String(rhs)) => return lhs.cmp(rhs),
        (Bson::DateTime(lhs), Bson::DateTime(rhs)) => return lhs.cmp(rhs),
        _ => {}
    };

    // Lastly, fall back on just the types themselves
    bson_type_sort_priority(lhs).cmp(&bson_type_sort_priority(rhs))
}

pub(crate) fn read_compressed_file<P: AsRef<Path>>(
    path: P,
    compression: Compression,
) -> Result<Vec<u8>, String> {
    match compression {
        Compression::None => {
            let mut file = match fs::File::open(path) {
                Ok(file) => file,
                Err(_) => return Err("Failed to open file".to_string()),
            };
            let mut contents = Vec::new();
            if let Err(_) = std::io::Read::read_to_end(&mut file, &mut contents) {
                return Err("Failed to read file".to_string());
            }
            Ok(contents)
        }

        Compression::Zip => {
            // Read the zip-compressed file
            let file = match fs::File::open(path) {
                Ok(file) => file,
                Err(_) => return Err("Failed to open file".to_string()),
            };

            let mut archive = match zip::ZipArchive::new(file) {
                Ok(archive) => archive,
                Err(_) => return Err("Failed to read zip archive".to_string()),
            };

            if archive.len() != 1 {
                return Err("Zip compressed archives must contain exactly one file".to_string());
            }

            let mut zip_file = match archive.by_index(0) {
                Ok(zip_file) => zip_file,
                Err(_) => return Err("Failed to read file from zip archive".to_string()),
            };

            let mut contents = Vec::new();
            if let Err(_) = std::io::Read::read_to_end(&mut zip_file, &mut contents) {
                return Err("Failed to read file contents from zip archive".to_string());
            }

            Ok(contents)
        }
    }
}