better_auth_core/plugin_runtime/fields/
models.rs1use 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 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 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}