Skip to main content

sea_orm/entity/
active_enum.rs

1use crate::{ColIdx, ColumnDef, DbErr, Iterable, QueryResult, TryFromU64, TryGetError, TryGetable};
2use sea_query::{DynIden, Expr, ExprTrait, Nullable, SimpleExpr, Value, ValueType};
3
4/// A Rust representation of enum defined in database.
5///
6/// # Implementations
7///
8/// You can implement [ActiveEnum] manually by hand or use the derive macro [DeriveActiveEnum](sea_orm_macros::DeriveActiveEnum).
9///
10/// # Examples
11///
12/// Implementing it manually versus using the derive macro [DeriveActiveEnum](sea_orm_macros::DeriveActiveEnum).
13///
14/// > See [DeriveActiveEnum](sea_orm_macros::DeriveActiveEnum) for the full specification of macro attributes.
15///
16/// ```rust
17/// use sea_orm::entity::prelude::*;
18///
19/// // Using the derive macro
20/// #[derive(Debug, PartialEq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
21/// #[sea_orm(
22///     rs_type = "String",
23///     db_type = "String(StringLen::N(1))",
24///     enum_name = "category"
25/// )]
26/// pub enum DeriveCategory {
27///     #[sea_orm(string_value = "B")]
28///     Big,
29///     #[sea_orm(string_value = "S")]
30///     Small,
31/// }
32///
33/// // Implementing it manually
34/// #[derive(Debug, PartialEq, EnumIter)]
35/// pub enum Category {
36///     Big,
37///     Small,
38/// }
39///
40/// #[derive(Debug, DeriveIden)]
41/// pub struct CategoryEnum;
42///
43/// impl ActiveEnum for Category {
44///     // The macro attribute `rs_type` is being pasted here
45///     type Value = String;
46///
47///     type ValueVec = Vec<String>;
48///
49///     // Will be atomically generated by `DeriveActiveEnum`
50///     fn name() -> DynIden {
51///         SeaRc::new(CategoryEnum)
52///     }
53///
54///     // Will be atomically generated by `DeriveActiveEnum`
55///     fn to_value(&self) -> Self::Value {
56///         match self {
57///             Self::Big => "B",
58///             Self::Small => "S",
59///         }
60///         .to_owned()
61///     }
62///
63///     // Will be atomically generated by `DeriveActiveEnum`
64///     fn try_from_value(v: &Self::Value) -> Result<Self, DbErr> {
65///         match v.as_ref() {
66///             "B" => Ok(Self::Big),
67///             "S" => Ok(Self::Small),
68///             _ => Err(DbErr::Type(format!(
69///                 "unexpected value for Category enum: {}",
70///                 v
71///             ))),
72///         }
73///     }
74///
75///     fn db_type() -> ColumnDef {
76///         // The macro attribute `db_type` is being pasted here
77///         ColumnType::String(StringLen::N(1)).def()
78///     }
79/// }
80/// ```
81///
82/// Using [ActiveEnum] on Model.
83///
84/// ```
85/// use sea_orm::entity::prelude::*;
86///
87/// // Define the `Category` active enum
88/// #[derive(Debug, Clone, PartialEq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
89/// #[sea_orm(rs_type = "String", db_type = "String(StringLen::N(1))")]
90/// pub enum Category {
91///     #[sea_orm(string_value = "B")]
92///     Big,
93///     #[sea_orm(string_value = "S")]
94///     Small,
95/// }
96///
97/// #[derive(Clone, Debug, PartialEq, DeriveEntityModel)]
98/// #[sea_orm(table_name = "active_enum")]
99/// pub struct Model {
100///     #[sea_orm(primary_key)]
101///     pub id: i32,
102///     // Represents a db column using `Category` active enum
103///     pub category: Category,
104///     pub category_opt: Option<Category>,
105/// }
106///
107/// #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
108/// pub enum Relation {}
109///
110/// impl ActiveModelBehavior for ActiveModel {}
111/// ```
112pub trait ActiveEnum: Sized + Iterable {
113    /// Define the Rust type that each enum variant corresponds.
114    type Value: ActiveEnumValue;
115
116    /// This has no purpose. It will be removed in the next major version.
117    type ValueVec;
118
119    /// Get the name of enum
120    fn name() -> DynIden;
121
122    /// Convert enum variant into the corresponding value.
123    fn to_value(&self) -> Self::Value;
124
125    /// Try to convert the corresponding value into enum variant.
126    fn try_from_value(v: &Self::Value) -> Result<Self, DbErr>;
127
128    /// Get the database column definition of this active enum.
129    fn db_type() -> ColumnDef;
130
131    /// Convert an owned enum variant into the corresponding value.
132    fn into_value(self) -> Self::Value {
133        Self::to_value(&self)
134    }
135
136    /// Construct a enum expression with casting
137    fn as_enum(&self) -> SimpleExpr {
138        let value: Value = Self::to_value(self).into();
139        match value {
140            Value::Enum(_) => Expr::val(value),
141            _ => Expr::val(value).as_enum(Self::name()),
142        }
143    }
144
145    /// Get the name of all enum variants
146    fn values() -> Vec<Self::Value> {
147        Self::iter().map(Self::into_value).collect()
148    }
149}
150
151/// The Rust Value backing ActiveEnums
152pub trait ActiveEnumValue: Into<Value> + ValueType + Nullable + TryGetable {
153    /// For getting an array of enum. Postgres only
154    fn try_get_vec_by<I: ColIdx>(res: &QueryResult, index: I) -> Result<Vec<Self>, TryGetError>;
155}
156
157macro_rules! impl_active_enum_value {
158    ($type:ident) => {
159        impl ActiveEnumValue for $type {
160            fn try_get_vec_by<I: ColIdx>(
161                _res: &QueryResult,
162                _index: I,
163            ) -> Result<Vec<Self>, TryGetError> {
164                Err(TryGetError::DbErr(DbErr::BackendNotSupported {
165                    db: "Postgres",
166                    ctx: "ActiveEnumValue::try_get_vec_by",
167                }))
168            }
169        }
170    };
171}
172
173macro_rules! impl_active_enum_value_with_pg_array {
174    ($type:ident) => {
175        impl ActiveEnumValue for $type {
176            fn try_get_vec_by<I: ColIdx>(
177                _res: &QueryResult,
178                _index: I,
179            ) -> Result<Vec<Self>, TryGetError> {
180                #[cfg(feature = "postgres-array")]
181                {
182                    <Vec<Self>>::try_get_by(_res, _index)
183                }
184                #[cfg(not(feature = "postgres-array"))]
185                Err(TryGetError::DbErr(DbErr::BackendNotSupported {
186                    db: "Postgres",
187                    ctx: "ActiveEnumValue::try_get_vec_by (`postgres-array` not enabled)",
188                }))
189            }
190        }
191    };
192}
193
194impl_active_enum_value!(u8);
195impl_active_enum_value!(u16);
196impl_active_enum_value!(u32);
197impl_active_enum_value!(u64);
198impl_active_enum_value_with_pg_array!(String);
199impl_active_enum_value_with_pg_array!(i8);
200impl_active_enum_value_with_pg_array!(i16);
201impl_active_enum_value_with_pg_array!(i32);
202impl_active_enum_value_with_pg_array!(i64);
203
204impl TryGetable for sea_query::Enum {
205    fn try_get_by<I: ColIdx>(res: &QueryResult, idx: I) -> Result<Self, TryGetError> {
206        let value: String = <String as TryGetable>::try_get_by(res, idx)?;
207        Ok(Self {
208            // `DeriveActiveEnum` overwrites `type_name` when constructing values for queries, but we
209            // can't reliably recover the enum type name from `QueryResult`. Keeping it empty may
210            // still cause issues if this value is later reused to build SQL (e.g. missing casts).
211            type_name: "".into(),
212            value: value.into(),
213        })
214    }
215}
216
217impl ActiveEnumValue for sea_query::Enum {
218    fn try_get_vec_by<I: ColIdx>(_res: &QueryResult, _index: I) -> Result<Vec<Self>, TryGetError> {
219        #[cfg(feature = "postgres-array")]
220        {
221            let values: Vec<String> = <Vec<String> as TryGetable>::try_get_by(_res, _index)?;
222            Ok(values
223                .into_iter()
224                .map(|value| Self {
225                    // See comment in `TryGetable for sea_query::Enum` about empty `type_name`.
226                    type_name: "".into(),
227                    value: value.into(),
228                })
229                .collect())
230        }
231        #[cfg(not(feature = "postgres-array"))]
232        {
233            Err(TryGetError::DbErr(DbErr::BackendNotSupported {
234                db: "Postgres",
235                ctx: "ActiveEnumValue::try_get_vec_by (`postgres-array` not enabled)",
236            }))
237        }
238    }
239}
240
241impl<T> TryFromU64 for T
242where
243    T: ActiveEnum,
244{
245    fn try_from_u64(_: u64) -> Result<Self, DbErr> {
246        Err(DbErr::ConvertFromU64(
247            "Fail to construct ActiveEnum from a u64, if your primary key consist of a ActiveEnum field, its auto increment should be set to false.",
248        ))
249    }
250}
251
252#[cfg(test)]
253mod tests {
254    use crate as sea_orm;
255    use crate::{
256        error::*,
257        sea_query::{SeaRc, StringLen},
258        *,
259    };
260    use pretty_assertions::assert_eq;
261
262    #[test]
263    fn active_enum_string() {
264        #[derive(Debug, PartialEq, Eq, EnumIter)]
265        pub enum Category {
266            Big,
267            Small,
268        }
269
270        #[derive(Debug, DeriveIden)]
271        #[sea_orm(iden = "category")]
272        pub struct CategoryEnum;
273
274        impl ActiveEnum for Category {
275            type Value = String;
276
277            type ValueVec = Vec<String>;
278
279            fn name() -> DynIden {
280                SeaRc::new(CategoryEnum)
281            }
282
283            fn to_value(&self) -> Self::Value {
284                match self {
285                    Self::Big => "B",
286                    Self::Small => "S",
287                }
288                .to_owned()
289            }
290
291            fn try_from_value(v: &Self::Value) -> Result<Self, DbErr> {
292                match v.as_ref() {
293                    "B" => Ok(Self::Big),
294                    "S" => Ok(Self::Small),
295                    _ => Err(type_err(format!("unexpected value for Category enum: {v}"))),
296                }
297            }
298
299            fn db_type() -> ColumnDef {
300                ColumnType::String(StringLen::N(1)).def()
301            }
302        }
303
304        #[derive(Debug, PartialEq, Eq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
305        #[sea_orm(
306            rs_type = "String",
307            db_type = "String(StringLen::N(1))",
308            enum_name = "category"
309        )]
310        pub enum DeriveCategory {
311            #[sea_orm(string_value = "B")]
312            Big,
313            #[sea_orm(string_value = "S")]
314            Small,
315        }
316
317        assert_eq!(Category::Big.to_value(), "B".to_owned());
318        assert_eq!(Category::Small.to_value(), "S".to_owned());
319        assert_eq!(DeriveCategory::Big.to_value(), "B".to_owned());
320        assert_eq!(DeriveCategory::Small.to_value(), "S".to_owned());
321
322        assert_eq!(
323            Category::try_from_value(&"A".to_owned()).err(),
324            Some(type_err("unexpected value for Category enum: A"))
325        );
326        assert_eq!(
327            Category::try_from_value(&"B".to_owned()).ok(),
328            Some(Category::Big)
329        );
330        assert_eq!(
331            Category::try_from_value(&"S".to_owned()).ok(),
332            Some(Category::Small)
333        );
334        assert_eq!(
335            DeriveCategory::try_from_value(&"A".to_owned()).err(),
336            Some(type_err("unexpected value for DeriveCategory enum: A"))
337        );
338        assert_eq!(
339            DeriveCategory::try_from_value(&"B".to_owned()).ok(),
340            Some(DeriveCategory::Big)
341        );
342        assert_eq!(
343            DeriveCategory::try_from_value(&"S".to_owned()).ok(),
344            Some(DeriveCategory::Small)
345        );
346
347        assert_eq!(
348            Category::db_type(),
349            ColumnType::String(StringLen::N(1)).def()
350        );
351        assert_eq!(
352            DeriveCategory::db_type(),
353            ColumnType::String(StringLen::N(1)).def()
354        );
355
356        assert_eq!(
357            Category::name().to_string(),
358            DeriveCategory::name().to_string()
359        );
360        assert_eq!(Category::values(), DeriveCategory::values());
361
362        assert_eq!(format!("{}", DeriveCategory::Big), "Big");
363        assert_eq!(format!("{}", DeriveCategory::Small), "Small");
364    }
365
366    #[test]
367    fn active_enum_derive_signed_integers() {
368        macro_rules! test_num_value_int {
369            ($ident: ident, $rs_type: expr, $db_type: expr, $col_def: ident) => {
370                #[derive(Debug, PartialEq, Eq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
371                #[sea_orm(rs_type = $rs_type, db_type = $db_type)]
372                pub enum $ident {
373                    #[sea_orm(num_value = -10)]
374                    Negative,
375                    #[sea_orm(num_value = 1)]
376                    Big,
377                    #[sea_orm(num_value = 0)]
378                    Small,
379                }
380
381                test_int!($ident, $rs_type, $db_type, $col_def);
382            };
383        }
384
385        macro_rules! test_fallback_int {
386            ($ident: ident, $fallback_type: ident, $rs_type: expr, $db_type: expr, $col_def: ident) => {
387                #[derive(Debug, PartialEq, Eq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
388                #[sea_orm(rs_type = $rs_type, db_type = $db_type)]
389                #[repr(i32)]
390                pub enum $ident {
391                    Big = 1,
392                    Small = 0,
393                    Negative = -10,
394                }
395
396                test_int!($ident, $rs_type, $db_type, $col_def);
397            };
398        }
399
400        macro_rules! test_int {
401            ($ident: ident, $rs_type: expr, $db_type: expr, $col_def: ident) => {
402                assert_eq!($ident::Big.to_value(), 1);
403                assert_eq!($ident::Small.to_value(), 0);
404                assert_eq!($ident::Negative.to_value(), -10);
405
406                assert_eq!($ident::try_from_value(&1).ok(), Some($ident::Big));
407                assert_eq!($ident::try_from_value(&0).ok(), Some($ident::Small));
408                assert_eq!($ident::try_from_value(&-10).ok(), Some($ident::Negative));
409                assert_eq!(
410                    $ident::try_from_value(&2).err(),
411                    Some(type_err(format!(
412                        "unexpected value for {} enum: 2",
413                        stringify!($ident)
414                    )))
415                );
416
417                assert_eq!($ident::db_type(), ColumnType::$col_def.def());
418
419                assert_eq!(format!("{}", $ident::Big), "Big");
420                assert_eq!(format!("{}", $ident::Small), "Small");
421                assert_eq!(format!("{}", $ident::Negative), "Negative");
422            };
423        }
424
425        test_num_value_int!(I8, "i8", "TinyInteger", TinyInteger);
426        test_num_value_int!(I16, "i16", "SmallInteger", SmallInteger);
427        test_num_value_int!(I32, "i32", "Integer", Integer);
428        test_num_value_int!(I64, "i64", "BigInteger", BigInteger);
429
430        test_fallback_int!(I8Fallback, i8, "i8", "TinyInteger", TinyInteger);
431        test_fallback_int!(I16Fallback, i16, "i16", "SmallInteger", SmallInteger);
432        test_fallback_int!(I32Fallback, i32, "i32", "Integer", Integer);
433        test_fallback_int!(I64Fallback, i64, "i64", "BigInteger", BigInteger);
434    }
435
436    #[test]
437    fn active_enum_derive_unsigned_integers() {
438        macro_rules! test_num_value_uint {
439            ($ident: ident, $rs_type: expr, $db_type: expr, $col_def: ident) => {
440                #[derive(Debug, PartialEq, Eq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
441                #[sea_orm(rs_type = $rs_type, db_type = $db_type)]
442                pub enum $ident {
443                    #[sea_orm(num_value = 1)]
444                    Big,
445                    #[sea_orm(num_value = 0)]
446                    Small,
447                }
448
449                test_uint!($ident, $rs_type, $db_type, $col_def);
450            };
451        }
452
453        macro_rules! test_fallback_uint {
454            ($ident: ident, $fallback_type: ident, $rs_type: expr, $db_type: expr, $col_def: ident) => {
455                #[derive(Debug, PartialEq, Eq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
456                #[sea_orm(rs_type = $rs_type, db_type = $db_type)]
457                #[repr($fallback_type)]
458                pub enum $ident {
459                    Big = 1,
460                    Small = 0,
461                }
462
463                test_uint!($ident, $rs_type, $db_type, $col_def);
464            };
465        }
466
467        macro_rules! test_uint {
468            ($ident: ident, $rs_type: expr, $db_type: expr, $col_def: ident) => {
469                assert_eq!($ident::Big.to_value(), 1);
470                assert_eq!($ident::Small.to_value(), 0);
471
472                assert_eq!($ident::try_from_value(&1).ok(), Some($ident::Big));
473                assert_eq!($ident::try_from_value(&0).ok(), Some($ident::Small));
474                assert_eq!(
475                    $ident::try_from_value(&2).err(),
476                    Some(type_err(format!(
477                        "unexpected value for {} enum: 2",
478                        stringify!($ident)
479                    )))
480                );
481
482                assert_eq!($ident::db_type(), ColumnType::$col_def.def());
483
484                assert_eq!(format!("{}", $ident::Big), "Big");
485                assert_eq!(format!("{}", $ident::Small), "Small");
486            };
487        }
488
489        test_num_value_uint!(U8, "u8", "TinyInteger", TinyInteger);
490        test_num_value_uint!(U16, "u16", "SmallInteger", SmallInteger);
491        test_num_value_uint!(U32, "u32", "Integer", Integer);
492        test_num_value_uint!(U64, "u64", "BigInteger", BigInteger);
493
494        test_fallback_uint!(U8Fallback, u8, "u8", "TinyInteger", TinyInteger);
495        test_fallback_uint!(U16Fallback, u16, "u16", "SmallInteger", SmallInteger);
496        test_fallback_uint!(U32Fallback, u32, "u32", "Integer", Integer);
497        test_fallback_uint!(U64Fallback, u64, "u64", "BigInteger", BigInteger);
498    }
499
500    #[test]
501    fn escaped_non_uax31() {
502        #[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Copy)]
503        #[sea_orm(rs_type = "String", db_type = "Enum", enum_name = "pop_os_names_typos")]
504        pub enum PopOSTypos {
505            #[sea_orm(string_value = "Pop!_OS")]
506            PopOSCorrect,
507            #[sea_orm(string_value = "Pop\u{2757}_OS")]
508            PopOSEmoji,
509            #[sea_orm(string_value = "Pop!_操作系统")]
510            PopOSChinese,
511            #[sea_orm(string_value = "PopOS")]
512            PopOSASCIIOnly,
513            #[sea_orm(string_value = "Pop OS")]
514            PopOSASCIIOnlyWithSpace,
515            #[sea_orm(string_value = "Pop!OS")]
516            PopOSNoUnderscore,
517            #[sea_orm(string_value = "Pop_OS")]
518            PopOSNoExclaimation,
519            #[sea_orm(string_value = "!PopOS_")]
520            PopOSAllOverThePlace,
521            #[sea_orm(string_value = "Pop!_OS22.04LTS")]
522            PopOSWithVersion,
523            #[sea_orm(string_value = "22.04LTSPop!_OS")]
524            PopOSWithVersionPrefix,
525            #[sea_orm(string_value = "!_")]
526            PopOSJustTheSymbols,
527            #[sea_orm(string_value = "")]
528            Nothing,
529            // This WILL fail:
530            // Both PopOs and PopOS will create identifier "Popos"
531            // #[sea_orm(string_value = "PopOs")]
532            // PopOSLowerCase,
533        }
534        let values = [
535            "Pop!_OS",
536            "Pop\u{2757}_OS",
537            "Pop!_操作系统",
538            "PopOS",
539            "Pop OS",
540            "Pop!OS",
541            "Pop_OS",
542            "!PopOS_",
543            "Pop!_OS22.04LTS",
544            "22.04LTSPop!_OS",
545            "!_",
546            "",
547        ];
548        for (variant, val) in PopOSTypos::iter().zip(values) {
549            assert_eq!(variant.to_value(), val);
550            assert_eq!(PopOSTypos::try_from_value(&val.to_owned()), Ok(variant));
551        }
552
553        #[derive(Clone, Debug, PartialEq, EnumIter, DeriveActiveEnum, DeriveDisplay)]
554        #[sea_orm(
555            rs_type = "String",
556            db_type = "String(StringLen::None)",
557            enum_name = "conflicting_string_values"
558        )]
559        pub enum ConflictingStringValues {
560            #[sea_orm(string_value = "")]
561            Member1,
562            #[sea_orm(string_value = "$")]
563            Member2,
564            #[sea_orm(string_value = "$$")]
565            Member3,
566            #[sea_orm(string_value = "AB")]
567            Member4,
568            #[sea_orm(string_value = "A_B")]
569            Member5,
570            #[sea_orm(string_value = "A$B")]
571            Member6,
572            #[sea_orm(string_value = "0 123")]
573            Member7,
574        }
575        type EnumVariant = ConflictingStringValuesVariant;
576        assert_eq!(EnumVariant::__Empty.to_string(), "");
577        assert_eq!(EnumVariant::_0x24.to_string(), "$");
578        assert_eq!(EnumVariant::_0x240x24.to_string(), "$$");
579        assert_eq!(EnumVariant::Ab.to_string(), "AB");
580        assert_eq!(EnumVariant::A0x5Fb.to_string(), "A_B");
581        assert_eq!(EnumVariant::A0x24B.to_string(), "A$B");
582        assert_eq!(EnumVariant::_0x300x20123.to_string(), "0 123");
583    }
584
585    #[test]
586    fn test_derive_display() {
587        use crate::DeriveDisplay;
588
589        #[derive(DeriveDisplay)]
590        enum DisplayTea {
591            EverydayTea,
592            #[sea_orm(display_value = "Breakfast Tea")]
593            BreakfastTea,
594        }
595        assert_eq!(format!("{}", DisplayTea::EverydayTea), "EverydayTea");
596        assert_eq!(format!("{}", DisplayTea::BreakfastTea), "Breakfast Tea");
597    }
598}