use crate::engine::Engine;
use crate::error::{ErrorCode, McpError};
use crate::schema::{
apply_schema_override, infer_csv_schema, infer_json_schema, json_type_name,
widen_csv_numeric_columns, ColumnSchema,
};
use crate::stats::{IngestStats, StatsTimer};
use hyperdb_api::AsyncConnection;
use std::path::{Path, PathBuf};
fn canonicalize_for_copy(path: &Path) -> PathBuf {
let canonical = path.canonicalize().unwrap_or_else(|_| path.to_path_buf());
#[cfg(windows)]
{
if let Some(s) = canonical.to_str() {
let stripped = match s.strip_prefix(r"\\?\") {
Some(rest) if !rest.starts_with("UNC\\") => rest,
_ => s,
};
return PathBuf::from(stripped.replace('\\', "/"));
}
}
canonical
}
use serde_json::Value;
const SCHEMA_INFERENCE_MAX_BYTES: u64 = 64 * 1024 * 1024;
fn read_text_sample(path: impl AsRef<Path>, max_bytes: u64) -> std::io::Result<String> {
use std::io::Read;
let file = std::fs::File::open(path.as_ref())?;
let file_len = file.metadata()?.len();
if file_len <= max_bytes {
return std::fs::read_to_string(path.as_ref());
}
let mut reader = std::io::BufReader::new(file);
let cap = usize::try_from(max_bytes).unwrap_or(usize::MAX);
let mut buf = vec![0u8; cap];
reader.read_exact(&mut buf)?;
if let Some(pos) = buf.iter().rposition(|&b| b == b'\n') {
buf.truncate(pos + 1);
}
match String::from_utf8(buf) {
Ok(s) => Ok(s),
Err(e) => {
let valid_up_to = e.utf8_error().valid_up_to();
let mut bytes = e.into_bytes();
bytes.truncate(valid_up_to);
if let Some(pos) = bytes.iter().rposition(|&b| b == b'\n') {
bytes.truncate(pos + 1);
}
String::from_utf8(bytes)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
}
}
}
#[derive(Debug)]
pub struct IngestOptions {
pub table: String,
pub mode: String,
pub schema_override: Option<serde_json::Map<String, Value>>,
pub merge_key: Option<Vec<String>>,
pub target_db: Option<String>,
}
pub fn qualified_table(opts: &IngestOptions) -> String {
match &opts.target_db {
Some(db) => {
let esc_db = db.replace('"', "\"\"");
let esc_tbl = opts.table.replace('"', "\"\"");
format!("\"{esc_db}\".\"public\".\"{esc_tbl}\"")
}
None => format!("\"{}\"", opts.table.replace('"', "\"\"")),
}
}
#[derive(Debug)]
pub struct IngestResult {
pub rows: u64,
pub schema: Vec<ColumnSchema>,
pub stats: IngestStats,
}
struct TempTableGuard<'a> {
engine: &'a Engine,
name: String,
target_db: Option<String>,
armed: bool,
}
impl<'a> TempTableGuard<'a> {
fn new(engine: &'a Engine, name: String, target_db: Option<String>) -> Self {
Self {
engine,
name,
target_db,
armed: true,
}
}
fn disarm(&mut self) {
self.armed = false;
}
}
impl Drop for TempTableGuard<'_> {
fn drop(&mut self) {
if !self.armed {
return;
}
let panicking = std::thread::panicking();
let quoted = match &self.target_db {
Some(db) => format!(
"\"{}\".\"public\".\"{}\"",
db.replace('"', "\"\""),
self.name.replace('"', "\"\""),
),
None => format!("\"{}\"", self.name.replace('"', "\"\"")),
};
let drop_sql = format!("DROP TABLE IF EXISTS {quoted}");
match self.engine.execute_command(&drop_sql) {
Ok(_) => {
if panicking {
tracing::info!(
tmp = %self.name,
"merge_via_temp_table: dropped temp table during panic unwind"
);
}
}
Err(e) => {
tracing::warn!(
tmp = %self.name,
panicking,
error = %e,
"merge_via_temp_table: failed to drop temp table on guard exit"
);
}
}
}
}
pub fn merge_via_temp_table<F>(
engine: &Engine,
opts: &IngestOptions,
replace_load: F,
) -> Result<IngestResult, McpError>
where
F: FnOnce(&IngestOptions) -> Result<IngestResult, McpError>,
{
debug_assert_eq!(
opts.mode, "merge",
"merge_via_temp_table called with non-merge mode `{}`",
opts.mode
);
let keys = opts.merge_key.as_ref().ok_or_else(|| {
McpError::new(
ErrorCode::InvalidArgument,
"merge mode requires merge_key (a column name or list of column names)",
)
})?;
if keys.is_empty() || keys.iter().any(String::is_empty) {
return Err(McpError::new(
ErrorCode::InvalidArgument,
"merge_key must be a non-empty list of non-empty column names",
));
}
use std::sync::atomic::{AtomicU64, Ordering};
static MERGE_COUNTER: AtomicU64 = AtomicU64::new(0);
let counter = MERGE_COUNTER.fetch_add(1, Ordering::Relaxed);
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0, |d| d.as_nanos());
let safe_target: String = opts
.table
.chars()
.map(|c| {
if c.is_ascii_alphanumeric() || c == '_' {
c
} else {
'_'
}
})
.collect();
let tmp = format!(
"__hyperdb_merge_{safe_target}_{}_{nanos}_{counter}",
std::process::id(),
);
let tmp_opts = IngestOptions {
table: tmp.clone(),
mode: "replace".into(),
schema_override: opts.schema_override.clone(),
merge_key: None,
target_db: opts.target_db.clone(),
};
let tmp_result = replace_load(&tmp_opts)?;
let mut guard = TempTableGuard::new(engine, tmp.clone(), opts.target_db.clone());
if !engine.table_exists_in(opts.target_db.as_deref(), &tmp)? {
return Err(McpError::new(
ErrorCode::InternalError,
format!(
"merge: temp table '{tmp}' was not produced by the format-specific \
replace load — this is a contract violation in the per-format ingest path"
),
));
}
if !engine.table_exists_in(opts.target_db.as_deref(), &opts.table)? {
let qualified_tmp_opts = IngestOptions {
table: tmp.clone(),
mode: "replace".into(),
schema_override: None,
merge_key: None,
target_db: opts.target_db.clone(),
};
let quoted_tmp = qualified_table(&qualified_tmp_opts);
let escaped_new = opts.table.replace('"', "\"\"");
engine.execute_command(&format!(
"ALTER TABLE {quoted_tmp} RENAME TO \"{escaped_new}\""
))?;
guard.disarm();
return Ok(IngestResult {
rows: tmp_result.rows,
schema: tmp_result.schema,
stats: IngestStats {
operation: tmp_result.stats.operation,
rows: tmp_result.stats.rows,
elapsed_ms: tmp_result.stats.elapsed_ms,
bytes_read: tmp_result.stats.bytes_read,
bytes_stored: tmp_result.stats.bytes_stored,
schema_inference_ms: tmp_result.stats.schema_inference_ms,
table: opts.table.clone(),
file_format: tmp_result.stats.file_format,
warning: tmp_result.stats.warning,
schema_changed: true,
},
});
}
let target_cols = engine.column_metadata_in(opts.target_db.as_deref(), &opts.table)?;
let tmp_cols = engine.column_metadata_in(opts.target_db.as_deref(), &tmp)?;
for k in keys {
let in_target = target_cols.iter().find(|c| c.name == *k);
let in_tmp = tmp_cols.iter().find(|c| c.name == *k);
match (in_target, in_tmp) {
(None, _) => {
return Err(McpError::new(
ErrorCode::InvalidArgument,
format!(
"merge_key column '{k}' is not in target table '{}'",
opts.table
),
));
}
(_, None) => {
return Err(McpError::new(
ErrorCode::InvalidArgument,
format!(
"merge_key column '{k}' is not in incoming data \
(column missing from the file)"
),
));
}
(Some(t), Some(s)) if !types_compatible(&t.hyper_type, &s.hyper_type) => {
return Err(McpError::new(
ErrorCode::InvalidArgument,
format!(
"merge_key column '{k}' type mismatch: target is {} but \
incoming is {}. Use mode=replace or apply a schema override.",
t.hyper_type, s.hyper_type
),
));
}
_ => {}
}
}
for tc in &tmp_cols {
if let Some(target_c) = target_cols.iter().find(|c| c.name == tc.name) {
if !types_compatible(&target_c.hyper_type, &tc.hyper_type) {
return Err(McpError::new(
ErrorCode::InvalidArgument,
format!(
"Column '{}' type mismatch: target is {} but incoming is {}. \
Use mode=replace or apply a schema override.",
tc.name, target_c.hyper_type, tc.hyper_type
),
));
}
}
}
let new_cols: Vec<ColumnSchema> = tmp_cols
.iter()
.filter(|c| !target_cols.iter().any(|t| t.name == c.name))
.cloned()
.map(|mut c| {
c.nullable = true;
c
})
.collect();
let schema_changed = !new_cols.is_empty();
engine.alter_table_add_columns_in(opts.target_db.as_deref(), &opts.table, &new_cols)?;
let quoted_tgt = qualified_table(opts);
let qualified_tmp_opts = IngestOptions {
table: tmp.clone(),
mode: "replace".into(),
schema_override: None,
merge_key: None,
target_db: opts.target_db.clone(),
};
let quoted_tmp = qualified_table(&qualified_tmp_opts);
let key_eq = keys
.iter()
.map(|k| {
let qk = k.replace('"', "\"\"");
format!("t.\"{qk}\" = s.\"{qk}\"")
})
.collect::<Vec<_>>()
.join(" AND ");
let delete_sql = format!("DELETE FROM {quoted_tgt} t USING {quoted_tmp} s WHERE {key_eq}");
engine.execute_command(&delete_sql)?;
let cols_csv = tmp_cols
.iter()
.map(|c| format!("\"{}\"", c.name.replace('"', "\"\"")))
.collect::<Vec<_>>()
.join(", ");
let insert_sql =
format!("INSERT INTO {quoted_tgt} ({cols_csv}) SELECT {cols_csv} FROM {quoted_tmp}");
let inserted = engine.execute_command(&insert_sql)?;
let final_schema = engine.column_metadata_in(opts.target_db.as_deref(), &opts.table)?;
Ok(IngestResult {
rows: inserted,
schema: final_schema,
stats: IngestStats {
operation: tmp_result.stats.operation,
rows: inserted,
elapsed_ms: tmp_result.stats.elapsed_ms,
bytes_read: tmp_result.stats.bytes_read,
bytes_stored: tmp_result.stats.bytes_stored,
schema_inference_ms: tmp_result.stats.schema_inference_ms,
table: opts.table.clone(),
file_format: tmp_result.stats.file_format,
warning: tmp_result.stats.warning,
schema_changed,
},
})
}
fn types_compatible(a: &str, b: &str) -> bool {
if a == b {
return true;
}
match (
crate::schema::map_hyper_type(a),
crate::schema::map_hyper_type(b),
) {
(Some(sa), Some(sb)) => sa == sb,
_ => true,
}
}
pub fn ingest_json(
engine: &Engine,
json_str: &str,
opts: &IngestOptions,
) -> Result<IngestResult, McpError> {
if opts.mode == "merge" {
return merge_via_temp_table(engine, opts, |tmp_opts| {
ingest_json(engine, json_str, tmp_opts)
});
}
let timer = StatsTimer::start();
let bytes_read = json_str.len() as u64;
let schema_timer = StatsTimer::start();
let inferred = infer_json_schema(json_str)?;
let columns = match &opts.schema_override {
Some(s) => apply_schema_override(inferred, s)?,
None => inferred,
};
let schema_ms = schema_timer.elapsed_ms();
let array: Vec<serde_json::Map<String, Value>> = serde_json::from_str(json_str)
.map_err(|e| McpError::new(ErrorCode::SchemaMismatch, format!("Invalid JSON: {e}")))?;
if array.is_empty() {
return Err(McpError::new(ErrorCode::EmptyData, "JSON array is empty"));
}
let is_replace = opts.mode != "append";
let qualified = qualified_table(opts);
let row_count = engine.execute_in_transaction(|engine| {
engine.create_table_in(&opts.table, &columns, is_replace, opts.target_db.as_deref())?;
let mut row_count = 0u64;
let col_names: Vec<String> = columns.iter().map(|c| format!("\"{}\"", c.name)).collect();
for obj in &array {
let values: Vec<String> = columns
.iter()
.map(|col| match obj.get(&col.name) {
None | Some(Value::Null) => "NULL".to_string(),
Some(Value::Bool(b)) => b.to_string(),
Some(Value::Number(n)) => n.to_string(),
Some(Value::String(s)) => format!("'{}'", s.replace('\'', "''")),
Some(other) => format!("'{}'", other.to_string().replace('\'', "''")),
})
.collect();
let sql = format!(
"INSERT INTO {} ({}) VALUES ({})",
qualified,
col_names.join(", "),
values.join(", ")
);
engine.execute_command(&sql)?;
row_count += 1;
}
Ok(row_count)
})?;
let elapsed = timer.elapsed_ms();
let stats = IngestStats {
operation: "load_data".into(),
rows: row_count,
elapsed_ms: elapsed,
bytes_read,
bytes_stored: 0,
schema_inference_ms: Some(schema_ms),
table: opts.table.clone(),
file_format: Some("json".into()),
warning: if bytes_read > 50_000_000 {
Some("Large inline data. Consider using load_file for better performance.".into())
} else {
None
},
schema_changed: false,
};
Ok(IngestResult {
rows: row_count,
schema: columns,
stats,
})
}
pub fn ingest_csv(
engine: &Engine,
csv_text: &str,
opts: &IngestOptions,
) -> Result<IngestResult, McpError> {
if opts.mode == "merge" {
return merge_via_temp_table(engine, opts, |tmp_opts| {
ingest_csv(engine, csv_text, tmp_opts)
});
}
let timer = StatsTimer::start();
let bytes_read = csv_text.len() as u64;
let schema_timer = StatsTimer::start();
let mut inferred = infer_csv_schema(csv_text, true)?;
widen_csv_numeric_columns(csv_text.as_bytes(), true, &mut inferred)?;
let columns = match &opts.schema_override {
Some(s) => apply_schema_override(inferred, s)?,
None => inferred,
};
let schema_ms = schema_timer.elapsed_ms();
let temp_path = tempfile::Builder::new()
.prefix("hyperdb_mcp_csv_")
.suffix(".csv")
.tempfile()
.map_err(|e| {
McpError::new(
ErrorCode::InternalError,
format!("Failed to create temp CSV file: {e}"),
)
})?
.into_temp_path();
std::fs::write(&temp_path, csv_text).map_err(|e| {
McpError::new(
ErrorCode::InternalError,
format!("Failed to write temp CSV: {e}"),
)
})?;
let canonical_temp = canonicalize_for_copy(&temp_path);
let qualified = qualified_table(opts);
let copy_sql = format!(
"COPY {} FROM {} WITH (FORMAT csv, NULL '', DELIMITER ',', HEADER)",
qualified,
hyperdb_api::escape_string_literal(canonical_temp.to_str().unwrap_or(""))
);
let is_replace = opts.mode != "append";
let row_count = engine.execute_in_transaction(|engine| {
engine.create_table_in(&opts.table, &columns, is_replace, opts.target_db.as_deref())?;
engine.execute_command(©_sql)
});
drop(temp_path);
let row_count = row_count?;
let elapsed = timer.elapsed_ms();
let stats = IngestStats {
operation: "load_data".into(),
rows: row_count,
elapsed_ms: elapsed,
bytes_read,
bytes_stored: 0,
schema_inference_ms: Some(schema_ms),
table: opts.table.clone(),
file_format: Some("csv".into()),
warning: if bytes_read > 50_000_000 {
Some("Large inline data. Consider using load_file for better performance.".into())
} else {
None
},
schema_changed: false,
};
Ok(IngestResult {
rows: row_count,
schema: columns,
stats,
})
}
pub fn ingest_csv_file(
engine: &Engine,
path: &str,
opts: &IngestOptions,
) -> Result<IngestResult, McpError> {
if opts.mode == "merge" {
return merge_via_temp_table(engine, opts, |tmp_opts| {
ingest_csv_file(engine, path, tmp_opts)
});
}
let timer = StatsTimer::start();
let abs_path = std::path::Path::new(path);
if !abs_path.exists() {
return Err(McpError::new(
ErrorCode::FileNotFound,
format!("File not found: {path}"),
));
}
let file_size = std::fs::metadata(abs_path).map_or(0, |m| m.len());
let schema_timer = StatsTimer::start();
let sample = read_text_sample(abs_path, SCHEMA_INFERENCE_MAX_BYTES)
.map_err(|e| McpError::new(ErrorCode::FileNotFound, format!("Cannot read file: {e}")))?;
let mut inferred = infer_csv_schema(&sample, true)?;
widen_csv_numeric_columns(sample.as_bytes(), true, &mut inferred)?;
let columns = match &opts.schema_override {
Some(s) => apply_schema_override(inferred, s)?,
None => inferred,
};
let schema_ms = schema_timer.elapsed_ms();
let canonical = canonicalize_for_copy(abs_path);
let qualified = qualified_table(opts);
let copy_sql = format!(
"COPY {} FROM {} WITH (FORMAT csv, NULL '', DELIMITER ',', HEADER)",
qualified,
hyperdb_api::escape_string_literal(canonical.to_str().unwrap_or(""))
);
let is_replace = opts.mode != "append";
let row_count = engine.execute_in_transaction(|engine| {
engine.create_table_in(&opts.table, &columns, is_replace, opts.target_db.as_deref())?;
engine.execute_command(©_sql)
})?;
let elapsed = timer.elapsed_ms();
let stats = IngestStats {
operation: "load_file".into(),
rows: row_count,
elapsed_ms: elapsed,
bytes_read: file_size,
bytes_stored: 0,
schema_inference_ms: Some(schema_ms),
table: opts.table.clone(),
file_format: Some("csv".into()),
warning: None,
schema_changed: false,
};
Ok(IngestResult {
rows: row_count,
schema: columns,
stats,
})
}
pub async fn ingest_csv_file_async(
conn: &AsyncConnection,
path: &str,
opts: &IngestOptions,
) -> Result<IngestResult, McpError> {
let timer = StatsTimer::start();
let abs_path = std::path::Path::new(path);
if !abs_path.exists() {
return Err(McpError::new(
ErrorCode::FileNotFound,
format!("File not found: {path}"),
));
}
let file_size = std::fs::metadata(abs_path).map_or(0, |m| m.len());
let schema_timer = StatsTimer::start();
let path_owned = path.to_string();
let override_owned = opts.schema_override.clone();
let columns: Vec<ColumnSchema> = tokio::task::spawn_blocking(move || -> Result<_, McpError> {
let sample = read_text_sample(&path_owned, SCHEMA_INFERENCE_MAX_BYTES).map_err(|e| {
McpError::new(ErrorCode::FileNotFound, format!("Cannot read file: {e}"))
})?;
let mut inferred = infer_csv_schema(&sample, true)?;
widen_csv_numeric_columns(sample.as_bytes(), true, &mut inferred)?;
let columns = match &override_owned {
Some(s) => apply_schema_override(inferred, s)?,
None => inferred,
};
Ok(columns)
})
.await
.map_err(|e| McpError::new(ErrorCode::InternalError, format!("Task join error: {e}")))??;
let schema_ms = schema_timer.elapsed_ms();
let canonical = canonicalize_for_copy(abs_path);
let qualified = qualified_table(opts);
let copy_sql = format!(
"COPY {} FROM {} WITH (FORMAT csv, NULL '', DELIMITER ',', HEADER)",
qualified,
hyperdb_api::escape_string_literal(canonical.to_str().unwrap_or(""))
);
let is_replace = opts.mode != "append";
conn.begin_transaction().await.map_err(McpError::from)?;
let inner: Result<u64, McpError> = async {
create_table_async(
conn,
&opts.table,
&columns,
is_replace,
opts.target_db.as_deref(),
)
.await?;
conn.execute_command(©_sql)
.await
.map_err(McpError::from)
}
.await;
let row_count = match inner {
Ok(n) => {
conn.commit().await.map_err(McpError::from)?;
n
}
Err(e) => {
if let Err(rb) = conn.rollback().await {
tracing::warn!("rollback after error failed: {}", rb);
}
return Err(e);
}
};
let elapsed = timer.elapsed_ms();
let stats = IngestStats {
operation: "load_file".into(),
rows: row_count,
elapsed_ms: elapsed,
bytes_read: file_size,
bytes_stored: 0,
schema_inference_ms: Some(schema_ms),
table: opts.table.clone(),
file_format: Some("csv".into()),
warning: None,
schema_changed: false,
};
Ok(IngestResult {
rows: row_count,
schema: columns,
stats,
})
}
pub fn ingest_json_file(
engine: &Engine,
path: &str,
opts: &IngestOptions,
) -> Result<IngestResult, McpError> {
let abs_path = std::path::Path::new(path);
if !abs_path.exists() {
return Err(McpError::new(
ErrorCode::FileNotFound,
format!("File not found: {path}"),
));
}
let text = std::fs::read_to_string(abs_path)
.map_err(|e| McpError::new(ErrorCode::FileNotFound, format!("Cannot read file: {e}")))?;
let json_array_text = normalize_json_or_jsonl(&text)?;
let mut result = ingest_json(engine, &json_array_text, opts)?;
result.stats.operation = "load_file".into();
result.stats.bytes_read = std::fs::metadata(abs_path).map_or(0, |m| m.len());
result.stats.file_format = Some(
if text.trim_start().starts_with('[') {
"json"
} else {
"jsonl"
}
.into(),
);
Ok(result)
}
pub async fn ingest_json_async(
conn: &AsyncConnection,
json_str: &str,
opts: &IngestOptions,
) -> Result<IngestResult, McpError> {
let timer = StatsTimer::start();
let bytes_read = json_str.len() as u64;
let schema_timer = StatsTimer::start();
let inferred = infer_json_schema(json_str)?;
let columns = match &opts.schema_override {
Some(s) => apply_schema_override(inferred, s)?,
None => inferred,
};
let schema_ms = schema_timer.elapsed_ms();
let array: Vec<serde_json::Map<String, Value>> = serde_json::from_str(json_str)
.map_err(|e| McpError::new(ErrorCode::SchemaMismatch, format!("Invalid JSON: {e}")))?;
if array.is_empty() {
return Err(McpError::new(ErrorCode::EmptyData, "JSON array is empty"));
}
let is_replace = opts.mode != "append";
let qualified = qualified_table(opts);
conn.begin_transaction().await.map_err(McpError::from)?;
let inner: Result<u64, McpError> = async {
create_table_async(
conn,
&opts.table,
&columns,
is_replace,
opts.target_db.as_deref(),
)
.await?;
let mut row_count = 0u64;
let col_names: Vec<String> = columns.iter().map(|c| format!("\"{}\"", c.name)).collect();
for obj in &array {
let values: Vec<String> = columns
.iter()
.map(|col| match obj.get(&col.name) {
None | Some(Value::Null) => "NULL".to_string(),
Some(Value::Bool(b)) => b.to_string(),
Some(Value::Number(n)) => n.to_string(),
Some(Value::String(s)) => format!("'{}'", s.replace('\'', "''")),
Some(other) => format!("'{}'", other.to_string().replace('\'', "''")),
})
.collect();
let sql = format!(
"INSERT INTO {} ({}) VALUES ({})",
qualified,
col_names.join(", "),
values.join(", ")
);
conn.execute_command(&sql).await.map_err(McpError::from)?;
row_count += 1;
}
Ok(row_count)
}
.await;
let row_count = match inner {
Ok(n) => {
conn.commit().await.map_err(McpError::from)?;
n
}
Err(e) => {
if let Err(rb) = conn.rollback().await {
tracing::warn!("rollback after error failed: {}", rb);
}
return Err(e);
}
};
let elapsed = timer.elapsed_ms();
let stats = IngestStats {
operation: "load_data".into(),
rows: row_count,
elapsed_ms: elapsed,
bytes_read,
bytes_stored: 0,
schema_inference_ms: Some(schema_ms),
table: opts.table.clone(),
file_format: Some("json".into()),
warning: if bytes_read > 50_000_000 {
Some("Large inline data. Consider using load_file for better performance.".into())
} else {
None
},
schema_changed: false,
};
Ok(IngestResult {
rows: row_count,
schema: columns,
stats,
})
}
pub async fn ingest_json_file_async(
conn: &AsyncConnection,
path: &str,
opts: &IngestOptions,
) -> Result<IngestResult, McpError> {
let abs_path = std::path::Path::new(path);
if !abs_path.exists() {
return Err(McpError::new(
ErrorCode::FileNotFound,
format!("File not found: {path}"),
));
}
let text = std::fs::read_to_string(abs_path)
.map_err(|e| McpError::new(ErrorCode::FileNotFound, format!("Cannot read file: {e}")))?;
let json_array_text = normalize_json_or_jsonl(&text)?;
let mut result = ingest_json_async(conn, &json_array_text, opts).await?;
result.stats.operation = "load_file".into();
result.stats.bytes_read = std::fs::metadata(abs_path).map_or(0, |m| m.len());
result.stats.file_format = Some(
if text.trim_start().starts_with('[') {
"json"
} else {
"jsonl"
}
.into(),
);
Ok(result)
}
pub(crate) async fn create_table_async(
conn: &AsyncConnection,
table_name: &str,
columns: &[ColumnSchema],
replace: bool,
target_db: Option<&str>,
) -> Result<(), McpError> {
if columns.is_empty() {
return Err(McpError::new(
ErrorCode::EmptyData,
"No columns to create table from",
));
}
for col in columns {
if crate::schema::map_hyper_type(&col.hyper_type).is_none() {
return Err(McpError::new(
ErrorCode::SchemaMismatch,
format!(
"Unknown type '{}' for column '{}'",
col.hyper_type, col.name
),
));
}
}
let quoted_table = match target_db {
Some(db) => {
let esc_db = db.replace('"', "\"\"");
let esc_tbl = table_name.replace('"', "\"\"");
format!("\"{esc_db}\".\"public\".\"{esc_tbl}\"")
}
None => format!("\"{}\"", table_name.replace('"', "\"\"")),
};
if replace {
conn.execute_command(&format!("DROP TABLE IF EXISTS {quoted_table}"))
.await
.map_err(McpError::from)?;
}
let col_defs: Vec<String> = columns
.iter()
.map(|c| {
let nullable = if c.nullable { "" } else { " NOT NULL" };
format!(
"\"{}\" {}{}",
c.name.replace('"', "\"\""),
c.hyper_type,
nullable
)
})
.collect();
let create_sql = format!(
"CREATE TABLE IF NOT EXISTS {} ({})",
quoted_table,
col_defs.join(", ")
);
conn.execute_command(&create_sql)
.await
.map_err(McpError::from)?;
Ok(())
}
pub fn normalize_json_or_jsonl(text: &str) -> Result<String, McpError> {
let trimmed = text.trim_start();
if trimmed.starts_with('[') {
return Ok(text.to_string());
}
let mut objects: Vec<Value> = Vec::new();
for (idx, line) in text.lines().enumerate() {
let line = line.trim();
if line.is_empty() {
continue;
}
let value: Value = serde_json::from_str(line).map_err(|e| {
McpError::new(
ErrorCode::SchemaMismatch,
format!("Invalid JSON on line {}: {e}", idx + 1),
)
})?;
objects.push(value);
}
if objects.is_empty() {
return Err(McpError::new(
ErrorCode::EmptyData,
"JSON/JSONL file contained no records",
));
}
serde_json::to_string(&Value::Array(objects)).map_err(|e| {
McpError::new(
ErrorCode::InternalError,
format!("Failed to serialize JSONL as array: {e}"),
)
})
}
pub fn extract_json_path(raw_json: &str, path: &str) -> Result<String, McpError> {
let mut current: Value = serde_json::from_str(raw_json).map_err(|e| {
McpError::new(
ErrorCode::SchemaMismatch,
format!("json_extract_path: file is not valid JSON: {e}"),
)
})?;
let segments: Vec<&str> = path.split('.').collect();
let mut traversed: Vec<&str> = Vec::new();
for segment in &segments {
if let Value::String(s) = ¤t {
current = serde_json::from_str(s).map_err(|_| {
McpError::new(
ErrorCode::SchemaMismatch,
format!(
"json_extract_path '{}': at segment '{}' (after '{}'): \
value is a string but not valid JSON",
path,
segment,
traversed.join(".")
),
)
})?;
}
current = if let Ok(idx) = segment.parse::<usize>() {
match current {
Value::Array(mut arr) => {
if idx >= arr.len() {
return Err(McpError::new(
ErrorCode::SchemaMismatch,
format!(
"json_extract_path '{}': at segment '{}' (after '{}'): \
array index {} out of bounds (length {})",
path,
segment,
traversed.join("."),
idx,
arr.len()
),
));
}
arr.swap_remove(idx)
}
other => {
return Err(McpError::new(
ErrorCode::SchemaMismatch,
format!(
"json_extract_path '{}': at segment '{}' (after '{}'): \
expected array, found {}",
path,
segment,
traversed.join("."),
json_type_name(&other)
),
));
}
}
} else {
match current {
Value::Object(mut map) => match map.remove(*segment) {
Some(v) => v,
None => {
return Err(McpError::new(
ErrorCode::SchemaMismatch,
format!(
"json_extract_path '{}': at segment '{}' (after '{}'): \
key not found in object",
path,
segment,
traversed.join(".")
),
));
}
},
other => {
return Err(McpError::new(
ErrorCode::SchemaMismatch,
format!(
"json_extract_path '{}': at segment '{}' (after '{}'): \
expected object, found {}",
path,
segment,
traversed.join("."),
json_type_name(&other)
),
));
}
}
};
traversed.push(segment);
}
if let Value::String(s) = ¤t {
if let Ok(parsed) = serde_json::from_str::<Value>(s) {
current = parsed;
}
}
serde_json::to_string(¤t).map_err(|e| {
McpError::new(
ErrorCode::InternalError,
format!("json_extract_path: failed to serialize extracted value: {e}"),
)
})
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum InferredFileFormat {
Parquet,
ArrowIpc,
Json,
Csv,
}
#[must_use]
pub fn detect_file_format(path: &std::path::Path) -> InferredFileFormat {
let ext = path
.extension()
.and_then(|e| e.to_str())
.unwrap_or("")
.to_lowercase();
match ext.as_str() {
"parquet" | "pq" => return InferredFileFormat::Parquet,
"arrow" | "ipc" | "feather" => return InferredFileFormat::ArrowIpc,
"json" | "jsonl" | "ndjson" => return InferredFileFormat::Json,
_ => {}
}
use std::io::Read;
if let Ok(mut f) = std::fs::File::open(path) {
let mut buf = [0u8; 4096];
if let Ok(n) = f.read(&mut buf) {
for b in &buf[..n] {
match b {
b' ' | b'\t' | b'\n' | b'\r' | 0xEF | 0xBB | 0xBF => {}
b'[' | b'{' => return InferredFileFormat::Json,
_ => return InferredFileFormat::Csv,
}
}
}
}
InferredFileFormat::Csv
}
#[cfg(test)]
mod read_text_sample_tests {
use super::read_text_sample;
use std::io::Write;
fn write_temp(name: &str, contents: &[u8]) -> tempfile::TempPath {
let mut tmp = tempfile::NamedTempFile::new().expect("temp file");
tmp.write_all(contents).unwrap();
tmp.flush().unwrap();
let _ = name;
tmp.into_temp_path()
}
#[test]
fn small_file_returns_full_contents() {
let path = write_temp("small.csv", b"a,b,c\n1,2,3\n");
let s = read_text_sample(&path, 1024 * 1024).unwrap();
assert_eq!(s, "a,b,c\n1,2,3\n");
}
#[test]
fn truncates_at_last_newline_within_budget() {
use std::fmt::Write;
let mut data = String::new();
for i in 0..20 {
writeln!(&mut data, "row-{i}").unwrap();
}
let path = write_temp("rows.csv", data.as_bytes());
let s = read_text_sample(&path, 30).unwrap();
assert!(s.ends_with('\n'));
assert!(s.len() <= 30);
}
#[test]
fn handles_utf8_boundary_split() {
let prefix = b"row1\n".to_vec();
let mut data = prefix.clone();
data.extend_from_slice("🔥".as_bytes());
let path = write_temp("emoji.csv", &data);
let s = read_text_sample(&path, 8).unwrap();
assert!(s.ends_with('\n'));
assert_eq!(s, "row1\n");
}
#[test]
fn returns_full_buffer_when_no_newline_under_budget() {
let path = write_temp("noeol.csv", b"a,b,c");
let s = read_text_sample(&path, 1024).unwrap();
assert_eq!(s, "a,b,c");
}
#[test]
fn handles_no_newline_in_first_max_bytes() {
let data = vec![b'a'; 2048];
let path = write_temp("noeol-big.csv", &data);
let s = read_text_sample(&path, 100).unwrap();
assert_eq!(s.len(), 100);
assert!(s.chars().all(|c| c == 'a'));
}
}