rustrails_record/

querying.rs

use std::collections::HashMap;

use sea_orm::{ColumnTrait, DatabaseConnection, EntityTrait, FromQueryResult, Iterable};
use serde_json::{Value, json};

use crate::{Record, RecordError, Relation};

/// Sort direction for ordered record queries.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OrderDirection {
    /// Sort in ascending order.
    Asc,
    /// Sort in descending order.
    Desc,
}

/// Async query helpers for [`Record`] types. Use `Querying` for the sync API.
#[allow(async_fn_in_trait, dead_code)]
pub(crate) trait AsyncQuerying: Record {
    /// Finds a record by primary key or returns [`RecordError::NotFound`].
    async fn find(id: i64, db: &DatabaseConnection) -> Result<Self, RecordError>
    where
        <Self::Entity as EntityTrait>::Column: ColumnTrait + Iterable,
    {
        Self::find_by_id(id, db).await?.ok_or(RecordError::NotFound)
    }

    /// Finds a record by primary key and returns `None` when missing.
    async fn find_by_id(id: i64, db: &DatabaseConnection) -> Result<Option<Self>, RecordError>
    where
        <Self::Entity as EntityTrait>::Column: ColumnTrait + Iterable,
    {
        let mut conditions = HashMap::new();
        conditions.insert(Self::primary_key_name().to_owned(), json!(id));
        Self::r#where(conditions).first(db).await
    }

    /// Finds the first record matching the provided conditions.
    async fn find_by(
        conditions: HashMap<String, Value>,
        db: &DatabaseConnection,
    ) -> Result<Option<Self>, RecordError>
    where
        <Self::Entity as EntityTrait>::Column: ColumnTrait + Iterable,
    {
        Self::r#where(conditions).first(db).await
    }

    /// Finds the first record matching the provided conditions or returns [`RecordError::NotFound`].
    async fn find_by_bang(
        conditions: HashMap<String, Value>,
        db: &DatabaseConnection,
    ) -> Result<Self, RecordError>
    where
        <Self::Entity as EntityTrait>::Column: ColumnTrait + Iterable,
    {
        Self::find_by(conditions, db)
            .await?
            .ok_or(RecordError::NotFound)
    }

    /// Returns one record without imposing an explicit order.
    async fn take(db: &DatabaseConnection) -> Result<Option<Self>, RecordError>
    where
        <Self::Entity as EntityTrait>::Column: ColumnTrait + Iterable,
    {
        Relation::<Self>::new().first(db).await
    }

    /// Returns one record without imposing an explicit order or raises when none exist.
    async fn take_bang(db: &DatabaseConnection) -> Result<Self, RecordError>
    where
        <Self::Entity as EntityTrait>::Column: ColumnTrait + Iterable,
    {
        Self::take(db).await?.ok_or(RecordError::NotFound)
    }

    /// Returns the only matching row.
    async fn sole(db: &DatabaseConnection) -> Result<Self, RecordError>
    where
        <Self::Entity as EntityTrait>::Column: ColumnTrait + Iterable,
    {
        Relation::<Self>::new().sole(db).await
    }

    /// Returns the only matching row for the provided conditions.
    async fn find_sole_by(
        conditions: HashMap<String, Value>,
        db: &DatabaseConnection,
    ) -> Result<Self, RecordError>
    where
        <Self::Entity as EntityTrait>::Column: ColumnTrait + Iterable,
    {
        Self::r#where(conditions).sole(db).await
    }

    /// Plucks the requested column from matching rows.
    async fn pluck(column: &str, db: &DatabaseConnection) -> Result<Vec<Value>, RecordError>
    where
        Self: serde::Serialize,
        <Self::Entity as EntityTrait>::Column: ColumnTrait + Iterable,
    {
        Relation::<Self>::new().pluck(column, db).await
    }

