-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
emit an $addFields stage before $sort with safe aliases if necessary (#…
…109) * emit an $addFields stage before $sort with safe aliases if necessary * update changelog
- Loading branch information
Showing
10 changed files
with
217 additions
and
76 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,14 @@ | ||
| use mongodb::bson; | ||
| use serde::{Deserialize, Serialize}; | ||
|
|
||
| /// Wraps a BSON document that represents a set of sort criteria. A SortDocument value is intended | ||
| /// to be used as the argument to a $sort pipeline stage. | ||
| #[derive(Clone, Debug, PartialEq, Deserialize, Serialize)] | ||
| #[serde(transparent)] | ||
| pub struct SortDocument(pub bson::Document); | ||
|
|
||
| impl SortDocument { | ||
| pub fn from_doc(doc: bson::Document) -> Self { | ||
| SortDocument(doc) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,65 +1,176 @@ | ||
| use itertools::Itertools as _; | ||
| use mongodb::bson::{bson, Document}; | ||
| use std::{collections::BTreeMap, iter::once}; | ||
|
|
||
| use itertools::join; | ||
| use mongodb::bson::bson; | ||
| use ndc_models::OrderDirection; | ||
|
|
||
| use crate::{ | ||
| interface_types::MongoAgentError, | ||
| mongo_query_plan::{OrderBy, OrderByTarget}, | ||
| mongodb::sanitize::safe_name, | ||
| mongodb::{sanitize::escape_invalid_variable_chars, SortDocument, Stage}, | ||
| }; | ||
|
|
||
| pub fn make_sort(order_by: &OrderBy) -> Result<Document, MongoAgentError> { | ||
| use super::column_ref::ColumnRef; | ||
|
|
||
| /// In a [SortDocument] there is no way to reference field names that need to be escaped, such as | ||
| /// names that begin with dollar signs. To sort on such fields we need to insert an $addFields | ||
| /// stage _before_ the $sort stage to map safe aliases. | ||
| type RequiredAliases<'a> = BTreeMap<String, ColumnRef<'a>>; | ||
|
|
||
| type Result<T> = std::result::Result<T, MongoAgentError>; | ||
|
|
||
| pub fn make_sort_stages(order_by: &OrderBy) -> Result<Vec<Stage>> { | ||
| let (sort_document, required_aliases) = make_sort(order_by)?; | ||
| let mut stages = vec![]; | ||
|
|
||
| if !required_aliases.is_empty() { | ||
| let fields = required_aliases | ||
| .into_iter() | ||
| .map(|(alias, expression)| (alias, expression.into_aggregate_expression())) | ||
| .collect(); | ||
| let stage = Stage::AddFields(fields); | ||
| stages.push(stage); | ||
| } | ||
|
|
||
| let sort_stage = Stage::Sort(sort_document); | ||
| stages.push(sort_stage); | ||
|
|
||
| Ok(stages) | ||
| } | ||
|
|
||
| fn make_sort(order_by: &OrderBy) -> Result<(SortDocument, RequiredAliases<'_>)> { | ||
| let OrderBy { elements } = order_by; | ||
|
|
||
| elements | ||
| .clone() | ||
| let keys_directions_expressions: BTreeMap<String, (OrderDirection, Option<ColumnRef<'_>>)> = | ||
| elements | ||
| .iter() | ||
| .map(|obe| { | ||
| let col_ref = ColumnRef::from_order_by_target(&obe.target)?; | ||
| let (key, required_alias) = match col_ref { | ||
| ColumnRef::MatchKey(key) => (key.to_string(), None), | ||
| ref_expr => (safe_alias(&obe.target)?, Some(ref_expr)), | ||
| }; | ||
| Ok((key, (obe.order_direction, required_alias))) | ||
| }) | ||
| .collect::<Result<BTreeMap<_, _>>>()?; | ||
|
|
||
| let sort_document = keys_directions_expressions | ||
| .iter() | ||
| .map(|obe| { | ||
| let direction = match obe.clone().order_direction { | ||
| .map(|(key, (direction, _))| { | ||
| let direction_bson = match direction { | ||
| OrderDirection::Asc => bson!(1), | ||
| OrderDirection::Desc => bson!(-1), | ||
| }; | ||
| match &obe.target { | ||
| OrderByTarget::Column { | ||
| name, | ||
| field_path, | ||
| path, | ||
| } => Ok(( | ||
| column_ref_with_path(name, field_path.as_deref(), path)?, | ||
| direction, | ||
| )), | ||
| OrderByTarget::SingleColumnAggregate { | ||
| column: _, | ||
| function: _, | ||
| path: _, | ||
| result_type: _, | ||
| } => | ||
| // TODO: MDB-150 | ||
| { | ||
| Err(MongoAgentError::NotImplemented( | ||
| "ordering by single column aggregate".into(), | ||
| )) | ||
| } | ||
| OrderByTarget::StarCountAggregate { path: _ } => Err( | ||
| // TODO: MDB-151 | ||
| MongoAgentError::NotImplemented("ordering by star count aggregate".into()), | ||
| ), | ||
| } | ||
| (key.clone(), direction_bson) | ||
| }) | ||
| .collect() | ||
| .collect(); | ||
|
|
||
| let required_aliases = keys_directions_expressions | ||
| .into_iter() | ||
| .flat_map(|(key, (_, expr))| expr.map(|e| (key, e))) | ||
| .collect(); | ||
|
|
||
| Ok((SortDocument(sort_document), required_aliases)) | ||
| } | ||
|
|
||
| // TODO: MDB-159 Replace use of [safe_name] with [ColumnRef]. | ||
| fn column_ref_with_path( | ||
| name: &ndc_models::FieldName, | ||
| field_path: Option<&[ndc_models::FieldName]>, | ||
| relation_path: &[ndc_models::RelationshipName], | ||
| ) -> Result<String, MongoAgentError> { | ||
| relation_path | ||
| .iter() | ||
| .map(|n| n.as_str()) | ||
| .chain(std::iter::once(name.as_str())) | ||
| .chain(field_path.into_iter().flatten().map(|n| n.as_str())) | ||
| .map(safe_name) | ||
| .process_results(|mut iter| iter.join(".")) | ||
| fn safe_alias(target: &OrderByTarget) -> Result<String> { | ||
| match target { | ||
| ndc_query_plan::OrderByTarget::Column { | ||
| name, | ||
| field_path, | ||
| path, | ||
| } => { | ||
| let name_and_path = once("__sort_key_") | ||
| .chain(path.iter().map(|n| n.as_str())) | ||
| .chain([name.as_str()]) | ||
| .chain( | ||
| field_path | ||
| .iter() | ||
| .flatten() | ||
| .map(|field_name| field_name.as_str()), | ||
| ); | ||
| let combine_all_elements_into_one_name = join(name_and_path, "_"); | ||
| Ok(escape_invalid_variable_chars( | ||
| &combine_all_elements_into_one_name, | ||
| )) | ||
| } | ||
| ndc_query_plan::OrderByTarget::SingleColumnAggregate { .. } => { | ||
| // TODO: ENG-1011 | ||
| Err(MongoAgentError::NotImplemented( | ||
| "ordering by single column aggregate".into(), | ||
| )) | ||
| } | ||
| ndc_query_plan::OrderByTarget::StarCountAggregate { .. } => { | ||
| // TODO: ENG-1010 | ||
| Err(MongoAgentError::NotImplemented( | ||
| "ordering by star count aggregate".into(), | ||
| )) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod tests { | ||
| use mongodb::bson::doc; | ||
| use ndc_models::{FieldName, OrderDirection}; | ||
| use ndc_query_plan::OrderByElement; | ||
| use pretty_assertions::assert_eq; | ||
|
|
||
| use crate::{mongo_query_plan::OrderBy, mongodb::SortDocument, query::column_ref::ColumnRef}; | ||
|
|
||
| use super::make_sort; | ||
|
|
||
| #[test] | ||
| fn escapes_field_names() -> anyhow::Result<()> { | ||
| let order_by = OrderBy { | ||
| elements: vec![OrderByElement { | ||
| order_direction: OrderDirection::Asc, | ||
| target: ndc_query_plan::OrderByTarget::Column { | ||
| name: "$schema".into(), | ||
| field_path: Default::default(), | ||
| path: Default::default(), | ||
| }, | ||
| }], | ||
| }; | ||
| let path: [FieldName; 1] = ["$schema".into()]; | ||
|
|
||
| let actual = make_sort(&order_by)?; | ||
| let expected_sort_doc = SortDocument(doc! { | ||
| "__sort_key__·24schema": 1 | ||
| }); | ||
| let expected_aliases = [( | ||
| "__sort_key__·24schema".into(), | ||
| ColumnRef::from_field_path(path.iter()), | ||
| )] | ||
| .into(); | ||
| assert_eq!(actual, (expected_sort_doc, expected_aliases)); | ||
| Ok(()) | ||
| } | ||
|
|
||
| #[test] | ||
| fn escapes_nested_field_names() -> anyhow::Result<()> { | ||
| let order_by = OrderBy { | ||
| elements: vec![OrderByElement { | ||
| order_direction: OrderDirection::Asc, | ||
| target: ndc_query_plan::OrderByTarget::Column { | ||
| name: "configuration".into(), | ||
| field_path: Some(vec!["$schema".into()]), | ||
| path: Default::default(), | ||
| }, | ||
| }], | ||
| }; | ||
| let path: [FieldName; 2] = ["configuration".into(), "$schema".into()]; | ||
|
|
||
| let actual = make_sort(&order_by)?; | ||
| let expected_sort_doc = SortDocument(doc! { | ||
| "__sort_key__configuration_·24schema": 1 | ||
| }); | ||
| let expected_aliases = [( | ||
| "__sort_key__configuration_·24schema".into(), | ||
| ColumnRef::from_field_path(path.iter()), | ||
| )] | ||
| .into(); | ||
| assert_eq!(actual, (expected_sort_doc, expected_aliases)); | ||
| Ok(()) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.