rqb 0.1.7

sqlx-first typed Postgres query builder.
Documentation
use super::*;

/// Editable single-row INSERT returned by [`insert()`].
///
/// Assignments stay editable through WITH, RETURNING and ON CONFLICT.
/// Batch, SELECT and DEFAULT VALUES select a complete body and return [`Insert`],
/// which has no additive row setters. Convert with `.into()` when a helper
/// needs a completed `Insert` or [`Stmt`].
#[derive(Clone, Debug)]
#[must_use]
pub struct InsertRow(pub(super) Insert);

impl InsertRow {
    pub(super) fn new(target: impl Into<Source>) -> Self {
        Self(Insert {
            ctes: Vec::new(),
            target: target.into(),
            body: InsertBody::Values(Vec::new()),
            conflict: None,
            returning: Vec::new(),
        })
    }

    /// Adds one column assignment. If the same database column was assigned
    /// earlier, this assignment replaces the earlier value.
    ///
    /// This makes it safe to layer server-owned values around a DTO mapping:
    /// call `values(&dto)` for request-owned fields and use `set(...)` for
    /// generated IDs, tenant IDs, status defaults, or explicit overrides.
    #[inline]
    pub fn set(mut self, assignment: Assignment) -> Self {
        push_assignment(self.values_mut(), assignment);
        self
    }

    /// Adds multiple column assignments. Later assignments for the same
    /// database column replace earlier values.
    pub fn set_many(mut self, assignments: impl IntoAssignments) -> Self {
        extend_assignments(self.values_mut(), assignments.into_assignments());
        self
    }

    /// Adds one assignment only when `condition` is true.
    #[inline]
    pub fn set_if(self, condition: bool, assignment: Assignment) -> Self {
        if condition {
            self.set(assignment)
        } else {
            self
        }
    }

    /// Adds one assignment built from an optional value.
    pub fn set_option<T>(self, value: Option<T>, f: impl FnOnce(T) -> Assignment) -> Self {
        match value {
            Some(value) => self.set(f(value)),
            None => self,
        }
    }

    /// Adds assignments produced by an [`Insertable`] DTO.
    ///
    /// The assignments participate in the same replacement rules as `set`.
    /// Call `values(&dto)` first, then `set(...)` for IDs, tenant fields, or
    /// other values owned by the server.
    pub fn values(mut self, values: impl Insertable) -> Self {
        extend_assignments(self.values_mut(), values.insert_assignments());
        self
    }

    /// Adds many [`Insertable`] DTO rows through `INSERT ... SELECT` over an
    /// inline `VALUES` source.
    ///
    /// Every row must produce the same insert fields in the same order. This is
    /// usually true for DTOs without `#[rqb(skip_none)]`; if optional insert
    /// fields can be omitted per row, normalize the input first or build an
    /// explicit `values_source(...)`.
    /// The row builder must not already contain single-row assignments.
    ///
    /// For an upsert, use `do_update_excluded` or `set_excluded`: PostgreSQL
    /// exposes incoming values as `EXCLUDED` in the conflict action.
    pub fn values_many<I, R>(self, rows: I) -> Result<Insert>
    where
        I: IntoIterator<Item = R>,
        R: Insertable,
    {
        if !matches!(&self.0.body, InsertBody::Values(assignments) if assignments.is_empty()) {
            return Err(Error::InvalidInsertShape {
                message: "batch insert cannot be combined with existing single-row assignments",
            });
        }

        let mut columns = None::<Vec<Meta>>;
        let mut values: Vec<Vec<ValueExpr>> = Vec::new();
        for row in rows {
            let assignments = normalized_assignments(row.insert_assignments());
            if assignments.is_empty() {
                return Err(Error::InvalidInsertShape {
                    message: "batch insert rows must contain at least one assignment",
                });
            }
            match &columns {
                Some(columns) if !same_batch_fields(columns, &assignments) => {
                    return Err(Error::InvalidInsertShape {
                        message: "batch insert rows must use the same fields in the same order",
                    });
                }
                Some(_) => {}
                None => {
                    columns = Some(
                        assignments
                            .iter()
                            .map(|assignment| assignment.field)
                            .collect(),
                    );
                }
            }
            values.push(
                assignments
                    .into_iter()
                    .map(|assignment| match assignment.value {
                        // VALUES resolves all-NULL columns before target coercion.
                        AssignmentValue::Expr(ValueExpr::Null) => {
                            Ok(ValueExpr::Null.cast(assignment.field.pg))
                        }
                        AssignmentValue::Expr(expr) => Ok(expr),
                        AssignmentValue::Default => Err(Error::InvalidInsertShape {
                            message: "batch insert rows cannot use DEFAULT assignments",
                        }),
                    })
                    .collect::<Result<Vec<_>>>()?,
            );
        }

        let Some(columns) = columns else {
            return Err(Error::InvalidInsertShape {
                message: "batch insert requires at least one row",
            });
        };
        Ok(self.from_select_all(crate::values_source(values, "incoming", columns)))
    }

    /// Replaces the row body with an INSERT ... SELECT.
    pub fn from_select(self, columns: impl IntoFieldMetas, select: Select) -> Insert {
        self.0.from_select(columns, select)
    }

    /// Replaces the row body with all exposed source fields.
    pub fn from_select_all(self, source: impl Into<Source>) -> Insert {
        self.0.from_select_all(source)
    }

