bestool-postgres 1.1.0

use fraction::ToPrimitive;

/// Convert a PostgreSQL row column to a JSON value
///
/// This function handles common PostgreSQL types and converts them to appropriate
/// serde_json::Value representations. Unsupported types are converted to JSON strings
/// using their text representation.
pub fn postgres_to_json_value(row: &tokio_postgres::Row, idx: usize) -> serde_json::Value {
	use tokio_postgres::types::Type;

	let column = &row.columns()[idx];
	match column.type_() {
		&Type::BOOL => row
			.try_get::<_, Option<bool>>(idx)
			.ok()
			.flatten()
			.map(serde_json::Value::Bool)
			.unwrap_or(serde_json::Value::Null),
		&Type::INT2 => row
			.try_get::<_, Option<i16>>(idx)
			.ok()
			.flatten()
			.map(|v| serde_json::Value::Number(v.into()))
			.unwrap_or(serde_json::Value::Null),
		&Type::INT4 => row
			.try_get::<_, Option<i32>>(idx)
			.ok()
			.flatten()
			.map(|v| serde_json::Value::Number(v.into()))
			.unwrap_or(serde_json::Value::Null),
		&Type::INT8 => row
			.try_get::<_, Option<i64>>(idx)
			.ok()
			.flatten()
			.map(|v| serde_json::Value::Number(v.into()))
			.unwrap_or(serde_json::Value::Null),
		&Type::FLOAT4 => row
			.try_get::<_, Option<f32>>(idx)
			.ok()
			.flatten()
			.and_then(|v| serde_json::Number::from_f64(v as f64))
			.map(serde_json::Value::Number)
			.unwrap_or(serde_json::Value::Null),
		&Type::FLOAT8 => row
			.try_get::<_, Option<f64>>(idx)
			.ok()
			.flatten()
			.and_then(serde_json::Number::from_f64)
			.map(serde_json::Value::Number)
			.unwrap_or(serde_json::Value::Null),
		&Type::NUMERIC => row
			.try_get::<_, Option<fraction::Decimal>>(idx)
			.ok()
			.flatten()
			.and_then(|v| {
				// Convert Decimal to f64, then to JSON number
				let float_val = v.to_f64()?;
				serde_json::Number::from_f64(float_val)
			})
			.map(serde_json::Value::Number)
			.unwrap_or(serde_json::Value::Null),
		&Type::TEXT | &Type::VARCHAR | &Type::BPCHAR | &Type::NAME => row
			.try_get::<_, Option<String>>(idx)
			.ok()
			.flatten()
			.map(serde_json::Value::String)
			.unwrap_or(serde_json::Value::Null),
		&Type::JSON | &Type::JSONB => {
			let val: Option<serde_json::Value> = row.get(idx);
			val.unwrap_or(serde_json::Value::Null)
		}
		&Type::TIMESTAMP => row
			.try_get::<_, Option<jiff::civil::DateTime>>(idx)
			.ok()
			.flatten()
			.map(|dt| serde_json::Value::String(dt.to_string()))
			.unwrap_or(serde_json::Value::Null),
		&Type::TIMESTAMPTZ => row
			.try_get::<_, Option<jiff::Timestamp>>(idx)
			.ok()
			.flatten()
			.map(|ts| serde_json::Value::String(ts.to_string()))
			.unwrap_or(serde_json::Value::Null),
		&Type::DATE => row
			.try_get::<_, Option<jiff::civil::Date>>(idx)
			.ok()
			.flatten()
			.map(|d| serde_json::Value::String(d.to_string()))
			.unwrap_or(serde_json::Value::Null),
		&Type::TIME => row
			.try_get::<_, Option<jiff::civil::Time>>(idx)
			.ok()
			.flatten()
			.map(|t| serde_json::Value::String(t.to_string()))
			.unwrap_or(serde_json::Value::Null),
		&Type::BYTEA => row
			.try_get::<_, Option<Vec<u8>>>(idx)
			.ok()
			.flatten()
			.map(|v| serde_json::Value::String(format!("\\x{}", hex::encode(v))))
			.unwrap_or(serde_json::Value::Null),
		// Array types
		&Type::TEXT_ARRAY | &Type::VARCHAR_ARRAY => row
			.try_get::<_, Option<Vec<String>>>(idx)
			.ok()
			.flatten()
			.map(|v| {
				serde_json::Value::Array(v.into_iter().map(serde_json::Value::String).collect())
			})
			.unwrap_or(serde_json::Value::Null),
		&Type::INT2_ARRAY => row
			.try_get::<_, Option<Vec<i16>>>(idx)
			.ok()
			.flatten()
			.map(|v| {
				serde_json::Value::Array(
					v.into_iter()
						.map(|n| serde_json::Value::Number(n.into()))
						.collect(),
				)
			})
			.unwrap_or(serde_json::Value::Null),
		&Type::INT4_ARRAY => row
			.try_get::<_, Option<Vec<i32>>>(idx)
			.ok()
			.flatten()
			.map(|v| {
				serde_json::Value::Array(
					v.into_iter()
						.map(|n| serde_json::Value::Number(n.into()))
						.collect(),
				)
			})
			.unwrap_or(serde_json::Value::Null),
		&Type::INT8_ARRAY => row
			.try_get::<_, Option<Vec<i64>>>(idx)
			.ok()
			.flatten()
			.map(|v| {
				serde_json::Value::Array(
					v.into_iter()
						.map(|n| serde_json::Value::Number(n.into()))
						.collect(),
				)
			})
			.unwrap_or(serde_json::Value::Null),
		&Type::FLOAT4_ARRAY => row
			.try_get::<_, Option<Vec<f32>>>(idx)
			.ok()
			.flatten()
			.map(|v| {
				serde_json::Value::Array(
					v.into_iter()
						.filter_map(|n| {
							serde_json::Number::from_f64(n as f64).map(serde_json::Value::Number)
						})
						.collect(),
				)
			})
			.unwrap_or(serde_json::Value::Null),
		&Type::FLOAT8_ARRAY => row
			.try_get::<_, Option<Vec<f64>>>(idx)
			.ok()
			.flatten()
			.map(|v| {
				serde_json::Value::Array(
					v.into_iter()
						.filter_map(|n| {
							serde_json::Number::from_f64(n).map(serde_json::Value::Number)
						})
						.collect(),
				)
			})
			.unwrap_or(serde_json::Value::Null),
		&Type::BOOL_ARRAY => row
			.try_get::<_, Option<Vec<bool>>>(idx)
			.ok()
			.flatten()
			.map(|v| serde_json::Value::Array(v.into_iter().map(serde_json::Value::Bool).collect()))
			.unwrap_or(serde_json::Value::Null),
		// For unknown types, try to get as string
		_ => row
			.try_get::<_, Option<String>>(idx)
			.ok()
			.flatten()
			.map(serde_json::Value::String)
			.unwrap_or(serde_json::Value::Null),
	}
}

