use std::collections::HashSet;
use crate::translation::projection_builder::ProjectionBuilder;
use crate::translation::PendingQuery;
use hamelin_lib::antlr::hamelinparser::{
DropCommandContext, DropCommandContextAttrs, SelectionContextAttrs,
};
use hamelin_lib::err::{TranslationError, TranslationErrors};
use hamelin_lib::sql::expression::identifier::HamelinIdentifier;
use hamelin_lib::sql::expression::identifier::{Identifier, SimpleIdentifier};
use hamelin_lib::sql::expression::literal::ColumnReference;
pub fn translate(
ctx: &DropCommandContext<'static>,
previous: &PendingQuery,
) -> Result<PendingQuery, TranslationErrors> {
let columns_to_drop = ctx
.selection_all()
.into_iter()
.flat_map(|sctx| sctx.identifier())
.map(|i| HamelinIdentifier::new(i))
.collect::<Vec<_>>();
previous.env.check_column_references(&columns_to_drop[..])?;
let rendered_columns_to_drop = columns_to_drop
.into_iter()
.map(|c| c.to_sql().and_then(|c| Ok(ColumnReference::new(c))))
.collect::<Result<Vec<_>, _>>()?;
let env = previous
.env
.clone()
.drop(rendered_columns_to_drop.clone())
.map_err(|e| TranslationError::wrap_box(ctx, e.into()).single())?;
let simple_columns_to_drop = rendered_columns_to_drop
.iter()
.flat_map(|cr| match cr.identifier.clone() {
Identifier::Simple(s) => Some(s),
_ => None,
})
.collect::<Vec<_>>();
let complex_columns_to_drop = rendered_columns_to_drop
.iter()
.filter(|cr| match cr.identifier {
Identifier::Simple(_) => false,
Identifier::Compound(_) => true,
})
.collect::<Vec<_>>();
if !complex_columns_to_drop.is_empty() {
let modified = complex_columns_to_drop
.iter()
.map(|cr| cr.identifier.first())
.collect::<HashSet<_>>();
let mut projection_builder = ProjectionBuilder::default();
for (ident, typ) in env.fields.fields.iter() {
if modified.contains(ident) {
projection_builder.initialize_key(
ident.clone(),
typ.clone()
.try_unwrap_struct()
.expect("ported a class cast"),
);
} else {
projection_builder.bind_column_reference(ident.clone(), typ.clone());
}
}
Ok(PendingQuery::new(
previous.query.replace_projections(
projection_builder
.build_projections()
.map_err(|e| TranslationError::wrap_box(ctx, e.into()).single())?,
),
env,
))
} else {
let idents: Vec<SimpleIdentifier> = simple_columns_to_drop
.into_iter()
.map(|s| s.into())
.collect::<Vec<_>>();
Ok(PendingQuery::new(
previous.query.remove_simple_projections(&idents[..]),
env,
))
}
}