Skip to main content

better_auth_core/plugin_runtime/fields/
models.rs

1use super::ModelFields;
2use crate::{
3    AuthError, AuthResult,
4    store::schema::{EntityRole, resolve_field_name},
5    user_fields::UserConfig,
6};
7
8impl ModelFields {
9    pub(crate) fn register_custom_model(
10        &mut self,
11        name: String,
12        model_name: Option<&str>,
13        fields: UserConfig,
14    ) -> AuthResult<()> {
15        if matches!(
16            name.as_str(),
17            "user"
18                | "session"
19                | "account"
20                | "verification"
21                | "organization"
22                | "member"
23                | "invitation"
24                | "team"
25                | "teamMember"
26                | "organizationRole"
27                | "apikey"
28                | "deviceCode"
29                | "passkey"
30                | "twoFactor"
31                | "jwks"
32                | "walletAddress"
33                | "rateLimit"
34        ) {
35            return Err(AuthError::config(format!(
36                "Model \"{name}\" is reserved for a native model and cannot be declared as a custom model"
37            )));
38        }
39        let physical_name = resolve_field_name(model_name, &name).to_owned();
40        let declaration = self.custom_models.entry(name).or_default();
41        declaration.0 = physical_name;
42        if let Some(fields) = fields.additional_fields {
43            declaration.1.fields_mut().extend(fields);
44        }
45        Ok(())
46    }
47
48    /// Read custom logical names, physical names, and merged field policies in declaration order.
49    /// Storage adapters retain responsibility for custom model operations and migrations.
50    pub fn custom_models(&self) -> impl Iterator<Item = (&str, &str, &UserConfig)> {
51        self.custom_models
52            .iter()
53            .map(|(name, (physical, fields))| (name.as_str(), physical.as_str(), fields))
54    }
55
56    pub(crate) fn resolve_model_name(
57        &self,
58        candidate: &str,
59        table_matches: &impl Fn(EntityRole, &str) -> bool,
60    ) -> AuthResult<&str> {
61        let models = self.schema_models.as_deref().unwrap_or(&[
62            (EntityRole::User, "user"),
63            (EntityRole::Session, "session"),
64            (EntityRole::Account, "account"),
65            (EntityRole::Verification, "verification"),
66        ]);
67        // Every logical name precedes physical aliases, including aliases of native models.
68        models
69            .iter()
70            .find(|(_, name)| *name == candidate)
71            .map(|(_, name)| *name)
72            .or_else(|| {
73                self.custom_models
74                    .get_key_value(candidate)
75                    .map(|(name, _)| name.as_str())
76            })
77            .or_else(|| {
78                models
79                    .iter()
80                    .find(|(role, _)| table_matches(*role, candidate))
81                    .map(|(_, name)| *name)
82            })
83            .or_else(|| {
84                self.custom_models()
85                    .find(|(_, physical, _)| *physical == candidate)
86                    .map(|(name, _, _)| name)
87            })
88            .filter(|name| !name.is_empty())
89            .ok_or_else(|| AuthError::config(format!("Model \"{candidate}\" not found in schema")))
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use crate::user_fields::UserFieldConfig;
97
98    #[test]
99    fn custom_model_declarations_merge_without_losing_schema_position() {
100        let mut schema = ModelFields::default();
101        let fields = |names: &[&str], required| UserConfig {
102            additional_fields: Some(
103                names
104                    .iter()
105                    .map(|name| {
106                        (
107                            (*name).into(),
108                            UserFieldConfig {
109                                required: Some(required),
110                                ..Default::default()
111                            },
112                        )
113                    })
114                    .collect(),
115            ),
116        };
117        schema
118            .register_custom_model(
119                "badge".into(),
120                Some("old_badges"),
121                fields(&["label", "rank"], false),
122            )
123            .unwrap();
124        schema
125            .register_custom_model("other".into(), Some("shared"), UserConfig::default())
126            .unwrap();
127        schema
128            .register_custom_model(
129                "badge".into(),
130                Some("shared"),
131                fields(&["label", "enabled"], true),
132            )
133            .unwrap();
134        let schema = schema.fresh_runtime();
135        assert_eq!(
136            schema
137                .custom_models()
138                .map(|(name, _, _)| name)
139                .collect::<Vec<_>>(),
140            ["badge", "other"]
141        );
142        let (_, physical, fields) = schema.custom_models().next().unwrap();
143        assert_eq!(physical, "shared");
144        assert_eq!(
145            fields
146                .fields()
147                .keys()
148                .map(String::as_str)
149                .collect::<Vec<_>>(),
150            ["label", "rank", "enabled"]
151        );
152        assert_eq!(fields.fields().get("label").unwrap().required, Some(true));
153        assert_eq!(fields.fields().get("rank").unwrap().required, Some(false));
154        assert_eq!(
155            schema.resolve_model_name("shared", &|_, _| false).unwrap(),
156            "badge"
157        );
158        assert!(
159            schema
160                .resolve_model_name("old_badges", &|_, _| false)
161                .is_err()
162        );
163    }
164
165    #[test]
166    fn custom_model_names_precede_aliases_and_reserve_native_roles() {
167        let mut schema = ModelFields::default();
168        schema
169            .register_custom_model("badge".into(), Some("user"), UserConfig::default())
170            .unwrap();
171        let aliases = |role, name: &str| role == EntityRole::User && name == "badge";
172        assert_eq!(
173            schema.resolve_model_name("badge", &aliases).unwrap(),
174            "badge"
175        );
176        assert_eq!(schema.resolve_model_name("user", &aliases).unwrap(), "user");
177        for name in [
178            "user",
179            "passkey",
180            "member",
181            "invitation",
182            "teamMember",
183            "organizationRole",
184            "rateLimit",
185        ] {
186            let error = schema
187                .register_custom_model(name.into(), None, UserConfig::default())
188                .unwrap_err();
189            assert_eq!(
190                error.instrumentation_message(),
191                format!(
192                    "Model \"{name}\" is reserved for a native model and cannot be declared as a custom model"
193                )
194            );
195        }
196        assert_eq!(schema.custom_models().count(), 1);
197        schema
198            .register_custom_model("badge".into(), Some(""), UserConfig::default())
199            .unwrap();
200        assert_eq!(schema.custom_models().next().unwrap().1, "badge");
201        assert_eq!(
202            schema
203                .resolve_model_name("missingBadge", &aliases)
204                .unwrap_err()
205                .instrumentation_message(),
206            "Model \"missingBadge\" not found in schema"
207        );
208    }
209}