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