Function try_extract_sql_columns

prqlc/prqlc/src/semantic/lowering.rs:1045–1133

Attempts to extract column names from an S-String to avoid wildcards when possible

(
    columns: Vec<RelationColumn>,
    items: &[InterpolateItem<rq::Expr>],
)

/// Attempts to extract column names from an S-String to avoid wildcards when possible
fn try_extract_sql_columns(
    columns: Vec<RelationColumn>,
    items: &[InterpolateItem<rq::Expr>],
) -> Vec<RelationColumn> {
    use sqlparser::ast;

    let mut has_wildcard = false;

    let sql_columns = items
        .iter()
        .map(|item| match item {
            InterpolateItem::String(s) => {
                let sql_ast =
                    sqlparser::parser::Parser::parse_sql(&sqlparser::dialect::GenericDialect {}, s)
                        .map_err(|err| format!("could not parse {item:?}: {err:?}"))?;
                if sql_ast.len() != 1 {
                    return Err(format!(
                        "expected exactly one statement, got {}",
                        sql_ast.len()
                    ));
                }

                let statement = sql_ast.into_iter().next().unwrap();

                if let sqlparser::ast::Statement::Query(query) = statement {
                    if let sqlparser::ast::SetExpr::Select(select_stmt) = *query.body {
                        select_stmt
                            .projection
                            .into_iter()
                            .map(|expr| match expr {
                                ast::SelectItem::UnnamedExpr(expr) => {
                                    if let ast::Expr::Identifier(ast::Ident { value, .. }) = expr {
                                        Ok(value)
                                    } else {
                                        Err(format!("Only Idents are supported, got {expr:?}"))
                                    }
                                }
                                ast::SelectItem::ExprWithAlias { alias, .. } => Ok(alias.value), // Store alias
                                ast::SelectItem::ExprWithAliases { .. } => Err(
                                    "Multi-alias projection (`expr AS (a, b)`) is not supported"
                                        .into(),
                                ),
                                ast::SelectItem::QualifiedWildcard(_, _)
                                | ast::SelectItem::Wildcard(_) => {
                                    has_wildcard = true;
                                    Err("columns contain a wildcard".into())
                                }
                            })
                            .collect::<Result<Vec<String>, String>>()
                    } else {
                        Err(format!("not a SELECT statement: {query:?}"))
                    }
                } else {
                    Err(format!("not a Query: {statement:?}"))
                }
            }
            InterpolateItem::Expr { .. } => Err(format!(
                "could not extract columns from item {item:?}: not a string"