diff --git a/pgml-sdks/pgml/src/lib.rs b/pgml-sdks/pgml/src/lib.rs index c32a4e5bb..87b99657c 100644 --- a/pgml-sdks/pgml/src/lib.rs +++ b/pgml-sdks/pgml/src/lib.rs @@ -1584,7 +1584,6 @@ mod tests { "nested_number": { "number": 3 }, - "tie": 2, }) .into(), @@ -1646,6 +1645,26 @@ mod tests { .collect::>(), vec![1, 2, 3] ); + let documents = collection + .get_documents(Some(json!({"order_by": { "COLUMN_id": "desc"}}).into())) + .await?; + assert_eq!( + documents + .iter() + .map(|d| d["row_id"].as_i64().unwrap()) + .collect::>(), + vec![3, 2, 1] + ); + let documents = collection + .get_documents(Some(json!({"order_by": { "COLUMN_id": "asc"}}).into())) + .await?; + assert_eq!( + documents + .iter() + .map(|d| d["row_id"].as_i64().unwrap()) + .collect::>(), + vec![1, 2, 3] + ); collection.archive().await?; Ok(()) } diff --git a/pgml-sdks/pgml/src/order_by_builder.rs b/pgml-sdks/pgml/src/order_by_builder.rs index 4198612af..4c3cd4269 100644 --- a/pgml-sdks/pgml/src/order_by_builder.rs +++ b/pgml-sdks/pgml/src/order_by_builder.rs @@ -7,6 +7,14 @@ pub(crate) struct OrderByBuilder<'a> { column_name: &'a str, } +fn str_to_order(order: &str) -> anyhow::Result { + match order { + "asc" | "ASC" => Ok(Order::Asc), + "desc" | "DESC" => Ok(Order::Desc), + _ => anyhow::bail!("Invalid `order_by`. Please refer to examples in the documentation for correct `order_by` syntax"), + } +} + fn build_recursive_access(key: &str, value: &serde_json::Value) -> anyhow::Result<(String, Order)> { if value.is_object() { let (new_key, new_value) = value @@ -14,19 +22,15 @@ fn build_recursive_access(key: &str, value: &serde_json::Value) -> anyhow::Resul .unwrap() .iter() .next() - .context("Invalid order by")?; + .context("Invalid `order_by`. Please refer to examples in the documentation for correct `order_by` syntax")?; let (path, order) = build_recursive_access(new_key, new_value)?; let path = format!("{},{}", key, path); Ok((path, order)) } else if value.is_string() { - let order = match value.as_str().unwrap() { - "asc" | "ASC" => Order::Asc, - "desc" | "DESC" => Order::Desc, - _ => return Err(anyhow::anyhow!("Invalid order by")), - }; + let order = str_to_order(value.as_str().unwrap())?; Ok((key.to_string(), order)) } else { - Err(anyhow::anyhow!("Invalid order by")) + Err(anyhow::anyhow!("Invalid `order_by`. Please refer to examples in the documentation for correct `order_by` syntax")) } } @@ -42,17 +46,22 @@ impl<'a> OrderByBuilder<'a> { pub fn build(self) -> anyhow::Result> { self.filter .as_object() - .context("Invalid order by")? + .context("`order_by` must be an object")? .iter() .map(|(k, v)| { - if let Ok((path, order)) = build_recursive_access(k, v) { + if k.starts_with("COLUMN_") { + Ok(( + Expr::cust(k.replace("COLUMN_", "")), + str_to_order(v.as_str().context("Invalid `order_by`. Please refer to examples in the documentation for correct `order_by` syntax")?)?, + )) + } else if let Ok((path, order)) = build_recursive_access(k, v) { let expr = Expr::cust(format!( "\"{}\".\"{}\"#>'{{{}}}'", self.table_name, self.column_name, path )); Ok((expr, order)) } else { - Err(anyhow::anyhow!("Invalid order by")) + Err(anyhow::anyhow!("Invalid `order_by`. Please refer to examples in the documentation for correct `order_by` syntax")) } }) .collect()