Skip to main content

sea_orm/entity/
identity.rs

1use crate::{ColumnTrait, EntityTrait, IdenStatic};
2use sea_query::{Alias, DynIden, Iden, IntoIden, SeaRc};
3use std::{borrow::Cow, fmt::Write};
4
5/// A one-or-many column identifier — the abstraction SeaORM uses to refer
6/// to either a single column or a composite (e.g. a composite primary key
7/// or foreign key).
8///
9/// Specialized variants exist for the common arities (1/2/3 columns) to
10/// avoid heap allocation; longer composites fall back to [`Identity::Many`].
11#[derive(Debug, Clone, PartialEq, Eq, Hash)]
12pub enum Identity {
13    /// A single column.
14    Unary(DynIden),
15    /// Two columns (composite).
16    Binary(DynIden, DynIden),
17    /// Three columns (composite).
18    Ternary(DynIden, DynIden, DynIden),
19    /// Four or more columns (composite).
20    Many(Vec<DynIden>),
21}
22
23impl Identity {
24    /// Number of columns in this identifier.
25    pub fn arity(&self) -> usize {
26        match self {
27            Self::Unary(_) => 1,
28            Self::Binary(_, _) => 2,
29            Self::Ternary(_, _, _) => 3,
30            Self::Many(vec) => vec.len(),
31        }
32    }
33
34    /// Iterate over each column iden in order.
35    pub fn iter(&self) -> BorrowedIdentityIter<'_> {
36        BorrowedIdentityIter {
37            identity: self,
38            index: 0,
39        }
40    }
41
42    /// `true` if `col` is one of the columns making up this identifier.
43    pub fn contains(&self, col: &DynIden) -> bool {
44        self.iter().any(|c| c == col)
45    }
46
47    /// `true` if every column of `other` also appears in `self`.
48    pub fn fully_contains(&self, other: &Identity) -> bool {
49        for col in other.iter() {
50            if !self.contains(col) {
51                return false;
52            }
53        }
54        true
55    }
56}
57
58impl IntoIterator for Identity {
59    type Item = DynIden;
60    type IntoIter = OwnedIdentityIter;
61
62    fn into_iter(self) -> Self::IntoIter {
63        OwnedIdentityIter {
64            identity: self,
65            index: 0,
66        }
67    }
68}
69
70impl Iden for Identity {
71    fn quoted(&self) -> Cow<'static, str> {
72        match self {
73            Identity::Unary(iden) => iden.inner(),
74            Identity::Binary(iden1, iden2) => Cow::Owned(format!("{iden1}{iden2}")),
75            Identity::Ternary(iden1, iden2, iden3) => Cow::Owned(format!("{iden1}{iden2}{iden3}")),
76            Identity::Many(vec) => {
77                let mut s = String::new();
78                for iden in vec.iter() {
79                    write!(&mut s, "{iden}").expect("Infallible");
80                }
81                Cow::Owned(s)
82            }
83        }
84    }
85
86    fn to_string(&self) -> String {
87        match self.quoted() {
88            Cow::Borrowed(s) => s.to_owned(),
89            Cow::Owned(s) => s,
90        }
91    }
92
93    fn unquoted(&self) -> &str {
94        panic!("Should not call this")
95    }
96}
97
98/// Borrowing iterator over the columns of an [`Identity`].
99#[derive(Debug)]
100pub struct BorrowedIdentityIter<'a> {
101    identity: &'a Identity,
102    index: usize,
103}
104
105/// Owning iterator over the columns of an [`Identity`]
106/// (returned by `IntoIterator for Identity`).
107#[derive(Debug)]
108pub struct OwnedIdentityIter {
109    identity: Identity,
110    index: usize,
111}
112
113impl<'a> Iterator for BorrowedIdentityIter<'a> {
114    type Item = &'a DynIden;
115
116    fn next(&mut self) -> Option<Self::Item> {
117        let result = match self.identity {
118            Identity::Unary(iden1) => {
119                if self.index == 0 {
120                    Some(iden1)
121                } else {
122                    None
123                }
124            }
125            Identity::Binary(iden1, iden2) => match self.index {
126                0 => Some(iden1),
127                1 => Some(iden2),
128                _ => None,
129            },
130            Identity::Ternary(iden1, iden2, iden3) => match self.index {
131                0 => Some(iden1),
132                1 => Some(iden2),
133                2 => Some(iden3),
134                _ => None,
135            },
136            Identity::Many(vec) => vec.get(self.index),
137        };
138        self.index += 1;
139        result
140    }
141}
142
143impl Iterator for OwnedIdentityIter {
144    type Item = DynIden;
145
146    fn next(&mut self) -> Option<Self::Item> {
147        let result = match &self.identity {
148            Identity::Unary(iden1) => {
149                if self.index == 0 {
150                    Some(iden1.clone())
151                } else {
152                    None
153                }
154            }
155            Identity::Binary(iden1, iden2) => match self.index {
156                0 => Some(iden1.clone()),
157                1 => Some(iden2.clone()),
158                _ => None,
159            },
160            Identity::Ternary(iden1, iden2, iden3) => match self.index {
161                0 => Some(iden1.clone()),
162                1 => Some(iden2.clone()),
163                2 => Some(iden3.clone()),
164                _ => None,
165            },
166            Identity::Many(vec) => vec.get(self.index).cloned(),
167        };
168        self.index += 1;
169        result
170    }
171}
172
173/// Conversion into an [`Identity`]. Implemented for `&str`/`String`, a single
174/// column iden, and tuples of column idens (for composites up to 12 columns).
175pub trait IntoIdentity {
176    /// Build the [`Identity`].
177    fn into_identity(self) -> Identity;
178}
179
180/// Conversion into an [`Identity`] whose columns are guaranteed to belong to
181/// entity `E`. Used by [`RelationBuilder::from`](crate::RelationBuilder::from) /
182/// [`to`](crate::RelationBuilder::to) so the type system enforces that
183/// relation columns reference the right table.
184pub trait IdentityOf<E>
185where
186    E: EntityTrait,
187{
188    /// Build the [`Identity`].
189    fn identity_of(self) -> Identity;
190}
191
192impl IntoIdentity for Identity {
193    fn into_identity(self) -> Identity {
194        self
195    }
196}
197
198impl IntoIdentity for String {
199    fn into_identity(self) -> Identity {
200        self.as_str().into_identity()
201    }
202}
203
204impl IntoIdentity for &str {
205    fn into_identity(self) -> Identity {
206        Identity::Unary(SeaRc::new(Alias::new(self)))
207    }
208}
209
210impl<T> IntoIdentity for T
211where
212    T: IdenStatic,
213{
214    fn into_identity(self) -> Identity {
215        Identity::Unary(self.into_iden())
216    }
217}
218
219impl<T, C> IntoIdentity for (T, C)
220where
221    T: IdenStatic,
222    C: IdenStatic,
223{
224    fn into_identity(self) -> Identity {
225        Identity::Binary(self.0.into_iden(), self.1.into_iden())
226    }
227}
228
229impl<T, C, R> IntoIdentity for (T, C, R)
230where
231    T: IdenStatic,
232    C: IdenStatic,
233    R: IdenStatic,
234{
235    fn into_identity(self) -> Identity {
236        Identity::Ternary(self.0.into_iden(), self.1.into_iden(), self.2.into_iden())
237    }
238}
239
240macro_rules! impl_into_identity {
241    ( $($T:ident : $N:tt),+ $(,)? ) => {
242        impl< $($T),+ > IntoIdentity for ( $($T),+ )
243        where
244            $($T: IdenStatic),+
245        {
246            fn into_identity(self) -> Identity {
247                Identity::Many(vec![
248                    $(self.$N.into_iden()),+
249                ])
250            }
251        }
252    };
253}
254
255#[rustfmt::skip]
256mod impl_into_identity {
257    use super::*;
258
259    impl_into_identity!(T0:0, T1:1, T2:2, T3:3);
260    impl_into_identity!(T0:0, T1:1, T2:2, T3:3, T4:4);
261    impl_into_identity!(T0:0, T1:1, T2:2, T3:3, T4:4, T5:5);
262    impl_into_identity!(T0:0, T1:1, T2:2, T3:3, T4:4, T5:5, T6:6);
263    impl_into_identity!(T0:0, T1:1, T2:2, T3:3, T4:4, T5:5, T6:6, T7:7);
264    impl_into_identity!(T0:0, T1:1, T2:2, T3:3, T4:4, T5:5, T6:6, T7:7, T8:8);
265    impl_into_identity!(T0:0, T1:1, T2:2, T3:3, T4:4, T5:5, T6:6, T7:7, T8:8, T9:9);
266    impl_into_identity!(T0:0, T1:1, T2:2, T3:3, T4:4, T5:5, T6:6, T7:7, T8:8, T9:9, T10:10);
267    impl_into_identity!(T0:0, T1:1, T2:2, T3:3, T4:4, T5:5, T6:6, T7:7, T8:8, T9:9, T10:10, T11:11);
268}
269
270impl<E, C> IdentityOf<E> for C
271where
272    E: EntityTrait<Column = C>,
273    C: ColumnTrait,
274{
275    fn identity_of(self) -> Identity {
276        self.into_identity()
277    }
278}
279
280macro_rules! impl_identity_of {
281    ( $($T:ident),+ $(,)? ) => {
282        impl<E, C> IdentityOf<E> for ( $($T),+ )
283        where
284            E: EntityTrait<Column = C>,
285            C: ColumnTrait,
286        {
287            fn identity_of(self) -> Identity {
288                self.into_identity()
289            }
290        }
291    };
292}
293
294#[rustfmt::skip]
295mod impl_identity_of {
296    use super::*;
297
298    impl_identity_of!(C, C);
299    impl_identity_of!(C, C, C);
300    impl_identity_of!(C, C, C, C);
301    impl_identity_of!(C, C, C, C, C);
302    impl_identity_of!(C, C, C, C, C, C);
303    impl_identity_of!(C, C, C, C, C, C, C);
304    impl_identity_of!(C, C, C, C, C, C, C, C);
305    impl_identity_of!(C, C, C, C, C, C, C, C, C);
306    impl_identity_of!(C, C, C, C, C, C, C, C, C, C);
307    impl_identity_of!(C, C, C, C, C, C, C, C, C, C, C);
308    impl_identity_of!(C, C, C, C, C, C, C, C, C, C, C, C);
309}
310
311#[cfg(test)]
312mod test {
313    use super::*;
314
315    #[test]
316    fn test_identity_contains() {
317        let abc = Identity::Ternary("a".into(), "b".into(), "c".into());
318        let a = Identity::Unary("a".into());
319        let ab = Identity::Binary("a".into(), "b".into());
320        let bc = Identity::Binary("b".into(), "c".into());
321        let d = Identity::Unary("d".into());
322        let bcd = Identity::Ternary("b".into(), "c".into(), "d".into());
323
324        assert!(abc.contains(&"a".into()));
325        assert!(abc.contains(&"b".into()));
326        assert!(abc.contains(&"c".into()));
327        assert!(!abc.contains(&"d".into()));
328
329        assert!(abc.fully_contains(&a));
330        assert!(abc.fully_contains(&ab));
331        assert!(abc.fully_contains(&bc));
332        assert!(!abc.fully_contains(&d));
333        assert!(!abc.fully_contains(&bcd));
334    }
335}