use crate::error::{ErrorCode, McpError};
use hyperdb_api::{SqlType, TableDefinition};
use serde_json::Value;
use std::collections::{BTreeMap, BTreeSet};
#[derive(Debug, Clone)]
pub struct ColumnSchema {
pub name: String,
pub hyper_type: String,
pub nullable: bool,
}
pub fn build_table_def(
table_name: &str,
columns: &[ColumnSchema],
) -> Result<TableDefinition, McpError> {
if columns.is_empty() {
return Err(McpError::new(
ErrorCode::EmptyData,
"No columns to create table from",
));
}
let mut def = TableDefinition::new(table_name);
for col in columns {
let sql_type = map_hyper_type(&col.hyper_type).ok_or_else(|| {
McpError::new(
ErrorCode::SchemaMismatch,
format!("Unknown type: {}", col.hyper_type),
)
})?;
if col.nullable {
def = def.add_nullable_column(&col.name, sql_type);
} else {
def = def.add_required_column(&col.name, sql_type);
}
}
Ok(def)
}
#[must_use]
pub fn map_hyper_type(type_name: &str) -> Option<SqlType> {
let upper = type_name.trim().to_uppercase();
match upper.as_str() {
"SMALLINT" | "INT2" => Some(SqlType::small_int()),
"INT" | "INTEGER" | "INT4" => Some(SqlType::int()),
"BIGINT" | "INT8" => Some(SqlType::big_int()),
"FLOAT" | "FLOAT4" | "REAL" => Some(SqlType::double()),
"DOUBLE" | "DOUBLE PRECISION" | "FLOAT8" => Some(SqlType::double()),
"TEXT" | "VARCHAR" | "STRING" => Some(SqlType::text()),
"BOOL" | "BOOLEAN" => Some(SqlType::bool()),
"DATE" => Some(SqlType::date()),
"TIME" => Some(SqlType::time()),
"TIMESTAMP" => Some(SqlType::timestamp()),
"TIMESTAMPTZ" | "TIMESTAMP WITH TIME ZONE" => Some(SqlType::timestamp_tz()),
"BYTEA" | "BYTES" => Some(SqlType::bytes()),
_ if upper.starts_with("NUMERIC") => {
if let Some(inner) = upper
.strip_prefix("NUMERIC(")
.and_then(|s| s.strip_suffix(')'))
{
let parts: Vec<&str> = inner.split(',').collect();
let precision = parts
.first()
.and_then(|p| p.trim().parse().ok())
.unwrap_or(38);
let scale = parts
.get(1)
.and_then(|s| s.trim().parse().ok())
.unwrap_or(0);
Some(SqlType::numeric(precision, scale))
} else {
Some(SqlType::numeric(38, 0))
}
}
_ => None,
}
}
#[derive(Debug, Clone, PartialEq)]
enum InferredType {
Null,
Bool,
Int,
BigInt,
Double,
Date,
Timestamp,
Text,
}
pub fn infer_json_schema(json_str: &str) -> Result<Vec<ColumnSchema>, McpError> {
let array: Vec<serde_json::Map<String, Value>> =
serde_json::from_str(json_str).map_err(|e| {
McpError::new(
ErrorCode::SchemaMismatch,
format!("Invalid JSON array: {e}"),
)
})?;
if array.is_empty() {
return Err(McpError::new(ErrorCode::EmptyData, "JSON array is empty"));
}
let mut all_keys = BTreeSet::new();
for obj in &array {
for key in obj.keys() {
all_keys.insert(key.clone());
}
}
let total_rows = array.len();
let mut col_types: BTreeMap<String, Vec<InferredType>> = BTreeMap::new();
let mut col_present: BTreeMap<String, usize> = BTreeMap::new();
for key in &all_keys {
col_types.insert(key.clone(), Vec::new());
col_present.insert(key.clone(), 0);
}
for obj in &array {
for key in &all_keys {
match obj.get(key.as_str()) {
None => {}
Some(Value::Null) => {
col_types.get_mut(key).unwrap().push(InferredType::Null);
*col_present.get_mut(key).unwrap() += 1;
}
Some(val) => {
col_types
.get_mut(key)
.unwrap()
.push(infer_json_value_type(val));
*col_present.get_mut(key).unwrap() += 1;
}
}
}
}
let mut columns = Vec::new();
for key in &all_keys {
let types = &col_types[key];
let present_count = col_present[key];
let nullable = present_count < total_rows || types.contains(&InferredType::Null);
let resolved = resolve_type(types);
columns.push(ColumnSchema {
name: key.clone(),
hyper_type: inferred_to_hyper_name(&resolved),
nullable,
});
}
Ok(columns)
}
fn infer_json_value_type(val: &Value) -> InferredType {
match val {
Value::Null => InferredType::Null,
Value::Bool(_) => InferredType::Bool,
Value::Number(n) => {
if let Some(i) = n.as_i64() {
if i32::try_from(i).is_ok() {
InferredType::Int
} else {
InferredType::BigInt
}
} else {
InferredType::Double
}
}
Value::String(s) => infer_string_type(s),
_ => InferredType::Text,
}
}
fn infer_string_type(s: &str) -> InferredType {
if s.len() == 10
&& s.chars().nth(4) == Some('-')
&& s.chars().nth(7) == Some('-')
&& s[0..4].parse::<u16>().is_ok()
&& s[5..7].parse::<u8>().is_ok()
&& s[8..10].parse::<u8>().is_ok()
{
return InferredType::Date;
}
if s.len() >= 19
&& s.chars().nth(10) == Some('T')
&& s[0..10].contains('-')
&& s[11..].contains(':')
{
return InferredType::Timestamp;
}
InferredType::Text
}
fn resolve_type(types: &[InferredType]) -> InferredType {
let non_null: Vec<&InferredType> = types.iter().filter(|t| **t != InferredType::Null).collect();
if non_null.is_empty() {
return InferredType::Text; }
let first = non_null[0];
let all_same = non_null.iter().all(|t| *t == first);
if all_same {
return first.clone();
}
let all_numeric = non_null.iter().all(|t| {
matches!(
t,
InferredType::Int | InferredType::BigInt | InferredType::Double
)
});
if all_numeric {
if non_null.iter().any(|t| **t == InferredType::Double) {
return InferredType::Double;
}
if non_null.iter().any(|t| **t == InferredType::BigInt) {
return InferredType::BigInt;
}
return InferredType::Int;
}
InferredType::Text
}
fn inferred_to_hyper_name(t: &InferredType) -> String {
match t {
InferredType::Null | InferredType::Text => "TEXT".into(),
InferredType::Bool => "BOOL".into(),
InferredType::Int => "INT".into(),
InferredType::BigInt => "BIGINT".into(),
InferredType::Double => "DOUBLE PRECISION".into(),
InferredType::Date => "DATE".into(),
InferredType::Timestamp => "TIMESTAMP".into(),
}
}
pub fn infer_csv_schema(csv_text: &str, has_header: bool) -> Result<Vec<ColumnSchema>, McpError> {
let mut reader = csv::ReaderBuilder::new()
.has_headers(has_header)
.from_reader(csv_text.as_bytes());
let headers: Vec<String> = if has_header {
reader
.headers()
.map_err(|e| {
McpError::new(ErrorCode::SchemaMismatch, format!("CSV header error: {e}"))
})?
.iter()
.map(std::string::ToString::to_string)
.collect()
} else {
let first = reader.records().next();
match first {
Some(Ok(ref rec)) => (0..rec.len()).map(|i| format!("col_{i}")).collect(),
_ => return Err(McpError::new(ErrorCode::EmptyData, "CSV has no data rows")),
}
};
if headers.is_empty() {
return Err(McpError::new(ErrorCode::EmptyData, "CSV has no columns"));
}
let num_cols = headers.len();
let mut col_types: Vec<Vec<InferredType>> = vec![Vec::new(); num_cols];
let mut sample_reader = csv::ReaderBuilder::new()
.has_headers(has_header)
.from_reader(csv_text.as_bytes());
for (row_idx, result) in sample_reader.records().enumerate() {
if row_idx >= 1000 {
break;
}
let record = result.map_err(|e| {
McpError::new(
ErrorCode::SchemaMismatch,
format!("CSV parse error at row {}: {e}", row_idx + 1),
)
})?;
for (col_idx, field) in record.iter().enumerate() {
if col_idx < num_cols {
col_types[col_idx].push(infer_csv_field_type(field));
}
}
}
let columns: Vec<ColumnSchema> = headers
.into_iter()
.enumerate()
.map(|(i, name)| {
let resolved = resolve_type(&col_types[i]);
ColumnSchema {
name,
hyper_type: inferred_to_hyper_name(&resolved),
nullable: true, }
})
.collect();
Ok(columns)
}
pub fn widen_csv_numeric_columns<R: std::io::Read>(
reader: R,
has_header: bool,
columns: &mut [ColumnSchema],
) -> Result<(), McpError> {
let candidate_idxs: Vec<usize> = columns
.iter()
.enumerate()
.filter(|(_, c)| {
matches!(
c.hyper_type.as_str(),
"INT" | "INTEGER" | "BIGINT" | "DOUBLE PRECISION"
)
})
.map(|(i, _)| i)
.collect();
if candidate_idxs.is_empty() {
return Ok(());
}
#[derive(Default)]
struct ColStats {
min: Option<i128>,
max: Option<i128>,
has_decimal: bool,
overflow_i128: bool,
}
let mut stats: std::collections::HashMap<usize, ColStats> = candidate_idxs
.iter()
.map(|i| (*i, ColStats::default()))
.collect();
let mut rdr = csv::ReaderBuilder::new()
.has_headers(has_header)
.from_reader(reader);
for (row_idx, result) in rdr.records().enumerate() {
let record = result.map_err(|e| {
McpError::new(
ErrorCode::SchemaMismatch,
format!("CSV parse error at row {}: {e}", row_idx + 1),
)
})?;
for &col_idx in &candidate_idxs {
let Some(field) = record.get(col_idx) else {
continue;
};
let trimmed = field.trim();
if trimmed.is_empty()
|| trimmed.eq_ignore_ascii_case("null")
|| trimmed.eq_ignore_ascii_case("na")
{
continue;
}
let s = stats.get_mut(&col_idx).expect("preallocated");
if trimmed.contains('.') || trimmed.contains('e') || trimmed.contains('E') {
s.has_decimal = true;
continue;
}
match trimmed.parse::<i128>() {
Ok(n) => {
s.min = Some(s.min.map_or(n, |m| m.min(n)));
s.max = Some(s.max.map_or(n, |m| m.max(n)));
}
Err(_) => {
s.overflow_i128 = true;
}
}
}
}
for (&col_idx, s) in &stats {
let col = &mut columns[col_idx];
let i32_range = i128::from(i32::MIN)..=i128::from(i32::MAX);
let i64_range = i128::from(i64::MIN)..=i128::from(i64::MAX);
match col.hyper_type.as_str() {
"INT" | "INTEGER" => {
if s.has_decimal {
col.hyper_type = "DOUBLE PRECISION".into();
} else if s.overflow_i128
|| !s.min.map_or(true, |m| i64_range.contains(&m))
|| !s.max.map_or(true, |m| i64_range.contains(&m))
{
col.hyper_type = "NUMERIC(38,0)".into();
} else if !s.min.map_or(true, |m| i32_range.contains(&m))
|| !s.max.map_or(true, |m| i32_range.contains(&m))
{
col.hyper_type = "BIGINT".into();
}
}
"BIGINT" => {
if s.has_decimal {
col.hyper_type = "DOUBLE PRECISION".into();
} else if s.overflow_i128
|| !s.min.map_or(true, |m| i64_range.contains(&m))
|| !s.max.map_or(true, |m| i64_range.contains(&m))
{
col.hyper_type = "NUMERIC(38,0)".into();
}
}
_ => {}
}
}
Ok(())
}
fn infer_csv_field_type(field: &str) -> InferredType {
let trimmed = field.trim();
if trimmed.is_empty()
|| trimmed.eq_ignore_ascii_case("null")
|| trimmed.eq_ignore_ascii_case("na")
{
return InferredType::Null;
}
if trimmed.eq_ignore_ascii_case("true") || trimmed.eq_ignore_ascii_case("false") {
return InferredType::Bool;
}
if let Ok(i) = trimmed.parse::<i64>() {
if i32::try_from(i).is_ok() {
return InferredType::Int;
}
return InferredType::BigInt;
}
if trimmed.parse::<f64>().is_ok() {
return InferredType::Double;
}
infer_string_type(trimmed)
}
pub fn parse_schema_override(
schema: &serde_json::Map<String, Value>,
) -> Result<Vec<ColumnSchema>, McpError> {
let mut columns = Vec::new();
for (name, type_val) in schema {
let type_name = type_val.as_str().ok_or_else(|| {
McpError::new(
ErrorCode::SchemaMismatch,
format!("Schema type for '{name}' must be a string"),
)
})?;
if map_hyper_type(type_name).is_none() {
return Err(McpError::new(
ErrorCode::SchemaMismatch,
format!("Unknown type '{type_name}' for column '{name}'"),
));
}
columns.push(ColumnSchema {
name: name.clone(),
hyper_type: type_name.to_uppercase(),
nullable: true,
});
}
Ok(columns)
}
pub fn normalize_schema_param(
schema: Option<&Value>,
) -> Result<Option<serde_json::Map<String, Value>>, McpError> {
let Some(v) = schema else {
return Ok(None);
};
match v {
Value::Null => Ok(None),
Value::Object(m) => Ok(Some(m.clone())),
Value::String(s) => {
let trimmed = s.trim();
if trimmed.is_empty() {
return Ok(None);
}
let parsed: Value = serde_json::from_str(trimmed).map_err(|e| {
McpError::new(
ErrorCode::SchemaMismatch,
format!(
"`schema` parameter is a string but not valid JSON: {e}. \
Expected an object like {{\"col\": \"TEXT\"}}."
),
)
})?;
match parsed {
Value::Object(m) => Ok(Some(m)),
Value::Null => Ok(None),
other => Err(McpError::new(
ErrorCode::SchemaMismatch,
format!(
"`schema` parameter must be a JSON object mapping column names \
to type strings, got {}.",
json_type_name(&other)
),
)),
}
}
other => Err(McpError::new(
ErrorCode::SchemaMismatch,
format!(
"`schema` parameter must be a JSON object mapping column names to type \
strings, got {}.",
json_type_name(other)
),
)),
}
}
pub(crate) fn json_type_name(v: &Value) -> &'static str {
match v {
Value::Null => "null",
Value::Bool(_) => "boolean",
Value::Number(_) => "number",
Value::String(_) => "string",
Value::Array(_) => "array",
Value::Object(_) => "object",
}
}
pub fn apply_schema_override(
mut inferred: Vec<ColumnSchema>,
override_map: &serde_json::Map<String, Value>,
) -> Result<Vec<ColumnSchema>, McpError> {
let known: std::collections::HashSet<&str> = inferred.iter().map(|c| c.name.as_str()).collect();
for (name, type_val) in override_map {
if !known.contains(name.as_str()) {
let real: Vec<&str> = inferred.iter().map(|c| c.name.as_str()).collect();
return Err(McpError::new(
ErrorCode::SchemaMismatch,
format!("Override key '{name}' does not match any column. Known columns: {real:?}"),
));
}
let type_name = type_val.as_str().ok_or_else(|| {
McpError::new(
ErrorCode::SchemaMismatch,
format!("Schema type for '{name}' must be a string"),
)
})?;
if map_hyper_type(type_name).is_none() {
return Err(McpError::new(
ErrorCode::SchemaMismatch,
format!("Unknown type '{type_name}' for column '{name}'"),
));
}
}
for col in &mut inferred {
if let Some(v) = override_map.get(&col.name).and_then(|v| v.as_str()) {
col.hyper_type = v.trim().to_uppercase();
}
}
Ok(inferred)
}