pub fn get_value(
	row: &tokio_postgres::Row,
	column_index: usize,
	unprintable_columns: &[usize],
) -> String {
	if !unprintable_columns.contains(&column_index) {
		return format_value(row, column_index);
	}

	// For unprintable columns without async context, show a placeholder
	// The actual text casting happens in the display layer which is async
	"(binary data)".to_string()
}

pub fn format_value(row: &tokio_postgres::Row, i: usize) -> String {
	// Check for void type first
	let column = row.columns().get(i);
	if let Some(col) = column
		&& col.type_().name() == "void"
	{
		return "(void)".to_string();
	}

	// Try numeric type with fraction crate
	if let Ok(v) = row.try_get::<_, fraction::Decimal>(i) {
		v.to_string()
	} else if let Ok(v) = row.try_get::<_, String>(i) {
		v
	} else if let Ok(v) = row.try_get::<_, i16>(i) {
		v.to_string()
	} else if let Ok(v) = row.try_get::<_, i32>(i) {
		v.to_string()
	} else if let Ok(v) = row.try_get::<_, i64>(i) {
		v.to_string()
	} else if let Ok(v) = row.try_get::<_, f32>(i) {
		format!("{}", v)
	} else if let Ok(v) = row.try_get::<_, f64>(i) {
		format!("{}", v)
	} else if let Ok(v) = row.try_get::<_, bool>(i) {
		v.to_string()
	} else if let Ok(v) = row.try_get::<_, Vec<u8>>(i) {
		format!("\\x{encoded}", encoded = hex::encode(v))
	} else if let Ok(v) = row.try_get::<_, jiff::Timestamp>(i) {
		v.to_string()
	} else if let Ok(v) = row.try_get::<_, jiff::civil::Date>(i) {
		v.to_string()
	} else if let Ok(v) = row.try_get::<_, jiff::civil::Time>(i) {
		v.to_string()
	} else if let Ok(v) = row.try_get::<_, jiff::civil::DateTime>(i) {
		v.to_string()
	} else if let Ok(v) = row.try_get::<_, serde_json::Value>(i) {
		v.to_string()
	} else if let Ok(v) = row.try_get::<_, Vec<String>>(i) {
		format!("{{{}}}", v.join(","))
	} else if let Ok(v) = row.try_get::<_, Vec<i32>>(i) {
		format!(
			"{{{}}}",
			v.iter()
				.map(|x| x.to_string())
				.collect::<Vec<_>>()
				.join(",")
		)
	} else if let Ok(v) = row.try_get::<_, Vec<i64>>(i) {
		format!(
			"{{{}}}",
			v.iter()
				.map(|x| x.to_string())
				.collect::<Vec<_>>()
				.join(",")
		)
	} else if let Ok(v) = row.try_get::<_, Vec<f32>>(i) {
		format!(
			"{{{}}}",
			v.iter()
				.map(|x| x.to_string())
				.collect::<Vec<_>>()
				.join(",")
		)
	} else if let Ok(v) = row.try_get::<_, Vec<f64>>(i) {
		format!(
			"{{{}}}",
			v.iter()
				.map(|x| x.to_string())
				.collect::<Vec<_>>()
				.join(",")
		)
	} else if let Ok(v) = row.try_get::<_, Vec<bool>>(i) {
		format!(
			"{{{}}}",
			v.iter()
				.map(|x| x.to_string())
				.collect::<Vec<_>>()
				.join(",")
		)
	} else {
		// Try to get as string - many types can be retrieved as text
		match row.try_get::<_, String>(i) {
			Ok(v) => v,
			Err(_) => match row.try_get::<_, Option<String>>(i) {
				Ok(None) => "NULL".to_string(),
				Ok(Some(v)) => v,
				Err(_) => "NULL".to_string(),
			},
		}
	}
}

