hamelin_legacy 0.3.9

use hamelin_lib::{
    antlr::hamelinparser::{UnnestCommandContext, UnnestCommandContextAttrs},
    err::{TranslationError, TranslationErrors},
    sql::{
        expression::{identifier::SimpleIdentifier, literal::BooleanLiteral},
        query::{
            projection::{Binding, ColumnProjection},
            Join, JoinClause, JoinType, SQLQuery, SubQuery, TableAlias, TableExpression,
            TableFunctionApply,
        },
    },
    types::Type,
};

use crate::env::Environment;
use crate::translation::PendingQuery;
use crate::{
    ast::{expression::HamelinExpression, pipeline::HamelinPipeline},
    translation::sql_query_helpers::prepend_projections,
};

pub fn translate(
    ctx: &UnnestCommandContext<'static>,
    pipeline: &HamelinPipeline,
    previous_command: &PendingQuery,
) -> Result<PendingQuery, TranslationErrors> {
    let expression_ctx = TranslationErrors::expect(ctx, ctx.expression())?;
    let expression = HamelinExpression::new(
        expression_ctx.clone(),
        pipeline
            .context
            .default_expression_translation_context(&previous_command.env),
    );
    let expression_translation = expression.translate()?;

    let left_alias = SimpleIdentifier::new("_left");
    let right_alias = SimpleIdentifier::new("_right");

    let left_sql = previous_command.query.clone();
    let left_query = SubQuery::new(left_sql.clone().into()).alias(left_alias.clone().into());

    match expression_translation.typ {
        Type::Array(array) => match array.element_type.as_ref() {
            Type::Struct(srct) => {
                let right_env = Environment::new(srct.clone());
                let merged_env = previous_command
                    .env
                    .clone()
                    .merge_prepend_overwrite(&right_env)
                    .map_err(|e| TranslationError::wrap(ctx, e))?;

                let unnest = TableFunctionApply::new("unnest".to_string())
                    .with_argument(expression_translation.sql.clone())
                    .with_alias(TableAlias::new(right_alias.clone().into()));

                let from: TableExpression = if left_sql == SQLQuery::default() {
                    unnest.into()
                } else {
                    Join::new(left_query.into())
                        .with_clause(JoinClause {
                            table: unnest.into(),
                            join_type: JoinType::LEFT,
                            condition: Some(BooleanLiteral::new(true).into()),
                        })
                        .into()
                };

                let query = SQLQuery::default().from(from).select(
                    merged_env
                        .get_column_projections()
                        .into_iter()
                        .map(|cp| {
                            if right_env.lookup(&cp.identifier).is_ok() {
                                ColumnProjection::new(
                                    cp.identifier
                                        .prefix_from_simples(&[right_alias.clone()])
                                        .into(),
                                )
                                .into()
                            } else {
                                ColumnProjection::new(
                                    cp.identifier
                                        .prefix_from_simples(&[left_alias.clone()])
                                        .into(),
                                )
                                .into()
                            }
                        })
                        .collect(),
                );

                Ok(PendingQuery::new(query, merged_env))
            }
            t => {
                return Err(TranslationError::msg(
                    expression_ctx.as_ref(),
                    &format!(
                    "Unsupported type {} for unnest command. Must be struct or array-of-struct.",
                    t
                ),
                )
                .single())
            }
        },
        Type::Struct(srct) => {
            let right_env = Environment::new(srct.clone());
            let prev_sql = previous_command.query.clone();
            let prev_env = previous_command.env.clone();
            let merged_env = prev_env
                .merge_prepend_overwrite(&right_env)
                .map_err(|e| TranslationError::wrap(ctx, e))?;

            let query = prepend_projections(
                &prev_sql,
                srct.fields
                    .iter()
                    .map(|(k, _)| {
                        Binding::new(
                            k.clone(),
                            expression_translation.sql.clone().dot(k.clone().into()),
                        )
                        .into()
                    })
                    .collect(),
                &merged_env,
            );

            Ok(PendingQuery::new(query, merged_env))
        }
        t => {
            return Err(TranslationError::msg(
                expression_ctx.as_ref(),
                &format!(
                    "Unsupported type {} for unnest command. Must be struct or array-of-struct.",
                    t
                ),
            )
            .single())
        }
    }
}