1use crate::{
10 assignment::AssignmentContext,
11 ast::{ColumnDef, Expr, TableHierarchy},
12 schema::SchemaExpressionCatalog,
13 type_resolution::FunctionTypeResolver,
14 ResultRow, RowSchema, SQLError, SQLParam,
15};
16use std::cmp::Ordering;
17use uqa_core::Value;
18
19pub trait PartitionCatalog {
21 fn try_table_hierarchy(&self, table: &str) -> Result<TableHierarchy, String>;
22 fn direct_hierarchy_children(&self, parent: &str) -> Result<Vec<String>, SQLError>;
23 fn try_resolve_table_name(&self, name: &str) -> Result<Option<String>, String>;
24 fn try_describe_table(&self, table: &str) -> Result<Option<Vec<ColumnDef>>, String>;
25 fn try_table_object_id(&self, table: &str) -> Result<Option<[u8; 16]>, String>;
27}
28
29pub trait PartitionExpressions {
31 fn expression_text(&self, expression: &Expr) -> Result<String, SQLError>;
33 fn evaluate_bound(&self, expression: &Expr, params: &[SQLParam]) -> Result<Value, SQLError>;
34 fn evaluate_row(
35 &self,
36 expression: &Expr,
37 row: &ResultRow,
38 schema: &RowSchema,
39 params: &[SQLParam],
40 ) -> Result<Value, SQLError>;
41}
42
43#[derive(Clone, Copy)]
44pub struct PartitionContext<'a> {
45 pub catalog: &'a dyn PartitionCatalog,
46 pub expressions: &'a dyn PartitionExpressions,
47 pub types: &'a dyn FunctionTypeResolver,
48 pub assignment: &'a dyn AssignmentContext,
50 pub schema: &'a dyn SchemaExpressionCatalog,
52}
53
54mod admission;
55mod bounds;
56mod datum_text;
57mod hash;
58mod key;
59pub use key::key_type as partition_key_type;
60
61pub use admission::validate_new_partition_bound;
62pub use bounds::transform_partition_bound;
63pub use datum_text::{partition_datum_text, range_bound_text, stored_datum};
64
65pub fn validate_hash_partition_spec(
66 context: &PartitionContext<'_>,
67 spec: &crate::ast::PartitionSpec,
68 columns: &[crate::ast::ColumnDef],
69) -> Result<(), SQLError> {
70 hash::validate_partition_spec(context.types, spec, columns)
71}
72
73pub fn prospective_partition_bound_accepts_document(
77 context: &PartitionContext<'_>,
78 parent: &str,
79 bound: &crate::ast::PartitionBound,
80 document: &ResultRow,
81) -> Result<bool, SQLError> {
82 let hierarchy = context
83 .catalog
84 .try_table_hierarchy(parent)
85 .map_err(|error| SQLError::Internal(format!("read parent partition metadata: {error}")))?;
86 let spec = hierarchy
87 .partition_spec
88 .as_ref()
89 .ok_or_else(|| SQLError::Routine {
90 sqlstate: "42809".into(),
91 message: format!("relation \"{parent}\" is not partitioned"),
92 })?;
93 let key = routing_key(context, parent, spec, document, &[])?;
94 if !matches!(bound, crate::ast::PartitionBound::Default) {
95 return partition_bound_matches(context, bound, &key, &[]);
96 }
97 for sibling in context.catalog.direct_hierarchy_children(parent)? {
98 let sibling_hierarchy = context
99 .catalog
100 .try_table_hierarchy(&sibling)
101 .map_err(|error| SQLError::Internal(format!("read child partition: {error}")))?;
102 let Some(sibling_bound) = sibling_hierarchy.partition_bound.as_ref() else {
103 continue;
104 };
105 if matches!(sibling_bound, crate::ast::PartitionBound::Default) {
106 continue;
107 }
108 if partition_bound_matches(context, sibling_bound, &key, &[])? {
109 return Ok(false);
110 }
111 }
112 Ok(true)
113}
114
115pub fn partition_constraint_accepts_document(
119 context: &PartitionContext<'_>,
120 table: &str,
121 spec: &crate::ast::PartitionSpec,
122 bound: &crate::ast::PartitionBound,
123 document: &ResultRow,
124) -> Result<bool, SQLError> {
125 let key = routing_key(context, table, spec, document, &[])?;
126 partition_bound_matches(context, bound, &key, &[])
127}
128
129struct RoutingKey {
131 values: Vec<Value>,
132 hash: Option<u64>,
133}
134
135fn routing_key(
136 context: &PartitionContext<'_>,
137 table: &str,
138 spec: &crate::ast::PartitionSpec,
139 document: &ResultRow,
140 params: &[SQLParam],
141) -> Result<RoutingKey, SQLError> {
142 let (values, definitions) =
143 evaluate_partition_keys(context, table, &spec.keys, document, params)?;
144 let hash = (spec.strategy == crate::ast::PartitionStrategy::Hash)
145 .then(|| hash::row_hash(context, spec, &definitions, &values))
146 .transpose()?;
147 Ok(RoutingKey { values, hash })
148}
149
150#[derive(Debug, Clone, PartialEq)]
152pub enum PartitionRejection {
153 Constraint { relation: String },
155 NoPartition { relation: String, keys: Vec<Value> },
157}
158
159#[derive(Debug, Clone, PartialEq)]
161pub enum PartitionRoute {
162 Target(String),
164 Rejected(PartitionRejection),
166}
167
168pub fn route_partition_insert(
170 context: &PartitionContext<'_>,
171 requested_table: &str,
172 document: &ResultRow,
173 params: &[SQLParam],
174 include_descendants: bool,
175) -> Result<PartitionRoute, SQLError> {
176 let table = context
177 .catalog
178 .try_resolve_table_name(requested_table)
179 .map_err(|error| SQLError::Internal(format!("resolve INSERT table: {error}")))?
180 .ok_or_else(|| SQLError::UnknownTable(requested_table.to_string()))?;
181 let hierarchy = context
182 .catalog
183 .try_table_hierarchy(&table)
184 .map_err(|error| SQLError::Internal(format!("read partition metadata: {error}")))?;
185 let Some(spec) = hierarchy.partition_spec.as_ref() else {
186 return Ok(PartitionRoute::Target(table));
187 };
188 if !include_descendants {
189 return Err(SQLError::Routine {
190 sqlstate: "42809".into(),
191 message: format!("cannot insert into partitioned table \"{requested_table}\""),
192 });
193 }
194 if !partition_constraint_accepts_row(context, &table, document, params)? {
195 return Ok(PartitionRoute::Rejected(PartitionRejection::Constraint {
196 relation: table,
197 }));
198 }
199 route_partition_tree(context, &table, spec, document, params)
200}
201
202pub fn partition_constraint_accepts_row(
204 context: &PartitionContext<'_>,
205 table: &str,
206 document: &ResultRow,
207 params: &[SQLParam],
208) -> Result<bool, SQLError> {
209 let mut child = table.to_string();
210 let mut visited = std::collections::BTreeSet::new();
211 loop {
212 if !visited.insert(child.clone()) {
213 return Err(SQLError::Internal(format!(
214 "partition hierarchy cycle reaches `{child}`"
215 )));
216 }
217 let hierarchy = context
218 .catalog
219 .try_table_hierarchy(&child)
220 .map_err(|error| SQLError::Internal(format!("read partition metadata: {error}")))?;
221 if hierarchy.partition_bound.is_none() {
222 return Ok(true);
223 }
224 let parent = hierarchy.parents.first().ok_or_else(|| {
225 SQLError::Internal(format!("partition `{child}` has no parent relation"))
226 })?;
227 let parent_hierarchy = context
228 .catalog
229 .try_table_hierarchy(parent)
230 .map_err(|error| {
231 SQLError::Internal(format!("read parent partition metadata: {error}"))
232 })?;
233 let spec = parent_hierarchy.partition_spec.as_ref().ok_or_else(|| {
234 SQLError::Internal(format!("partition parent `{parent}` has no partition key"))
235 })?;
236 match select_direct_partition(context, parent, spec, document, params)? {
237 DirectPartition::Child(selected) if selected == child => {}
238 DirectPartition::Child(_) | DirectPartition::None { .. } => return Ok(false),
239 }
240 child.clone_from(parent);
241 }
242}
243
244fn route_partition_tree(
245 context: &PartitionContext<'_>,
246 table: &str,
247 spec: &crate::ast::PartitionSpec,
248 document: &ResultRow,
249 params: &[SQLParam],
250) -> Result<PartitionRoute, SQLError> {
251 let child = match select_direct_partition(context, table, spec, document, params)? {
252 DirectPartition::Child(child) => child,
253 DirectPartition::None { keys } => {
254 return Ok(PartitionRoute::Rejected(PartitionRejection::NoPartition {
255 relation: table.to_string(),
256 keys,
257 }))
258 }
259 };
260 let hierarchy = context
261 .catalog
262 .try_table_hierarchy(&child)
263 .map_err(|error| SQLError::Internal(format!("read partition metadata: {error}")))?;
264 match hierarchy.partition_spec.as_ref() {
265 Some(spec) => route_partition_tree(context, &child, spec, document, params),
266 None => Ok(PartitionRoute::Target(child)),
267 }
268}
269
270enum DirectPartition {
272 Child(String),
273 None { keys: Vec<Value> },
274}
275
276fn select_direct_partition(
277 context: &PartitionContext<'_>,
278 parent: &str,
279 spec: &crate::ast::PartitionSpec,
280 document: &ResultRow,
281 params: &[SQLParam],
282) -> Result<DirectPartition, SQLError> {
283 let key = routing_key(context, parent, spec, document, params)?;
284 Ok(match select_partition(context, parent, &key, params)? {
285 Some(child) => DirectPartition::Child(child),
286 None => DirectPartition::None { keys: key.values },
287 })
288}
289
290fn select_partition(
291 context: &PartitionContext<'_>,
292 parent: &str,
293 key: &RoutingKey,
294 params: &[SQLParam],
295) -> Result<Option<String>, SQLError> {
296 let mut default = None;
297 for child in context.catalog.direct_hierarchy_children(parent)? {
298 let child_hierarchy = context
299 .catalog
300 .try_table_hierarchy(&child)
301 .map_err(|error| SQLError::Internal(format!("read child partition: {error}")))?;
302 let Some(bound) = child_hierarchy.partition_bound.as_ref() else {
303 continue;
304 };
305 if matches!(bound, crate::ast::PartitionBound::Default) {
306 if default.replace(child).is_some() {
307 return Err(SQLError::Internal(format!(
308 "partitioned table `{parent}` has more than one default partition"
309 )));
310 }
311 continue;
312 }
313 if partition_bound_matches(context, bound, key, params)? {
314 return Ok(Some(child));
315 }
316 }
317 Ok(default)
318}
319
320fn evaluate_partition_keys(
321 context: &PartitionContext<'_>,
322 table: &str,
323 expressions: &[crate::ast::Expr],
324 document: &ResultRow,
325 params: &[SQLParam],
326) -> Result<(Vec<Value>, Vec<crate::ast::ColumnDef>), SQLError> {
327 let definitions = context
328 .catalog
329 .try_describe_table(table)
330 .map_err(|error| SQLError::Internal(format!("read partition row type: {error}")))?
331 .ok_or_else(|| SQLError::UnknownTable(table.to_string()))?;
332 let schema = crate::RowSchema::with_types(
333 definitions
334 .iter()
335 .map(|definition| definition.name.clone())
336 .collect(),
337 definitions
338 .iter()
339 .map(|definition| Some(definition.ty.clone()))
340 .collect(),
341 );
342 let values = expressions
343 .iter()
344 .map(|expression| {
345 context
346 .expressions
347 .evaluate_row(expression, document, &schema, params)
348 })
349 .collect::<Result<Vec<_>, _>>()?;
350 Ok((values, definitions))
351}
352
353fn partition_bound_matches(
354 context: &PartitionContext<'_>,
355 bound: &crate::ast::PartitionBound,
356 key: &RoutingKey,
357 params: &[SQLParam],
358) -> Result<bool, SQLError> {
359 use crate::ast::PartitionBound;
360 let keys = key.values.as_slice();
361 match bound {
362 PartitionBound::Default => Ok(true),
363 PartitionBound::List(values) => {
364 let [key] = keys else {
365 return Err(SQLError::Internal(
366 "LIST partition has more than one partition key".into(),
367 ));
368 };
369 for expression in values {
370 if context
371 .expressions
372 .evaluate_bound(expression, params)?
373 .cmp(key)
374 == Ordering::Equal
375 {
376 return Ok(true);
377 }
378 }
379 Ok(false)
380 }
381 PartitionBound::Range { lower, upper } => {
382 if keys.iter().any(|value| matches!(value, Value::Null)) {
383 return Ok(false);
384 }
385 Ok(
386 compare_key_to_bound(context, keys, lower, params)? != Ordering::Less
387 && compare_key_to_bound(context, keys, upper, params)? == Ordering::Less,
388 )
389 }
390 PartitionBound::Hash { modulus, remainder } => hash::bound_matches(
391 key.hash.ok_or_else(|| {
392 SQLError::Internal("HASH partition bound has no computed row hash".into())
393 })?,
394 *modulus,
395 *remainder,
396 ),
397 }
398}
399
400fn compare_key_to_bound(
401 context: &PartitionContext<'_>,
402 keys: &[Value],
403 bound: &[crate::ast::PartitionRangeDatum],
404 params: &[SQLParam],
405) -> Result<Ordering, SQLError> {
406 if keys.len() != bound.len() {
407 return Err(SQLError::Internal(format!(
408 "partition key width {} differs from bound width {}",
409 keys.len(),
410 bound.len()
411 )));
412 }
413 for (key, datum) in keys.iter().zip(bound) {
414 let ordering = match datum {
415 crate::ast::PartitionRangeDatum::MinValue => Ordering::Greater,
416 crate::ast::PartitionRangeDatum::MaxValue => Ordering::Less,
417 crate::ast::PartitionRangeDatum::Value(expression) => {
418 key.cmp(&context.expressions.evaluate_bound(expression, params)?)
419 }
420 };
421 if ordering != Ordering::Equal {
422 return Ok(ordering);
423 }
424 }
425 Ok(Ordering::Equal)
426}
427
428mod identity;
429mod order;
430pub use order::partition_bound_order;
431mod tree;
432pub use identity::{
433 foreign_key_scan_tables, partition_ancestor_tables, partition_hierarchy_root,
434 partition_identity_owner,
435};
436pub use tree::{partition_tree, PartitionTreeNode};