pub fn can_print(row: &tokio_postgres::Row, i: usize) -> bool {
	// Check for void type
	let column = row.columns().get(i);
	if let Some(col) = column
		&& col.type_().name() == "void"
	{
		return true;
	}

	if row.try_get::<_, fraction::Decimal>(i).is_ok()
		|| row.try_get::<_, String>(i).is_ok()
		|| row.try_get::<_, i16>(i).is_ok()
		|| row.try_get::<_, i32>(i).is_ok()
		|| row.try_get::<_, i64>(i).is_ok()
		|| row.try_get::<_, f32>(i).is_ok()
		|| row.try_get::<_, f64>(i).is_ok()
		|| row.try_get::<_, bool>(i).is_ok()
		|| row.try_get::<_, Vec<u8>>(i).is_ok()
		|| row.try_get::<_, jiff::Timestamp>(i).is_ok()
		|| row.try_get::<_, jiff::civil::Date>(i).is_ok()
		|| row.try_get::<_, jiff::civil::Time>(i).is_ok()
		|| row.try_get::<_, jiff::civil::DateTime>(i).is_ok()
		|| row.try_get::<_, serde_json::Value>(i).is_ok()
		|| row.try_get::<_, Vec<String>>(i).is_ok()
		|| row.try_get::<_, Vec<i32>>(i).is_ok()
		|| row.try_get::<_, Vec<i64>>(i).is_ok()
		|| row.try_get::<_, Vec<f32>>(i).is_ok()
		|| row.try_get::<_, Vec<f64>>(i).is_ok()
		|| row.try_get::<_, Vec<bool>>(i).is_ok()
	{
		return true;
	}

	matches!(row.try_get::<_, Option<String>>(i), Ok(None))
}

