Skip to main content

datafusion_odata/
collection.rs

1use datafusion::prelude::*;
2
3use crate::{error::ODataError, filter::ODataFilter};
4
5///////////////////////////////////////////////////////////////////////////////
6
7#[derive(Debug, serde::Deserialize)]
8pub struct QueryParamsRaw {
9    #[serde(rename = "$select")]
10    pub select: Option<String>,
11    #[serde(rename = "$orderby")]
12    pub order_by: Option<String>,
13    #[serde(rename = "$skip")]
14    pub skip: Option<u64>,
15    #[serde(rename = "$top")]
16    pub top: Option<u64>,
17    #[serde(rename = "$filter")]
18    pub filter: Option<ODataFilter>,
19}
20
21///////////////////////////////////////////////////////////////////////////////
22
23impl QueryParamsRaw {
24    pub fn decode(self) -> Result<QueryParams, ODataError> {
25        let select = self.select.unwrap_or_default();
26        let mut select: Vec<_> = select.split(',').map(|s| s.to_string()).collect();
27        select.retain(|i| !i.is_empty());
28
29        let order_by_s = self.order_by.unwrap_or_default();
30        let mut order_by_s: Vec<_> = order_by_s.split(',').collect();
31        order_by_s.retain(|i| !i.is_empty());
32
33        let mut order_by = Vec::new();
34        for el in order_by_s {
35            let (cname, asc) = if let Some(cname) = el.strip_suffix(" asc") {
36                (cname, true)
37            } else if let Some(cname) = el.strip_suffix(" desc") {
38                (cname, false)
39            } else {
40                (el, true)
41            };
42            order_by.push((cname.to_string(), asc));
43        }
44
45        let skip = self.skip.map(|v| v as usize);
46        let top = self.top.map(|v| v as usize);
47
48        Ok(QueryParams {
49            select,
50            order_by,
51            skip,
52            top,
53            filter: self.filter.map(Into::into),
54        })
55    }
56}
57
58///////////////////////////////////////////////////////////////////////////////
59
60#[derive(Debug)]
61pub struct QueryParams {
62    /// Column names
63    pub select: Vec<String>,
64    /// Tuples (column_name, ascending)
65    pub order_by: Vec<(String, bool)>,
66    /// Number of records to skip
67    pub skip: Option<usize>,
68    /// Maximum number of records to return
69    pub top: Option<usize>,
70    /// Filter a collection of resources   
71    pub filter: Option<Expr>,
72}
73
74///////////////////////////////////////////////////////////////////////////////
75
76impl QueryParams {
77    pub fn apply(
78        self,
79        df: DataFrame,
80        addr: &CollectionAddr,
81        key_column: &str,
82        key_column_alias: &str,
83        default_rows: usize,
84        max_rows: usize,
85    ) -> datafusion::error::Result<DataFrame> {
86        // Add key column as alias
87        let df = df.with_column(key_column_alias, col(key_column))?;
88
89        // Select desired columns
90        let df = if self.select.is_empty() {
91            df
92        } else {
93            let mut select: Vec<_> = self.select.iter().map(String::as_str).collect();
94            select.push(key_column_alias);
95            df.select_columns(&select)?
96        };
97
98        // If queried by key - ignore the rest
99        if let Some(key) = &addr.key {
100            return df.filter(col(key_column_alias).eq(lit(key.clone())));
101        }
102
103        let df = match self.filter {
104            Some(filter) => df.filter(filter)?,
105            None => df,
106        };
107
108        // Order by
109        let df = if self.order_by.is_empty() {
110            df
111        } else {
112            df.sort(
113                self.order_by
114                    .into_iter()
115                    .map(|(c, asc)| col(c).sort(asc, true))
116                    .collect(),
117            )?
118        };
119
120        // Skip / limit
121        df.limit(
122            self.skip.unwrap_or(0),
123            Some(std::cmp::min(self.top.unwrap_or(default_rows), max_rows)),
124        )
125    }
126}
127
128///////////////////////////////////////////////////////////////////////////////
129
130#[derive(Debug, Clone, PartialEq, Eq)]
131pub struct CollectionAddr {
132    pub name: String,
133    pub key: Option<String>,
134}
135
136impl CollectionAddr {
137    pub fn decode(collection_path_element: &str) -> Option<Self> {
138        let re = regex::Regex::new(r#"^(?<name>[A-Za-z0-9._-]+)(\((?<key>[^)]+)\))?$"#).unwrap();
139        let c = re.captures(collection_path_element)?;
140
141        let name = c.name("name")?.as_str().to_string();
142        let key = c.name("key").map(|m| m.as_str().to_string());
143
144        Some(Self { name, key })
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use crate::collection::CollectionAddr;
151
152    #[test]
153    fn test_collection_addr_decode() {
154        assert_eq!(
155            CollectionAddr::decode("coll"),
156            Some(CollectionAddr {
157                name: "coll".to_string(),
158                key: None,
159            })
160        );
161
162        assert_eq!(
163            CollectionAddr::decode("Coll123"),
164            Some(CollectionAddr {
165                name: "Coll123".to_string(),
166                key: None,
167            })
168        );
169
170        assert_eq!(
171            CollectionAddr::decode("Coll.x_12-3"),
172            Some(CollectionAddr {
173                name: "Coll.x_12-3".to_string(),
174                key: None,
175            })
176        );
177
178        assert_eq!(
179            CollectionAddr::decode("Coll(123)"),
180            Some(CollectionAddr {
181                name: "Coll".to_string(),
182                key: Some("123".to_string()),
183            })
184        );
185
186        assert_eq!(
187            CollectionAddr::decode("Coll('key')"),
188            Some(CollectionAddr {
189                name: "Coll".to_string(),
190                key: Some("'key'".to_string()),
191            })
192        );
193    }
194}