rainier_orm/

repo.rs

//! Generic CRUD over any [`Entity`], plus a keyset [`Cursor`] for walking
//! large tables without holding them in memory.
//!
//! Each function builds a `sea_query` statement from the entity's metadata,
//! renders it for the executor's dialect, and runs it. No per-entity SQL is
//! written anywhere — this module *is* the repository for every derived
//! struct.
//!
//! # Every future here must be `Send`
//!
//! `sea_query`'s statement builders hold `Rc<dyn Iden>`, which is `!Send`. A
//! value that is merely **in scope** across an `.await` is captured by the
//! generated future — even if it is never read again — so leaving a statement
//! alive across the executor call makes the whole future `!Send`.
//!
//! That is not a cosmetic problem: a `!Send` future cannot be `tokio::spawn`ed,
//! so a multi-threaded web server could not call any of these from a request
//! handler.
//!
//! So every function here follows the same shape — build and render the
//! statement inside a block that **ends before the await**, and carry only the
//! rendered `String` and `Vec<Value>` across it:
//!
//! ```ignore
//! let (route, sql, params) = {
//!     let mut stmt = Query::select();
//!     // …
//!     let (sql, params) = exec.dialect().build_query(&stmt);
//!     (route, sql, into_vec(params))
//! };                                    // <- stmt dropped here
//! let rows = exec.fetch_all_routed(route, &sql, params).await?;
//! ```
//!
//! `tests/send_futures.rs` asserts this at compile time for every public async
//! API, so a regression fails the build rather than surfacing as a confusing
//! error in a downstream crate.

use crate::route::route_for;
use crate::{Entity, Executor, Result, ShardRoute};
use sea_query::{Alias, Expr, OnConflict, Order, Query, SimpleExpr, Value};

fn col(name: &str) -> Alias {
    Alias::new(name)
}

