From 85eadd38d222b07621ff5b5ca55b650295bbe682 Mon Sep 17 00:00:00 2001
From: Chris Tsang <tyt2y7@gmail.com>
Date: Tue, 28 Sep 2021 13:28:43 +0800
Subject: [PATCH] Change mock connector behavior to mock Postgres regardless of
 feature flags

---
 src/entity/base_entity.rs |  4 ++--
 src/executor/insert.rs    | 12 +++++-------
 src/executor/paginator.rs | 11 ++++-------
 3 files changed, 11 insertions(+), 16 deletions(-)

diff --git a/src/entity/base_entity.rs b/src/entity/base_entity.rs
index 7ba1e965..764f2524 100644
--- a/src/entity/base_entity.rs
+++ b/src/entity/base_entity.rs
@@ -293,7 +293,7 @@ pub trait EntityTrait: EntityName {
     /// assert_eq!(
     ///     db.into_transaction_log(),
     ///     vec![Transaction::from_sql_and_values(
-    ///         DbBackend::Postgres, r#"INSERT INTO "cake" ("name") VALUES ($1)"#, vec!["Apple Pie".into()]
+    ///         DbBackend::Postgres, r#"INSERT INTO "cake" ("name") VALUES ($1) RETURNING "id""#, vec!["Apple Pie".into()]
     ///     )]);
     /// ```
     fn insert<A>(model: A) -> Insert<A>
@@ -344,7 +344,7 @@ pub trait EntityTrait: EntityName {
     /// assert_eq!(
     ///     db.into_transaction_log(),
     ///     vec![Transaction::from_sql_and_values(
-    ///         DbBackend::Postgres, r#"INSERT INTO "cake" ("name") VALUES ($1), ($2)"#,
+    ///         DbBackend::Postgres, r#"INSERT INTO "cake" ("name") VALUES ($1), ($2) RETURNING "id""#,
     ///         vec!["Apple Pie".into(), "Orange Scone".into()]
     ///     )]);
     /// ```
diff --git a/src/executor/insert.rs b/src/executor/insert.rs
index 1f4936ba..d580f110 100644
--- a/src/executor/insert.rs
+++ b/src/executor/insert.rs
@@ -1,5 +1,5 @@
 use crate::{
-    error::*, ActiveModelTrait, DatabaseConnection, EntityTrait, Insert, PrimaryKeyTrait,
+    error::*, ActiveModelTrait, DatabaseConnection, DbBackend, EntityTrait, Insert, PrimaryKeyTrait,
     Statement, TryFromU64,
 };
 use sea_query::InsertStatement;
@@ -37,8 +37,7 @@ where
         // TODO: extract primary key's value from query
         // so that self is dropped before entering await
         let mut query = self.query;
-        #[cfg(feature = "sqlx-postgres")]
-        if let DatabaseConnection::SqlxPostgresPoolConnection(_) = db {
+        if db.get_database_backend() == DbBackend::Postgres {
             use crate::{sea_query::Query, Iterable};
             if <A::Entity as EntityTrait>::PrimaryKey::iter().count() > 0 {
                 query.returning(
@@ -86,14 +85,13 @@ where
 {
     type PrimaryKey<A> = <<A as ActiveModelTrait>::Entity as EntityTrait>::PrimaryKey;
     type ValueTypeOf<A> = <PrimaryKey<A> as PrimaryKeyTrait>::ValueType;
-    let last_insert_id = match db {
-        #[cfg(feature = "sqlx-postgres")]
-        DatabaseConnection::SqlxPostgresPoolConnection(conn) => {
+    let last_insert_id = match db.get_database_backend() {
+        DbBackend::Postgres => {
             use crate::{sea_query::Iden, Iterable};
             let cols = PrimaryKey::<A>::iter()
                 .map(|col| col.to_string())
                 .collect::<Vec<_>>();
-            let res = conn.query_one(statement).await?.unwrap();
+            let res = db.query_one(statement).await?.unwrap();
             res.try_get_many("", cols.as_ref()).unwrap_or_default()
         }
         _ => {
diff --git a/src/executor/paginator.rs b/src/executor/paginator.rs
index 24822111..608d9dc1 100644
--- a/src/executor/paginator.rs
+++ b/src/executor/paginator.rs
@@ -1,4 +1,4 @@
-use crate::{error::*, DatabaseConnection, SelectorTrait};
+use crate::{error::*, DatabaseConnection, DbBackend, SelectorTrait};
 use async_stream::stream;
 use futures::Stream;
 use sea_query::{Alias, Expr, SelectStatement};
@@ -63,11 +63,8 @@ where
             Some(res) => res,
             None => return Ok(0),
         };
-        let num_items = match self.db {
-            #[cfg(feature = "sqlx-postgres")]
-            DatabaseConnection::SqlxPostgresPoolConnection(_) => {
-                result.try_get::<i64>("", "num_items")? as usize
-            }
+        let num_items = match builder {
+            DbBackend::Postgres => result.try_get::<i64>("", "num_items")? as usize,
             _ => result.try_get::<i32>("", "num_items")? as usize,
         };
         Ok(num_items)
@@ -192,7 +189,7 @@ mod tests {
         (db, vec![page1, page2, page3])
     }
 
-    fn setup_num_items() -> (DatabaseConnection, i32) {
+    fn setup_num_items() -> (DatabaseConnection, i64) {
         let num_items = 3;
         let db = MockDatabase::new(DbBackend::Postgres)
             .append_query_results(vec![vec![maplit::btreemap! {