    /// Replaces the row body with PostgreSQL DEFAULT VALUES.
    pub fn default_values(self) -> Insert {
        self.0.default_values()
    }

    /// Validates the insert without rendering SQL.
    pub fn validate(&self) -> Result<()> {
        self.0.validate()
    }

    /// Validates and renders parameterized SQL without consuming the row builder.
    pub fn build(&self) -> Result<crate::BuiltQuery> {
        self.0.build()
    }

    pub(crate) fn check_returning(&self) -> Result<()> {
        self.0.check_returning()
    }

    fn values_mut(&mut self) -> &mut Vec<Assignment> {
        let InsertBody::Values(assignments) = &mut self.0.body else {
            unreachable!("InsertRow only constructs and preserves the Values body")
        };
        assignments
    }
}

impl From<InsertRow> for Insert {
    fn from(row: InsertRow) -> Self {
        row.0
    }
}

impl Insert {
    /// Uses a select statement as the insert source.
    ///
    /// `columns` owns the complete target-column list for the insert-select
    /// body. This keeps the AST from carrying pending target columns before a
    /// select source exists.
    pub fn from_select(mut self, columns: impl IntoFieldMetas, select: Select) -> Self {
        self.body = InsertBody::Select {
            columns: columns.into_field_metas(),
            select: Box::new(select),
        };
        self
    }

    /// Uses every exposed field from `source` as both target columns and the
    /// `INSERT ... SELECT` projection.
    ///
    /// This is the compact path for bulk loads from `values_source(...)` or
    /// staging sources whose exposed metadata already matches the insert target.
    pub fn from_select_all(mut self, source: impl Into<Source>) -> Self {
        let source = source.into();
        let qualifier = source.explicit_alias().map(str::to_owned);
        let mut columns = Vec::new();
        let mut projection = Vec::new();
        source.for_each_field(|field| {
            push_column(&mut columns, *field);
            projection.push(SelectItem {
                expr: ValueExpr::field(*field, qualifier.clone()),
                alias: None,
            });
        });
        let mut select = Select::from(source);
        select.projection = projection;
        self.body = InsertBody::Select {
            columns,
            select: Box::new(select),
        };
        self
    }

    /// Uses PostgreSQL `DEFAULT VALUES` instead of explicit insert values.
    ///
    /// This is for rows where every target column should be populated by its
    /// database default or remain nullable. It can still be combined with
    /// `RETURNING` and `ON CONFLICT`. A later `from_select(...)` or
    /// `from_select_all(...)` explicitly replaces this body.
    #[inline]
    pub fn default_values(mut self) -> Self {
        self.body = InsertBody::DefaultValues;
        self
    }
}

macro_rules! impl_insert_clauses {
    ($ty:ty $(, $field:tt)?) => {
        impl $ty {
            /// Adds a CTE to the insert statement.
            #[inline]
            pub fn with(mut self, cte: Cte) -> Self {
                self $(.$field)?.ctes.push(cte);
                self
            }

            /// Starts an `ON CONFLICT (columns...)` clause.
            /// Completing a later conflict clause replaces the previous clause.
            ///
            /// Use this for column/index targets. For a named database constraint, use
            /// [`Insert::on_conflict_constraint`] with a generated constraint constant.
            pub fn on_conflict(self, fields: impl ConflictFields) -> ColumnConflictBuilder<Self> {
                let mut target_fields = Vec::with_capacity(fields.conflict_field_count());
                fields.push_conflict_fields(&mut target_fields);
                ColumnConflictBuilder {
                    insert: self,
                    fields: target_fields,
                    predicate: None,
                }
            }

            /// Starts an `ON CONFLICT ON CONSTRAINT` clause.
            /// Completing a later conflict clause replaces the previous clause.
            ///
            /// Generated schema exposes unique constraint names under
            /// `relation::constraints`.
            pub fn on_conflict_constraint(
                self,
                constraint: impl Into<String>,
            ) -> ConstraintConflictBuilder<Self> {
                ConstraintConflictBuilder {
                    insert: self,
                    constraint: constraint.into(),
                }
            }

            /// Adds one or more fields to `RETURNING`.
            ///
            /// Use a field for one returned column or tuple syntax for a small
            /// heterogeneous field list: `.returning((users::ID, users::EMAIL))`.
            pub fn returning(mut self, fields: impl IntoColumns) -> Self {
                self $(.$field)?.returning.extend(fields.into_columns().items);
                self
            }

            /// Adds an aliased expression to `RETURNING`.
            pub fn returning_as(mut self, expr: impl Into<ValueExpr>, alias: impl Into<String>) -> Self {
                self $(.$field)?.returning.push(SelectItem {
                    expr: expr.into(),
                    alias: Some(alias.into()),
                });
                self
            }

            /// Replaces `RETURNING` with every field exposed by the target source.
            #[inline]
            pub fn returning_all(mut self) -> Self {
                self $(.$field)?.returning.clear();
                push_returning_fields(&self $(.$field)?.target, &mut self $(.$field)?.returning);
                self
            }

        }
    };
}

impl_insert_clauses!(Insert);
impl_insert_clauses!(InsertRow, 0);

fn same_batch_fields(columns: &[Meta], assignments: &[Assignment]) -> bool {
    columns.len() == assignments.len()
        && columns
            .iter()
            .zip(assignments)
            .all(|(column, assignment)| column.db == assignment.field.db)
}