1use std::collections::HashMap;
16use std::marker::PhantomData;
17
18use sz_orm_core::{
19 dialect::get_dialect,
20 repository::{PageResult, RepositoryError, WhereCondition, WhereOp},
21 DbError, DbType, Model, ModelExt, Pool, QueryBuilder, Value,
22};
23
24struct SqlRepoModel;
25
26impl Model for SqlRepoModel {
27 type PrimaryKey = i64;
28 fn table_name() -> &'static str {
29 "__sql_repo__"
30 }
31 fn pk(&self) -> Self::PrimaryKey {
32 0
33 }
34 fn set_pk(&mut self, _pk: Self::PrimaryKey) {}
35}
36
37pub struct SqlRepository<E: Model + ModelExt + Default + Clone + Send + Sync + 'static> {
60 pool: Pool,
61 dialect: DbType,
62 _marker: PhantomData<E>,
63}
64
65impl<E: Model + ModelExt + Default + Clone + Send + Sync + 'static> SqlRepository<E> {
66 pub fn new(pool: Pool, dialect: DbType) -> Self {
68 Self {
69 pool,
70 dialect,
71 _marker: PhantomData,
72 }
73 }
74
75 fn table() -> &'static str {
77 E::table_name()
78 }
79
80 fn entity_to_row(entity: &E) -> HashMap<String, Value> {
82 let mut row = HashMap::new();
83 for col in E::fillable() {
84 if let Some(val) = entity.get_column_value(col) {
85 row.insert(col.to_string(), val);
86 }
87 }
88 row
89 }
90
91 fn entities_to_rows(entities: &[E]) -> Vec<HashMap<String, Value>> {
93 entities.iter().map(Self::entity_to_row).collect()
94 }
95
96 fn row_to_entity(row: HashMap<String, Value>) -> E {
98 let mut entity = E::default();
99 entity.from_value(row);
100 entity
101 }
102
103 fn db_err(e: DbError) -> RepositoryError {
105 RepositoryError::DatabaseError(e.to_string())
106 }
107
108 fn pool_err(e: sz_orm_core::PoolError) -> RepositoryError {
110 RepositoryError::DatabaseError(e.to_string())
111 }
112
113 fn apply_conditions(
115 qb: QueryBuilder<SqlRepoModel>,
116 conditions: &[WhereCondition],
117 ) -> QueryBuilder<SqlRepoModel> {
118 let mut qb = qb;
119 for cond in conditions {
120 qb = match cond.op {
121 WhereOp::Eq => qb.where_eq(&cond.field, cond.value.clone()),
122 WhereOp::Ne => qb.where_ne(&cond.field, cond.value.clone()),
123 WhereOp::Gt => qb.where_gt(&cond.field, cond.value.clone()),
124 WhereOp::Ge => qb.where_ge(&cond.field, cond.value.clone()),
125 WhereOp::Lt => qb.where_lt(&cond.field, cond.value.clone()),
126 WhereOp::Le => qb.where_le(&cond.field, cond.value.clone()),
127 WhereOp::Like => qb.where_like(&cond.field, cond.value.clone()),
128 WhereOp::In => qb.where_in(&cond.field, cond.extra_values.clone()),
129 WhereOp::NotIn => qb.where_not_in(&cond.field, cond.extra_values.clone()),
130 WhereOp::IsNull => qb.where_null(&cond.field),
131 WhereOp::IsNotNull => qb.where_not_null(&cond.field),
132 WhereOp::Between => {
133 if cond.extra_values.len() == 1 {
134 qb.where_between(
135 &cond.field,
136 cond.value.clone(),
137 cond.extra_values[0].clone(),
138 )
139 } else {
140 qb
141 }
142 }
143 };
144 }
145 qb
146 }
147
148 pub async fn find_all(&self) -> Result<Vec<E>, RepositoryError> {
150 let dialect = get_dialect(self.dialect).map_err(Self::db_err)?;
151 let qb = QueryBuilder::<SqlRepoModel>::new(dialect).table(Self::table().to_string());
152 let (sql, params) = qb.build_select();
153 let mut conn = self.pool.acquire().await.map_err(Self::pool_err)?;
154 let rows = conn
155 .query_with_params(&sql, ¶ms)
156 .await
157 .map_err(Self::db_err)?;
158 Ok(rows.into_iter().map(Self::row_to_entity).collect())
159 }
160
161 pub async fn find_by(&self, conditions: &[WhereCondition]) -> Result<Vec<E>, RepositoryError> {
163 let dialect = get_dialect(self.dialect).map_err(Self::db_err)?;
164 let qb = QueryBuilder::<SqlRepoModel>::new(dialect).table(Self::table().to_string());
165 let qb = Self::apply_conditions(qb, conditions);
166 let (sql, params) = qb.build_select();
167 let mut conn = self.pool.acquire().await.map_err(Self::pool_err)?;
168 let rows = conn
169 .query_with_params(&sql, ¶ms)
170 .await
171 .map_err(Self::db_err)?;
172 Ok(rows.into_iter().map(Self::row_to_entity).collect())
173 }
174
175 pub async fn save(&self, entity: &mut E) -> Result<u64, RepositoryError> {
177 let fillable = E::fillable();
178 if fillable.is_empty() {
179 return Err(RepositoryError::InvalidEntity(
180 "fillable() 为空,无法构造 INSERT".to_string(),
181 ));
182 }
183 let row = Self::entity_to_row(entity);
184 let dialect = get_dialect(self.dialect).map_err(Self::db_err)?;
185 let qb = QueryBuilder::<SqlRepoModel>::new(dialect).table(Self::table().to_string());
186 let (sql, params) = qb.build_insert_with_params(&row);
187 let mut conn = self.pool.acquire().await.map_err(Self::pool_err)?;
188 conn.execute_with_params(&sql, ¶ms)
189 .await
190 .map_err(Self::db_err)
191 }
192
193 pub async fn count(&self) -> Result<u64, RepositoryError> {
195 let dialect = get_dialect(self.dialect).map_err(Self::db_err)?;
196 let qb = QueryBuilder::<SqlRepoModel>::new(dialect).table(Self::table().to_string());
197 let sql = qb.build_count();
198 let mut conn = self.pool.acquire().await.map_err(Self::pool_err)?;
199 let rows = conn
200 .query_with_params(&sql, &[])
201 .await
202 .map_err(Self::db_err)?;
203 Self::extract_count(rows)
204 }
205
206 fn extract_count(rows: Vec<HashMap<String, Value>>) -> Result<u64, RepositoryError> {
208 if let Some(row) = rows.into_iter().next() {
209 if let Some(val) = row.get("total") {
210 return Ok(Self::value_to_u64(val));
211 }
212 }
213 Ok(0)
214 }
215
216 fn value_to_u64(val: &Value) -> u64 {
218 match val {
219 Value::I64(n) => *n as u64,
220 Value::I32(n) => *n as u64,
221 Value::U64(n) => *n,
222 Value::U32(n) => *n as u64,
223 _ => 0,
224 }
225 }
226
227 pub async fn batch_create(&self, entities: &[E]) -> Result<u64, RepositoryError> {
234 if entities.is_empty() {
235 return Ok(0);
236 }
237 let rows = Self::entities_to_rows(entities);
238 let dialect = get_dialect(self.dialect).map_err(Self::db_err)?;
239 let qb = QueryBuilder::<SqlRepoModel>::new(dialect).table(Self::table().to_string());
240 let (sql, params) = qb.build_batch_insert_with_params(&rows);
241 if sql.is_empty() {
242 return Ok(0);
243 }
244 let mut conn = self.pool.acquire().await.map_err(Self::pool_err)?;
245 conn.execute_with_params(&sql, ¶ms)
246 .await
247 .map_err(Self::db_err)
248 }
249
250 pub async fn batch_upsert(
258 &self,
259 entities: &[E],
260 conflict_cols: &[&str],
261 update_cols: &[&str],
262 ) -> Result<u64, RepositoryError> {
263 if entities.is_empty() {
264 return Err(RepositoryError::InvalidParams(
265 "batch_upsert: entities cannot be empty".to_string(),
266 ));
267 }
268 let rows = Self::entities_to_rows(entities);
269 let dialect = get_dialect(self.dialect).map_err(Self::db_err)?;
270 let qb = QueryBuilder::<SqlRepoModel>::new(dialect).table(Self::table().to_string());
271 let (sql, params) = qb
272 .build_batch_upsert_with_params(&rows, conflict_cols, update_cols)
273 .map_err(Self::db_err)?;
274 let mut conn = self.pool.acquire().await.map_err(Self::pool_err)?;
275 conn.execute_with_params(&sql, ¶ms)
276 .await
277 .map_err(Self::db_err)
278 }
279
280 pub async fn paginate_sql(
287 &self,
288 page: u64,
289 page_size: u64,
290 ) -> Result<PageResult<E>, RepositoryError> {
291 if page == 0 || page_size == 0 {
292 return Err(RepositoryError::InvalidParams(
293 "page must be >= 1 and page_size must be >= 1".to_string(),
294 ));
295 }
296 let offset = (page.saturating_sub(1)) * page_size;
297
298 let dialect = get_dialect(self.dialect).map_err(Self::db_err)?;
299 let qb = QueryBuilder::<SqlRepoModel>::new(dialect.clone_box())
300 .table(Self::table().to_string())
301 .limit(page_size as usize)
302 .offset(offset as usize);
303 let (sql, params) = qb.build_select();
304
305 let count_sql = QueryBuilder::<SqlRepoModel>::new(dialect)
306 .table(Self::table().to_string())
307 .build_count();
308
309 let mut conn = self.pool.acquire().await.map_err(Self::pool_err)?;
310
311 let rows = conn
312 .query_with_params(&sql, ¶ms)
313 .await
314 .map_err(Self::db_err)?;
315 let items: Vec<E> = rows.into_iter().map(Self::row_to_entity).collect();
316
317 let count_rows = conn
318 .query_with_params(&count_sql, &[])
319 .await
320 .map_err(Self::db_err)?;
321 let total = Self::extract_count(count_rows)?;
322
323 Ok(PageResult::new(items, total, page, page_size))
324 }
325}
326
327#[cfg(test)]
328mod tests {
329 use super::*;
330
331 #[derive(Debug, Clone, Default, PartialEq)]
332 struct TestEntity {
333 id: i64,
334 name: String,
335 }
336
337 impl Model for TestEntity {
338 type PrimaryKey = i64;
339 fn table_name() -> &'static str {
340 "test_entities"
341 }
342 fn pk(&self) -> Self::PrimaryKey {
343 self.id
344 }
345 fn set_pk(&mut self, pk: Self::PrimaryKey) {
346 self.id = pk;
347 }
348 }
349
350 impl ModelExt for TestEntity {
351 fn columns() -> Vec<&'static str> {
352 vec!["id", "name"]
353 }
354 fn fillable() -> Vec<&'static str> {
355 vec!["id", "name"]
356 }
357 fn get_column_value(&self, column: &str) -> Option<Value> {
358 match column {
359 "id" => Some(Value::I64(self.id)),
360 "name" => Some(Value::String(self.name.clone())),
361 _ => None,
362 }
363 }
364 fn from_value(&mut self, map: HashMap<String, Value>) {
365 if let Some(Value::I64(id)) = map.get("id") {
366 self.id = *id;
367 }
368 if let Some(Value::String(name)) = map.get("name") {
369 self.name = name.clone();
370 }
371 }
372 }
373
374 fn make_entity(id: i64, name: &str) -> TestEntity {
375 TestEntity {
376 id,
377 name: name.to_string(),
378 }
379 }
380
381 #[test]
382 fn test_sql_repo_entity_to_row() {
383 let entity = make_entity(1, "alice");
384 let row = SqlRepository::<TestEntity>::entity_to_row(&entity);
385 assert_eq!(row.get("id"), Some(&Value::I64(1)));
386 assert_eq!(row.get("name"), Some(&Value::String("alice".to_string())));
387 }
388
389 #[test]
390 fn test_sql_repo_entities_to_rows_multi_row() {
391 let entities = vec![
392 make_entity(1, "a"),
393 make_entity(2, "b"),
394 make_entity(3, "c"),
395 ];
396 let rows = SqlRepository::<TestEntity>::entities_to_rows(&entities);
397 assert_eq!(rows.len(), 3);
398 assert_eq!(rows[0].get("id"), Some(&Value::I64(1)));
399 assert_eq!(rows[2].get("name"), Some(&Value::String("c".to_string())));
400 }
401
402 #[test]
403 fn test_sql_repo_row_to_entity() {
404 let mut row = HashMap::new();
405 row.insert("id".to_string(), Value::I64(42));
406 row.insert("name".to_string(), Value::String("bob".to_string()));
407 let entity = SqlRepository::<TestEntity>::row_to_entity(row);
408 assert_eq!(entity.id, 42);
409 assert_eq!(entity.name, "bob");
410 }
411
412 #[test]
413 fn test_sql_repo_table_name() {
414 assert_eq!(SqlRepository::<TestEntity>::table(), "test_entities");
415 }
416
417 #[test]
418 fn test_sql_repo_batch_create_empty_returns_zero() {
419 let entities: Vec<TestEntity> = vec![];
420 let rows = SqlRepository::<TestEntity>::entities_to_rows(&entities);
421 assert!(rows.is_empty());
422 }
423
424 #[test]
425 fn test_sql_repo_batch_create_non_empty_rows() {
426 let entities = vec![
427 make_entity(1, "a"),
428 make_entity(2, "b"),
429 make_entity(3, "c"),
430 ];
431 let rows = SqlRepository::<TestEntity>::entities_to_rows(&entities);
432 assert_eq!(rows.len(), 3);
433 assert!(rows[0].contains_key("id"));
434 assert!(rows[0].contains_key("name"));
435 }
436
437 #[test]
438 fn test_sql_repo_default_entity() {
439 let entity = TestEntity::default();
440 assert_eq!(entity.id, 0);
441 assert_eq!(entity.name, "");
442 }
443
444 #[test]
445 fn test_sql_repo_parametrized_values_not_inlined() {
446 let entity = make_entity(1, "'; DROP TABLE users; --");
447 let row = SqlRepository::<TestEntity>::entity_to_row(&entity);
448 let name_val = row.get("name").unwrap();
449 assert!(matches!(name_val, Value::String(s) if s.contains("DROP TABLE")));
450 }
451
452 #[test]
453 fn test_sql_repo_db_error_mapping() {
454 let db_err = DbError::Internal("constraint violation".to_string());
455 let repo_err = SqlRepository::<TestEntity>::db_err(db_err);
456 assert!(matches!(repo_err, RepositoryError::DatabaseError(_)));
457 }
458
459 #[test]
460 fn test_sql_repo_apply_conditions_eq() {
461 let dialect = get_dialect(DbType::MySQL).unwrap();
462 let qb = QueryBuilder::<SqlRepoModel>::new(dialect).table("test".to_string());
463 let conditions = vec![WhereCondition::new("id", WhereOp::Eq, Value::I64(1))];
464 let qb = SqlRepository::<TestEntity>::apply_conditions(qb, &conditions);
465 let (sql, params) = qb.build_select();
466 assert!(sql.contains("WHERE") || sql.contains("where"));
467 assert!(!params.is_empty());
468 }
469
470 #[test]
471 fn test_sql_repo_apply_conditions_in() {
472 let dialect = get_dialect(DbType::MySQL).unwrap();
473 let qb = QueryBuilder::<SqlRepoModel>::new(dialect).table("test".to_string());
474 let conditions = vec![WhereCondition::in_op(
475 "id",
476 WhereOp::In,
477 vec![Value::I64(1), Value::I64(2), Value::I64(3)],
478 )];
479 let qb = SqlRepository::<TestEntity>::apply_conditions(qb, &conditions);
480 let (sql, _) = qb.build_select();
481 assert!(sql.contains("IN") || sql.contains("in"));
482 }
483
484 #[test]
485 fn test_sql_repo_apply_conditions_empty() {
486 let dialect = get_dialect(DbType::MySQL).unwrap();
487 let qb = QueryBuilder::<SqlRepoModel>::new(dialect).table("test".to_string());
488 let qb = SqlRepository::<TestEntity>::apply_conditions(qb, &[]);
489 let (sql, params) = qb.build_select();
490 assert!(!sql.is_empty());
491 assert!(params.is_empty());
492 }
493
494 #[test]
495 fn test_sql_repo_apply_conditions_multiple() {
496 let dialect = get_dialect(DbType::MySQL).unwrap();
497 let qb = QueryBuilder::<SqlRepoModel>::new(dialect).table("test".to_string());
498 let conditions = vec![
499 WhereCondition::new("age", WhereOp::Ge, Value::I64(18)),
500 WhereCondition::new("name", WhereOp::Like, Value::String("%kw%".to_string())),
501 ];
502 let qb = SqlRepository::<TestEntity>::apply_conditions(qb, &conditions);
503 let (sql, params) = qb.build_select();
504 assert!(!sql.is_empty());
505 assert_eq!(params.len(), 2);
506 }
507
508 #[test]
509 fn test_sql_repo_paginate_sql_invalid_params_validation() {
510 let valid = |p: u64, ps: u64| p != 0 && ps != 0;
511 assert!(!valid(0, 20));
512 assert!(!valid(1, 0));
513 assert!(!valid(0, 0));
514 assert!(valid(1, 20));
515 assert!(valid(2, 10));
516 }
517
518 #[test]
519 fn test_sql_repo_value_to_u64() {
520 assert_eq!(
521 SqlRepository::<TestEntity>::value_to_u64(&Value::I64(42)),
522 42
523 );
524 assert_eq!(
525 SqlRepository::<TestEntity>::value_to_u64(&Value::I32(42)),
526 42
527 );
528 assert_eq!(
529 SqlRepository::<TestEntity>::value_to_u64(&Value::U64(42)),
530 42
531 );
532 assert_eq!(SqlRepository::<TestEntity>::value_to_u64(&Value::Null), 0);
533 assert_eq!(
534 SqlRepository::<TestEntity>::value_to_u64(&Value::Bool(true)),
535 0
536 );
537 }
538
539 #[test]
540 fn test_sql_repo_extract_count() {
541 let mut row = HashMap::new();
542 row.insert("total".to_string(), Value::I64(100));
543 let rows = vec![row];
544 assert_eq!(
545 SqlRepository::<TestEntity>::extract_count(rows).unwrap(),
546 100
547 );
548
549 let empty_rows: Vec<HashMap<String, Value>> = vec![];
550 assert_eq!(
551 SqlRepository::<TestEntity>::extract_count(empty_rows).unwrap(),
552 0
553 );
554 }
555
556 #[test]
557 fn test_sql_repo_batch_create_generates_multi_value_insert() {
558 let entities = vec![
559 make_entity(1, "a"),
560 make_entity(2, "b"),
561 make_entity(3, "c"),
562 ];
563 let rows = SqlRepository::<TestEntity>::entities_to_rows(&entities);
564 let dialect = get_dialect(DbType::MySQL).unwrap();
565 let qb = QueryBuilder::<SqlRepoModel>::new(dialect).table("test_entities".to_string());
566 let (sql, params) = qb.build_batch_insert_with_params(&rows);
567 assert!(sql.contains("INSERT INTO"));
568 assert!(sql.contains("VALUES"));
569 assert!(sql.contains("?"));
570 assert_eq!(params.len(), 6);
571 }
572
573 #[test]
574 fn test_sql_repo_batch_upsert_generates_on_conflict() {
575 let entities = vec![make_entity(1, "a"), make_entity(2, "b")];
576 let rows = SqlRepository::<TestEntity>::entities_to_rows(&entities);
577 let dialect = get_dialect(DbType::MySQL).unwrap();
578 let qb = QueryBuilder::<SqlRepoModel>::new(dialect).table("test_entities".to_string());
579 let (sql, _) = qb
580 .build_batch_upsert_with_params(&rows, &["id"], &["name"])
581 .unwrap();
582 assert!(sql.contains("INSERT INTO"));
583 assert!(sql.contains("ON DUPLICATE KEY UPDATE") || sql.contains("ON CONFLICT"));
584 }
585
586 #[test]
587 fn test_sql_repo_batch_upsert_pg_generates_on_conflict() {
588 let entities = vec![make_entity(1, "a"), make_entity(2, "b")];
589 let rows = SqlRepository::<TestEntity>::entities_to_rows(&entities);
590 let dialect = get_dialect(DbType::PostgreSQL).unwrap();
591 let qb = QueryBuilder::<SqlRepoModel>::new(dialect).table("test_entities".to_string());
592 let (sql, _) = qb
593 .build_batch_upsert_with_params(&rows, &["id"], &["name"])
594 .unwrap();
595 assert!(sql.contains("ON CONFLICT"));
596 }
597
598 #[test]
599 fn test_sql_repo_paginate_sql_builds_limit_offset() {
600 let dialect = get_dialect(DbType::MySQL).unwrap();
601 let qb = QueryBuilder::<SqlRepoModel>::new(dialect)
602 .table("test_entities".to_string())
603 .limit(20)
604 .offset(0);
605 let (sql, _) = qb.build_select();
606 assert!(sql.contains("LIMIT") || sql.contains("limit"));
607 }
608
609 #[test]
610 fn test_sql_repo_paginate_sql_builds_count() {
611 let dialect = get_dialect(DbType::MySQL).unwrap();
612 let qb = QueryBuilder::<SqlRepoModel>::new(dialect).table("test_entities".to_string());
613 let sql = qb.build_count();
614 assert!(sql.contains("COUNT"));
615 assert!(sql.contains("test_entities"));
616 }
617}