oxisql_parse/

dml.rs

//! DML planning helpers.
//!
//! Converts sqlparser DML [`Statement`]s into a high-level [`DmlPlan`] that
//! the rest of OxiSQL can reason about without depending on sqlparser's AST
//! directly.
//!
//! # Supported statements
//!
//! | SQL form | [`DmlPlan`] variant |
//! |---|---|
//! | `INSERT INTO t SELECT …` | [`DmlPlan::InsertSelect`] |
//! | `INSERT INTO t … ON CONFLICT … DO UPDATE` | [`DmlPlan::Upsert`] |
//! | `UPDATE t SET … WHERE …` | [`DmlPlan::Update`] |
//!
//! All other statement types return `None` from [`plan_dml`].

use crate::{plan_query, LogicalPlan};
use sqlparser::ast::{
    OnConflictAction, OnInsert, SetExpr, Statement, TableFactor, TableWithJoins,
    UpdateTableFromKind,
};

// ── Public types ─────────────────────────────────────────────────────────────

/// A high-level DML plan derived from a SQL statement.
#[derive(Debug, Clone)]
pub enum DmlPlan {
    /// `INSERT INTO table SELECT …`
    InsertSelect {
        /// Target table name.
        table: String,
        /// Optional column list from `INSERT INTO t (col1, col2, …)`.
        columns: Vec<String>,
        /// Logical plan for the `SELECT` sub-query.
        source: Box<LogicalPlan>,
    },
    /// `INSERT … ON CONFLICT … DO UPDATE SET …` (PostgreSQL / SQLite UPSERT).
    Upsert {
        /// Target table name.
        table: String,
        /// Column list from `INSERT INTO t (col1, …)`.
        columns: Vec<String>,
        /// Value rows as string expressions.
        values: Vec<Vec<String>>,
        /// The conflict target column(s), joined with `", "`.
        conflict_column: String,
        /// `(column, expression)` pairs from the `DO UPDATE SET` clause.
        update_exprs: Vec<(String, String)>,
    },
    /// `UPDATE t SET col = expr [FROM t2] [WHERE pred]`.
    Update {
        /// Target table name.
        table: String,
        /// `(column, expression)` pairs from the `SET` clause.
        set_exprs: Vec<(String, String)>,
        /// Optional `WHERE` predicate rendered as a string.
        predicate: Option<String>,
        /// Optional source table for Postgres-style `UPDATE … FROM t2 …`.
        from_table: Option<String>,
    },
}

// ── Public API ────────────────────────────────────────────────────────────────

/// Convert a sqlparser [`Statement`] into a [`DmlPlan`].
///
/// Returns `None` when the statement is not a recognised DML form (e.g. a
/// plain `SELECT`).
pub fn plan_dml(stmt: &Statement) -> Option<DmlPlan> {
    match stmt {
        Statement::Insert(insert) => plan_insert(insert),
        Statement::Update(update) => plan_update(update),
        _ => None,
    }
}

// ── Private helpers ───────────────────────────────────────────────────────────

fn plan_insert(insert: &sqlparser::ast::Insert) -> Option<DmlPlan> {
    let table = insert.table.to_string();

    // Collect column names.
    let columns: Vec<String> = insert
        .columns
        .iter()
        .filter_map(|col| {
            col.0
                .last()
                .and_then(|p| p.as_ident())
                .map(|id| id.value.clone())
        })
        .collect();

    // Check for ON CONFLICT clause first → Upsert.
    if let Some(OnInsert::OnConflict(on_conflict)) = &insert.on {
        if let OnConflictAction::DoUpdate(do_update) = &on_conflict.action {
            // Conflict target columns.
            let conflict_column = on_conflict
                .conflict_target
                .as_ref()
                .map(|ct| match ct {
                    sqlparser::ast::ConflictTarget::Columns(cols) => cols
                        .iter()
                        .map(|id| id.value.clone())
                        .collect::<Vec<_>>()
                        .join(", "),
                    sqlparser::ast::ConflictTarget::OnConstraint(name) => name.to_string(),
                })
                .unwrap_or_default();

            // SET expressions from DO UPDATE.
            let update_exprs: Vec<(String, String)> = do_update
                .assignments
                .iter()
                .map(|a| (a.target.to_string(), a.value.to_string()))
                .collect();

            // Collect VALUES rows as strings.
            let values = collect_values_rows(&insert.source);

            return Some(DmlPlan::Upsert {
                table,
                columns,
                values,
                conflict_column,
                update_exprs,
            });
        }
    }

    // INSERT … SELECT: source is a Query whose body is *not* Values.
    if let Some(src) = &insert.source {
        if !matches!(*src.body, SetExpr::Values(_)) {
            // Plan the SELECT sub-query.
            let source_plan = plan_query(src).ok()?;
            return Some(DmlPlan::InsertSelect {
                table,
                columns,
                source: Box::new(source_plan),
            });
        }
    }

    // INSERT … VALUES that carries an ON CONFLICT we already handled above,
    // but plain VALUES without ON CONFLICT is not a DML plan variant here.
    None
}

