From 38ace07e83e5de5b2399fe9726ec0d264fdcc03f Mon Sep 17 00:00:00 2001
From: Billy Chan <ccw.billy.123@gmail.com>
Date: Sat, 14 Aug 2021 17:00:03 +0800
Subject: [PATCH] Schema prefix in table join

---
 src/entity/relation.rs              | 14 ++++++-------
 src/query/helper.rs                 | 18 ++++++++++++++---
 src/query/join.rs                   | 31 +++++++++++++++++++++++------
 src/tests_cfg/cake_filling_price.rs |  4 ++++
 4 files changed, 51 insertions(+), 16 deletions(-)

diff --git a/src/entity/relation.rs b/src/entity/relation.rs
index 7d2059e2..955660e2 100644
--- a/src/entity/relation.rs
+++ b/src/entity/relation.rs
@@ -1,6 +1,6 @@
 use crate::{EntityTrait, Identity, IdentityOf, Iterable, QuerySelect, Select};
 use core::marker::PhantomData;
-use sea_query::{DynIden, IntoIden, JoinType};
+use sea_query::{JoinType, TableRef};
 use std::fmt::Debug;
 
 #[derive(Clone, Debug)]
@@ -30,8 +30,8 @@ where
 
 pub struct RelationDef {
     pub rel_type: RelationType,
-    pub from_tbl: DynIden,
-    pub to_tbl: DynIden,
+    pub from_tbl: TableRef,
+    pub to_tbl: TableRef,
     pub from_col: Identity,
     pub to_col: Identity,
 }
@@ -43,8 +43,8 @@ where
 {
     entities: PhantomData<(E, R)>,
     rel_type: RelationType,
-    from_tbl: DynIden,
-    to_tbl: DynIden,
+    from_tbl: TableRef,
+    to_tbl: TableRef,
     from_col: Option<Identity>,
     to_col: Option<Identity>,
 }
@@ -71,8 +71,8 @@ where
         Self {
             entities: PhantomData,
             rel_type,
-            from_tbl: from.into_iden(),
-            to_tbl: to.into_iden(),
+            from_tbl: from.table_ref(),
+            to_tbl: to.table_ref(),
             from_col: None,
             to_col: None,
         }
diff --git a/src/query/helper.rs b/src/query/helper.rs
index 9691e1ab..6ade581a 100644
--- a/src/query/helper.rs
+++ b/src/query/helper.rs
@@ -2,7 +2,9 @@ use crate::{
     ColumnTrait, EntityTrait, Identity, IntoSimpleExpr, Iterable, ModelTrait, PrimaryKeyToColumn,
     RelationDef,
 };
-use sea_query::{Alias, Expr, IntoCondition, SeaRc, SelectExpr, SelectStatement, SimpleExpr};
+use sea_query::{
+    Alias, Expr, IntoCondition, SeaRc, SelectExpr, SelectStatement, SimpleExpr, TableRef,
+};
 pub use sea_query::{Condition, ConditionalStatement, DynIden, JoinType, Order, OrderedStatement};
 
 // LINT: when the column does not appear in tables selected from
@@ -269,8 +271,8 @@ pub trait QueryFilter: Sized {
 }
 
 fn join_condition(rel: RelationDef) -> SimpleExpr {
-    let from_tbl = rel.from_tbl.clone();
-    let to_tbl = rel.to_tbl.clone();
+    let from_tbl = unpack_table_ref(&rel.from_tbl);
+    let to_tbl = unpack_table_ref(&rel.to_tbl);
     let owner_keys = rel.from_col;
     let foreign_keys = rel.to_col;
 
@@ -292,3 +294,13 @@ fn join_condition(rel: RelationDef) -> SimpleExpr {
         _ => panic!("Owner key and foreign key mismatch"),
     }
 }
+
+fn unpack_table_ref(table_ref: &TableRef) -> DynIden {
+    match table_ref {
+        TableRef::Table(tbl) => SeaRc::clone(tbl),
+        TableRef::SchemaTable(_, tbl) => SeaRc::clone(tbl),
+        TableRef::TableAlias(tbl, _) => SeaRc::clone(tbl),
+        TableRef::SchemaTableAlias(_, tbl, _) => SeaRc::clone(tbl),
+        TableRef::SubQuery(_, tbl) => SeaRc::clone(tbl),
+    }
+}
diff --git a/src/query/join.rs b/src/query/join.rs
index d6e00301..72726d14 100644
--- a/src/query/join.rs
+++ b/src/query/join.rs
@@ -190,13 +190,32 @@ mod tests {
         let find_cake_filling_price: Select<cake_filling_price::Entity> =
             cake_filling::Entity::find_related();
         assert_eq!(
-            find_cake_filling_price.build(DbBackend::MySql).to_string(),
+            find_cake_filling_price.build(DbBackend::Postgres).to_string(),
             [
-                "SELECT `cake_filling_price`.`cake_id`, `cake_filling_price`.`filling_id`, `cake_filling_price`.`price`",
-                "FROM `cake_filling_price`",
-                "INNER JOIN `cake_filling` ON",
-                "(`cake_filling`.`cake_id` = `cake_filling_price`.`cake_id`) AND",
-                "(`cake_filling`.`filling_id` = `cake_filling_price`.`filling_id`)",
+                r#"SELECT "cake_filling_price"."cake_id", "cake_filling_price"."filling_id", "cake_filling_price"."price""#,
+                r#"FROM "public"."cake_filling_price""#,
+                r#"INNER JOIN "cake_filling" ON"#,
+                r#"("cake_filling"."cake_id" = "cake_filling_price"."cake_id") AND"#,
+                r#"("cake_filling"."filling_id" = "cake_filling_price"."filling_id")"#,
+            ]
+            .join(" ")
+        );
+    }
+
+    #[test]
+    fn join_9() {
+        use crate::{Related, Select};
+
+        let find_cake_filling: Select<cake_filling::Entity> =
+            cake_filling_price::Entity::find_related();
+        assert_eq!(
+            find_cake_filling.build(DbBackend::Postgres).to_string(),
+            [
+                r#"SELECT "cake_filling"."cake_id", "cake_filling"."filling_id""#,
+                r#"FROM "cake_filling""#,
+                r#"INNER JOIN "public"."cake_filling_price" ON"#,
+                r#"("cake_filling_price"."cake_id" = "cake_filling"."cake_id") AND"#,
+                r#"("cake_filling_price"."filling_id" = "cake_filling"."filling_id")"#,
             ]
             .join(" ")
         );
diff --git a/src/tests_cfg/cake_filling_price.rs b/src/tests_cfg/cake_filling_price.rs
index eae07871..c0bcbea1 100644
--- a/src/tests_cfg/cake_filling_price.rs
+++ b/src/tests_cfg/cake_filling_price.rs
@@ -5,6 +5,10 @@ use crate::entity::prelude::*;
 pub struct Entity;
 
 impl EntityName for Entity {
+    fn schema_name(&self) -> Option<&str> {
+        Some("public")
+    }
+
     fn table_name(&self) -> &str {
         "cake_filling_price"
     }