use scythe_backend::manifest::BackendManifest;
use scythe_backend::naming::{
enum_type_name, enum_variant_name, fn_name, row_struct_name, to_pascal_case, to_snake_case,
};
use scythe_backend::types::resolve_type;
use std::fmt::Write;
use scythe_core::analyzer::{AnalyzedQuery, CompositeInfo, EnumInfo};
use scythe_core::errors::{ErrorCode, ScytheError};
use scythe_core::parser::QueryCommand;
use crate::backend_trait::{CodegenBackend, ResolvedColumn, ResolvedParam};
use crate::singularize;
const DEFAULT_MANIFEST_TOML: &str = include_str!("../../manifests/rust-tiberius.toml");
pub struct RustTiberiusBackend {
manifest: BackendManifest,
}
impl RustTiberiusBackend {
pub fn new(engine: &str) -> Result<Self, ScytheError> {
match engine {
"mssql" => {}
_ => {
return Err(ScytheError::new(
ErrorCode::InternalError,
format!("rust-tiberius only supports MSSQL, got engine '{}'", engine),
));
}
}
let manifest = super::parse_manifest(DEFAULT_MANIFEST_TOML)?;
Ok(Self { manifest })
}
fn tiberius_field_init(col: &ResolvedColumn) -> String {
if col.nullable {
if col.neutral_type == "string" {
format!(
"{}: row.try_get::<&str, _>(\"{}\")?.map(|s| s.to_string())",
col.field_name, col.name
)
} else {
format!("{}: row.try_get(\"{}\")?", col.field_name, col.name)
}
} else if col.neutral_type == "string" {
format!(
"{}: row.try_get::<&str, _>(\"{name}\")?.ok_or_else(|| tiberius::error::Error::Protocol(\"unexpected NULL for non-nullable column '{name}'\".into()))?.to_string()",
col.field_name,
name = col.name
)
} else {
format!(
"{}: row.try_get(\"{name}\")?.ok_or_else(|| tiberius::error::Error::Protocol(\"unexpected NULL for non-nullable column '{name}'\".into()))?",
col.field_name,
name = col.name
)
}
}
fn tiberius_decode_col(col: &ResolvedColumn) -> String {
if col.nullable {
if col.neutral_type == "string" {
format!("row.try_get::<&str, _>(\"{}\")?.map(|s| s.to_string())", col.name)
} else {
format!("row.try_get(\"{}\")?", col.name)
}
} else if col.neutral_type == "string" {
format!(
"row.try_get::<&str, _>(\"{name}\")?.ok_or_else(|| tiberius::error::Error::Protocol(\"unexpected NULL for non-nullable column '{name}'\".into()))?.to_string()",
name = col.name
)
} else {
format!(
"row.try_get(\"{name}\")?.ok_or_else(|| tiberius::error::Error::Protocol(\"unexpected NULL for non-nullable column '{name}'\".into()))?",
name = col.name
)
}
}
}
impl CodegenBackend for RustTiberiusBackend {
fn name(&self) -> &str {
"rust-tiberius"
}
fn manifest(&self) -> &scythe_backend::manifest::BackendManifest {
&self.manifest
}
fn supported_engines(&self) -> &[&str] {
&["mssql"]
}
fn file_header(&self) -> String {
"// Auto-generated by scythe. Do not edit.\n#![allow(dead_code, unused_imports, clippy::all)]".to_string()
}
fn generate_row_struct(&self, query_name: &str, columns: &[ResolvedColumn]) -> Result<String, ScytheError> {
let struct_name = row_struct_name(query_name, &self.manifest.naming);
let mut out = String::new();
let _ = writeln!(out, "#[derive(Debug, Clone)]");
let _ = writeln!(out, "pub struct {} {{", struct_name);
for col in columns {
let _ = writeln!(out, " pub {}: {},", col.field_name, col.full_type);
}
let _ = writeln!(out, "}}");
let _ = writeln!(out);
let _ = writeln!(out, "impl {} {{", struct_name);
let _ = writeln!(
out,
" pub fn from_row(row: &tiberius::Row) -> Result<Self, tiberius::error::Error> {{"
);
let _ = writeln!(out, " Ok(Self {{");
for col in columns {
if col.nullable {
if col.neutral_type == "string" {
let _ = writeln!(
out,
" {}: row.try_get::<&str, _>(\"{}\")?.map(|s| s.to_string()),",
col.field_name, col.name
);
} else {
let _ = writeln!(out, " {}: row.try_get(\"{}\")?,", col.field_name, col.name);
}
} else {
if col.neutral_type == "string" {
let _ = writeln!(
out,
" {}: row.try_get::<&str, _>(\"{}\")?.ok_or_else(|| tiberius::error::Error::Protocol(\"unexpected NULL for non-nullable column '{}'\".into()))?.to_string(),",
col.field_name, col.name, col.name
);
} else {
let _ = writeln!(
out,
" {}: row.try_get(\"{}\")?.ok_or_else(|| tiberius::error::Error::Protocol(\"unexpected NULL for non-nullable column '{}'\".into()))?,",
col.field_name, col.name, col.name
);
}
}
}
let _ = writeln!(out, " }})");
let _ = writeln!(out, " }}");
let _ = write!(out, "}}");
Ok(out)
}
fn generate_model_struct(&self, table_name: &str, columns: &[ResolvedColumn]) -> Result<String, ScytheError> {
let singular = singularize(table_name);
let name = to_pascal_case(&singular).into_owned();
self.generate_row_struct(&name, columns)
}
fn generate_query_fn(
&self,
analyzed: &AnalyzedQuery,
struct_name: &str,
_columns: &[ResolvedColumn],
params: &[ResolvedParam],
) -> Result<String, ScytheError> {
let func_name = fn_name(&analyzed.name, &self.manifest.naming);
let mut out = String::new();
if let Some(ref msg) = analyzed.deprecated {
let _ = writeln!(out, "#[deprecated(note = \"{}\")]", msg);
}
let sql = super::rewrite_pg_placeholders(
&super::clean_sql_with_optional(&analyzed.sql, &analyzed.optional_params, &analyzed.params),
|n| format!("@p{n}"),
);
let mut param_parts: Vec<String> =
vec!["client: &mut tiberius::Client<tokio_util::compat::Compat<tokio::net::TcpStream>>".to_string()];
for param in params {
param_parts.push(format!("{}: {}", param.field_name, param.borrowed_type));
}
if matches!(analyzed.command, QueryCommand::Batch) {
let batch_fn_name = format!("{}_batch", func_name);
if params.len() > 1 {
let params_struct_name = format!("{}BatchParams", struct_name);
let _ = writeln!(out, "#[derive(Debug, Clone)]");
let _ = writeln!(out, "pub struct {} {{", params_struct_name);
for param in params {
let _ = writeln!(out, " pub {}: {},", param.field_name, param.full_type);
}
let _ = writeln!(out, "}}");
let _ = writeln!(out);
let _ = writeln!(
out,
"pub async fn {}(client: &mut tiberius::Client<tokio_util::compat::Compat<tokio::net::TcpStream>>, items: &[{}]) -> Result<(), tiberius::error::Error> {{",
batch_fn_name, params_struct_name
);
let _ = writeln!(out, " for item in items {{");
let bind_args: Vec<String> = params
.iter()
.map(|p| format!("&item.{} as &dyn tiberius::ToSql", p.field_name))
.collect();
let _ = writeln!(
out,
" client.execute(r#\"{}\"#, &[{}]).await?;",
sql,
bind_args.join(", ")
);
let _ = writeln!(out, " }}");
let _ = writeln!(out, " Ok(())");
} else if params.len() == 1 {
let _ = writeln!(
out,
"pub async fn {}(client: &mut tiberius::Client<tokio_util::compat::Compat<tokio::net::TcpStream>>, items: &[{}]) -> Result<(), tiberius::error::Error> {{",
batch_fn_name, params[0].full_type
);
let _ = writeln!(out, " for item in items {{");
let _ = writeln!(out, " client.execute(r#\"{}\"#, &[item]).await?;", sql);
let _ = writeln!(out, " }}");
let _ = writeln!(out, " Ok(())");
} else {
let _ = writeln!(
out,
"pub async fn {}(client: &mut tiberius::Client<tokio_util::compat::Compat<tokio::net::TcpStream>>, count: usize) -> Result<(), tiberius::error::Error> {{",
batch_fn_name
);
let _ = writeln!(out, " for _ in 0..count {{");
let _ = writeln!(out, " client.execute(r#\"{}\"#, &[]).await?;", sql);
let _ = writeln!(out, " }}");
let _ = writeln!(out, " Ok(())");
}
let _ = write!(out, "}}");
return Ok(out);
}
let return_type = match &analyzed.command {
QueryCommand::One | QueryCommand::Opt => struct_name.to_string(),
QueryCommand::Many => format!("Vec<{}>", struct_name),
QueryCommand::Exec => "()".to_string(),
QueryCommand::ExecResult | QueryCommand::ExecRows => "u64".to_string(),
QueryCommand::Batch => unreachable!(),
QueryCommand::Grouped => {
unreachable!("Grouped command is routed through generate_grouped_query_fn, not generate_query_fn")
}
};
let _ = writeln!(
out,
"pub async fn {}({}) -> Result<{}, tiberius::error::Error> {{",
func_name,
param_parts.join(", "),
return_type
);
let param_slice: String = if params.is_empty() {
"&[]".to_string()
} else {
let param_refs: Vec<String> = params
.iter()
.map(|p| {
let bt = &p.borrowed_type;
if bt == "&str" || bt == "Option<&str>" {
if bt == "Option<&str>" {
format!("&{}.map(|s| s.to_string()) as &dyn tiberius::ToSql", p.field_name)
} else {
format!("&{}.to_string() as &dyn tiberius::ToSql", p.field_name)
}
} else if bt.starts_with('&') {
format!("{} as &dyn tiberius::ToSql", p.field_name)
} else {
format!("&{} as &dyn tiberius::ToSql", p.field_name)
}
})
.collect();
format!("&[{}]", param_refs.join(", "))
};
match &analyzed.command {
QueryCommand::One | QueryCommand::Opt => {
let _ = writeln!(
out,
" let stream = client.query(r#\"{}\"#, {}).await?;",
sql, param_slice
);
let _ = writeln!(
out,
" let row = stream.into_row().await?.expect(\"expected one row\");"
);
let _ = writeln!(out, " Ok({}::from_row(&row)?)", struct_name);
}
QueryCommand::Many => {
let _ = writeln!(
out,
" let stream = client.query(r#\"{}\"#, {}).await?;",
sql, param_slice
);
let _ = writeln!(out, " let rows = stream.into_first_result().await?;");
let _ = writeln!(out, " rows.iter().map({}::from_row).collect()", struct_name);
}
QueryCommand::Exec => {
let _ = writeln!(out, " client.execute(r#\"{}\"#, {}).await?;", sql, param_slice);
let _ = writeln!(out, " Ok(())");
}
QueryCommand::ExecResult | QueryCommand::ExecRows => {
let _ = writeln!(
out,
" let result = client.execute(r#\"{}\"#, {}).await?;",
sql, param_slice
);
let _ = writeln!(out, " Ok(result.total())");
}
QueryCommand::Batch => unreachable!(),
QueryCommand::Grouped => {
unreachable!("Grouped command is routed through generate_grouped_query_fn, not generate_query_fn")
}
}
let _ = write!(out, "}}");
Ok(out)
}
fn generate_enum_def(&self, enum_info: &EnumInfo) -> Result<String, ScytheError> {
let type_name = enum_type_name(&enum_info.sql_name, &self.manifest.naming);
let mut out = String::with_capacity(512);
let _ = writeln!(out, "#[derive(Debug, Clone, PartialEq, Eq)]");
let _ = writeln!(out, "pub enum {} {{", type_name);
for value in &enum_info.values {
let variant = enum_variant_name(value, &self.manifest.naming);
let _ = writeln!(out, " {},", variant);
}
let _ = writeln!(out, "}}");
let _ = writeln!(out);
let _ = writeln!(out, "impl std::fmt::Display for {} {{", type_name);
let _ = writeln!(
out,
" fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {{"
);
let _ = writeln!(out, " match self {{");
for value in &enum_info.values {
let variant = enum_variant_name(value, &self.manifest.naming);
let _ = writeln!(
out,
" {}::{} => write!(f, \"{}\"),",
type_name, variant, value
);
}
let _ = writeln!(out, " }}");
let _ = writeln!(out, " }}");
let _ = writeln!(out, "}}");
let _ = writeln!(out);
let _ = writeln!(out, "impl std::str::FromStr for {} {{", type_name);
let _ = writeln!(out, " type Err = String;");
let _ = writeln!(out, " fn from_str(s: &str) -> Result<Self, Self::Err> {{");
let _ = writeln!(out, " match s {{");
for value in &enum_info.values {
let variant = enum_variant_name(value, &self.manifest.naming);
let _ = writeln!(out, " \"{}\" => Ok({}::{}),", value, type_name, variant);
}
let _ = writeln!(out, " _ => Err(format!(\"unknown variant: {{}}\", s)),");
let _ = writeln!(out, " }}");
let _ = writeln!(out, " }}");
let _ = write!(out, "}}");
Ok(out)
}
fn generate_grouped_structs(
&self,
parent_struct_name: &str,
child_struct_name: &str,
parent_columns: &[crate::backend_trait::ResolvedColumn],
child_columns: &[crate::backend_trait::ResolvedColumn],
_key_column: &str,
) -> Result<String, ScytheError> {
let mut out = String::new();
let _ = writeln!(out, "#[derive(Debug, Clone)]");
let _ = writeln!(out, "pub struct {} {{", child_struct_name);
for col in child_columns {
let _ = writeln!(out, " pub {}: {},", col.field_name, col.full_type);
}
let _ = writeln!(out, "}}");
let _ = writeln!(out);
let _ = writeln!(out, "impl {} {{", child_struct_name);
let _ = writeln!(
out,
" pub fn from_row(row: &tiberius::Row) -> Result<Self, tiberius::error::Error> {{"
);
let _ = writeln!(out, " Ok(Self {{");
for col in child_columns {
let _ = writeln!(out, " {},", Self::tiberius_field_init(col));
}
let _ = writeln!(out, " }})");
let _ = writeln!(out, " }}");
let _ = writeln!(out, "}}");
let _ = writeln!(out);
let _ = writeln!(out, "#[derive(Debug, Clone)]");
let _ = writeln!(out, "pub struct {} {{", parent_struct_name);
for col in parent_columns {
let _ = writeln!(out, " pub {}: {},", col.field_name, col.full_type);
}
let _ = writeln!(out, " pub children: Vec<{}>,", child_struct_name);
let _ = write!(out, "}}");
Ok(out)
}
fn generate_grouped_query_fn(
&self,
request: &crate::backend_trait::GroupedQueryFn<'_>,
) -> Result<String, ScytheError> {
let analyzed = request.analyzed;
let parent_struct_name = request.parent_struct_name;
let child_struct_name = request.child_struct_name;
let parent_columns = request.parent_columns;
let params = request.params;
let key_column = request.key_column;
let func_name = fn_name(&analyzed.name, &self.manifest.naming);
let key_field = to_snake_case(key_column);
let mut out = String::new();
if let Some(ref msg) = analyzed.deprecated {
let _ = writeln!(out, "#[deprecated(note = \"{msg}\")]");
}
let mut param_parts: Vec<String> =
vec!["client: &mut tiberius::Client<tokio_util::compat::Compat<tokio::net::TcpStream>>".to_string()];
for param in params {
param_parts.push(format!("{}: {}", param.field_name, param.borrowed_type));
}
let sql = super::rewrite_pg_placeholders(
&super::clean_sql_with_optional(&analyzed.sql, &analyzed.optional_params, &analyzed.params),
|n| format!("@p{n}"),
);
let param_slice: String = if params.is_empty() {
"&[]".to_string()
} else {
let refs: Vec<String> = params
.iter()
.map(|p| {
let bt = &p.borrowed_type;
if bt == "&str" || bt == "Option<&str>" {
if bt == "Option<&str>" {
format!("&{}.map(|s| s.to_string()) as &dyn tiberius::ToSql", p.field_name)
} else {
format!("&{}.to_string() as &dyn tiberius::ToSql", p.field_name)
}
} else if bt.starts_with('&') {
format!("{} as &dyn tiberius::ToSql", p.field_name)
} else {
format!("&{} as &dyn tiberius::ToSql", p.field_name)
}
})
.collect();
format!("&[{}]", refs.join(", "))
};
let key_col_type = parent_columns
.iter()
.find(|c| c.name == key_column || c.field_name == key_field)
.map(|c| c.full_type.as_str())
.unwrap_or("i64");
let _ = writeln!(
out,
"pub async fn {}({}) -> Result<Vec<{}>, tiberius::error::Error> {{",
func_name,
param_parts.join(", "),
parent_struct_name
);
let _ = writeln!(
out,
" let stream = client.query(r#\"{}\"#, {}).await?;",
sql, param_slice
);
let _ = writeln!(out, " let rows = stream.into_first_result().await?;");
let _ = writeln!(out, " let mut result: Vec<{}> = Vec::new();", parent_struct_name);
let _ = writeln!(out, " for row in &rows {{");
for col in parent_columns {
let decode = Self::tiberius_decode_col(col);
let _ = writeln!(out, " let {}: {} = {decode};", col.field_name, col.full_type);
}
let _ = writeln!(out, " let key: {key_col_type} = {key_field}.clone();");
let _ = writeln!(out, " let child = {}::from_row(row)?;", child_struct_name);
let _ = writeln!(
out,
" if let Some(parent) = result.iter_mut().rev().find(|p| p.{key_field} == key) {{"
);
let _ = writeln!(out, " parent.children.push(child);");
let _ = writeln!(out, " }} else {{");
let _ = writeln!(out, " result.push({} {{", parent_struct_name);
for col in parent_columns {
let _ = writeln!(out, " {},", col.field_name);
}
let _ = writeln!(out, " children: vec![child],");
let _ = writeln!(out, " }});");
let _ = writeln!(out, " }}");
let _ = writeln!(out, " }}");
let _ = writeln!(out, " Ok(result)");
let _ = write!(out, "}}");
Ok(out)
}
fn generate_composite_def(&self, composite: &CompositeInfo) -> Result<String, ScytheError> {
let struct_name = to_pascal_case(&composite.sql_name).into_owned();
let mut out = String::new();
let _ = writeln!(out, "#[derive(Debug, Clone)]");
let _ = writeln!(out, "pub struct {} {{", struct_name);
for field in &composite.fields {
let rust_type = resolve_type(&field.neutral_type, &self.manifest, false)
.map(|t| t.into_owned())
.map_err(|e| {
ScytheError::new(ErrorCode::InternalError, format!("composite field type error: {}", e))
})?;
let _ = writeln!(out, " pub {}: {},", to_snake_case(&field.name), rust_type);
}
let _ = write!(out, "}}");
Ok(out)
}
}
#[cfg(test)]
mod tests {
use scythe_core::analyzer::{AnalyzedColumn, AnalyzedQuery, GroupByConfig};
use scythe_core::parser::QueryCommand;
use super::RustTiberiusBackend;
use crate::generate_with_backend;
fn make_grouped_query() -> AnalyzedQuery {
let parent_cols = vec![
AnalyzedColumn {
name: "id".to_string(),
neutral_type: "int32".to_string(),
nullable: false,
..Default::default()
},
AnalyzedColumn {
name: "name".to_string(),
neutral_type: "string".to_string(),
nullable: false,
..Default::default()
},
AnalyzedColumn {
name: "email".to_string(),
neutral_type: "string".to_string(),
nullable: false,
..Default::default()
},
];
let child_cols = vec![
AnalyzedColumn {
name: "order_id".to_string(),
neutral_type: "int32".to_string(),
nullable: false,
..Default::default()
},
AnalyzedColumn {
name: "total".to_string(),
neutral_type: "decimal".to_string(),
nullable: true,
..Default::default()
},
AnalyzedColumn {
name: "order_date".to_string(),
neutral_type: "datetime".to_string(),
nullable: false,
..Default::default()
},
];
let all_cols = [parent_cols.clone(), child_cols.clone()].concat();
AnalyzedQuery {
name: "GetUsersWithOrders".to_string(),
command: QueryCommand::Grouped,
sql: "-- @name GetUsersWithOrders\n-- @returns :grouped\n-- @group_by users.id\n\
SELECT u.id, u.name, u.email, o.id AS order_id, o.total, o.created_at AS order_date\n\
FROM users u\n\
JOIN orders o ON o.user_id = u.id"
.to_string(),
columns: all_cols,
params: vec![],
deprecated: None,
source_table: None,
composites: vec![],
enums: vec![],
optional_params: vec![],
group_by: Some(GroupByConfig {
table: "users".to_string(),
key_column: "id".to_string(),
parent_columns: parent_cols,
child_columns: child_cols,
}),
custom: vec![],
}
}
#[test]
fn test_grouped_tiberius_structs() {
let backend = RustTiberiusBackend::new("mssql").unwrap();
let query = make_grouped_query();
let result = generate_with_backend(&query, &backend).unwrap();
let row_struct = result.row_struct.as_deref().unwrap();
assert!(
row_struct.contains("pub struct GetUsersWithOrdersChildRow"),
"missing child struct; got:\n{row_struct}"
);
assert!(
row_struct.contains("pub order_id: i32"),
"child struct missing order_id; got:\n{row_struct}"
);
assert!(
row_struct.contains("pub struct GetUsersWithOrdersRow"),
"missing parent struct; got:\n{row_struct}"
);
assert!(
row_struct.contains("pub id: i32"),
"parent struct missing id; got:\n{row_struct}"
);
assert!(
row_struct.contains("pub name: String"),
"parent struct missing name; got:\n{row_struct}"
);
assert!(
row_struct.contains("pub children: Vec<GetUsersWithOrdersChildRow>"),
"parent struct missing children field; got:\n{row_struct}"
);
let child_pos = row_struct.find("GetUsersWithOrdersChildRow").unwrap();
let parent_pos = row_struct.find("pub struct GetUsersWithOrdersRow").unwrap();
assert!(child_pos < parent_pos, "child struct must appear before parent struct");
assert!(
row_struct.contains("fn from_row(row: &tiberius::Row)"),
"child struct should include from_row; got:\n{row_struct}"
);
assert!(result.model_struct.is_none(), "grouped must not produce a model_struct");
}
#[test]
fn test_grouped_tiberius_query_fn() {
let backend = RustTiberiusBackend::new("mssql").unwrap();
let query = make_grouped_query();
let result = generate_with_backend(&query, &backend).unwrap();
let query_fn = result.query_fn.as_deref().unwrap();
assert!(
query_fn.contains("pub async fn get_users_with_orders("),
"missing fn; got:\n{query_fn}"
);
assert!(
query_fn.contains("-> Result<Vec<GetUsersWithOrdersRow>, tiberius::error::Error>"),
"wrong return type; got:\n{query_fn}"
);
assert!(
query_fn.contains("into_first_result()"),
"fn must use into_first_result; got:\n{query_fn}"
);
assert!(
query_fn.contains("GetUsersWithOrdersChildRow::from_row(row)?"),
"fn must call child from_row; got:\n{query_fn}"
);
assert!(
query_fn.contains("children: vec![child]"),
"fn must initialize children vec; got:\n{query_fn}"
);
assert!(
query_fn.contains("parent.children.push(child)"),
"fn must fold child into existing parent; got:\n{query_fn}"
);
assert!(
query_fn.contains("Ok(result)"),
"fn must return result; got:\n{query_fn}"
);
}
}