hamelin_legacy 0.3.8

//! Extension methods for ProjectionBuilder that depend on hamelin_legacy types.
//!
//! The core ProjectionBuilder lives in hamelin_lib::sql::projection_builder.
//! These extensions add methods that depend on HamelinExpression and Environment.

use anyhow::{anyhow, bail};

use hamelin_lib::sql::expression::identifier::{CompoundIdentifier, Identifier};
use hamelin_lib::sql::expression::literal::ColumnReference;
use hamelin_lib::types::struct_type::Struct;
use hamelin_lib::types::Type;

// Re-export ProjectionBuilder for backwards compatibility
pub use hamelin_lib::sql::projection_builder::ProjectionBuilder;

use crate::ast::expression::HamelinExpression;
use crate::env::Environment;

/// Extension trait for ProjectionBuilder with legacy-specific methods.
pub trait ProjectionBuilderExt {
    /// Construct a new projection builder with items initialized to reference the given environment.
    /// The returned ProjectionBuilder would be useful to deeply select all columns out of the given
    /// environment. The concept "deep," here, means that it refers to nested columns explicitly with
    /// bindings in the projection builder. That way if you want to override just one of them, you can.
    fn deep_initialize_from_environment(env: &Environment) -> ProjectionBuilder;

    /// Construct a new projection builder with items initialized to reference the given environment.
    /// The returned ProjectionBuilder would be useful to just select all the first-level fields out
    /// of the given environment. This is useful for "nesting" operations, where you know you will
    /// never need to reach inside and mess with anything.
    fn shallow_initialize_from_environment(env: &Environment) -> ProjectionBuilder;

    /// Construct a projection builder from a struct expression. This will make
    /// the expression for every leaf in the struct expression a dereference
    /// of the original struct expression, no matter how nested the type of
    /// the struct expression is.
    fn from_struct_expression(expression: HamelinExpression) -> anyhow::Result<ProjectionBuilder>;

    /// Bind a new column value in the list of projections. Initialize structs to refer to existing columns.
    fn intialize_struct_reference_then_bind(
        &mut self,
        id: Identifier,
        value: hamelin_lib::sql::expression::SQLExpression,
        type_: Type,
        environment: &Environment,
    );

    /// Build an Environment from this ProjectionBuilder.
    fn build_environment(self) -> Environment;
}

impl ProjectionBuilderExt for ProjectionBuilder {
    fn deep_initialize_from_environment(env: &Environment) -> ProjectionBuilder {
        ProjectionBuilder::deep_initialize_from_struct(None, env.fields.clone())
    }

    fn shallow_initialize_from_environment(env: &Environment) -> ProjectionBuilder {
        env.fields
            .fields
            .iter()
            .fold(ProjectionBuilder::default(), |pb, (id, typ)| {
                pb.with_binding(
                    id.clone().into(),
                    ColumnReference::new(id.clone().into()).into(),
                    typ.clone(),
                )
            })
    }

    fn from_struct_expression(expression: HamelinExpression) -> anyhow::Result<ProjectionBuilder> {
        from_struct_expression_helper(None, expression)
    }

    fn intialize_struct_reference_then_bind(
        &mut self,
        id: Identifier,
        value: hamelin_lib::sql::expression::SQLExpression,
        type_: Type,
        environment: &Environment,
    ) {
        if let Ok(binding) = environment.lookup(&id.first().clone().into()) {
            if let Type::Struct(struct_type) = binding {
                if !self.is_present(&id.first().clone().into()) {
                    self.initialize_key(id.first().clone().into(), struct_type);
                }
            }
        }

        self.bind(id, value, type_);
    }

    fn build_environment(self) -> Environment {
        Environment::new(self.build_hamelin_type())
    }
}

fn from_struct_expression_helper(
    name: Option<Identifier>,
    expression: HamelinExpression,
) -> anyhow::Result<ProjectionBuilder> {
    let result = dereference_struct_field(name.clone(), expression.clone())?;
    let mut builder = ProjectionBuilder::default();

    for (id, typ) in result.fields.into_iter() {
        let root_ident: Identifier = name
            .clone()
            .map(|n| CompoundIdentifier::from_idents(&[n, id.clone().into()]).into())
            .unwrap_or(id.clone().into());
        if let Type::Struct(_) = typ {
            builder.bind(
                id.clone().into(),
                expression.translate()?.sql.dot(root_ident.clone()),
                typ.clone(),
            );
            // Recurse for nested struct
            let nested = from_struct_expression_helper(Some(root_ident), expression.clone())?;
            for (nested_id, nested_typ) in nested.build_hamelin_type().fields.into_iter() {
                let compound =
                    CompoundIdentifier::from_idents(&[id.clone().into(), nested_id.into()]);
                if let Type::Struct(_) = nested_typ {
                    // Already handled by recursion
                } else {
                    builder.bind(
                        compound.into(),
                        expression.translate()?.sql.dot(
                            name.clone()
                                .map(|n| {
                                    CompoundIdentifier::from_idents(&[n, id.clone().into()]).into()
                                })
                                .unwrap_or(id.clone().into()),
                        ),
                        nested_typ,
                    );
                }
            }
        } else {
            builder.bind(id.into(), expression.translate()?.sql.dot(root_ident), typ);
        }
    }

    Ok(builder)
}

/// Dereference a struct field from a struct expression by traversing the struct type of the
/// expression using the identifier list.
fn dereference_struct_field(
    id: Option<Identifier>,
    expression: HamelinExpression,
) -> anyhow::Result<Struct> {
    let initial_struct = expression
        .translate()
        .map_err(|e| anyhow!(e))
        .and_then(|t| match t.typ {
            Type::Struct(s) => Ok(s),
            _ => bail!("The expression {} is not a struct", expression.text()),
        });

    id.map(|i| i.simples())
        .unwrap_or(vec![])
        .iter()
        .fold(initial_struct, |acc, ident| match acc {
            Ok(Struct { fields }) => fields
                .get(ident)
                .cloned()
                .ok_or(anyhow!("field {} not found in {:#?}", ident, fields))
                .and_then(|t| match t {
                    Type::Struct(s) => Ok(s),
                    _ => Err(anyhow!("field {} in {:#?} is not a struct", ident, fields)),
                }),
            Err(e) => Err(e),
        })
}