fn plan_update(update: &sqlparser::ast::Update) -> Option<DmlPlan> {
    // Extract table name from the TableWithJoins relation.
    let table = match &update.table.relation {
        sqlparser::ast::TableFactor::Table { name, .. } => name.to_string(),
        other => other.to_string(),
    };

    let set_exprs: Vec<(String, String)> = update
        .assignments
        .iter()
        .map(|a| (a.target.to_string(), a.value.to_string()))
        .collect();

    let predicate = update.selection.as_ref().map(|e| e.to_string());

    // Postgres-style `UPDATE … FROM t2 …` — extract the source table name.
    let from_table = update.from.as_ref().and_then(|kind| {
        let twjs: &[TableWithJoins] = match kind {
            UpdateTableFromKind::BeforeSet(v) | UpdateTableFromKind::AfterSet(v) => v,
        };
        twjs.first().map(|twj| match &twj.relation {
            TableFactor::Table { name, .. } => name.to_string(),
            other => other.to_string(),
        })
    });

    Some(DmlPlan::Update {
        table,
        set_exprs,
        predicate,
        from_table,
    })
}

/// Extract value rows from an INSERT source query.
///
/// Returns an empty `Vec` when the source is absent or is not a `VALUES`
/// clause.
fn collect_values_rows(source: &Option<Box<sqlparser::ast::Query>>) -> Vec<Vec<String>> {
    source
        .as_ref()
        .and_then(|q| {
            if let SetExpr::Values(vals) = &*q.body {
                Some(
                    vals.rows
                        .iter()
                        .map(|row| row.iter().map(|e| e.to_string()).collect())
                        .collect(),
                )
            } else {
                None
            }
        })
        .unwrap_or_default()
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parse_one;

    /// `INSERT INTO archive SELECT * FROM orders WHERE status = 'done'`
    /// should produce `DmlPlan::InsertSelect` targeting "archive".
    #[test]
    fn test_plan_insert_select() {
        let stmt = parse_one("INSERT INTO archive SELECT * FROM orders WHERE status = 'done'")
            .expect("parse");

        let plan = plan_dml(&stmt).expect("should produce a DmlPlan");

        match plan {
            DmlPlan::InsertSelect { table, .. } => {
                assert_eq!(table, "archive");
            }
            other => panic!("expected InsertSelect, got {other:?}"),
        }
    }

    /// `UPDATE users SET active = false WHERE last_login < '2020-01-01'`
    /// should produce `DmlPlan::Update`.
    #[test]
    fn test_plan_update() {
        let stmt = parse_one("UPDATE users SET active = false WHERE last_login < '2020-01-01'")
            .expect("parse");

        let plan = plan_dml(&stmt).expect("should produce a DmlPlan");

        match plan {
            DmlPlan::Update {
                table,
                predicate,
                from_table,
                ..
            } => {
                assert_eq!(table, "users");
                assert!(predicate.is_some(), "expected a WHERE predicate");
                assert!(
                    from_table.is_none(),
                    "plain UPDATE should have no from_table"
                );
            }
            other => panic!("expected Update, got {other:?}"),
        }
    }

    /// `UPDATE … FROM` — Postgres-style join-based update.
    ///
    /// Verified with the PostgreSQL dialect which parses `UPDATE … FROM`.
    /// The generic dialect does not support this syntax, so we fall back
    /// gracefully if that path is hit.
    #[test]
    fn test_plan_update_from() {
        // Try with the Postgres dialect first — it supports UPDATE … FROM.
        let sql = "UPDATE orders SET status = 'done' FROM users WHERE orders.user_id = users.id";
        let result = crate::parse_postgres(sql);
        match result {
            Ok(stmts) => {
                if let Some(stmt) = stmts.into_iter().next() {
                    let plan = plan_dml(&stmt);
                    match plan {
                        Some(DmlPlan::Update {
                            table, from_table, ..
                        }) => {
                            assert_eq!(table, "orders");
                            // Postgres dialect populates from_table for UPDATE … FROM.
                            assert_eq!(
                                from_table.as_deref(),
                                Some("users"),
                                "from_table should be 'users', got {:?}",
                                from_table
                            );
                        }
                        other => {
                            // Parsing succeeded but plan_dml returned None or a
                            // different variant — acceptable if dialect did something
                            // unexpected.
                            let _ = other;
                        }
                    }
                }
            }
            Err(_) => {
                // If even the Postgres dialect rejects this SQL, document the
                // limitation and pass.
            }
        }
    }

    /// A plain `SELECT` statement should return `None` from `plan_dml`.
    #[test]
    fn test_plan_dml_returns_none_for_select() {
        let stmt = parse_one("SELECT id FROM users").expect("parse");
        let result = plan_dml(&stmt);
        assert!(result.is_none(), "expected None for SELECT, got {result:?}");
    }

    /// `INSERT … ON CONFLICT … DO UPDATE` (UPSERT) is recognised.
    #[test]
    fn test_plan_upsert() {
        let stmt = parse_one(
            "INSERT INTO users (id, name) VALUES (1, 'alice') \
             ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name",
        )
        .expect("parse");

        let plan = plan_dml(&stmt).expect("should produce a DmlPlan");

        match plan {
            DmlPlan::Upsert {
                table,
                conflict_column,
                update_exprs,
                ..
            } => {
                assert_eq!(table, "users");
                assert_eq!(conflict_column, "id");
                assert!(
                    !update_exprs.is_empty(),
                    "expected at least one SET expression"
                );
            }
            other => panic!("expected Upsert, got {other:?}"),
        }
    }
}