qusql-type 0.4.0

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use qusql_parse::{InsertReplaceType, Statement, WithBlock};

use alloc::vec::Vec;

use crate::{
    schema::{Column, Schema},
    type_call::type_call,
    type_delete::type_delete,
    type_insert_replace::{AutoIncrementId, type_insert_replace},
    type_lock::{type_lock, type_unlock},
    type_select::{SelectType, type_compound_query},
    type_set::type_set,
    type_truncate::type_truncate,
    type_update::type_update,
    typer::Typer,
};

pub(crate) enum InnerStatementType<'a> {
    Select(SelectType<'a>),
    Delete {
        returning: Option<SelectType<'a>>,
    },
    Insert {
        auto_increment_id: AutoIncrementId,
        returning: Option<SelectType<'a>>,
    },
    Update {
        returning: Option<SelectType<'a>>,
    },
    Replace {
        returning: Option<SelectType<'a>>,
    },
    Truncate,
    Call,
    Transaction,
    Set,
    Lock,
    Invalid,
}

fn type_with_query<'a>(
    typer: &mut Typer<'a, '_>,
    with_blocks: &[WithBlock<'a>],
    inner: &Statement<'a>,
) -> InnerStatementType<'a> {
    if let Some((block, rem_blocks)) = with_blocks.split_first() {
        let r = type_statement(typer, &block.statement);
        let s = match r {
            InnerStatementType::Select(v) => Some(v),
            InnerStatementType::Delete { returning } => returning,
            InnerStatementType::Insert { returning, .. } => returning,
            InnerStatementType::Update { returning } => returning,
            _ => None,
        };
        if let Some(s) = s {
            let mut columns = Vec::new();
            for c in s.columns {
                if let Some(name) = &c.name {
                    columns.push(Column {
                        identifier: name.clone(),
                        type_: c.type_,
                        default: false,
                        auto_increment: false,
                        as_: None,
                        generated: false,
                    });
                }
            }
            let schema: Schema<'a> = Schema {
                identifier_span: block.identifier.span.clone(),
                columns,
                view: true,
            };

            let mut schemas = typer.with_schemas.clone();
            schemas.insert(block.identifier.as_str(), &schema);
            let mut child = typer.with_schemas(schemas);
            let result = type_with_query(&mut child, rem_blocks, inner);
            // Propagate any argument constraints discovered inside the CTE scope back
            // to the outer typer so they are visible to the caller.
            typer.arg_types = child.arg_types;
            result
        } else {
            type_with_query(typer, rem_blocks, inner)
        }
    } else {
        type_statement(typer, inner)
    }
}

pub(crate) fn type_statement<'a>(
    typer: &mut Typer<'a, '_>,
    statement: &Statement<'a>,
) -> InnerStatementType<'a> {
    match statement {
        Statement::Select(s) => InnerStatementType::Select(crate::type_select::type_select(
            typer,
            s,
            typer.options.warn_duplicate_column_in_select,
        )),
        Statement::Delete(d) => {
            let returning = type_delete(typer, d);
            InnerStatementType::Delete { returning }
        }
        Statement::InsertReplace(ior) => {
            let (auto_increment_id, returning) = type_insert_replace(typer, ior);
            match &ior.type_ {
                InsertReplaceType::Insert(_) => InnerStatementType::Insert {
                    auto_increment_id,
                    returning,
                },
                InsertReplaceType::Replace(_) => InnerStatementType::Replace { returning },
            }
        }
        Statement::Update(u) => {
            let returning = type_update(typer, u);
            InnerStatementType::Update { returning }
        }
        Statement::CompoundQuery(u) => InnerStatementType::Select(type_compound_query(typer, u)),
        Statement::WithQuery(w) => type_with_query(typer, &w.with_blocks, &w.statement),
        Statement::TruncateTable(t) => {
            type_truncate(typer, t);
            InnerStatementType::Truncate
        }
        Statement::Call(c) => {
            type_call(typer, c);
            InnerStatementType::Call
        }
        Statement::Begin(_)
        | Statement::Commit(_)
        | Statement::StartTransaction(_)
        | Statement::End(_) => InnerStatementType::Transaction,
        Statement::Set(s) => {
            type_set(typer, s);
            InnerStatementType::Set
        }
        Statement::Lock(l) => {
            type_lock(typer, l);
            InnerStatementType::Lock
        }
        Statement::Unlock(u) => {
            type_unlock(typer, u);
            InnerStatementType::Lock
        }
        s => {
            typer.issues.err("Cannot type statement of this type", s);
            InnerStatementType::Invalid
        }
    }
}