From e76cbb9fe19be1d89e2c4d540785d71deaf1484f Mon Sep 17 00:00:00 2001
From: Billy Chan <ccw.billy.123@gmail.com>
Date: Sun, 16 Oct 2022 19:02:48 +0800
Subject: [PATCH] Add `into_model` & `into_json` for `Cursor` (#1112)

---
 src/executor/cursor.rs | 29 +++++++++++++++++++++++++++++
 tests/cursor_tests.rs  | 36 +++++++++++++++++++++++++++++++++++-
 2 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/src/executor/cursor.rs b/src/executor/cursor.rs
index 264b7dc6..e690747c 100644
--- a/src/executor/cursor.rs
+++ b/src/executor/cursor.rs
@@ -8,6 +8,9 @@ use sea_query::{
 };
 use std::marker::PhantomData;
 
+#[cfg(feature = "with-json")]
+use crate::JsonValue;
+
 /// Cursor pagination
 #[derive(Debug, Clone)]
 pub struct Cursor<S>
@@ -140,6 +143,32 @@ where
         }
         Ok(buffer)
     }
+
+    /// Construct a [Cursor] that fetch any custom struct
+    pub fn into_model<M>(self) -> Cursor<SelectModel<M>>
+    where
+        M: FromQueryResult,
+    {
+        Cursor {
+            query: self.query,
+            table: self.table,
+            order_columns: self.order_columns,
+            last: self.last,
+            phantom: PhantomData,
+        }
+    }
+
+    /// Construct a [Cursor] that fetch JSON value
+    #[cfg(feature = "with-json")]
+    pub fn into_json(self) -> Cursor<SelectModel<JsonValue>> {
+        Cursor {
+            query: self.query,
+            table: self.table,
+            order_columns: self.order_columns,
+            last: self.last,
+            phantom: PhantomData,
+        }
+    }
 }
 
 impl<S> QueryOrder for Cursor<S>
diff --git a/tests/cursor_tests.rs b/tests/cursor_tests.rs
index 19790205..e6feef9f 100644
--- a/tests/cursor_tests.rs
+++ b/tests/cursor_tests.rs
@@ -2,7 +2,8 @@ pub mod common;
 
 pub use common::{features::*, setup::*, TestContext};
 use pretty_assertions::assert_eq;
-use sea_orm::entity::prelude::*;
+use sea_orm::{entity::prelude::*, FromQueryResult};
+use serde_json::json;
 
 #[sea_orm_macros::test]
 #[cfg(any(
@@ -199,5 +200,38 @@ pub async fn cursor_pagination(db: &DatabaseConnection) -> Result<(), DbErr> {
         vec![Model { id: 6 }, Model { id: 7 }]
     );
 
+    // Fetch custom struct
+
+    #[derive(FromQueryResult, Debug, PartialEq)]
+    struct Row {
+        id: i32,
+    }
+
+    let mut cursor = cursor.into_model::<Row>();
+
+    assert_eq!(
+        cursor.first(2).all(db).await?,
+        vec![Row { id: 6 }, Row { id: 7 }]
+    );
+
+    assert_eq!(
+        cursor.first(3).all(db).await?,
+        vec![Row { id: 6 }, Row { id: 7 }]
+    );
+
+    // Fetch JSON value
+
+    let mut cursor = cursor.into_json();
+
+    assert_eq!(
+        cursor.first(2).all(db).await?,
+        vec![json!({ "id": 6 }), json!({ "id": 7 })]
+    );
+
+    assert_eq!(
+        cursor.first(3).all(db).await?,
+        vec![json!({ "id": 6 }), json!({ "id": 7 })]
+    );
+
     Ok(())
 }