1use crate::{
10 ast::{ColumnDef, Expr, TableHierarchy},
11 type_resolution::FunctionTypeResolver,
12 ResultRow, RowSchema, SQLError, SQLParam,
13};
14use std::cmp::Ordering;
15use uqa_core::Value;
16
17pub trait PartitionCatalog {
19 fn try_table_hierarchy(&self, table: &str) -> Result<TableHierarchy, String>;
20 fn direct_hierarchy_children(&self, parent: &str) -> Result<Vec<String>, SQLError>;
21 fn try_resolve_table_name(&self, name: &str) -> Result<Option<String>, String>;
22 fn try_describe_table(&self, table: &str) -> Result<Option<Vec<ColumnDef>>, String>;
23 fn try_table_object_id(&self, table: &str) -> Result<Option<[u8; 16]>, String>;
25}
26
27pub trait PartitionExpressions {
29 fn evaluate_bound(&self, expression: &Expr, params: &[SQLParam]) -> Result<Value, SQLError>;
30 fn evaluate_row(
31 &self,
32 expression: &Expr,
33 row: &ResultRow,
34 schema: &RowSchema,
35 params: &[SQLParam],
36 ) -> Result<Value, SQLError>;
37}
38
39#[derive(Clone, Copy)]
40pub struct PartitionContext<'a> {
41 pub catalog: &'a dyn PartitionCatalog,
42 pub expressions: &'a dyn PartitionExpressions,
43 pub types: &'a dyn FunctionTypeResolver,
44}
45
46mod hash;
47
48pub fn validate_hash_partition_spec(
49 context: &PartitionContext<'_>,
50 spec: &crate::ast::PartitionSpec,
51 columns: &[crate::ast::ColumnDef],
52) -> Result<(), SQLError> {
53 hash::validate_partition_spec(context.types, spec, columns)
54}
55
56pub fn validate_new_partition_bound(
57 context: &PartitionContext<'_>,
58 parent: &str,
59 bound: &crate::ast::PartitionBound,
60) -> Result<(), SQLError> {
61 let hierarchy = context
62 .catalog
63 .try_table_hierarchy(parent)
64 .map_err(|error| SQLError::Internal(format!("read parent partition metadata: {error}")))?;
65 let spec = hierarchy
66 .partition_spec
67 .as_ref()
68 .ok_or_else(|| SQLError::Routine {
69 sqlstate: "42809".into(),
70 message: format!("relation \"{parent}\" is not partitioned"),
71 })?;
72 validate_partition_bound_width(spec, bound)?;
73 if let crate::ast::PartitionBound::Hash { modulus, remainder } = bound {
74 hash::validate_bound(*modulus, *remainder)?;
75 let mut existing_moduli = Vec::new();
76 for sibling in context.catalog.direct_hierarchy_children(parent)? {
77 let sibling_hierarchy = context
78 .catalog
79 .try_table_hierarchy(&sibling)
80 .map_err(|error| SQLError::Internal(format!("read sibling partition: {error}")))?;
81 match sibling_hierarchy.partition_bound.as_ref() {
82 Some(crate::ast::PartitionBound::Hash { modulus, remainder }) => {
83 hash::validate_bound(*modulus, *remainder)?;
84 existing_moduli.push(*modulus);
85 }
86 Some(crate::ast::PartitionBound::Default) => {
87 return Err(SQLError::Internal(format!(
88 "HASH-partitioned table `{parent}` has a default partition"
89 )))
90 }
91 Some(_) => {
92 return Err(SQLError::Internal(
93 "partition siblings use different bound strategies".into(),
94 ))
95 }
96 None => {}
97 }
98 }
99 hash::validate_modulus_chain(*modulus, existing_moduli)?;
100 }
101 if let crate::ast::PartitionBound::Range { lower, upper } = bound {
102 if compare_partition_points(context, lower, upper)? != Ordering::Less {
103 return Err(invalid_partition_bound(
104 "empty range bound specified for partition",
105 ));
106 }
107 }
108 for sibling in context.catalog.direct_hierarchy_children(parent)? {
109 let sibling_hierarchy = context
110 .catalog
111 .try_table_hierarchy(&sibling)
112 .map_err(|error| SQLError::Internal(format!("read sibling partition: {error}")))?;
113 let Some(sibling_bound) = sibling_hierarchy.partition_bound.as_ref() else {
114 continue;
115 };
116 if partition_bounds_overlap(context, bound, sibling_bound)? {
117 return Err(invalid_partition_bound(format!(
118 "partition would overlap partition \"{sibling}\""
119 )));
120 }
121 }
122 Ok(())
123}
124
125pub fn prospective_partition_bound_accepts_document(
129 context: &PartitionContext<'_>,
130 parent: &str,
131 bound: &crate::ast::PartitionBound,
132 document: &ResultRow,
133) -> Result<bool, SQLError> {
134 let hierarchy = context
135 .catalog
136 .try_table_hierarchy(parent)
137 .map_err(|error| SQLError::Internal(format!("read parent partition metadata: {error}")))?;
138 let spec = hierarchy
139 .partition_spec
140 .as_ref()
141 .ok_or_else(|| SQLError::Routine {
142 sqlstate: "42809".into(),
143 message: format!("relation \"{parent}\" is not partitioned"),
144 })?;
145 let (keys, row_hash) = partition_key_values_and_hash(context, parent, spec, document)?;
146 if !matches!(bound, crate::ast::PartitionBound::Default) {
147 return partition_bound_matches(context, bound, &keys, &[], row_hash);
148 }
149 for sibling in context.catalog.direct_hierarchy_children(parent)? {
150 let sibling_hierarchy = context
151 .catalog
152 .try_table_hierarchy(&sibling)
153 .map_err(|error| SQLError::Internal(format!("read child partition: {error}")))?;
154 let Some(sibling_bound) = sibling_hierarchy.partition_bound.as_ref() else {
155 continue;
156 };
157 if matches!(sibling_bound, crate::ast::PartitionBound::Default) {
158 continue;
159 }
160 if partition_bound_matches(context, sibling_bound, &keys, &[], row_hash)? {
161 return Ok(false);
162 }
163 }
164 Ok(true)
165}
166
167pub fn partition_constraint_accepts_document(
171 context: &PartitionContext<'_>,
172 table: &str,
173 spec: &crate::ast::PartitionSpec,
174 bound: &crate::ast::PartitionBound,
175 document: &ResultRow,
176) -> Result<bool, SQLError> {
177 let (keys, row_hash) = partition_key_values_and_hash(context, table, spec, document)?;
178 partition_bound_matches(context, bound, &keys, &[], row_hash)
179}
180
181fn partition_key_values_and_hash(
182 context: &PartitionContext<'_>,
183 table: &str,
184 spec: &crate::ast::PartitionSpec,
185 document: &ResultRow,
186) -> Result<(Vec<Value>, Option<u64>), SQLError> {
187 let (keys, definitions) = evaluate_partition_keys(context, table, &spec.keys, document, &[])?;
188 let row_hash = (spec.strategy == crate::ast::PartitionStrategy::Hash)
189 .then(|| hash::row_hash(context.types, spec, &definitions, &keys))
190 .transpose()?;
191 Ok((keys, row_hash))
192}
193
194fn validate_partition_bound_width(
195 spec: &crate::ast::PartitionSpec,
196 bound: &crate::ast::PartitionBound,
197) -> Result<(), SQLError> {
198 use crate::ast::{PartitionBound, PartitionStrategy};
199 match (spec.strategy, bound) {
200 (_, PartitionBound::Default) => Ok(()),
201 (PartitionStrategy::List, PartitionBound::List(_)) if spec.keys.len() != 1 => {
202 Err(invalid_partition_bound(
203 "cannot use list partition bounds with more than one partition key",
204 ))
205 }
206 (PartitionStrategy::List, PartitionBound::List(_)) => Ok(()),
207 (PartitionStrategy::Range, PartitionBound::Range { lower, upper })
208 if lower.len() != spec.keys.len() || upper.len() != spec.keys.len() =>
209 {
210 Err(invalid_partition_bound(
211 "partition bound has the wrong number of columns",
212 ))
213 }
214 (PartitionStrategy::Range, PartitionBound::Range { .. })
215 | (PartitionStrategy::Hash, PartitionBound::Hash { .. }) => Ok(()),
216 (strategy, _) => Err(invalid_partition_bound(format!(
217 "invalid bound specification for a {} partitioned table",
218 match strategy {
219 PartitionStrategy::List => "list",
220 PartitionStrategy::Range => "range",
221 PartitionStrategy::Hash => "hash",
222 }
223 ))),
224 }
225}
226
227fn partition_bounds_overlap(
228 context: &PartitionContext<'_>,
229 left: &crate::ast::PartitionBound,
230 right: &crate::ast::PartitionBound,
231) -> Result<bool, SQLError> {
232 use crate::ast::PartitionBound;
233 match (left, right) {
234 (PartitionBound::Default, PartitionBound::Default) => Ok(true),
235 (PartitionBound::Default, _) | (_, PartitionBound::Default) => Ok(false),
236 (PartitionBound::List(left), PartitionBound::List(right)) => {
237 let left = evaluate_bound_values(context, left)?;
238 let right = evaluate_bound_values(context, right)?;
239 Ok(left.iter().any(|value| right.contains(value)))
240 }
241 (
242 PartitionBound::Range {
243 lower: left_lower,
244 upper: left_upper,
245 },
246 PartitionBound::Range {
247 lower: right_lower,
248 upper: right_upper,
249 },
250 ) => Ok(
251 compare_partition_points(context, left_lower, right_upper)? == Ordering::Less
252 && compare_partition_points(context, right_lower, left_upper)? == Ordering::Less,
253 ),
254 (
255 PartitionBound::Hash {
256 modulus: left_modulus,
257 remainder: left_remainder,
258 },
259 PartitionBound::Hash {
260 modulus: right_modulus,
261 remainder: right_remainder,
262 },
263 ) => hash::bounds_overlap(
264 *left_modulus,
265 *left_remainder,
266 *right_modulus,
267 *right_remainder,
268 ),
269 _ => Err(SQLError::Internal(
270 "partition siblings use different bound strategies".into(),
271 )),
272 }
273}
274
275fn evaluate_bound_values(
276 context: &PartitionContext<'_>,
277 expressions: &[crate::ast::Expr],
278) -> Result<Vec<Value>, SQLError> {
279 expressions
280 .iter()
281 .map(|expression| context.expressions.evaluate_bound(expression, &[]))
282 .collect()
283}
284
285fn compare_partition_points(
286 context: &PartitionContext<'_>,
287 left: &[crate::ast::PartitionRangeDatum],
288 right: &[crate::ast::PartitionRangeDatum],
289) -> Result<Ordering, SQLError> {
290 if left.len() != right.len() {
291 return Err(invalid_partition_bound(
292 "partition range points have different widths",
293 ));
294 }
295 for (left, right) in left.iter().zip(right) {
296 let ordering = match (left, right) {
297 (
298 crate::ast::PartitionRangeDatum::MinValue,
299 crate::ast::PartitionRangeDatum::MinValue,
300 )
301 | (
302 crate::ast::PartitionRangeDatum::MaxValue,
303 crate::ast::PartitionRangeDatum::MaxValue,
304 ) => Ordering::Equal,
305 (crate::ast::PartitionRangeDatum::MinValue, _)
306 | (_, crate::ast::PartitionRangeDatum::MaxValue) => Ordering::Less,
307 (crate::ast::PartitionRangeDatum::MaxValue, _)
308 | (_, crate::ast::PartitionRangeDatum::MinValue) => Ordering::Greater,
309 (
310 crate::ast::PartitionRangeDatum::Value(left),
311 crate::ast::PartitionRangeDatum::Value(right),
312 ) => context
313 .expressions
314 .evaluate_bound(left, &[])?
315 .cmp(&context.expressions.evaluate_bound(right, &[])?),
316 };
317 if ordering != Ordering::Equal {
318 return Ok(ordering);
319 }
320 }
321 Ok(Ordering::Equal)
322}
323
324fn invalid_partition_bound(message: impl Into<String>) -> SQLError {
325 SQLError::Routine {
326 sqlstate: "42P17".into(),
327 message: message.into(),
328 }
329}
330
331#[derive(Debug, Clone, PartialEq)]
333pub enum PartitionRejection {
334 Constraint { relation: String },
336 NoPartition { relation: String, keys: Vec<Value> },
338}
339
340#[derive(Debug, Clone, PartialEq)]
342pub enum PartitionRoute {
343 Target(String),
345 Rejected(PartitionRejection),
347}
348
349pub fn route_partition_insert(
351 context: &PartitionContext<'_>,
352 requested_table: &str,
353 document: &ResultRow,
354 params: &[SQLParam],
355 include_descendants: bool,
356) -> Result<PartitionRoute, SQLError> {
357 let table = context
358 .catalog
359 .try_resolve_table_name(requested_table)
360 .map_err(|error| SQLError::Internal(format!("resolve INSERT table: {error}")))?
361 .ok_or_else(|| SQLError::UnknownTable(requested_table.to_string()))?;
362 let hierarchy = context
363 .catalog
364 .try_table_hierarchy(&table)
365 .map_err(|error| SQLError::Internal(format!("read partition metadata: {error}")))?;
366 let Some(spec) = hierarchy.partition_spec.as_ref() else {
367 return Ok(PartitionRoute::Target(table));
368 };
369 if !include_descendants {
370 return Err(SQLError::Routine {
371 sqlstate: "42809".into(),
372 message: format!("cannot insert into partitioned table \"{requested_table}\""),
373 });
374 }
375 if !partition_constraint_accepts_row(context, &table, document, params)? {
376 return Ok(PartitionRoute::Rejected(PartitionRejection::Constraint {
377 relation: table,
378 }));
379 }
380 route_partition_tree(context, &table, spec, document, params)
381}
382
383pub fn partition_constraint_accepts_row(
385 context: &PartitionContext<'_>,
386 table: &str,
387 document: &ResultRow,
388 params: &[SQLParam],
389) -> Result<bool, SQLError> {
390 let mut child = table.to_string();
391 let mut visited = std::collections::BTreeSet::new();
392 loop {
393 if !visited.insert(child.clone()) {
394 return Err(SQLError::Internal(format!(
395 "partition hierarchy cycle reaches `{child}`"
396 )));
397 }
398 let hierarchy = context
399 .catalog
400 .try_table_hierarchy(&child)
401 .map_err(|error| SQLError::Internal(format!("read partition metadata: {error}")))?;
402 if hierarchy.partition_bound.is_none() {
403 return Ok(true);
404 }
405 let parent = hierarchy.parents.first().ok_or_else(|| {
406 SQLError::Internal(format!("partition `{child}` has no parent relation"))
407 })?;
408 let parent_hierarchy = context
409 .catalog
410 .try_table_hierarchy(parent)
411 .map_err(|error| {
412 SQLError::Internal(format!("read parent partition metadata: {error}"))
413 })?;
414 let spec = parent_hierarchy.partition_spec.as_ref().ok_or_else(|| {
415 SQLError::Internal(format!("partition parent `{parent}` has no partition key"))
416 })?;
417 match select_direct_partition(context, parent, spec, document, params)? {
418 DirectPartition::Child(selected) if selected == child => {}
419 DirectPartition::Child(_) | DirectPartition::None { .. } => return Ok(false),
420 }
421 child.clone_from(parent);
422 }
423}
424
425fn route_partition_tree(
426 context: &PartitionContext<'_>,
427 table: &str,
428 spec: &crate::ast::PartitionSpec,
429 document: &ResultRow,
430 params: &[SQLParam],
431) -> Result<PartitionRoute, SQLError> {
432 let child = match select_direct_partition(context, table, spec, document, params)? {
433 DirectPartition::Child(child) => child,
434 DirectPartition::None { keys } => {
435 return Ok(PartitionRoute::Rejected(PartitionRejection::NoPartition {
436 relation: table.to_string(),
437 keys,
438 }))
439 }
440 };
441 let hierarchy = context
442 .catalog
443 .try_table_hierarchy(&child)
444 .map_err(|error| SQLError::Internal(format!("read partition metadata: {error}")))?;
445 match hierarchy.partition_spec.as_ref() {
446 Some(spec) => route_partition_tree(context, &child, spec, document, params),
447 None => Ok(PartitionRoute::Target(child)),
448 }
449}
450
451enum DirectPartition {
453 Child(String),
454 None { keys: Vec<Value> },
455}
456
457fn select_direct_partition(
458 context: &PartitionContext<'_>,
459 parent: &str,
460 spec: &crate::ast::PartitionSpec,
461 document: &ResultRow,
462 params: &[SQLParam],
463) -> Result<DirectPartition, SQLError> {
464 let (keys, definitions) =
465 evaluate_partition_keys(context, parent, &spec.keys, document, params)?;
466 let row_hash = (spec.strategy == crate::ast::PartitionStrategy::Hash)
467 .then(|| hash::row_hash(context.types, spec, &definitions, &keys))
468 .transpose()?;
469 let mut default = None;
470 for child in context.catalog.direct_hierarchy_children(parent)? {
471 let child_hierarchy = context
472 .catalog
473 .try_table_hierarchy(&child)
474 .map_err(|error| SQLError::Internal(format!("read child partition: {error}")))?;
475 let Some(bound) = child_hierarchy.partition_bound.as_ref() else {
476 continue;
477 };
478 if matches!(bound, crate::ast::PartitionBound::Default) {
479 if default.replace(child).is_some() {
480 return Err(SQLError::Internal(format!(
481 "partitioned table `{parent}` has more than one default partition"
482 )));
483 }
484 continue;
485 }
486 if partition_bound_matches(context, bound, &keys, params, row_hash)? {
487 return Ok(DirectPartition::Child(child));
488 }
489 }
490 Ok(default.map_or(DirectPartition::None { keys }, DirectPartition::Child))
491}
492
493fn evaluate_partition_keys(
494 context: &PartitionContext<'_>,
495 table: &str,
496 expressions: &[crate::ast::Expr],
497 document: &ResultRow,
498 params: &[SQLParam],
499) -> Result<(Vec<Value>, Vec<crate::ast::ColumnDef>), SQLError> {
500 let definitions = context
501 .catalog
502 .try_describe_table(table)
503 .map_err(|error| SQLError::Internal(format!("read partition row type: {error}")))?
504 .ok_or_else(|| SQLError::UnknownTable(table.to_string()))?;
505 let schema = crate::RowSchema::with_types(
506 definitions
507 .iter()
508 .map(|definition| definition.name.clone())
509 .collect(),
510 definitions
511 .iter()
512 .map(|definition| Some(definition.ty.clone()))
513 .collect(),
514 );
515 let values = expressions
516 .iter()
517 .map(|expression| {
518 context
519 .expressions
520 .evaluate_row(expression, document, &schema, params)
521 })
522 .collect::<Result<Vec<_>, _>>()?;
523 Ok((values, definitions))
524}
525
526fn partition_bound_matches(
527 context: &PartitionContext<'_>,
528 bound: &crate::ast::PartitionBound,
529 keys: &[Value],
530 params: &[SQLParam],
531 row_hash: Option<u64>,
532) -> Result<bool, SQLError> {
533 use crate::ast::PartitionBound;
534 match bound {
535 PartitionBound::Default => Ok(true),
536 PartitionBound::List(values) => {
537 let [key] = keys else {
538 return Err(SQLError::Internal(
539 "LIST partition has more than one partition key".into(),
540 ));
541 };
542 for expression in values {
543 if context.expressions.evaluate_bound(expression, params)? == *key {
544 return Ok(true);
545 }
546 }
547 Ok(false)
548 }
549 PartitionBound::Range { lower, upper } => {
550 if keys.iter().any(|value| matches!(value, Value::Null)) {
551 return Ok(false);
552 }
553 Ok(
554 compare_key_to_bound(context, keys, lower, params)? != Ordering::Less
555 && compare_key_to_bound(context, keys, upper, params)? == Ordering::Less,
556 )
557 }
558 PartitionBound::Hash { modulus, remainder } => hash::bound_matches(
559 row_hash.ok_or_else(|| {
560 SQLError::Internal("HASH partition bound has no computed row hash".into())
561 })?,
562 *modulus,
563 *remainder,
564 ),
565 }
566}
567
568fn compare_key_to_bound(
569 context: &PartitionContext<'_>,
570 keys: &[Value],
571 bound: &[crate::ast::PartitionRangeDatum],
572 params: &[SQLParam],
573) -> Result<Ordering, SQLError> {
574 if keys.len() != bound.len() {
575 return Err(SQLError::Internal(format!(
576 "partition key width {} differs from bound width {}",
577 keys.len(),
578 bound.len()
579 )));
580 }
581 for (key, datum) in keys.iter().zip(bound) {
582 let ordering = match datum {
583 crate::ast::PartitionRangeDatum::MinValue => Ordering::Greater,
584 crate::ast::PartitionRangeDatum::MaxValue => Ordering::Less,
585 crate::ast::PartitionRangeDatum::Value(expression) => {
586 key.cmp(&context.expressions.evaluate_bound(expression, params)?)
587 }
588 };
589 if ordering != Ordering::Equal {
590 return Ok(ordering);
591 }
592 }
593 Ok(Ordering::Equal)
594}
595
596mod identity;
597mod order;
598pub use order::partition_bound_order;
599mod tree;
600pub use identity::{
601 foreign_key_scan_tables, partition_ancestor_tables, partition_hierarchy_root,
602 partition_identity_owner,
603};
604pub use tree::{partition_tree, PartitionTreeNode};