use std::collections::{HashMap, HashSet};
#[derive(Debug, Clone)]
pub struct SelectSpec {
pub model_name: String,
pub fields: FieldSelection,
pub relations: HashMap<String, SelectSpec>,
}
impl SelectSpec {
pub fn new(model_name: impl Into<String>) -> Self {
Self {
model_name: model_name.into(),
fields: FieldSelection::All,
relations: HashMap::new(),
}
}
pub fn all(model_name: impl Into<String>) -> Self {
Self {
model_name: model_name.into(),
fields: FieldSelection::All,
relations: HashMap::new(),
}
}
pub fn only(
model_name: impl Into<String>,
fields: impl IntoIterator<Item = impl Into<String>>,
) -> Self {
Self {
model_name: model_name.into(),
fields: FieldSelection::Only(fields.into_iter().map(Into::into).collect()),
relations: HashMap::new(),
}
}
pub fn except(
model_name: impl Into<String>,
fields: impl IntoIterator<Item = impl Into<String>>,
) -> Self {
Self {
model_name: model_name.into(),
fields: FieldSelection::Except(fields.into_iter().map(Into::into).collect()),
relations: HashMap::new(),
}
}
pub fn field(mut self, name: impl Into<String>) -> Self {
match &mut self.fields {
FieldSelection::All => {
self.fields = FieldSelection::Only(HashSet::from([name.into()]));
}
FieldSelection::Only(fields) => {
fields.insert(name.into());
}
FieldSelection::Except(fields) => {
fields.remove(&name.into());
}
}
self
}
pub fn fields(mut self, names: impl IntoIterator<Item = impl Into<String>>) -> Self {
for name in names {
self = self.field(name);
}
self
}
pub fn relation(mut self, name: impl Into<String>, select: SelectSpec) -> Self {
self.relations.insert(name.into(), select);
self
}
pub fn is_field_selected(&self, field: &str) -> bool {
self.fields.includes(field)
}
pub fn selected_fields(&self) -> Option<&HashSet<String>> {
match &self.fields {
FieldSelection::Only(fields) => Some(fields),
_ => None,
}
}
pub fn excluded_fields(&self) -> Option<&HashSet<String>> {
match &self.fields {
FieldSelection::Except(fields) => Some(fields),
_ => None,
}
}
pub fn is_all(&self) -> bool {
matches!(self.fields, FieldSelection::All)
}
pub fn to_sql_columns(&self, all_columns: &[&str], table_alias: Option<&str>) -> String {
let columns: Vec<_> = match &self.fields {
FieldSelection::All => all_columns.iter().map(|&s| s.to_string()).collect(),
FieldSelection::Only(fields) => all_columns
.iter()
.filter(|&c| fields.contains(*c))
.map(|&s| s.to_string())
.collect(),
FieldSelection::Except(fields) => all_columns
.iter()
.filter(|&c| !fields.contains(*c))
.map(|&s| s.to_string())
.collect(),
};
match table_alias {
Some(alias) => columns
.into_iter()
.map(|c| format!("{}.{}", alias, c))
.collect::<Vec<_>>()
.join(", "),
None => columns.join(", "),
}
}
}
#[derive(Debug, Clone, Default)]
pub enum FieldSelection {
#[default]
All,
Only(HashSet<String>),
Except(HashSet<String>),
}
impl FieldSelection {
pub fn includes(&self, field: &str) -> bool {
match self {
Self::All => true,
Self::Only(fields) => fields.contains(field),
Self::Except(fields) => !fields.contains(field),
}
}
pub fn is_all(&self) -> bool {
matches!(self, Self::All)
}
}
pub fn select(model: impl Into<String>) -> SelectSpec {
SelectSpec::new(model)
}
pub fn select_only(
model: impl Into<String>,
fields: impl IntoIterator<Item = impl Into<String>>,
) -> SelectSpec {
SelectSpec::only(model, fields)
}
pub fn select_except(
model: impl Into<String>,
fields: impl IntoIterator<Item = impl Into<String>>,
) -> SelectSpec {
SelectSpec::except(model, fields)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_select_spec_all() {
let spec = SelectSpec::all("User");
assert!(spec.is_all());
assert!(spec.is_field_selected("id"));
assert!(spec.is_field_selected("email"));
}
#[test]
fn test_select_spec_only() {
let spec = SelectSpec::only("User", ["id", "email"]);
assert!(!spec.is_all());
assert!(spec.is_field_selected("id"));
assert!(spec.is_field_selected("email"));
assert!(!spec.is_field_selected("password"));
}
#[test]
fn test_select_spec_except() {
let spec = SelectSpec::except("User", ["password"]);
assert!(!spec.is_all());
assert!(spec.is_field_selected("id"));
assert!(!spec.is_field_selected("password"));
}
#[test]
fn test_select_spec_with_relation() {
let spec = SelectSpec::only("User", ["id", "name"])
.relation("posts", SelectSpec::only("Post", ["id", "title"]));
assert!(spec.relations.contains_key("posts"));
}
#[test]
fn test_to_sql_columns() {
let spec = SelectSpec::only("User", ["id", "email"]);
let columns = spec.to_sql_columns(&["id", "email", "name", "password"], None);
assert!(columns.contains("id"));
assert!(columns.contains("email"));
assert!(!columns.contains("password"));
}
#[test]
fn test_to_sql_columns_with_alias() {
let spec = SelectSpec::only("User", ["id", "email"]);
let columns = spec.to_sql_columns(&["id", "email"], Some("u"));
assert!(columns.contains("u.id"));
assert!(columns.contains("u.email"));
}
}