/// The shard route for an entity instance: the first of its `(column, value)`
/// pairs that names a shard-encoded column. `Global` if none (a global entity,
/// or a single-database backend ignores it anyway).
fn route_from_pairs<E: Entity>(pairs: &[(&'static str, Value)]) -> ShardRoute {
    if E::shard_columns().is_empty() {
        return ShardRoute::Global;
    }
    for (c, v) in pairs {
        let r = route_for::<E>(c, v);
        if !matches!(r, ShardRoute::Global) {
            return r;
        }
    }
    ShardRoute::Global
}

/// If `E`'s primary key is itself a shard key, is not auto-increment, and is
/// unset (0), mint a shard-encoded id for it from another shard-key column's
/// value (the owning `user_id`) via the executor — transparent id assignment
/// for the sharded data tier. Returns the new id (caller writes it into the pk
/// pair). `None` for global entities, auto-increment keys, an already-set key,
/// or an executor that doesn't allocate (single-DB).
fn assign_shard_pk<E: Entity, X: Executor>(exec: &X, pairs: &[(&str, Value)]) -> Option<u64> {
    let pk = E::primary_key();
    if !E::shard_columns().contains(&pk) {
        return None;
    }
    let pk_spec = E::columns().iter().find(|c| c.name == pk)?;
    if pk_spec.auto_increment {
        return None;
    }
    // pk currently unset?
    let pk_val = pairs.iter().find(|(c, _)| *c == pk).map(|(_, v)| v)?;
    if crate::route::routing_key(pk_val) != Some(0) {
        return None;
    }
    // derive the shard from another shard-key column (the user id)
    let (_, sk_val) = pairs.iter().find(|(c, v)| {
        *c != pk && E::shard_columns().contains(c) && crate::route::routing_key(v).is_some()
    })?;
    let shard_key = crate::route::routing_key(sk_val)?;
    exec.allocate_id(shard_key)
}

/// Start a fluent [`Query`](crate::Query) over `E` — `where_eq` / `where_like`
/// / `where_in` / joins / `first_or_create` / etc. Sugar for
/// `Query::<E>::new()`.
pub fn query<E: Entity>() -> crate::Query<E> {
    crate::Query::new()
}

/// Insert `entity`. Returns the autogenerated key
/// ([`last_insert_id`](crate::ExecOutcome::last_insert_id)) when the table has an auto-increment
/// primary key, else `0`.
pub async fn insert<E, X>(exec: &X, entity: &E) -> Result<i64>
where
    E: Entity,
    X: Executor,
{
    // Everything `sea_query` is built *and dropped* inside this block, so only
    // the rendered `String` and `Vec<Value>` survive to the await. See the
    // module docs on why that matters.
    let (route, sql, params, assigned) = {
        let mut pairs = entity.insert_values();
        // Transparent shard-encoded pk: the connector mints it (see `allocate_id`).
        let assigned = assign_shard_pk::<E, X>(exec, &pairs);
        if let Some(id) = assigned {
            let pk = E::primary_key();
            for (c, v) in pairs.iter_mut() {
                if *c == pk {
                    *v = id.into();
                }
            }
        }
        let route = route_from_pairs::<E>(&pairs);
        let mut stmt = Query::insert();
        stmt.into_table(col(E::table()));
        stmt.columns(pairs.iter().map(|(c, _)| col(c)));
        stmt.values_panic(pairs.into_iter().map(|(_, v)| SimpleExpr::from(Expr::val(v))));

        let (sql, params) = exec.dialect().build_query(&stmt);
        (route, sql, into_vec(params), assigned)
    };

    let outcome = exec.execute_routed(route, &sql, params).await?;
    // The assigned id is the truth when we minted one (the DB has no
    // auto-increment to report for an app-assigned key).
    Ok(assigned.map(|i| i as i64).unwrap_or(outcome.last_insert_id))
}

/// Insert `entity`, or on a conflict with `conflict_columns` update
/// `update_columns` to the inserted values (`INSERT … ON CONFLICT … DO UPDATE`
/// / MySQL `ON DUPLICATE KEY UPDATE`, rendered per dialect by sea-query). An
/// empty `update_columns` makes it insert-or-ignore (a no-op self-update,
/// portable across dialects). Returns rows affected.
pub async fn upsert<E, X>(
    exec: &X,
    entity: &E,
    conflict_columns: &[&str],
    update_columns: &[&str],
) -> Result<u64>
where
    E: Entity,
    X: Executor,
{
    let (route, sql, params) = {
        let pairs = entity.insert_values();
        let route = route_from_pairs::<E>(&pairs);
        let mut stmt = Query::insert();
        stmt.into_table(col(E::table()));
        stmt.columns(pairs.iter().map(|(c, _)| col(c)));
        stmt.values_panic(pairs.into_iter().map(|(_, v)| SimpleExpr::from(Expr::val(v))));

        let mut on_conflict = OnConflict::columns(conflict_columns.iter().map(|c| col(c)));
        if !update_columns.is_empty() {
            on_conflict.update_columns(update_columns.iter().map(|c| col(c)));
        } else if let Some(c) = conflict_columns.first() {
            // Insert-or-ignore. MySQL has no portable `DO NOTHING` — sea-query
            // renders `do_nothing()` as invalid `IGNORE` syntax there — so emit a
            // no-op self-update on a conflict column instead. `col = col` is valid on
            // every dialect and leaves the existing row untouched.
            on_conflict.value(col(c), Expr::col(col(c)));
        } else {
            on_conflict.do_nothing();
        }
        stmt.on_conflict(on_conflict);

        let (sql, params) = exec.dialect().build_query(&stmt);
        (route, sql, into_vec(params))
    };

    Ok(exec.execute_routed(route, &sql, params).await?.rows_affected)
}

/// Find the row whose primary key equals `pk`, decoded into `E`. `None` if no
/// row matches.
pub async fn find_by_pk<E, X, V>(exec: &X, pk: V) -> Result<Option<E>>
where
    E: Entity,
    X: Executor,
    V: Into<Value>,
{
    let (route, sql, params) = {
        let pk_val: Value = pk.into();
        let route = route_for::<E>(E::primary_key(), &pk_val);
        let mut stmt = Query::select();
        stmt.from(col(E::table()))
            .columns(E::columns().iter().map(|c| col(c.name)))
            .and_where(Expr::col(col(E::primary_key())).eq(pk_val))
            .limit(1);

        let (sql, params) = exec.dialect().build_query(&stmt);
        (route, sql, into_vec(params))
    };

    let rows = exec.fetch_all_routed(route, &sql, params).await?;
    match rows.first() {
        Some(row) => Ok(Some(E::from_row(row.as_ref())?)),
        None => Ok(None),
    }
}

/// Fetch every row of `E` where `column` equals `value`. This is how you
/// traverse a relationship — e.g. a user's posts is
/// `find_by::<Post, _, _>(db, "author_id", user_id)`. There is no generated
/// navigation method; traversal is an explicit query against the flat FK
/// column, because the related rows may live in a different backend entirely
/// (a cross-backend JOIN is impossible) and even within one database this
/// crate doesn't replace the query planner.
pub async fn find_by<E, X, V>(exec: &X, column: &str, value: V) -> Result<Vec<E>>
where
    E: Entity,
    X: Executor,
    V: Into<Value>,
{
    let (route, sql, params) = {
        let val: Value = value.into();
        let route = route_for::<E>(column, &val);
        let mut stmt = Query::select();
        stmt.from(col(E::table()))
            .columns(E::columns().iter().map(|c| col(c.name)))
            .and_where(Expr::col(col(column)).eq(val));

        let (sql, params) = exec.dialect().build_query(&stmt);
        (route, sql, into_vec(params))
    };

    let rows = exec.fetch_all_routed(route, &sql, params).await?;
    rows.iter().map(|r| E::from_row(r.as_ref())).collect()
}

/// Like [`find_by`] but for a column you expect to be unique — returns the
/// first match, or `None`.
pub async fn find_one_by<E, X, V>(exec: &X, column: &str, value: V) -> Result<Option<E>>
where
    E: Entity,
    X: Executor,
    V: Into<Value>,
{
    let (route, sql, params) = {
        let val: Value = value.into();
        let route = route_for::<E>(column, &val);
        let mut stmt = Query::select();
        stmt.from(col(E::table()))
            .columns(E::columns().iter().map(|c| col(c.name)))
            .and_where(Expr::col(col(column)).eq(val))
            .limit(1);

        let (sql, params) = exec.dialect().build_query(&stmt);
        (route, sql, into_vec(params))
    };

    let rows = exec.fetch_all_routed(route, &sql, params).await?;
    match rows.first() {
        Some(row) => Ok(Some(E::from_row(row.as_ref())?)),
        None => Ok(None),
    }
}

/// Fetch every row of `E`'s table. Convenient for small tables; prefer
/// [`cursor`] for large ones. Routes `Global`: on a sharded entity this hits
/// only the global/directory database, so use it for global tables — a
/// cross-shard scan is an explicit fan-out, not a single `all`.
pub async fn all<E, X>(exec: &X) -> Result<Vec<E>>
where
    E: Entity,
    X: Executor,
{
    let (sql, params) = {
        let mut stmt = Query::select();
        stmt.from(col(E::table())).columns(E::columns().iter().map(|c| col(c.name)));

        let (sql, params) = exec.dialect().build_query(&stmt);
        (sql, into_vec(params))
    };

    let rows = exec.fetch_all_routed(ShardRoute::Global, &sql, params).await?;
    rows.iter().map(|r| E::from_row(r.as_ref())).collect()
}

/// Update every non-primary-key column of `entity`, matched by its primary
/// key. Returns the number of rows affected (0 or 1 for a unique key).
pub async fn update<E, X>(exec: &X, entity: &E) -> Result<u64>
where
    E: Entity,
    X: Executor,
{
    let (route, sql, params) = {
        let pk_val = entity.pk_value();
        let route = route_for::<E>(E::primary_key(), &pk_val);
        let mut stmt = Query::update();
        stmt.table(col(E::table()));
        for (c, v) in entity.update_values() {
            stmt.value(col(c), v);
        }
        stmt.and_where(Expr::col(col(E::primary_key())).eq(pk_val));

        let (sql, params) = exec.dialect().build_query(&stmt);
        (route, sql, into_vec(params))
    };

    let outcome = exec.execute_routed(route, &sql, params).await?;
    Ok(outcome.rows_affected)
}

/// Delete the row whose primary key equals `pk`. Returns rows affected.
pub async fn delete_by_pk<E, X, V>(exec: &X, pk: V) -> Result<u64>
where
    E: Entity,
    X: Executor,
    V: Into<Value>,
{
    let (route, sql, params) = {
        let pk_val: Value = pk.into();
        let route = route_for::<E>(E::primary_key(), &pk_val);
        let mut stmt = Query::delete();
        stmt.from_table(col(E::table())).and_where(Expr::col(col(E::primary_key())).eq(pk_val));

        let (sql, params) = exec.dialect().build_query(&stmt);
        (route, sql, into_vec(params))
    };

    let outcome = exec.execute_routed(route, &sql, params).await?;
    Ok(outcome.rows_affected)
}

/// Open a [`Cursor`] over `E`'s whole table, walking it in primary-key order
/// in pages of `page_size`.
pub fn cursor<E, X>(exec: &X, page_size: u64) -> Cursor<'_, E, X>
where
    E: Entity,
    X: Executor,
{
    Cursor::new(exec, page_size)
}

/// A forward-only keyset cursor for streaming large tables.
///
/// Rather than `OFFSET` (which the database must scan past, getting slower
/// deeper in) or loading the whole table, the cursor remembers the last
/// primary key it saw and asks for `WHERE pk > ? ORDER BY pk LIMIT page_size`.
/// Every page is an index seek, so cost stays flat however deep you walk —
/// and only one page is in memory at a time. This is the only iteration
/// strategy that is also correct on D1-over-HTTP, which cannot hold a
/// streaming server-side cursor.
///
/// ```ignore
/// let mut cur = repo::cursor::<User, _>(&exec, 500);
/// while let Some(page) = cur.next_page().await? {
///     for user in page { /* ... */ }
/// }
/// ```
pub struct Cursor<'a, E, X> {
    exec: &'a X,
    page_size: u64,
    after: Option<Value>,
    done: bool,
    _entity: core::marker::PhantomData<E>,
}

impl<'a, E, X> Cursor<'a, E, X>
where
    E: Entity,
    X: Executor,
{
    fn new(exec: &'a X, page_size: u64) -> Self {
        Self {
            exec,
            page_size: page_size.max(1),
            after: None,
            done: false,
            _entity: core::marker::PhantomData,
        }
    }

    /// Resume from a known primary key — the next page starts strictly after
    /// `pk`. Lets a walk be checkpointed and continued across requests (e.g.
    /// one D1 HTTP call per invocation).
    pub fn after(mut self, pk: impl Into<Value>) -> Self {
        self.after = Some(pk.into());
        self
    }

    /// Fetch the next page. Returns `None` once the table is exhausted. A
    /// final short page (fewer than `page_size` rows) marks the end, so the
    /// following call returns `None` without a wasted round-trip.
    pub async fn next_page(&mut self) -> Result<Option<Vec<E>>> {
        if self.done {
            return Ok(None);
        }

        let pk = E::primary_key();
        let (route, sql, params) = {
            let mut stmt = Query::select();
            stmt.from(col(E::table()))
                .columns(E::columns().iter().map(|c| col(c.name)))
                .order_by(col(pk), Order::Asc)
                .limit(self.page_size);
            if let Some(after) = &self.after {
                stmt.and_where(Expr::col(col(pk)).gt(after.clone()));
            }

            // Route by the keyset position when it's a shard-encoded pk (walks one
            // shard); otherwise `Global` (a single/global table).
            let route = match &self.after {
                Some(after) => route_for::<E>(pk, after),
                None => ShardRoute::Global,
            };
            let (sql, params) = self.exec.dialect().build_query(&stmt);
            (route, sql, into_vec(params))
        };

        let rows = self.exec.fetch_all_routed(route, &sql, params).await?;
        if rows.is_empty() {
            self.done = true;
            return Ok(None);
        }

        let page: Vec<E> = rows.iter().map(|r| E::from_row(r.as_ref())).collect::<Result<_>>()?;

        // Advance the keyset to the last row's primary key for the next page.
        if let Some(last) = rows.last() {
            self.after = Some(pk_from_row::<E>(last.as_ref())?);
        }
        if (rows.len() as u64) < self.page_size {
            self.done = true;
        }
        Ok(Some(page))
    }

    /// The primary key the next page will start strictly after, if a page has
    /// been read. Persist this to resume the walk later via [`after`].
    ///
    /// [`after`]: Cursor::after
    pub fn checkpoint(&self) -> Option<&Value> {
        self.after.as_ref()
    }
}

/// Read just the primary-key column out of a row, by its declared
/// [`ColumnType`](crate::ColumnType), to advance the keyset.
fn pk_from_row<E: Entity>(row: &dyn crate::Row) -> Result<Value> {
    use crate::ColumnType;
    let pk = E::primary_key();
    let spec = E::columns()
        .iter()
        .find(|c| c.name == pk)
        .ok_or_else(|| crate::Error::msg("primary key column not in metadata"))?;
    let null = || crate::Error::msg(format!("primary key `{pk}` was NULL"));
    Ok(match spec.ty {
        ColumnType::Bool => row.get_bool(pk)?.ok_or_else(null)?.into(),
        ColumnType::Int => row.get_i32(pk)?.ok_or_else(null)?.into(),
        ColumnType::BigInt => row.get_i64(pk)?.ok_or_else(null)?.into(),
        ColumnType::Uint => row.get_u32(pk)?.ok_or_else(null)?.into(),
        ColumnType::BigUint => row.get_u64(pk)?.ok_or_else(null)?.into(),
        ColumnType::Double => row.get_f64(pk)?.ok_or_else(null)?.into(),
        ColumnType::Text => row.get_string(pk)?.ok_or_else(null)?.into(),
        ColumnType::Binary => row.get_bytes(pk)?.ok_or_else(null)?.into(),
        ColumnType::Timestamp => row.get_datetime(pk)?.ok_or_else(null)?.into(),
        ColumnType::Date => row.get_naive_date(pk)?.ok_or_else(null)?.into(),
    })
}

/// `sea_query::Values` (newtype over `Vec<Value>`) → the `Vec<Value>` the
/// executor binds.
fn into_vec(values: sea_query::Values) -> Vec<Value> {
    values.0
}