use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::OnceLock;
use anyhow::Result;
use indicatif::{ProgressBar, ProgressStyle};
use velesdb_core::{Database, Point, VectorCollection};
static CONFIG_PATH: OnceLock<Option<PathBuf>> = OnceLock::new();
pub fn set_config_path(path: Option<PathBuf>) {
let _ = CONFIG_PATH.set(path);
}
pub fn open_database(path: &Path) -> Result<Database> {
let config_path = CONFIG_PATH
.get()
.and_then(Option::as_ref)
.map(PathBuf::as_path);
open_database_with_config(path, config_path)
}
pub fn open_database_with_config(path: &Path, config_path: Option<&Path>) -> Result<Database> {
open_database_with_default_file(path, config_path, Path::new(DEFAULT_CONFIG_FILE))
}
const DEFAULT_CONFIG_FILE: &str = "velesdb.toml";
#[derive(Debug, PartialEq, Eq)]
enum ConfigSource<'a> {
Explicit(&'a Path),
Implicit(&'a Path),
CoreDefaults,
}
impl<'a> ConfigSource<'a> {
fn resolve(config_path: Option<&'a Path>, default_file: &'a Path) -> Result<Self> {
match config_path {
Some(cfg) if cfg.exists() => Ok(Self::Explicit(cfg)),
Some(cfg) => anyhow::bail!("config file not found: {}", cfg.display()),
None if default_file.exists() => Ok(Self::Implicit(default_file)),
None => Ok(Self::CoreDefaults),
}
}
fn load(&self) -> Result<velesdb_core::config::VelesConfig> {
let (file, origin) = match *self {
Self::Explicit(file) => (file, String::new()),
Self::Implicit(file) => (
file,
" (read implicitly from the working directory because no --config or \
VELESDB_CONFIG was given; pass `--config <file>` to use another file)"
.to_owned(),
),
Self::CoreDefaults => {
return velesdb_core::config::VelesConfig::load_engine_only_from_env()
.map_err(|e| anyhow::anyhow!("failed to load VelesDB config: {e}"));
}
};
velesdb_core::config::VelesConfig::load_from_path_engine_only(file).map_err(|e| {
anyhow::anyhow!(
"failed to load VelesDB config from {}{origin}: {e}",
file.display()
)
})
}
}
fn open_database_with_default_file(
path: &Path,
config_path: Option<&Path>,
default_file: &Path,
) -> Result<Database> {
let source = ConfigSource::resolve(config_path, default_file)?;
Ok(Database::open_with_config(path, source.load()?)?)
}
pub struct BatchImporter<'a> {
collection: &'a VectorCollection,
batch: Vec<Point>,
batch_size: usize,
pub stats: ImportAccumulator,
}
#[derive(Debug, Default)]
pub struct ImportAccumulator {
pub imported: usize,
pub errors: usize,
}
impl<'a> BatchImporter<'a> {
pub fn new(collection: &'a VectorCollection, batch_size: usize) -> Self {
Self {
collection,
batch: Vec::with_capacity(batch_size),
batch_size,
stats: ImportAccumulator::default(),
}
}
pub fn push(&mut self, point: Point) -> Result<()> {
self.batch.push(point);
self.stats.imported += 1;
if self.batch.len() >= self.batch_size {
self.collection.upsert_bulk(&self.batch)?;
self.batch.clear();
}
Ok(())
}
pub fn record_error(&mut self) {
self.stats.errors += 1;
}
pub fn flush(self) -> Result<ImportAccumulator> {
if !self.batch.is_empty() {
self.collection.upsert_bulk(&self.batch)?;
}
Ok(self.stats)
}
}
#[must_use]
pub fn create_progress_bar(total: usize, show: bool) -> ProgressBar {
if show {
let pb = ProgressBar::new(total as u64);
if let Ok(style) = ProgressStyle::default_bar().template(
"{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} ({eta})",
) {
pb.set_style(style.progress_chars("#>-"));
}
pb
} else {
ProgressBar::hidden()
}
}
pub fn set_import_message(progress: &ProgressBar, total: usize, file_size: u64, show: bool) {
if show {
#[allow(clippy::cast_precision_loss)]
let size_mb = file_size as f64 / (1024.0 * 1024.0);
progress.set_message(format!("Importing {total} vectors ({size_mb:.1} MB)"));
}
}
pub fn point_payload_to_row(
id: u64,
payload: &Option<serde_json::Value>,
) -> HashMap<String, serde_json::Value> {
let mut row = HashMap::new();
row.insert("id".to_string(), serde_json::json!(id));
if let Some(serde_json::Value::Object(map)) = payload {
for (k, v) in map {
row.insert(k.clone(), v.clone());
}
}
row
}
pub fn point_payload_to_browse_row(
id: u64,
payload: &Option<serde_json::Value>,
) -> HashMap<String, serde_json::Value> {
let mut row = HashMap::new();
row.insert("id".to_string(), serde_json::json!(id));
if let Some(serde_json::Value::Object(map)) = payload {
for (k, v) in map {
row.insert(k.clone(), truncate_display_value(v));
}
}
row
}
fn truncate_display_value(value: &serde_json::Value) -> serde_json::Value {
match value {
serde_json::Value::String(s) if s.len() > 50 => {
let truncated: String = s.chars().take(47).collect();
serde_json::json!(format!("{truncated}..."))
}
other => other.clone(),
}
}
pub fn print_json(data: &serde_json::Value) -> Result<()> {
println!("{}", serde_json::to_string_pretty(data)?);
Ok(())
}
pub fn point_to_export_record(
id: u64,
vector: Option<&[f32]>,
payload: &Option<serde_json::Value>,
) -> serde_json::Value {
let mut record = serde_json::Map::new();
record.insert("id".to_string(), serde_json::json!(id));
if let Some(v) = vector {
record.insert("vector".to_string(), serde_json::json!(v));
}
if let Some(p) = payload {
record.insert("payload".to_string(), p.clone());
}
serde_json::Value::Object(record)
}
pub fn write_export_file(records: &[serde_json::Value], filename: &str) -> Result<(), String> {
let json_str = serde_json::to_string_pretty(records)
.map_err(|e| format!("Failed to serialize records: {e}"))?;
std::fs::write(filename, json_str).map_err(|e| format!("Failed to write file: {e}"))?;
Ok(())
}
#[cfg(test)]
#[path = "helpers_tests.rs"]
mod helpers_tests;