1const GROUP_BY_STATUS: &str = "GROUP BY ExecutionStatus";
14
15pub fn count_query(filter: &str) -> String {
20 let filter = strip_clause(filter, "group by");
21 let filter = strip_clause(&filter, "order by");
22 let filter = filter.trim();
23 if filter.is_empty() {
24 GROUP_BY_STATUS.to_string()
25 } else {
26 format!("{filter} {GROUP_BY_STATUS}")
27 }
28}
29
30fn strip_clause(query: &str, keyword: &str) -> String {
36 match find_clause(query, keyword) {
37 Some(at) => query[..at].trim_end().to_string(),
38 None => query.trim_end().to_string(),
39 }
40}
41
42fn find_clause(query: &str, keyword: &str) -> Option<usize> {
46 let words: Vec<&str> = keyword.split_whitespace().collect();
47 let bytes = query.as_bytes();
48 let mut in_quote = false;
49 let mut found = None;
50 let mut i = 0;
51
52 while i < bytes.len() {
53 if bytes[i] == b'\'' {
54 in_quote = !in_quote;
55 i += 1;
56 continue;
57 }
58 if in_quote {
59 i += 1;
60 continue;
61 }
62 let at_boundary = i == 0 || !is_word_byte(bytes[i - 1]);
64 if at_boundary && let Some(end) = match_words(bytes, i, &words) {
65 found = Some(i);
66 i = end;
67 continue;
68 }
69 i += 1;
70 }
71 found
72}
73
74fn match_words(bytes: &[u8], start: usize, words: &[&str]) -> Option<usize> {
76 let mut i = start;
77 for (n, word) in words.iter().enumerate() {
78 if n > 0 {
79 let ws = i;
80 while i < bytes.len() && bytes[i].is_ascii_whitespace() {
81 i += 1;
82 }
83 if i == ws {
84 return None;
85 }
86 }
87 let end = i + word.len();
88 if end > bytes.len() || !bytes[i..end].eq_ignore_ascii_case(word.as_bytes()) {
89 return None;
90 }
91 i = end;
92 }
93 if i < bytes.len() && is_word_byte(bytes[i]) {
95 return None;
96 }
97 Some(i)
98}
99
100fn is_word_byte(b: u8) -> bool {
101 b.is_ascii_alphanumeric() || b == b'_'
102}
103
104#[cfg(test)]
105mod tests {
106 use super::*;
107
108 #[test]
109 fn an_empty_filter_still_groups() {
110 assert_eq!(count_query(""), "GROUP BY ExecutionStatus");
111 assert_eq!(count_query(" "), "GROUP BY ExecutionStatus");
112 }
113
114 #[test]
115 fn a_filter_is_kept_and_grouped() {
116 assert_eq!(
117 count_query("WorkflowType = 'Foo'"),
118 "WorkflowType = 'Foo' GROUP BY ExecutionStatus"
119 );
120 }
121
122 #[test]
123 fn order_by_is_stripped_because_count_rejects_it() {
124 assert_eq!(
125 count_query("WorkflowType = 'Foo' ORDER BY StartTime DESC"),
126 "WorkflowType = 'Foo' GROUP BY ExecutionStatus"
127 );
128 assert_eq!(
130 count_query("WorkflowType = 'Foo' order by StartTime"),
131 "WorkflowType = 'Foo' GROUP BY ExecutionStatus"
132 );
133 }
134
135 #[test]
136 fn an_existing_group_by_is_replaced_not_appended() {
137 assert_eq!(
138 count_query("WorkflowType = 'Foo' GROUP BY WorkflowType"),
139 "WorkflowType = 'Foo' GROUP BY ExecutionStatus"
140 );
141 }
142
143 #[test]
144 fn both_clauses_are_stripped_together() {
145 assert_eq!(
146 count_query("A = 1 GROUP BY B ORDER BY C"),
147 "A = 1 GROUP BY ExecutionStatus"
148 );
149 }
150
151 #[test]
152 fn a_quoted_value_is_not_mistaken_for_a_clause() {
153 assert_eq!(
156 count_query("WorkflowId = 'daily order by region'"),
157 "WorkflowId = 'daily order by region' GROUP BY ExecutionStatus"
158 );
159 }
160
161 #[test]
162 fn a_keyword_inside_an_identifier_is_not_a_clause() {
163 assert_eq!(
164 count_query("MyOrder By_Field = 1"),
165 "MyOrder By_Field = 1 GROUP BY ExecutionStatus"
166 );
167 }
168}