    /// Picks the first value for the requested column.
    async fn pick(column: &str, db: &DatabaseConnection) -> Result<Option<Value>, RecordError>
    where
        Self: serde::Serialize,
        <Self::Entity as EntityTrait>::Column: ColumnTrait + Iterable,
    {
        Relation::<Self>::new().pick(column, db).await
    }

    /// Returns all primary key values.
    async fn ids(db: &DatabaseConnection) -> Result<Vec<i64>, RecordError>
    where
        Self: serde::Serialize,
        <Self::Entity as EntityTrait>::Column: ColumnTrait + Iterable,
    {
        Relation::<Self>::new().ids(db).await
    }

    /// Loads all records for the entity.
    async fn all(db: &DatabaseConnection) -> Result<Vec<Self>, RecordError>
    where
        <Self::Entity as EntityTrait>::Column: ColumnTrait + Iterable,
    {
        Relation::<Self>::new().load(db).await
    }

    /// Loads the first record ordered by primary key ascending.
    async fn first(db: &DatabaseConnection) -> Result<Option<Self>, RecordError>
    where
        <Self::Entity as EntityTrait>::Column: ColumnTrait + Iterable,
    {
        Self::order(Self::primary_key_name(), OrderDirection::Asc)
            .first(db)
            .await
    }

    /// Loads the last record ordered by primary key descending.
    async fn last(db: &DatabaseConnection) -> Result<Option<Self>, RecordError>
    where
        <Self::Entity as EntityTrait>::Column: ColumnTrait + Iterable,
    {
        Self::order(Self::primary_key_name(), OrderDirection::Desc)
            .first(db)
            .await
    }

    /// Counts all rows for the entity.
    async fn count(db: &DatabaseConnection) -> Result<u64, RecordError>
    where
        <Self::Entity as EntityTrait>::Column: ColumnTrait + Iterable,
        <Self::Entity as EntityTrait>::Model: FromQueryResult + Send + Sync,
    {
        Relation::<Self>::new().count(db).await
    }

    /// Returns `true` when any record matches the provided conditions.
    async fn exists_with_conditions(
        conditions: HashMap<String, Value>,
        db: &DatabaseConnection,
    ) -> Result<bool, RecordError>
    where
        <Self::Entity as EntityTrait>::Column: ColumnTrait + Iterable,
    {
        Self::r#where(conditions).exists(db).await
    }

    /// Returns `true` when any record matches the provided conditions.
    async fn exists(
        conditions: HashMap<String, Value>,
        db: &DatabaseConnection,
    ) -> Result<bool, RecordError>
    where
        <Self::Entity as EntityTrait>::Column: ColumnTrait + Iterable,
    {
        Self::exists_with_conditions(conditions, db).await
    }

