1use std::marker::PhantomData;
2
3use sea_orm::{
4 entity::prelude::*, ColumnTrait, Condition, ConnectionTrait, DbBackend, EntityTrait,
5 QueryFilter, QueryOrder,
6};
7
8use sea_orm::sea_query::Expr;
9
10use crate::config::{ClosureTreeConfig, OrderStrategy};
11use crate::error::ClosureTreeError;
12use crate::lock::LockedTransaction;
13use crate::traits::ClosureTreeModel;
14
15#[derive(Debug, Default)]
17pub struct ClosureTreeRepository<M>
18where
19 M: ClosureTreeModel,
20{
21 _marker: PhantomData<M>,
22}
23
24impl<M> ClosureTreeRepository<M>
25where
26 M: ClosureTreeModel,
27{
28 pub fn new() -> Self {
29 Self {
30 _marker: PhantomData,
31 }
32 }
33
34 fn config(&self) -> &'static ClosureTreeConfig {
35 M::closure_tree_config()
36 }
37
38 fn ensure_postgres(conn: &impl ConnectionTrait) -> Result<(), ClosureTreeError> {
39 if conn.get_database_backend() == DbBackend::Postgres {
40 Ok(())
41 } else {
42 Err(ClosureTreeError::UnsupportedBackend)
43 }
44 }
45
46 pub async fn parent(
47 &self,
48 conn: &DatabaseConnection,
49 model: &M,
50 ) -> Result<Option<M>, ClosureTreeError> {
51 Self::ensure_postgres(conn)?;
52 match model.parent_id() {
53 Some(parent_id) => {
54 let parent = M::Entity::find()
55 .filter(M::id_column().eq(M::id_to_value(&parent_id)))
56 .one(conn)
57 .await?;
58 Ok(parent)
59 }
60 None => Ok(None),
61 }
62 }
63
64 pub async fn children(
65 &self,
66 conn: &DatabaseConnection,
67 model: &M,
68 ) -> Result<Vec<M>, ClosureTreeError> {
69 Self::ensure_postgres(conn)?;
70 let id = model.id();
71 let parent_value = M::id_to_value(&id);
72 let mut query = M::Entity::find().filter(M::parent_column().eq(parent_value));
73 if let Some(OrderStrategy::NumericColumn { column }) = self.config().order_strategy() {
74 query = query.order_by_asc(Expr::cust(column.clone()));
75 }
76 query = query.order_by_asc(M::name_column());
77 let rows = query.all(conn).await?;
78 Ok(rows)
79 }
80
81 pub async fn roots(&self, conn: &DatabaseConnection) -> Result<Vec<M>, ClosureTreeError> {
82 Self::ensure_postgres(conn)?;
83 let rows = M::Entity::find()
84 .filter(M::parent_column().is_null())
85 .order_by_asc(M::name_column())
86 .all(conn)
87 .await?;
88 Ok(rows)
89 }
90
91 pub async fn descendants(
92 &self,
93 conn: &DatabaseConnection,
94 model: &M,
95 ) -> Result<Vec<M>, ClosureTreeError> {
96 Self::ensure_postgres(conn)?;
97 let rows = self.descendants_with_conn(conn, &model.id(), true).await?;
98 Ok(rows)
99 }
100
101 pub async fn self_and_descendants(
102 &self,
103 conn: &DatabaseConnection,
104 model: &M,
105 ) -> Result<Vec<M>, ClosureTreeError> {
106 Self::ensure_postgres(conn)?;
107 let mut nodes = Vec::with_capacity(1);
108 nodes.push(model.clone());
109 let mut descendants = self.descendants_with_conn(conn, &model.id(), true).await?;
110 nodes.append(&mut descendants);
111 Ok(nodes)
112 }
113
114 pub async fn find_by_path<S: AsRef<str>>(
115 &self,
116 conn: &DatabaseConnection,
117 segments: &[S],
118 ) -> Result<Option<M>, ClosureTreeError> {
119 Self::ensure_postgres(conn)?;
120 self.find_by_path_on(conn, segments).await
121 }
122
123 pub async fn find_or_create_by_path<S: AsRef<str>>(
124 &self,
125 conn: &DatabaseConnection,
126 segments: &[S],
127 ) -> Result<M, ClosureTreeError> {
128 Self::ensure_postgres(conn)?;
129
130 if segments.is_empty() {
131 return Err(ClosureTreeError::EmptyPath);
132 }
133
134 let strategy = self.config().advisory_lock_strategy().clone();
135 let guard = LockedTransaction::acquire(&strategy, conn).await?;
136 self.find_or_create_with_guard(guard, segments).await
137 }
138
139 async fn find_or_create_with_guard<S: AsRef<str>>(
140 &self,
141 guard: LockedTransaction,
142 segments: &[S],
143 ) -> Result<M, ClosureTreeError> {
144 let result = self
145 .find_or_create_by_path_on(guard.connection(), segments)
146 .await;
147
148 match result {
149 Ok(model) => {
150 guard.commit().await?;
151 Ok(model)
152 }
153 Err(err) => {
154 let _ = guard.rollback().await;
155 Err(err)
156 }
157 }
158 }
159
160 async fn find_by_path_on<S: AsRef<str>, C: ConnectionTrait>(
161 &self,
162 conn: &C,
163 segments: &[S],
164 ) -> Result<Option<M>, ClosureTreeError> {
165 if segments.is_empty() {
166 return Ok(None);
167 }
168
169 let mut current_parent: Option<M::Id> = None;
170 let mut current: Option<M> = None;
171
172 for segment in segments {
173 let name = segment.as_ref();
174 let node = self
175 .find_child_by_name(conn, current_parent.as_ref(), name)
176 .await?;
177
178 match node {
179 Some(model) => {
180 current_parent = Some(model.id());
181 current = Some(model);
182 }
183 None => return Ok(None),
184 }
185 }
186
187 Ok(current)
188 }
189
190 async fn find_or_create_by_path_on<S: AsRef<str>, C: ConnectionTrait>(
191 &self,
192 conn: &C,
193 segments: &[S],
194 ) -> Result<M, ClosureTreeError> {
195 let mut current_parent: Option<M::Id> = None;
196 let mut current: Option<M> = None;
197
198 for segment in segments {
199 let name = segment.as_ref();
200 match self
201 .find_child_by_name(conn, current_parent.as_ref(), name)
202 .await?
203 {
204 Some(model) => {
205 current_parent = Some(model.id());
206 current = Some(model);
207 }
208 None => {
209 let created = self
210 .insert_child(conn, current_parent.as_ref(), name)
211 .await?;
212 current_parent = Some(created.id());
213 current = Some(created);
214 }
215 }
216 }
217
218 current.ok_or_else(|| ClosureTreeError::invariant("path segments produced no model"))
219 }
220
221 async fn insert_child<C: ConnectionTrait>(
222 &self,
223 conn: &C,
224 parent_id: Option<&M::Id>,
225 name: &str,
226 ) -> Result<M, ClosureTreeError> {
227 let mut active = M::ActiveModel::default();
228 M::set_parent(&mut active, parent_id.cloned());
229 M::set_name(&mut active, name);
230
231 let model = active.insert(conn).await?;
232 self.insert_hierarchy_rows(conn, &model, parent_id).await?;
233 Ok(model)
234 }
235
236 async fn insert_hierarchy_rows<C: ConnectionTrait>(
237 &self,
238 conn: &C,
239 model: &M,
240 parent_id: Option<&M::Id>,
241 ) -> Result<(), ClosureTreeError> {
242 let mut rows = Vec::new();
243 let model_id = model.id();
244
245 rows.push(M::hierarchy_build_row(
246 model_id.clone(),
247 model_id.clone(),
248 0,
249 ));
250
251 if let Some(parent_id) = parent_id {
252 let ancestors = M::HierarchyEntity::find()
253 .filter(M::hierarchy_descendant_column().eq(M::hierarchy_id_to_value(parent_id)))
254 .all(conn)
255 .await?;
256
257 for ancestor in ancestors {
258 let ancestor_id = M::hierarchy_model_ancestor(&ancestor);
259 let generations = M::hierarchy_model_generations(&ancestor) + 1;
260 rows.push(M::hierarchy_build_row(
261 ancestor_id,
262 model_id.clone(),
263 generations,
264 ));
265 }
266 }
267
268 M::HierarchyEntity::insert_many(rows).exec(conn).await?;
269 Ok(())
270 }
271
272 async fn find_child_by_name<C: ConnectionTrait>(
273 &self,
274 conn: &C,
275 parent_id: Option<&M::Id>,
276 name: &str,
277 ) -> Result<Option<M>, ClosureTreeError> {
278 let mut condition = Condition::all().add(M::name_column().eq(name));
279
280 if let Some(parent_id) = parent_id {
281 condition = condition.add(M::parent_column().eq(M::id_to_value(parent_id)));
282 } else {
283 condition = condition.add(M::parent_column().is_null());
284 }
285
286 let model = M::Entity::find().filter(condition).one(conn).await?;
287 Ok(model)
288 }
289
290 async fn descendants_with_conn<C: ConnectionTrait>(
291 &self,
292 conn: &C,
293 ancestor_id: &M::Id,
294 exclude_root: bool,
295 ) -> Result<Vec<M>, ClosureTreeError> {
296 let mut query = M::HierarchyEntity::find()
297 .filter(M::hierarchy_ancestor_column().eq(M::hierarchy_id_to_value(ancestor_id)));
298
299 if exclude_root {
300 query = query.filter(M::hierarchy_generations_column().gt(0));
301 }
302
303 let rows = query.all(conn).await?;
304
305 let mut descendant_ids = Vec::with_capacity(rows.len());
306 for hierarchy in rows {
307 let descendant = M::hierarchy_model_descendant(&hierarchy);
308 descendant_ids.push(descendant);
309 }
310
311 if descendant_ids.is_empty() {
312 return Ok(Vec::new());
313 }
314
315 let values = descendant_ids
316 .iter()
317 .map(|id| M::id_to_value(id))
318 .collect::<Vec<_>>();
319
320 let mut query = M::Entity::find().filter(M::id_column().is_in(values));
321 if let Some(OrderStrategy::NumericColumn { column }) = self.config().order_strategy() {
322 query = query.order_by_asc(Expr::cust(column.clone()));
323 }
324 query = query.order_by_asc(M::name_column());
325
326 let models = query.all(conn).await?;
327 Ok(models)
328 }
329}