reinhardt_db/orm/
composite_pk.rs

1//! Composite primary key support for database models
2//!
3//! This module provides functionality for defining and working with composite primary keys,
4//! which consist of multiple fields combined to form a unique identifier for a database record.
5
6use super::constraints::Constraint;
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9
10/// Error types for composite primary key operations
11#[non_exhaustive]
12#[derive(Debug, Clone, PartialEq)]
13pub enum CompositePkError {
14	/// Empty fields list provided
15	EmptyFields,
16	/// Missing required field value
17	MissingField(String),
18	/// Invalid field value type
19	InvalidFieldType {
20		/// The field name.
21		field: String,
22		/// The expected type name.
23		expected: String,
24	},
25	/// Duplicate field name in definition
26	DuplicateField(String),
27}
28
29impl std::fmt::Display for CompositePkError {
30	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
31		match self {
32			CompositePkError::EmptyFields => {
33				write!(f, "Composite primary key must have at least one field")
34			}
35			CompositePkError::MissingField(field) => write!(f, "Missing required field: {}", field),
36			CompositePkError::InvalidFieldType { field, expected } => {
37				write!(
38					f,
39					"Invalid type for field '{}': expected {}",
40					field, expected
41				)
42			}
43			CompositePkError::DuplicateField(field) => write!(f, "Duplicate field name: {}", field),
44		}
45	}
46}
47
48impl std::error::Error for CompositePkError {}
49
50/// Value types supported in composite primary keys
51#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
52#[serde(untagged)]
53pub enum PkValue {
54	/// String variant.
55	String(String),
56	/// Int variant.
57	Int(i64),
58	/// Uint variant.
59	Uint(u64),
60	/// Bool variant.
61	Bool(bool),
62}
63
64impl PkValue {
65	/// Convert the value to a string representation for SQL
66	///
67	/// # Examples
68	///
69	/// ```
70	/// use reinhardt_db::orm::composite_pk::PkValue;
71	///
72	/// let value = PkValue::Int(42);
73	/// assert_eq!(value.to_sql_string(), "42");
74	///
75	/// let value = PkValue::String("test".to_string());
76	/// assert_eq!(value.to_sql_string(), "'test'");
77	/// ```
78	pub fn to_sql_string(&self) -> String {
79		match self {
80			PkValue::String(s) => format!("'{}'", s.replace('\'', "''")),
81			PkValue::Int(i) => i.to_string(),
82			PkValue::Uint(u) => u.to_string(),
83			PkValue::Bool(b) => if *b { "TRUE" } else { "FALSE" }.to_string(),
84		}
85	}
86}
87
88// Conversion implementations for common types
89impl From<String> for PkValue {
90	fn from(value: String) -> Self {
91		PkValue::String(value)
92	}
93}
94
95impl From<&str> for PkValue {
96	fn from(value: &str) -> Self {
97		PkValue::String(value.to_string())
98	}
99}
100
101impl From<i32> for PkValue {
102	fn from(value: i32) -> Self {
103		PkValue::Int(value as i64)
104	}
105}
106
107impl From<i64> for PkValue {
108	fn from(value: i64) -> Self {
109		PkValue::Int(value)
110	}
111}
112
113impl From<u32> for PkValue {
114	fn from(value: u32) -> Self {
115		PkValue::Uint(value as u64)
116	}
117}
118
119impl From<u64> for PkValue {
120	fn from(value: u64) -> Self {
121		PkValue::Uint(value)
122	}
123}
124
125impl From<bool> for PkValue {
126	fn from(value: bool) -> Self {
127		PkValue::Bool(value)
128	}
129}
130
131impl From<&i32> for PkValue {
132	fn from(value: &i32) -> Self {
133		PkValue::Int(*value as i64)
134	}
135}
136
137impl From<&i64> for PkValue {
138	fn from(value: &i64) -> Self {
139		PkValue::Int(*value)
140	}
141}
142
143impl From<&u32> for PkValue {
144	fn from(value: &u32) -> Self {
145		PkValue::Uint(*value as u64)
146	}
147}
148
149impl From<&u64> for PkValue {
150	fn from(value: &u64) -> Self {
151		PkValue::Uint(*value)
152	}
153}
154
155impl From<&bool> for PkValue {
156	fn from(value: &bool) -> Self {
157		PkValue::Bool(*value)
158	}
159}
160
161impl From<&String> for PkValue {
162	fn from(value: &String) -> Self {
163		PkValue::String(value.clone())
164	}
165}
166
167/// Quote a SQL identifier using ANSI SQL double-quote escaping.
168///
169/// Embedded double quotes are escaped by doubling them.
170fn quote_identifier(ident: &str) -> String {
171	format!("\"{}\"", ident.replace('"', "\"\""))
172}
173
174/// Composite primary key definition consisting of multiple fields
175///
176/// A composite primary key is a primary key that spans multiple columns in a database table.
177/// This is useful when no single field can uniquely identify a record.
178#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
179pub struct CompositePrimaryKey {
180	/// Field names that compose the primary key
181	fields: Vec<String>,
182	/// Optional constraint name
183	name: Option<String>,
184}
185
186impl CompositePrimaryKey {
187	/// Create a new composite primary key with the specified fields
188	///
189	/// # Examples
190	///
191	/// ```
192	/// use reinhardt_db::orm::composite_pk::CompositePrimaryKey;
193	///
194	/// let pk = CompositePrimaryKey::new(vec!["user_id".to_string(), "role_id".to_string()]);
195	/// assert!(pk.is_ok());
196	/// assert_eq!(pk.unwrap().fields().len(), 2);
197	/// ```
198	pub fn new(fields: Vec<String>) -> Result<Self, CompositePkError> {
199		if fields.is_empty() {
200			return Err(CompositePkError::EmptyFields);
201		}
202
203		let mut seen = HashMap::new();
204		for field in &fields {
205			if seen.contains_key(field) {
206				return Err(CompositePkError::DuplicateField(field.clone()));
207			}
208			seen.insert(field.clone(), true);
209		}
210
211		Ok(Self { fields, name: None })
212	}
213
214	/// Create a composite primary key with a custom constraint name
215	///
216	/// # Examples
217	///
218	/// ```
219	/// use reinhardt_db::orm::composite_pk::CompositePrimaryKey;
220	///
221	/// let pk = CompositePrimaryKey::with_name(
222	///     vec!["user_id".to_string(), "group_id".to_string()],
223	///     "user_group_pk"
224	/// );
225	/// assert!(pk.is_ok());
226	/// ```
227	pub fn with_name(
228		fields: Vec<String>,
229		name: impl Into<String>,
230	) -> Result<Self, CompositePkError> {
231		let mut pk = Self::new(fields)?;
232		pk.name = Some(name.into());
233		Ok(pk)
234	}
235
236	/// Get the list of field names in the composite key
237	///
238	/// # Examples
239	///
240	/// ```
241	/// use reinhardt_db::orm::composite_pk::CompositePrimaryKey;
242	///
243	/// let pk = CompositePrimaryKey::new(vec!["a".to_string(), "b".to_string()]).unwrap();
244	/// assert_eq!(pk.fields(), &["a", "b"]);
245	/// ```
246	pub fn fields(&self) -> &[String] {
247		&self.fields
248	}
249
250	/// Get the constraint name if set
251	pub fn name(&self) -> Option<&str> {
252		self.name.as_deref()
253	}
254
255	/// Generate SQL PRIMARY KEY constraint definition
256	///
257	/// # Examples
258	///
259	/// ```
260	/// use reinhardt_db::orm::composite_pk::CompositePrimaryKey;
261	///
262	/// let pk = CompositePrimaryKey::new(vec!["user_id".to_string(), "order_id".to_string()]).unwrap();
263	/// let sql = pk.to_sql();
264	/// assert!(sql.contains("PRIMARY KEY"));
265	/// assert!(sql.contains("user_id"));
266	/// assert!(sql.contains("order_id"));
267	/// ```
268	pub fn to_sql(&self) -> String {
269		let fields: Vec<String> = self.fields.iter().map(|f| quote_identifier(f)).collect();
270		let fields = fields.join(", ");
271		if let Some(ref name) = self.name {
272			format!(
273				"CONSTRAINT {} PRIMARY KEY ({})",
274				quote_identifier(name),
275				fields
276			)
277		} else {
278			format!("PRIMARY KEY ({})", fields)
279		}
280	}
281
282	/// Validate that all required fields are present in the provided values
283	///
284	/// # Examples
285	///
286	/// ```
287	/// use reinhardt_db::orm::composite_pk::{CompositePrimaryKey, PkValue};
288	/// use std::collections::HashMap;
289	///
290	/// let pk = CompositePrimaryKey::new(vec!["id".to_string(), "type".to_string()]).unwrap();
291	/// let mut values = HashMap::new();
292	/// values.insert("id".to_string(), PkValue::Int(1));
293	/// values.insert("type".to_string(), PkValue::String("admin".to_string()));
294	///
295	/// assert!(pk.validate(&values).is_ok());
296	/// ```
297	pub fn validate(&self, values: &HashMap<String, PkValue>) -> Result<(), CompositePkError> {
298		for field in &self.fields {
299			if !values.contains_key(field) {
300				return Err(CompositePkError::MissingField(field.clone()));
301			}
302		}
303		Ok(())
304	}
305
306	/// Generate a WHERE clause for querying by this composite key
307	///
308	/// # Examples
309	///
310	/// ```
311	/// use reinhardt_db::orm::composite_pk::{CompositePrimaryKey, PkValue};
312	/// use std::collections::HashMap;
313	///
314	/// let pk = CompositePrimaryKey::new(vec!["user_id".to_string(), "role_id".to_string()]).unwrap();
315	/// let mut values = HashMap::new();
316	/// values.insert("user_id".to_string(), PkValue::Int(100));
317	/// values.insert("role_id".to_string(), PkValue::Int(5));
318	///
319	/// let where_clause = pk.to_where_clause(&values);
320	/// assert!(where_clause.is_ok());
321	/// assert!(where_clause.unwrap().contains("\"user_id\" = 100"));
322	/// ```
323	pub fn to_where_clause(
324		&self,
325		values: &HashMap<String, PkValue>,
326	) -> Result<String, CompositePkError> {
327		self.validate(values)?;
328
329		let conditions: Vec<String> = self
330			.fields
331			.iter()
332			.map(|field| {
333				let value = values.get(field).unwrap();
334				format!("{} = {}", quote_identifier(field), value.to_sql_string())
335			})
336			.collect();
337
338		Ok(conditions.join(" AND "))
339	}
340
341	/// Check if this composite key contains a specific field
342	///
343	/// # Examples
344	///
345	/// ```
346	/// use reinhardt_db::orm::composite_pk::CompositePrimaryKey;
347	///
348	/// let pk = CompositePrimaryKey::new(vec!["a".to_string(), "b".to_string()]).unwrap();
349	/// assert!(pk.contains_field("a"));
350	/// assert!(pk.contains_field("b"));
351	/// assert!(!pk.contains_field("c"));
352	/// ```
353	pub fn contains_field(&self, field: &str) -> bool {
354		self.fields.iter().any(|f| f == field)
355	}
356
357	/// Get the number of fields in this composite key
358	///
359	/// # Examples
360	///
361	/// ```
362	/// use reinhardt_db::orm::composite_pk::CompositePrimaryKey;
363	///
364	/// let pk = CompositePrimaryKey::new(vec!["a".to_string(), "b".to_string(), "c".to_string()]).unwrap();
365	/// assert_eq!(pk.field_count(), 3);
366	/// ```
367	pub fn field_count(&self) -> usize {
368		self.fields.len()
369	}
370}
371
372impl Constraint for CompositePrimaryKey {
373	fn to_sql(&self) -> String {
374		CompositePrimaryKey::to_sql(self)
375	}
376
377	fn name(&self) -> &str {
378		self.name.as_deref().unwrap_or("composite_pk")
379	}
380}
381
382#[cfg(test)]
383mod tests {
384	use super::*;
385	use rstest::rstest;
386
387	#[test]
388	fn test_composite_pk_new_valid() {
389		let pk = CompositePrimaryKey::new(vec!["user_id".to_string(), "role_id".to_string()]);
390		let pk = pk.unwrap();
391		assert_eq!(pk.fields().len(), 2);
392		assert_eq!(pk.fields()[0], "user_id");
393		assert_eq!(pk.fields()[1], "role_id");
394	}
395
396	#[test]
397	fn test_composite_pk_new_empty_fields() {
398		let pk = CompositePrimaryKey::new(vec![]);
399		assert!(pk.is_err());
400		assert_eq!(pk.unwrap_err(), CompositePkError::EmptyFields);
401	}
402
403	#[test]
404	fn test_composite_pk_duplicate_fields() {
405		let pk =
406			CompositePrimaryKey::new(vec!["id".to_string(), "type".to_string(), "id".to_string()]);
407		assert!(pk.is_err());
408		match pk.unwrap_err() {
409			CompositePkError::DuplicateField(field) => assert_eq!(field, "id"),
410			_ => panic!("Expected DuplicateField error"),
411		}
412	}
413
414	#[test]
415	fn test_composite_pk_with_name() {
416		let pk =
417			CompositePrimaryKey::with_name(vec!["a".to_string(), "b".to_string()], "custom_pk");
418		let pk = pk.unwrap();
419		assert_eq!(pk.name(), Some("custom_pk"));
420	}
421
422	#[test]
423	fn test_composite_pk_to_sql_without_name() {
424		let pk =
425			CompositePrimaryKey::new(vec!["user_id".to_string(), "order_id".to_string()]).unwrap();
426		let sql = pk.to_sql();
427		assert_eq!(sql, "PRIMARY KEY (\"user_id\", \"order_id\")");
428	}
429
430	#[test]
431	fn test_composite_pk_to_sql_with_name() {
432		let pk = CompositePrimaryKey::with_name(
433			vec!["user_id".to_string(), "order_id".to_string()],
434			"user_order_pk",
435		)
436		.unwrap();
437		let sql = pk.to_sql();
438		assert_eq!(
439			sql,
440			"CONSTRAINT \"user_order_pk\" PRIMARY KEY (\"user_id\", \"order_id\")"
441		);
442	}
443
444	#[test]
445	fn test_pk_value_to_sql_string() {
446		assert_eq!(PkValue::Int(42).to_sql_string(), "42");
447		assert_eq!(PkValue::Uint(100).to_sql_string(), "100");
448		assert_eq!(PkValue::Bool(true).to_sql_string(), "TRUE");
449		assert_eq!(PkValue::Bool(false).to_sql_string(), "FALSE");
450		assert_eq!(
451			PkValue::String("test".to_string()).to_sql_string(),
452			"'test'"
453		);
454	}
455
456	#[test]
457	fn test_pk_value_sql_string_escapes_quotes() {
458		let value = PkValue::String("O'Brien".to_string());
459		assert_eq!(value.to_sql_string(), "'O''Brien'");
460	}
461
462	#[test]
463	fn test_validate_success() {
464		let pk = CompositePrimaryKey::new(vec!["id".to_string(), "type".to_string()]).unwrap();
465		let mut values = HashMap::new();
466		values.insert("id".to_string(), PkValue::Int(1));
467		values.insert("type".to_string(), PkValue::String("admin".to_string()));
468
469		assert!(pk.validate(&values).is_ok());
470	}
471
472	#[test]
473	fn test_validate_missing_field() {
474		let pk = CompositePrimaryKey::new(vec!["id".to_string(), "type".to_string()]).unwrap();
475		let mut values = HashMap::new();
476		values.insert("id".to_string(), PkValue::Int(1));
477
478		let result = pk.validate(&values);
479		assert!(result.is_err());
480		match result.unwrap_err() {
481			CompositePkError::MissingField(field) => assert_eq!(field, "type"),
482			_ => panic!("Expected MissingField error"),
483		}
484	}
485
486	#[test]
487	fn test_to_where_clause_success() {
488		let pk =
489			CompositePrimaryKey::new(vec!["user_id".to_string(), "role_id".to_string()]).unwrap();
490		let mut values = HashMap::new();
491		values.insert("user_id".to_string(), PkValue::Int(100));
492		values.insert("role_id".to_string(), PkValue::Int(5));
493
494		let where_clause = pk.to_where_clause(&values);
495		let clause = where_clause.unwrap();
496		assert!(clause.contains("\"user_id\" = 100"));
497		assert!(clause.contains("\"role_id\" = 5"));
498		assert!(clause.contains(" AND "));
499	}
500
501	#[test]
502	fn test_to_where_clause_missing_field() {
503		let pk =
504			CompositePrimaryKey::new(vec!["user_id".to_string(), "role_id".to_string()]).unwrap();
505		let mut values = HashMap::new();
506		values.insert("user_id".to_string(), PkValue::Int(100));
507
508		let result = pk.to_where_clause(&values);
509		assert!(result.is_err());
510	}
511
512	#[test]
513	fn test_contains_field() {
514		let pk = CompositePrimaryKey::new(vec!["a".to_string(), "b".to_string()]).unwrap();
515		assert!(pk.contains_field("a"));
516		assert!(pk.contains_field("b"));
517		assert!(!pk.contains_field("c"));
518	}
519
520	#[test]
521	fn test_field_count() {
522		let pk = CompositePrimaryKey::new(vec!["a".to_string(), "b".to_string(), "c".to_string()])
523			.unwrap();
524		assert_eq!(pk.field_count(), 3);
525	}
526
527	#[test]
528	fn test_constraint_trait_implementation() {
529		let pk = CompositePrimaryKey::with_name(vec!["id".to_string()], "test_pk").unwrap();
530		assert_eq!(pk.name(), Some("test_pk"));
531		assert!(pk.to_sql().contains("PRIMARY KEY"));
532	}
533
534	#[test]
535	fn test_composite_pk_three_fields() {
536		let pk = CompositePrimaryKey::new(vec![
537			"org_id".to_string(),
538			"user_id".to_string(),
539			"project_id".to_string(),
540		])
541		.unwrap();
542		assert_eq!(pk.field_count(), 3);
543		let sql = pk.to_sql();
544		assert!(sql.contains("org_id"));
545		assert!(sql.contains("user_id"));
546		assert!(sql.contains("project_id"));
547	}
548
549	#[test]
550	fn test_where_clause_with_string_values() {
551		let pk = CompositePrimaryKey::new(vec!["country".to_string(), "city".to_string()]).unwrap();
552		let mut values = HashMap::new();
553		values.insert("country".to_string(), PkValue::String("USA".to_string()));
554		values.insert("city".to_string(), PkValue::String("New York".to_string()));
555
556		let where_clause = pk.to_where_clause(&values).unwrap();
557		assert!(where_clause.contains("\"country\" = 'USA'"));
558		assert!(where_clause.contains("\"city\" = 'New York'"));
559	}
560
561	#[test]
562	fn test_where_clause_with_mixed_types() {
563		let pk = CompositePrimaryKey::new(vec!["id".to_string(), "active".to_string()]).unwrap();
564		let mut values = HashMap::new();
565		values.insert("id".to_string(), PkValue::Uint(42));
566		values.insert("active".to_string(), PkValue::Bool(true));
567
568		let where_clause = pk.to_where_clause(&values).unwrap();
569		assert!(where_clause.contains("\"id\" = 42"));
570		assert!(where_clause.contains("\"active\" = TRUE"));
571	}
572
573	#[test]
574	fn test_error_display() {
575		let err = CompositePkError::EmptyFields;
576		assert_eq!(
577			err.to_string(),
578			"Composite primary key must have at least one field"
579		);
580
581		let err = CompositePkError::MissingField("user_id".to_string());
582		assert_eq!(err.to_string(), "Missing required field: user_id");
583
584		let err = CompositePkError::DuplicateField("id".to_string());
585		assert_eq!(err.to_string(), "Duplicate field name: id");
586	}
587
588	#[test]
589	fn test_pk_value_serialization() {
590		let value = PkValue::Int(42);
591		let serialized = serde_json::to_string(&value).unwrap();
592		assert_eq!(serialized, "42");
593
594		let value = PkValue::String("test".to_string());
595		let serialized = serde_json::to_string(&value).unwrap();
596		assert_eq!(serialized, "\"test\"");
597	}
598
599	#[rstest]
600	fn test_to_sql_quotes_normal_field_names() {
601		// Arrange
602		let pk =
603			CompositePrimaryKey::new(vec!["user_id".to_string(), "role_id".to_string()]).unwrap();
604
605		// Act
606		let sql = pk.to_sql();
607
608		// Assert
609		assert_eq!(sql, "PRIMARY KEY (\"user_id\", \"role_id\")");
610	}
611
612	#[rstest]
613	fn test_to_sql_escapes_special_characters_in_field_names() {
614		// Arrange
615		let pk =
616			CompositePrimaryKey::new(vec!["field\"; DROP TABLE users; --".to_string()]).unwrap();
617
618		// Act
619		let sql = pk.to_sql();
620
621		// Assert
622		// The embedded double quote is escaped by doubling, preventing identifier breakout
623		assert_eq!(sql, "PRIMARY KEY (\"field\"\"; DROP TABLE users; --\")");
624	}
625
626	#[rstest]
627	fn test_to_sql_escapes_special_characters_in_constraint_name() {
628		// Arrange
629		let pk =
630			CompositePrimaryKey::with_name(vec!["id".to_string()], "pk\"; DROP TABLE users; --")
631				.unwrap();
632
633		// Act
634		let sql = pk.to_sql();
635
636		// Assert
637		assert_eq!(
638			sql,
639			"CONSTRAINT \"pk\"\"; DROP TABLE users; --\" PRIMARY KEY (\"id\")"
640		);
641	}
642
643	#[rstest]
644	fn test_to_sql_escapes_embedded_double_quotes_in_field() {
645		// Arrange
646		let pk = CompositePrimaryKey::new(vec!["field\"name".to_string()]).unwrap();
647
648		// Act
649		let sql = pk.to_sql();
650
651		// Assert
652		assert_eq!(sql, "PRIMARY KEY (\"field\"\"name\")");
653	}
654
655	#[rstest]
656	fn test_where_clause_quotes_field_names_with_special_characters() {
657		// Arrange
658		let malicious_field = "id\"; DROP TABLE users; --".to_string();
659		let pk = CompositePrimaryKey::new(vec![malicious_field.clone()]).unwrap();
660		let mut values = HashMap::new();
661		values.insert(malicious_field, PkValue::Int(1));
662
663		// Act
664		let clause = pk.to_where_clause(&values).unwrap();
665
666		// Assert
667		assert_eq!(clause, "\"id\"\"; DROP TABLE users; --\" = 1");
668	}
669
670	#[rstest]
671	fn test_where_clause_quotes_normal_field_names() {
672		// Arrange
673		let pk = CompositePrimaryKey::new(vec!["user_id".to_string()]).unwrap();
674		let mut values = HashMap::new();
675		values.insert("user_id".to_string(), PkValue::Int(42));
676
677		// Act
678		let clause = pk.to_where_clause(&values).unwrap();
679
680		// Assert
681		assert_eq!(clause, "\"user_id\" = 42");
682	}
683}
reinhardt_db/orm/composite_pk.rs

reinhardt_db/orm/
composite_pk.rs
