From acdee4574411bed47dda123cf4397a9920f66e73 Mon Sep 17 00:00:00 2001
From: Chris Tsang <tyt2y7@gmail.com>
Date: Mon, 17 May 2021 23:28:58 +0800
Subject: [PATCH] Trait bound RelationBuilder

---
 src/entity/base.rs     | 10 +++++-----
 src/entity/relation.rs | 39 ++++++++++++++++++++++-----------------
 2 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/src/entity/base.rs b/src/entity/base.rs
index 3d690f57..c1d1dc90 100644
--- a/src/entity/base.rs
+++ b/src/entity/base.rs
@@ -2,7 +2,7 @@ use crate::{
     ColumnTrait, ModelTrait, PrimaryKeyOfModel, PrimaryKeyTrait, QueryHelper, RelationBuilder,
     RelationTrait, RelationType, Select,
 };
-use sea_query::{Iden, IntoIden, Value};
+use sea_query::{Iden, Value};
 use std::fmt::Debug;
 pub use strum::IntoEnumIterator as Iterable;
 
@@ -25,16 +25,16 @@ pub trait EntityTrait: EntityName {
         true
     }
 
-    fn has_one<E>(entity: E) -> RelationBuilder
+    fn has_one<R>(entity: R) -> RelationBuilder<Self, R>
     where
-        E: IntoIden,
+        R: EntityTrait,
     {
         RelationBuilder::new(RelationType::HasOne, Self::default(), entity)
     }
 
-    fn has_many<E>(entity: E) -> RelationBuilder
+    fn has_many<R>(entity: R) -> RelationBuilder<Self, R>
     where
-        E: IntoIden,
+        R: EntityTrait,
     {
         RelationBuilder::new(RelationType::HasMany, Self::default(), entity)
     }
diff --git a/src/entity/relation.rs b/src/entity/relation.rs
index a1529c2c..0694a88b 100644
--- a/src/entity/relation.rs
+++ b/src/entity/relation.rs
@@ -1,4 +1,5 @@
 use crate::{EntityTrait, Identity, IntoIdentity, Select};
+use core::marker::PhantomData;
 use sea_query::{Iden, IntoIden, JoinType};
 use std::fmt::Debug;
 use std::rc::Rc;
@@ -36,7 +37,12 @@ pub struct RelationDef {
     pub to_col: Identity,
 }
 
-pub struct RelationBuilder {
+pub struct RelationBuilder<E, R>
+where
+    E: EntityTrait,
+    R: EntityTrait,
+{
+    entities: PhantomData<(E, R)>,
     rel_type: RelationType,
     from_tbl: Rc<dyn Iden>,
     to_tbl: Rc<dyn Iden>,
@@ -44,13 +50,14 @@ pub struct RelationBuilder {
     to_col: Option<Identity>,
 }
 
-impl RelationBuilder {
-    pub(crate) fn new<E, T>(rel_type: RelationType, from: E, to: T) -> Self
-    where
-        E: IntoIden,
-        T: IntoIden,
-    {
+impl<E, R> RelationBuilder<E, R>
+where
+    E: EntityTrait,
+    R: EntityTrait,
+{
+    pub(crate) fn new(rel_type: RelationType, from: E, to: R) -> Self {
         Self {
+            entities: PhantomData,
             rel_type,
             from_tbl: from.into_iden(),
             to_tbl: to.into_iden(),
@@ -59,25 +66,23 @@ impl RelationBuilder {
         }
     }
 
-    pub fn from<T>(mut self, identifier: T) -> Self
-    where
-        T: IntoIdentity,
-    {
+    pub fn from(mut self, identifier: E::Column) -> Self {
         self.from_col = Some(identifier.into_identity());
         self
     }
 
-    pub fn to<T>(mut self, identifier: T) -> Self
-    where
-        T: IntoIdentity,
-    {
+    pub fn to(mut self, identifier: R::Column) -> Self {
         self.to_col = Some(identifier.into_identity());
         self
     }
 }
 
-impl From<RelationBuilder> for RelationDef {
-    fn from(b: RelationBuilder) -> Self {
+impl<E, R> From<RelationBuilder<E, R>> for RelationDef
+where
+    E: EntityTrait,
+    R: EntityTrait,
+{
+    fn from(b: RelationBuilder<E, R>) -> Self {
         RelationDef {
             rel_type: b.rel_type,
             from_tbl: b.from_tbl,