/// Determine whether the value at the given column is SQL NULL.
///
/// Type-aware: matches on the column type to pick the right `Option<T>` so it works for
/// non-text columns (where `Option<String>` would be a type mismatch rather than NULL).
pub fn is_null(row: &tokio_postgres::Row, i: usize) -> bool {
	use tokio_postgres::types::Type;

	macro_rules! check {
		($t:ty) => {
			matches!(row.try_get::<_, Option<$t>>(i), Ok(None))
		};
	}

	match row.columns().get(i).map(|c| c.type_()) {
		Some(&Type::BOOL) => check!(bool),
		Some(&Type::INT2) => check!(i16),
		Some(&Type::INT4) => check!(i32),
		Some(&Type::INT8) => check!(i64),
		Some(&Type::FLOAT4) => check!(f32),
		Some(&Type::FLOAT8) => check!(f64),
		Some(&Type::NUMERIC) => check!(fraction::Decimal),
		Some(&Type::BYTEA) => check!(Vec<u8>),
		Some(&Type::TIMESTAMP) => check!(jiff::civil::DateTime),
		Some(&Type::TIMESTAMPTZ) => check!(jiff::Timestamp),
		Some(&Type::DATE) => check!(jiff::civil::Date),
		Some(&Type::TIME) => check!(jiff::civil::Time),
		_ => {
			if let Ok(v) = row.try_get::<_, Option<String>>(i) {
				v.is_none()
			} else {
				matches!(row.try_get::<_, Option<Vec<u8>>>(i), Ok(None))
			}
		}
	}
}

/// Render a cell's already-stringified text as a SQL literal suitable for an INSERT.
///
/// Numeric and boolean types are emitted bare; everything else is single-quoted with
/// embedded quotes doubled. NULLs become the `NULL` keyword.
pub fn sql_quote(ty: &tokio_postgres::types::Type, text: &str, is_null: bool) -> String {
	use tokio_postgres::types::Type;

	if is_null {
		return "NULL".to_string();
	}

	match ty {
		&Type::INT2
		| &Type::INT4
		| &Type::INT8
		| &Type::FLOAT4
		| &Type::FLOAT8
		| &Type::NUMERIC
		| &Type::OID => text.to_string(),
		&Type::BOOL => {
			if text == "true" || text == "t" {
				"TRUE".to_string()
			} else {
				"FALSE".to_string()
			}
		}
		_ => format!("'{}'", text.replace('\'', "''")),
	}
}

#[cfg(test)]
mod tests {
	use super::*;

	#[tokio::test]
	async fn test_void_type_handling() {
		let connection_string =
			std::env::var("DATABASE_URL").expect("DATABASE_URL must be set for this test");

		let pool = crate::pool::create_pool(&connection_string, "test")
			.await
			.expect("Failed to create pool");

		let client = pool.get().await.expect("Failed to get connection");

		// Test void type - pg_sleep returns void
		let rows = client
			.query("SELECT pg_sleep(0)", &[])
			.await
			.expect("Query failed");

		assert_eq!(rows.len(), 1);
		let row = &rows[0];

		// Check that void type can be printed
		assert!(can_print(row, 0));

		// Check that void type is formatted as "(void)"
		let value = format_value(row, 0);
		assert_eq!(value, "(void)");
	}

