use crate::engine::Engine;
use crate::error::{ErrorCode, McpError};
use crate::ingest::{IngestOptions, IngestResult};
use crate::schema::ColumnSchema;
use crate::stats::{IngestStats, StatsTimer};
use arrow::datatypes::{DataType, Schema as ArrowSchema};
use arrow::record_batch::RecordBatch;
use hyperdb_api::AsyncConnection;
use std::path::Path;
fn arrow_type_to_hyper(dt: &DataType) -> String {
match dt {
DataType::Boolean => "BOOL".into(),
DataType::Int8 | DataType::Int16 => "SMALLINT".into(),
DataType::Int32 | DataType::UInt16 => "INT".into(),
DataType::Int64 | DataType::UInt32 => "BIGINT".into(),
DataType::UInt64 => "BIGINT".into(),
DataType::Float16 | DataType::Float32 => "DOUBLE PRECISION".into(),
DataType::Float64 => "DOUBLE PRECISION".into(),
DataType::Utf8 | DataType::LargeUtf8 => "TEXT".into(),
DataType::Binary | DataType::LargeBinary => "BYTEA".into(),
DataType::Date32 | DataType::Date64 => "DATE".into(),
DataType::Time32(_) | DataType::Time64(_) => "TIME".into(),
DataType::Timestamp(_, None) => "TIMESTAMP".into(),
DataType::Timestamp(_, Some(_)) => "TIMESTAMPTZ".into(),
DataType::Decimal128(p, s) | DataType::Decimal256(p, s) => {
format!("NUMERIC({p}, {s})")
}
_ => "TEXT".into(),
}
}
#[must_use]
pub fn arrow_schema_to_columns(schema: &ArrowSchema) -> Vec<ColumnSchema> {
schema
.fields()
.iter()
.map(|f| ColumnSchema {
name: f.name().clone(),
hyper_type: arrow_type_to_hyper(f.data_type()),
nullable: f.is_nullable(),
})
.collect()
}
fn record_batches_to_ipc_stream(batches: &[RecordBatch]) -> Result<Vec<u8>, McpError> {
let schema = batches
.first()
.map(arrow::array::RecordBatch::schema)
.ok_or_else(|| McpError::new(ErrorCode::EmptyData, "Arrow IPC file has no batches"))?;
let mut buf = Vec::new();
{
let mut writer =
arrow::ipc::writer::StreamWriter::try_new(&mut buf, &schema).map_err(|e| {
McpError::new(
ErrorCode::InternalError,
format!("Failed to create Arrow IPC StreamWriter: {e}"),
)
})?;
for batch in batches {
writer.write(batch).map_err(|e| {
McpError::new(
ErrorCode::InternalError,
format!("Failed to write Arrow batch: {e}"),
)
})?;
}
writer.finish().map_err(|e| {
McpError::new(
ErrorCode::InternalError,
format!("Failed to finish Arrow IPC stream: {e}"),
)
})?;
}
Ok(buf)
}
fn infer_parquet_schema(path: &str) -> Result<Vec<ColumnSchema>, McpError> {
let file = std::fs::File::open(path)
.map_err(|e| McpError::new(ErrorCode::FileNotFound, format!("Cannot open file: {e}")))?;
let reader = parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder::try_new(file)
.map_err(|e| {
McpError::new(
ErrorCode::UnsupportedFormat,
format!("Invalid Parquet file: {e}"),
)
})?;
let arrow_schema = reader.schema();
Ok(arrow_schema_to_columns(arrow_schema))
}
fn parquet_select_projection(inferred: &[ColumnSchema], final_columns: &[ColumnSchema]) -> String {
inferred
.iter()
.zip(final_columns.iter())
.map(|(orig, col)| {
let quoted = format!("\"{}\"", col.name.replace('"', "\"\""));
if orig.hyper_type == col.hyper_type {
quoted
} else {
format!("{quoted}::{ty} AS {quoted}", ty = col.hyper_type)
}
})
.collect::<Vec<_>>()
.join(", ")
}
fn build_parquet_ingest_sql(
table: &str,
path: &str,
projection: &str,
is_replace: bool,
target_db: Option<&str>,
) -> String {
let quoted_table = match target_db {
Some(db) => {
let esc_db = db.replace('"', "\"\"");
let esc_tbl = table.replace('"', "\"\"");
format!("\"{esc_db}\".\"public\".\"{esc_tbl}\"")
}
None => format!("\"{}\"", table.replace('"', "\"\"")),
};
let quoted_path = hyperdb_api::escape_string_literal(path);
if is_replace {
format!(
"CREATE TABLE {quoted_table} AS SELECT {projection} FROM external({quoted_path}, FORMAT => 'parquet')"
)
} else {
format!(
"INSERT INTO {quoted_table} SELECT {projection} FROM external({quoted_path}, FORMAT => 'parquet')"
)
}
}
fn resolve_parquet_path(path: &str) -> Result<(String, u64), McpError> {
let file_path = Path::new(path);
if !file_path.exists() {
return Err(McpError::new(
ErrorCode::FileNotFound,
format!("File not found: {path}"),
));
}
let absolute = std::fs::canonicalize(file_path)
.map_err(|e| {
McpError::new(
ErrorCode::FileNotFound,
format!("Cannot resolve path {path}: {e}"),
)
})?
.to_string_lossy()
.into_owned();
let file_size = std::fs::metadata(file_path).map_or(0, |m| m.len());
Ok((absolute, file_size))
}
fn count_rows_sync(engine: &Engine, table: &str) -> Result<u64, McpError> {
let quoted = format!("\"{}\"", table.replace('"', "\"\""));
let sql = format!("SELECT COUNT(*) FROM {quoted}");
let rows = engine.execute_query_to_json(&sql)?;
rows.first()
.and_then(|r| r.get("count"))
.and_then(serde_json::Value::as_u64)
.ok_or_else(|| {
McpError::new(
ErrorCode::InternalError,
"Could not read row count after parquet ingest",
)
})
}
async fn count_rows_async(conn: &AsyncConnection, table: &str) -> Result<u64, McpError> {
let quoted = format!("\"{}\"", table.replace('"', "\"\""));
let sql = format!("SELECT COUNT(*) FROM {quoted}");
let row = conn.fetch_one(&sql).await.map_err(McpError::from)?;
let count: i64 = row.get(0).ok_or_else(|| {
McpError::new(
ErrorCode::InternalError,
"Could not read row count after parquet ingest",
)
})?;
Ok(u64::try_from(count).unwrap_or(0))
}
pub fn ingest_parquet_file(
engine: &Engine,
path: &str,
opts: &IngestOptions,
) -> Result<IngestResult, McpError> {
if opts.mode == "merge" {
return crate::ingest::merge_via_temp_table(engine, opts, |tmp_opts| {
ingest_parquet_file(engine, path, tmp_opts)
});
}
let timer = StatsTimer::start();
let (absolute_path, file_size) = resolve_parquet_path(path)?;
let inferred = infer_parquet_schema(&absolute_path)?;
let final_columns = match &opts.schema_override {
Some(s) => crate::schema::apply_schema_override(inferred.clone(), s)?,
None => inferred.clone(),
};
let projection = parquet_select_projection(&inferred, &final_columns);
let is_replace = opts.mode != "append";
let sql = build_parquet_ingest_sql(
&opts.table,
&absolute_path,
&projection,
is_replace,
opts.target_db.as_deref(),
);
let affected = engine.execute_in_transaction(|engine| {
if is_replace {
let qualified = crate::ingest::qualified_table(opts);
engine.execute_command(&format!("DROP TABLE IF EXISTS {qualified}"))?;
}
engine.execute_command(&sql)
})?;
let row_count = if is_replace {
count_rows_sync(engine, &opts.table)?
} else {
affected
};
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(0),
table: opts.table.clone(),
file_format: Some("parquet".into()),
warning: None,
schema_changed: false,
};
Ok(IngestResult {
rows: row_count,
schema: final_columns,
stats,
})
}
pub async fn ingest_parquet_file_async(
conn: &AsyncConnection,
path: &str,
opts: &IngestOptions,
) -> Result<IngestResult, McpError> {
let timer = StatsTimer::start();
let (absolute_path, file_size) = resolve_parquet_path(path)?;
let path_for_infer = absolute_path.clone();
let override_owned = opts.schema_override.clone();
let (inferred, final_columns): (Vec<ColumnSchema>, Vec<ColumnSchema>) =
tokio::task::spawn_blocking(move || -> Result<_, McpError> {
let inferred = infer_parquet_schema(&path_for_infer)?;
let final_columns = match &override_owned {
Some(s) => crate::schema::apply_schema_override(inferred.clone(), s)?,
None => inferred.clone(),
};
Ok((inferred, final_columns))
})
.await
.map_err(|e| McpError::new(ErrorCode::InternalError, format!("Task join error: {e}")))??;
let projection = parquet_select_projection(&inferred, &final_columns);
let is_replace = opts.mode != "append";
let sql = build_parquet_ingest_sql(
&opts.table,
&absolute_path,
&projection,
is_replace,
opts.target_db.as_deref(),
);
conn.begin_transaction().await.map_err(McpError::from)?;
let result: Result<u64, McpError> = async {
if is_replace {
let qualified = crate::ingest::qualified_table(opts);
conn.execute_command(&format!("DROP TABLE IF EXISTS {qualified}"))
.await
.map_err(McpError::from)?;
}
conn.execute_command(&sql).await.map_err(McpError::from)
}
.await;
let affected = match result {
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 row_count = if is_replace {
count_rows_async(conn, &opts.table).await?
} else {
affected
};
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(0),
table: opts.table.clone(),
file_format: Some("parquet".into()),
warning: None,
schema_changed: false,
};
Ok(IngestResult {
rows: row_count,
schema: final_columns,
stats,
})
}
fn reject_ipc_schema_override(opts: &IngestOptions) -> Result<(), McpError> {
if opts.schema_override.is_some() {
return Err(McpError::new(
ErrorCode::SchemaMismatch,
"Schema overrides are not supported for Arrow IPC files. \
The embedded Arrow schema is authoritative on this path \
because the binary COPY protocol requires an exact type \
match between the file and the target table.",
));
}
Ok(())
}
const ARROW_IPC_FILE_MAGIC: &[u8] = b"ARROW1";
fn read_arrow_ipc_file(path: &str) -> Result<(Vec<ColumnSchema>, Vec<RecordBatch>), McpError> {
use std::io::Read;
let mut file = std::fs::File::open(path)
.map_err(|e| McpError::new(ErrorCode::FileNotFound, format!("Cannot open file: {e}")))?;
let mut magic = [0u8; 6];
let read = file.read(&mut magic).map_err(|e| {
McpError::new(
ErrorCode::UnsupportedFormat,
format!("Cannot read Arrow IPC header: {e}"),
)
})?;
use std::io::Seek;
file.rewind().map_err(|e| {
McpError::new(
ErrorCode::InternalError,
format!("Cannot rewind file handle: {e}"),
)
})?;
let is_file_format = read == 6 && magic == ARROW_IPC_FILE_MAGIC;
if is_file_format {
let reader = arrow::ipc::reader::FileReader::try_new(file, None).map_err(|e| {
McpError::new(
ErrorCode::UnsupportedFormat,
format!("Invalid Arrow IPC file: {e}"),
)
})?;
let inferred = arrow_schema_to_columns(&reader.schema());
let batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>().map_err(|e| {
McpError::new(
ErrorCode::InternalError,
format!("Arrow IPC read error: {e}"),
)
})?;
Ok((inferred, batches))
} else {
let reader = arrow::ipc::reader::StreamReader::try_new(file, None).map_err(|e| {
McpError::new(
ErrorCode::UnsupportedFormat,
format!("Invalid Arrow IPC stream: {e}"),
)
})?;
let inferred = arrow_schema_to_columns(&reader.schema());
let batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>().map_err(|e| {
McpError::new(
ErrorCode::InternalError,
format!("Arrow IPC read error: {e}"),
)
})?;
Ok((inferred, batches))
}
}
pub fn ingest_arrow_ipc_file(
engine: &Engine,
path: &str,
opts: &IngestOptions,
) -> Result<IngestResult, McpError> {
if opts.mode == "merge" {
return crate::ingest::merge_via_temp_table(engine, opts, |tmp_opts| {
ingest_arrow_ipc_file(engine, path, tmp_opts)
});
}
let timer = StatsTimer::start();
reject_ipc_schema_override(opts)?;
let file_path = Path::new(path);
if !file_path.exists() {
return Err(McpError::new(
ErrorCode::FileNotFound,
format!("File not found: {path}"),
));
}
let file_size = std::fs::metadata(file_path).map_or(0, |m| m.len());
let (columns, batches) = read_arrow_ipc_file(path)?;
let is_replace = opts.mode != "append";
let _search_guard = if let Some(ref db) = opts.target_db {
Some(engine.scoped_search_path(db)?)
} else {
None
};
let row_count = engine.execute_in_transaction(|engine| {
engine.create_table_in(&opts.table, &columns, is_replace, opts.target_db.as_deref())?;
let mut inserter =
hyperdb_api::ArrowInserter::from_table(engine.connection(), opts.table.as_str())
.map_err(McpError::from)?;
inserter
.insert_batches(batches.iter())
.map_err(McpError::from)?;
inserter.execute().map_err(McpError::from)
})?;
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(0),
table: opts.table.clone(),
file_format: Some("arrow_ipc".into()),
warning: None,
schema_changed: false,
};
Ok(IngestResult {
rows: row_count,
schema: columns,
stats,
})
}
pub async fn ingest_arrow_ipc_file_async(
conn: &AsyncConnection,
path: &str,
opts: &IngestOptions,
) -> Result<IngestResult, McpError> {
let timer = StatsTimer::start();
reject_ipc_schema_override(opts)?;
let file_path = Path::new(path);
if !file_path.exists() {
return Err(McpError::new(
ErrorCode::FileNotFound,
format!("File not found: {path}"),
));
}
let file_size = std::fs::metadata(file_path).map_or(0, |m| m.len());
let path_owned = path.to_string();
let (columns, ipc_stream): (Vec<ColumnSchema>, Vec<u8>) =
tokio::task::spawn_blocking(move || -> Result<_, McpError> {
let (columns, batches) = read_arrow_ipc_file(&path_owned)?;
let ipc_stream = record_batches_to_ipc_stream(&batches)?;
Ok((columns, ipc_stream))
})
.await
.map_err(|e| McpError::new(ErrorCode::InternalError, format!("Task join error: {e}")))??;
let is_replace = opts.mode != "append";
let table_def = crate::schema::build_table_def(&opts.table, &columns)?;
conn.begin_transaction().await.map_err(McpError::from)?;
let result: Result<u64, McpError> = async {
crate::ingest::create_table_async(
conn,
&opts.table,
&columns,
is_replace,
opts.target_db.as_deref(),
)
.await?;
let mut inserter =
hyperdb_api::AsyncArrowInserter::new(conn, &table_def).map_err(McpError::from)?;
inserter
.insert_data(&ipc_stream)
.await
.map_err(McpError::from)?;
inserter.execute().await.map_err(McpError::from)
}
.await;
let row_count = match result {
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(0),
table: opts.table.clone(),
file_format: Some("arrow_ipc".into()),
warning: None,
schema_changed: false,
};
Ok(IngestResult {
rows: row_count,
schema: columns,
stats,
})
}