    /// Starts a relation scoped by equality conditions.
    fn r#where(conditions: HashMap<String, Value>) -> Relation<Self> {
        Relation::new().r#where(conditions)
    }

    /// Starts a relation scoped by ordering.
    fn order(column: &str, dir: OrderDirection) -> Relation<Self> {
        Relation::new().order(column, dir)
    }

    /// Starts a relation scoped by a limit.
    fn limit(n: u64) -> Relation<Self> {
        Relation::new().limit(n)
    }

    /// Starts a relation scoped by an offset.
    fn offset(n: u64) -> Relation<Self> {
        Relation::new().offset(n)
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use sea_orm::{ActiveModelTrait, ActiveValue::Set};
    use serde_json::json;

    use super::{AsyncQuerying, OrderDirection};
    use crate::{
        RecordError,
        base::test_support::{TestUser, seed_users, setup_db, test_user},
    };

    #[tokio::test]
    async fn find_returns_matching_record() {
        let db = setup_db().await;
        seed_users(&db).await;

        let user = TestUser::find(2, &db)
            .await
            .expect("find should return a record");

        assert_eq!(user.name, "Bob");
    }

    #[tokio::test]
    async fn find_missing_returns_not_found() {
        let db = setup_db().await;

        let error = TestUser::find(404, &db)
            .await
            .expect_err("missing row should return an error");

        assert!(matches!(error, RecordError::NotFound));
    }

    #[tokio::test]
    async fn find_by_id_returns_none_when_missing() {
        let db = setup_db().await;

        let user = TestUser::find_by_id(404, &db)
            .await
            .expect("query should succeed");

        assert!(user.is_none());
    }

    #[tokio::test]
    async fn find_by_filters_by_conditions() {
        let db = setup_db().await;
        seed_users(&db).await;

        let mut conditions = HashMap::new();
        conditions.insert("email".to_owned(), json!("carol@example.com"));

        let user = TestUser::find_by(conditions, &db)
            .await
            .expect("query should succeed")
            .expect("row should exist");

        assert_eq!(user.name, "Carol");
    }

    #[tokio::test]
    async fn all_returns_every_row() {
        let db = setup_db().await;
        seed_users(&db).await;

        let users = TestUser::all(&db).await.expect("all should succeed");

        assert_eq!(users.len(), 3);
    }

    #[tokio::test]
    async fn first_returns_lowest_primary_key() {
        let db = setup_db().await;
        seed_users(&db).await;

        let user = TestUser::first(&db)
            .await
            .expect("query should succeed")
            .expect("row should exist");

        assert_eq!(user.name, "Alice");
    }

    #[tokio::test]
    async fn last_returns_highest_primary_key() {
        let db = setup_db().await;
        seed_users(&db).await;

        let user = TestUser::last(&db)
            .await
            .expect("query should succeed")
            .expect("row should exist");

        assert_eq!(user.name, "Carol");
    }

    #[tokio::test]
    async fn count_returns_row_total() {
        let db = setup_db().await;
        seed_users(&db).await;

        assert_eq!(TestUser::count(&db).await.expect("count should succeed"), 3);
    }

    #[tokio::test]
    async fn exists_returns_true_for_matching_conditions() {
        let db = setup_db().await;
        seed_users(&db).await;

        let mut conditions = HashMap::new();
        conditions.insert("name".to_owned(), json!("Bob"));

        assert!(
            TestUser::exists(conditions, &db)
                .await
                .expect("exists should succeed")
        );
    }

    #[tokio::test]
    async fn exists_returns_false_for_missing_conditions() {
        let db = setup_db().await;
        seed_users(&db).await;

        let mut conditions = HashMap::new();
        conditions.insert("name".to_owned(), json!("Nobody"));

        assert!(
            !TestUser::exists(conditions, &db)
                .await
                .expect("exists should succeed")
        );
    }

    #[tokio::test]
    async fn query_builder_helpers_start_relations() {
        let db = setup_db().await;
        seed_users(&db).await;

        let ordered = TestUser::order("id", OrderDirection::Desc)
            .first(&db)
            .await
            .expect("query should succeed")
            .expect("row should exist");
        let limited = TestUser::limit(2)
            .load(&db)
            .await
            .expect("limit should load");
        let offset = TestUser::offset(2)
            .order("id", OrderDirection::Asc)
            .load(&db)
            .await
            .expect("offset should load");

        assert_eq!(ordered.name, "Carol");
        assert_eq!(limited.len(), 2);
        assert_eq!(offset.len(), 1);
        assert_eq!(offset[0].name, "Carol");
    }
    #[tokio::test]
    async fn find_by_id_returns_matching_record_when_present() {
        let db = setup_db().await;
        seed_users(&db).await;

        let user = TestUser::find_by_id(3, &db)
            .await
            .expect("query should succeed")
            .expect("row should exist");

        assert_eq!(user.name, "Carol");
    }

    #[tokio::test]
    async fn find_by_returns_none_when_no_match_exists() {
        let db = setup_db().await;
        seed_users(&db).await;

        let user = TestUser::find_by(
            HashMap::from([("email".to_owned(), json!("missing@example.com"))]),
            &db,
        )
        .await
        .expect("query should succeed");

        assert!(user.is_none());
    }

    #[tokio::test]
    async fn find_by_with_multiple_conditions_requires_all_matches() {
        let db = setup_db().await;
        seed_users(&db).await;

        let user = TestUser::find_by(
            HashMap::from([
                ("name".to_owned(), json!("Bob")),
                ("email".to_owned(), json!("bob@example.com")),
            ]),
            &db,
        )
        .await
        .expect("query should succeed")
        .expect("row should exist");

        assert_eq!(user.id, Some(2));
    }

    #[tokio::test]
    async fn all_marks_loaded_rows_persisted() {
        let db = setup_db().await;
        seed_users(&db).await;

        let users = TestUser::all(&db).await.expect("all should succeed");

        assert!(
            users
                .iter()
                .all(|user| user.state == crate::RecordState::Persisted)
        );
    }

    #[tokio::test]
    async fn first_returns_none_when_table_is_empty() {
        let db = setup_db().await;

        let user = TestUser::first(&db).await.expect("query should succeed");

        assert!(user.is_none());
    }

    #[tokio::test]
    async fn last_returns_none_when_table_is_empty() {
        let db = setup_db().await;

        let user = TestUser::last(&db).await.expect("query should succeed");

        assert!(user.is_none());
    }

    #[tokio::test]
    async fn count_returns_zero_when_table_is_empty() {
        let db = setup_db().await;

        assert_eq!(TestUser::count(&db).await.expect("count should succeed"), 0);
    }

    #[tokio::test]
    async fn exists_with_empty_conditions_is_false_without_rows() {
        let db = setup_db().await;

        assert!(
            !TestUser::exists(HashMap::new(), &db)
                .await
                .expect("exists should succeed")
        );
    }

    #[tokio::test]
    async fn exists_with_empty_conditions_is_true_when_rows_exist() {
        let db = setup_db().await;
        seed_users(&db).await;

        assert!(
            TestUser::exists(HashMap::new(), &db)
                .await
                .expect("exists should succeed")
        );
    }

    #[tokio::test]
    async fn exists_with_multiple_conditions_requires_all_conditions() {
        let db = setup_db().await;
        seed_users(&db).await;

        assert!(
            !TestUser::exists(
                HashMap::from([
                    ("name".to_owned(), json!("Bob")),
                    ("email".to_owned(), json!("alice@example.com")),
                ]),
                &db,
            )
            .await
            .expect("exists should succeed")
        );
    }

    #[tokio::test]
    async fn where_helper_loads_matching_rows() {
        let db = setup_db().await;
        seed_users(&db).await;

        let users = TestUser::r#where(HashMap::from([("name".to_owned(), json!("Bob"))]))
            .load(&db)
            .await
            .expect("where relation should load");

        assert_eq!(users.len(), 1);
        assert_eq!(users[0].email, "bob@example.com");
    }

    #[tokio::test]
    async fn where_helper_with_multiple_conditions_loads_match() {
        let db = setup_db().await;
        seed_users(&db).await;

        let users = TestUser::r#where(HashMap::from([
            ("name".to_owned(), json!("Carol")),
            ("email".to_owned(), json!("carol@example.com")),
        ]))
        .load(&db)
        .await
        .expect("where relation should load");

        assert_eq!(users.len(), 1);
        assert_eq!(users[0].id, Some(3));
    }

    #[tokio::test]
    async fn where_helper_with_empty_conditions_loads_all_rows() {
        let db = setup_db().await;
        seed_users(&db).await;

        let users = TestUser::r#where(HashMap::new())
            .load(&db)
            .await
            .expect("where relation should load");

        assert_eq!(users.len(), 3);
    }

    #[tokio::test]
    async fn order_helper_descending_returns_expected_sequence() {
        let db = setup_db().await;
        seed_users(&db).await;

        let users = TestUser::order("id", OrderDirection::Desc)
            .load(&db)
            .await
            .expect("ordered relation should load");
        let names = users.into_iter().map(|user| user.name).collect::<Vec<_>>();

        assert_eq!(names, vec!["Carol", "Bob", "Alice"]);
    }

    #[tokio::test]
    async fn order_helper_ascending_returns_expected_sequence() {
        let db = setup_db().await;
        seed_users(&db).await;

        let users = TestUser::order("id", OrderDirection::Asc)
            .load(&db)
            .await
            .expect("ordered relation should load");
        let names = users.into_iter().map(|user| user.name).collect::<Vec<_>>();

        assert_eq!(names, vec!["Alice", "Bob", "Carol"]);
    }

    #[tokio::test]
    async fn limit_helper_zero_returns_no_rows() {
        let db = setup_db().await;
        seed_users(&db).await;

        let users = TestUser::limit(0)
            .load(&db)
            .await
            .expect("limit should load");

        assert!(users.is_empty());
    }

    #[tokio::test]
    async fn limit_helper_can_chain_with_order() {
        let db = setup_db().await;
        seed_users(&db).await;

        let users = TestUser::limit(2)
            .order("id", OrderDirection::Desc)
            .load(&db)
            .await
            .expect("relation should load");
        let names = users.into_iter().map(|user| user.name).collect::<Vec<_>>();

        assert_eq!(names, vec!["Carol", "Bob"]);
    }

    #[tokio::test]
    async fn offset_helper_beyond_row_count_returns_empty() {
        let db = setup_db().await;
        seed_users(&db).await;

        let users = TestUser::offset(10)
            .order("id", OrderDirection::Asc)
            .load(&db)
            .await
            .expect("relation should load");

        assert!(users.is_empty());
    }

    #[tokio::test]
    async fn offset_helper_can_chain_with_limit_and_order() {
        let db = setup_db().await;
        seed_users(&db).await;

        let users = TestUser::offset(1)
            .order("id", OrderDirection::Asc)
            .limit(1)
            .load(&db)
            .await
            .expect("relation should load");

        assert_eq!(users.len(), 1);
        assert_eq!(users[0].name, "Bob");
    }

    #[tokio::test]
    async fn find_by_id_returns_persisted_record() {
        let db = setup_db().await;
        seed_users(&db).await;

        let user = TestUser::find_by_id(1, &db)
            .await
            .expect("query should succeed")
            .expect("row should exist");

        assert_eq!(user.state, crate::RecordState::Persisted);
    }

    #[tokio::test]
    async fn find_by_bang_returns_matching_record() {
        let db = setup_db().await;
        seed_users(&db).await;

        let user = TestUser::find_by_bang(
            HashMap::from([("email".to_owned(), json!("bob@example.com"))]),
            &db,
        )
        .await
        .expect("row should exist");

        assert_eq!(user.name, "Bob");
    }

    #[tokio::test]
    async fn find_by_bang_returns_not_found_when_missing() {
        let db = setup_db().await;

        let error = TestUser::find_by_bang(
            HashMap::from([("email".to_owned(), json!("missing@example.com"))]),
            &db,
        )
        .await
        .expect_err("missing row should return an error");

        assert!(matches!(error, RecordError::NotFound));
    }

    #[tokio::test]
    async fn take_returns_first_available_row_without_order() {
        let db = setup_db().await;
        seed_users(&db).await;

        let user = TestUser::take(&db)
            .await
            .expect("take should succeed")
            .expect("row should exist");

        assert_eq!(user.id, Some(1));
    }

    #[tokio::test]
    async fn take_returns_none_when_table_is_empty() {
        let db = setup_db().await;

        let user = TestUser::take(&db).await.expect("take should succeed");

        assert!(user.is_none());
    }

    #[tokio::test]
    async fn take_bang_returns_not_found_when_table_is_empty() {
        let db = setup_db().await;

        let error = TestUser::take_bang(&db)
            .await
            .expect_err("empty tables should fail");

        assert!(matches!(error, RecordError::NotFound));
    }

    #[tokio::test]
    async fn sole_returns_only_row_when_table_has_one_record() {
        let db = setup_db().await;
        test_user::ActiveModel {
            name: Set("Solo".to_owned()),
            email: Set("solo@example.com".to_owned()),
            ..Default::default()
        }
        .insert(&db)
        .await
        .expect("fixture insert should succeed");

        let user = TestUser::sole(&db)
            .await
            .expect("sole should return the row");

        assert_eq!(user.name, "Solo");
    }

    #[tokio::test]
    async fn sole_returns_not_found_when_table_is_empty() {
        let db = setup_db().await;

        let error = TestUser::sole(&db)
            .await
            .expect_err("empty tables should fail");

        assert!(matches!(error, RecordError::NotFound));
    }

    #[tokio::test]
    async fn sole_returns_exceeded_when_multiple_rows_match() {
        let db = setup_db().await;
        seed_users(&db).await;

        let error = TestUser::sole(&db)
            .await
            .expect_err("multiple rows should fail sole");

        assert!(matches!(error, RecordError::SoleRecordExceeded));
    }

    #[tokio::test]
    async fn find_sole_by_returns_matching_row() {
        let db = setup_db().await;
        seed_users(&db).await;

        let user =
            TestUser::find_sole_by(HashMap::from([("name".to_owned(), json!("Alice"))]), &db)
                .await
                .expect("single matching row should exist");

        assert_eq!(user.email, "alice@example.com");
    }

    #[tokio::test]
    async fn find_sole_by_returns_exceeded_when_multiple_rows_match() {
        let db = setup_db().await;
        seed_users(&db).await;
        test_user::ActiveModel {
            name: Set("Bob".to_owned()),
            email: Set("bobby@example.com".to_owned()),
            ..Default::default()
        }
        .insert(&db)
        .await
        .expect("fixture insert should succeed");

        let error = TestUser::find_sole_by(HashMap::from([("name".to_owned(), json!("Bob"))]), &db)
            .await
            .expect_err("multiple matches should fail");

        assert!(matches!(error, RecordError::SoleRecordExceeded));
    }

    #[tokio::test]
    async fn pluck_returns_requested_column_values() {
        let db = setup_db().await;
        seed_users(&db).await;

        let values = TestUser::pluck("name", &db)
            .await
            .expect("pluck should succeed");

        assert_eq!(values, vec![json!("Alice"), json!("Bob"), json!("Carol")]);
    }

    #[tokio::test]
    async fn pick_returns_first_requested_column_value() {
        let db = setup_db().await;
        seed_users(&db).await;

        let value = TestUser::pick("email", &db)
            .await
            .expect("pick should succeed");

        assert_eq!(value, Some(json!("alice@example.com")));
    }

    #[tokio::test]
    async fn ids_returns_primary_keys() {
        let db = setup_db().await;
        seed_users(&db).await;

        let ids = TestUser::ids(&db).await.expect("ids should succeed");

        assert_eq!(ids, vec![1, 2, 3]);
    }

    #[tokio::test]
    async fn exists_with_conditions_matches_rows() {
        let db = setup_db().await;
        seed_users(&db).await;

        assert!(
            TestUser::exists_with_conditions(
                HashMap::from([("name".to_owned(), json!("Carol"))]),
                &db,
            )
            .await
            .expect("exists_with_conditions should succeed")
        );
    }

    #[tokio::test]
    async fn exists_with_conditions_returns_false_when_missing() {
        let db = setup_db().await;

        assert!(
            !TestUser::exists_with_conditions(
                HashMap::from([("name".to_owned(), json!("Nobody"))]),
                &db,
            )
            .await
            .expect("exists_with_conditions should succeed")
        );
    }
}