use crate::errors::Error;
use crate::output::print_json;
use crate::sqlite::Database;
use base64::Engine;
use serde::Deserialize;
use std::io::BufRead;
use std::path::Path;
use std::process::ExitCode;
use std::time::Duration;
use super::ImportResponse;
const EMBEDDING_DIMS: usize = 384;
const IMPORT_BUSY_TIMEOUT: Duration = Duration::from_secs(30);
#[derive(Deserialize)]
struct ImportRow {
id: String,
project_id: String,
content: String,
#[serde(default)]
metadata: Option<String>,
embedding: String,
created_at: String,
updated_at: String,
memory_type: String,
status: String,
#[serde(default)]
superseded_by: Option<String>,
retrieval_count: i64,
#[serde(default)]
last_retrieved_at: Option<String>,
}
fn read_source(source: &str) -> Result<String, Error> {
if source == "-" {
let stdin = std::io::stdin();
let mut reader = std::io::BufReader::new(stdin.lock());
let mut content = String::new();
let mut line = String::new();
while reader.read_line(&mut line).map_err(Error::from)? != 0 {
content.push_str(&line);
line.clear();
}
Ok(content)
} else {
std::fs::read_to_string(source).map_err(|e| {
Error::InvalidInput(format!("cannot read import source '{}': {}", source, e))
})
}
}
fn parse_row(line: &str, line_no: usize) -> Result<(ImportRow, Vec<u8>), Error> {
let line = line.trim_end_matches('\r');
if line.trim().is_empty() {
return Err(Error::InvalidInput(format!("line {}: empty row", line_no)));
}
let row: ImportRow = serde_json::from_str(line)
.map_err(|e| Error::InvalidInput(format!("line {}: malformed JSON: {}", line_no, e)))?;
if row.embedding.is_empty() {
return Err(Error::InvalidInput(format!(
"line {}: empty embedding (NULL/empty blobs are not restorable)",
line_no
)));
}
let blob = base64::engine::general_purpose::STANDARD
.decode(&row.embedding)
.map_err(|e| {
Error::InvalidInput(format!("line {}: invalid base64 embedding: {}", line_no, e))
})?;
if blob.len() != EMBEDDING_DIMS * 4 {
return Err(Error::InvalidInput(format!(
"line {}: wrong-dimension embedding: expected {} bytes ({}xf32), got {} bytes",
line_no,
EMBEDDING_DIMS * 4,
EMBEDDING_DIMS,
blob.len()
)));
}
Ok((row, blob))
}
fn validate_header(line: &str) -> Result<usize, Error> {
let line = line.trim_end_matches('\r');
let value: serde_json::Value = serde_json::from_str(line)
.map_err(|e| Error::InvalidInput(format!("line 1: malformed header JSON: {}", e)))?;
let obj = value
.as_object()
.ok_or_else(|| Error::InvalidInput("line 1: header must be a JSON object".to_string()))?;
if obj.get("type").and_then(|v| v.as_str()) != Some("export") {
return Err(Error::InvalidInput(
"line 1: not an export file (missing type=export)".to_string(),
));
}
if obj.get("version").is_none() {
return Err(Error::InvalidInput(
"line 1: header missing version".to_string(),
));
}
let rows = obj
.get("rows")
.and_then(|v| v.as_u64())
.ok_or_else(|| Error::InvalidInput("line 1: header missing rows".to_string()))?;
Ok(rows as usize)
}
pub(crate) fn run_import(db_path: &Path, source: &str) -> Result<ImportResponse, Error> {
let mut db = Database::open(db_path).map_err(|e| Error::Config(e.to_string()))?;
db.set_busy_timeout(IMPORT_BUSY_TIMEOUT)
.map_err(Error::from)?;
let skip_set: std::collections::HashSet<String> = db.existing_ids()?.into_iter().collect();
let content = read_source(source)?;
let lines: Vec<&str> = content.lines().collect();
if lines.is_empty() {
return Err(Error::InvalidInput(
"import source is empty (no header line)".to_string(),
));
}
validate_header(lines[0])?;
let mut parsed: Vec<(usize, ImportRow, Vec<u8>)> = Vec::new();
for (idx, line) in lines.iter().enumerate().skip(1) {
let line_no = idx + 1;
if line.trim().is_empty() {
continue;
}
let (row, blob) = parse_row(line, line_no)?;
parsed.push((line_no, row, blob));
}
let (mut inserted, mut skipped) = (0usize, 0usize);
let tx = db.connection().transaction()?;
for (line_no, row, blob) in parsed {
if skip_set.contains(&row.id) {
skipped += 1;
continue;
}
let inserted_here = tx.execute(
r#"
INSERT INTO memories (id, project_id, content, embedding, metadata, created_at, updated_at, type, status, superseded_by, retrieval_count, last_retrieved_at)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12)
"#,
rusqlite::params![
row.id,
row.project_id,
row.content,
blob,
row.metadata,
row.created_at,
row.updated_at,
row.memory_type,
row.status,
row.superseded_by,
row.retrieval_count,
row.last_retrieved_at,
],
)?;
if inserted_here == 0 {
return Err(Error::InvalidInput(format!(
"line {}: id {} already exists",
line_no, row.id
)));
}
inserted += 1;
}
tx.commit().map_err(Error::from)?;
Ok(ImportResponse { inserted, skipped })
}
pub fn handle_import(
db_path: &Path,
source: Option<&str>,
project_flag: Option<&str>,
json: bool,
) -> Result<ExitCode, Error> {
if let Some(project) = project_flag {
eprintln!(
"warning: --project ({project}) is ignored for import; import restores all projects in the file"
);
}
let source_str = source.unwrap_or("-");
match run_import(db_path, source_str) {
Ok(response) => {
if json {
print_json(&response);
} else {
println!(
"Import complete: {} inserted, {} skipped (already present)",
response.inserted, response.skipped
);
}
Ok(ExitCode::SUCCESS)
}
Err(e) => {
eprintln!("Import failed: {}", e);
Err(Error::InvalidInput(e.to_string()))
}
}
}