	#[tokio::test]
	async fn test_float_handling() {
		let connection_string =
			std::env::var("DATABASE_URL").expect("DATABASE_URL must be set for this test");

		let pool = crate::pool::create_pool(&connection_string, "test")
			.await
			.expect("Failed to create pool");

		let client = pool.get().await.expect("Failed to get connection");

		// Test float types
		let rows = client
			.query(
				"SELECT 3.14::real as float4, 2.718281828::double precision as float8",
				&[],
			)
			.await
			.expect("Query failed");

		assert_eq!(rows.len(), 1);
		let row = &rows[0];

		// Check that float types can be printed
		assert!(can_print(row, 0));
		assert!(can_print(row, 1));

		// Check that float types are formatted
		let value_f32 = format_value(row, 0);
		let value_f64 = format_value(row, 1);

		assert!(value_f32.contains("3.14"));
		assert!(value_f64.contains("2.718"));
	}

	#[tokio::test]
	async fn test_numeric_handling() {
		let connection_string =
			std::env::var("DATABASE_URL").expect("DATABASE_URL must be set for this test");

		let pool = crate::pool::create_pool(&connection_string, "test")
			.await
			.expect("Failed to create pool");

		let client = pool.get().await.expect("Failed to get connection");

		// Test numeric type
		let rows = client
			.query("SELECT 123.456::numeric as num", &[])
			.await
			.expect("Query failed");

		assert_eq!(rows.len(), 1);
		let row = &rows[0];

		// Numeric type should now be directly printable with fraction crate
		assert!(can_print(row, 0));

		// Check that the value can be formatted
		let value = format_value(row, 0);
		assert!(!value.is_empty());
		assert_ne!(value, "(error)");
		assert!(value.contains("123.456"));
	}

	#[tokio::test]
	async fn test_numeric_arithmetic_with_text_cast() {
		let connection_string =
			std::env::var("DATABASE_URL").expect("DATABASE_URL must be set for this test");

		let pool = crate::pool::create_pool(&connection_string, "test")
			.await
			.expect("Failed to create pool");

		let client = pool.get().await.expect("Failed to get connection");

		// Test numeric arithmetic with explicit text cast
		let rows = client
			.query("SELECT (12.34 + 37.28)::text as result", &[])
			.await
			.expect("Query failed");

		assert_eq!(rows.len(), 1);
		let row = &rows[0];

		// With text cast, it should be printable
		assert!(can_print(row, 0));

		// Should be able to format the result
		let value = format_value(row, 0);
		assert!(!value.is_empty());
		assert_ne!(value, "(error)");
		assert!(value.starts_with("49.6"));
	}

	#[tokio::test]
	async fn test_numeric_arithmetic_direct() {
		let connection_string =
			std::env::var("DATABASE_URL").expect("DATABASE_URL must be set for this test");

		let pool = crate::pool::create_pool(&connection_string, "test")
			.await
			.expect("Failed to create pool");

		let client = pool.get().await.expect("Failed to get connection");

		// Test numeric arithmetic (the original failing case) - now should work directly
		let rows = client
			.query("SELECT 12.34 + 37.28", &[])
			.await
			.expect("Query failed");

		assert_eq!(rows.len(), 1);
		let row = &rows[0];

		// With fraction crate, numeric should be directly printable
		assert!(can_print(row, 0));

		// Should be able to format the result
		let value = format_value(row, 0);
		assert!(!value.is_empty());
		assert_ne!(value, "(error)");
		assert!(value.starts_with("49.6"));
	}

	#[tokio::test]
	async fn test_numeric_arithmetic_question_column() {
		let connection_string =
			std::env::var("DATABASE_URL").expect("DATABASE_URL must be set for this test");

		let pool = crate::pool::create_pool(&connection_string, "test")
			.await
			.expect("Failed to create pool");

		let client = pool.get().await.expect("Failed to get connection");

		// Test numeric arithmetic with ?column? (this was failing in the REPL)
		let rows = client
			.query("SELECT 12.34 + 37.28", &[])
			.await
			.expect("Query failed");

		assert_eq!(rows.len(), 1);
		let row = &rows[0];

		// Verify the column name is ?column?
		assert_eq!(row.columns()[0].name(), "?column?");

		// With fraction crate, numeric should now be directly printable
		assert!(can_print(row, 0));

		// Should be able to format directly without text casting
		let value = format_value(row, 0);
		assert!(!value.is_empty());
		assert_ne!(value, "(error)");
		assert!(value.starts_with("49.6"));
	}
}