databend_common_ast/ast/statements/
workload.rs

1// Copyright 2021 Datafuse Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::BTreeMap;
16use std::fmt::Display;
17use std::fmt::Formatter;
18use std::time::Duration;
19
20use derive_visitor::Drive;
21use derive_visitor::DriveMut;
22
23use crate::ast::Identifier;
24
25#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
26pub struct ShowWorkloadGroupsStmt {}
27
28impl Display for ShowWorkloadGroupsStmt {
29    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
30        write!(f, "SHOW WORKLOAD GROUPS")
31    }
32}
33
34#[derive(Debug, Clone, Eq, PartialEq, serde::Serialize)]
35pub enum QuotaValueStmt {
36    Duration(Duration),
37    Percentage(usize),
38    Bytes(usize),
39    Number(usize),
40}
41
42impl QuotaValueStmt {
43    fn parse_percentage(v: &str) -> Option<QuotaValueStmt> {
44        let v = v.trim();
45        if v.is_empty() {
46            return None;
47        }
48
49        if v.ends_with('%') {
50            let num = v.trim_end_matches('%').trim();
51            if let Ok(value) = num.parse::<usize>() {
52                if value <= 100 {
53                    return Some(QuotaValueStmt::Percentage(value));
54                }
55            }
56        }
57        None
58    }
59
60    fn parse_human_size(v: &str) -> Option<QuotaValueStmt> {
61        let v = v.trim();
62        if v.is_empty() {
63            return None;
64        }
65
66        if v == "0" {
67            return Some(QuotaValueStmt::Bytes(0));
68        }
69
70        let (num_str, unit) = v.split_at(
71            v.find(|c: char| !c.is_ascii_digit() && c != '.')
72                .unwrap_or(v.len()),
73        );
74
75        let num = num_str.parse::<f64>().ok()?;
76        if num <= 0.0 {
77            return None;
78        }
79
80        let multiplier = match unit.trim().to_lowercase().as_str() {
81            "" | "b" => 1,
82            "k" | "kb" => 1024,
83            "m" | "mb" => 1024 * 1024,
84            "g" | "gb" => 1024 * 1024 * 1024,
85            _ => return None,
86        };
87
88        let bytes = (num * multiplier as f64).round() as usize;
89        Some(QuotaValueStmt::Bytes(bytes))
90    }
91
92    fn parse_human_timeout(v: &str) -> Option<QuotaValueStmt> {
93        let v = v.trim();
94        if v.is_empty() {
95            return None;
96        }
97
98        if v == "0" {
99            return Some(QuotaValueStmt::Duration(Duration::from_secs(0)));
100        }
101
102        let (num_str, unit) = v.split_at(
103            v.find(|c: char| !c.is_ascii_digit() && c != '.')
104                .unwrap_or(v.len()),
105        );
106
107        let num = num_str.parse::<f64>().ok()?;
108        if num <= 0.0 {
109            return None;
110        }
111
112        let duration = match unit.trim().to_lowercase().as_str() {
113            "s" | "sec" | "secs" => Duration::from_secs_f64(num),
114            "m" | "min" | "mins" => Duration::from_secs_f64(num * 60.0),
115            "h" | "hour" | "hours" => Duration::from_secs_f64(num * 3600.0),
116            "d" | "day" | "days" => Duration::from_secs_f64(num * 86400.0),
117            "ms" | "milli" | "millis" => Duration::from_millis(num as u64),
118            "" => Duration::from_secs_f64(num), // default to seconds if no unit
119            _ => return None,
120        };
121
122        Some(QuotaValueStmt::Duration(duration))
123    }
124
125    fn parse_number(v: &str) -> Option<QuotaValueStmt> {
126        v.trim().parse::<usize>().ok().map(QuotaValueStmt::Number)
127    }
128
129    pub fn new(key: &str, v: String) -> Result<QuotaValueStmt, &'static str> {
130        match key {
131            "cpu_quota" => Self::parse_percentage(&v).ok_or("Invalid CPU quota value, expected percentage (e.g. '50%') between 0-100"),
132            "memory_quota" => Self::parse_percentage(&v).or_else(|| Self::parse_human_size(&v)).ok_or("Invalid memory quota value, expected percentage (e.g. '50%') or size (e.g. '1GB', '512MB')"),
133            "query_timeout" => Self::parse_human_timeout(&v).ok_or("Invalid query timeout value, expected duration (e.g. '30s', '5min', '1h')"),
134            "max_concurrency" => Self::parse_number(&v).filter(|x| !matches!(x, QuotaValueStmt::Number(0))).ok_or("Invalid max concurrency value, expected positive integer"),
135            "query_queued_timeout" => Self::parse_human_timeout(&v).ok_or("Invalid queued query timeout value, expected duration (e.g. '30s', '5min', '1h')"),
136            "max_memory_usage_ratio" => Self::parse_percentage(&v).ok_or("Invalid max_memory_usage_ratio value, expected percentage (e.g. '50%') between 0-100"),
137            _ => Err("Unknown quota key"),
138        }
139    }
140}
141
142impl Display for QuotaValueStmt {
143    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
144        match self {
145            QuotaValueStmt::Percentage(v) => write!(f, "{}%", v),
146            QuotaValueStmt::Duration(v) => write!(f, "{:?}", v),
147            QuotaValueStmt::Bytes(v) => write!(f, "{}", v),
148            QuotaValueStmt::Number(v) => write!(f, "{}", v),
149        }
150    }
151}
152
153#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
154pub struct CreateWorkloadGroupStmt {
155    pub name: Identifier,
156    pub if_not_exists: bool,
157    #[drive(skip)]
158    pub quotas: BTreeMap<String, QuotaValueStmt>,
159}
160
161impl Display for CreateWorkloadGroupStmt {
162    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
163        write!(f, "CREATE WORKLOAD GROUP")?;
164
165        if self.if_not_exists {
166            write!(f, " IF NOT EXISTS")?;
167        }
168
169        write!(f, " {}", self.name)?;
170
171        if !self.quotas.is_empty() {
172            write!(f, " WITH ")?;
173
174            for (idx, (key, value)) in self.quotas.iter().enumerate() {
175                if idx != 0 {
176                    write!(f, ",")?;
177                }
178
179                write!(f, " {} = '{}'", key, value)?;
180            }
181        }
182
183        Ok(())
184    }
185}
186
187#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
188pub struct DropWorkloadGroupStmt {
189    pub name: Identifier,
190    pub if_exists: bool,
191}
192
193impl Display for DropWorkloadGroupStmt {
194    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
195        write!(f, "DROP WORKLOAD GROUP ")?;
196
197        if self.if_exists {
198            write!(f, " IF EXISTS")?;
199        }
200
201        write!(f, " {}", self.name)
202    }
203}
204
205#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
206pub struct RenameWorkloadGroupStmt {
207    pub name: Identifier,
208    pub new_name: Identifier,
209}
210
211impl Display for RenameWorkloadGroupStmt {
212    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
213        write!(
214            f,
215            "RENAME WORKLOAD GROUP {} TO {}",
216            self.name, self.new_name
217        )
218    }
219}
220
221#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
222pub struct SetWorkloadGroupQuotasStmt {
223    pub name: Identifier,
224    #[drive(skip)]
225    pub quotas: BTreeMap<String, QuotaValueStmt>,
226}
227
228impl Display for SetWorkloadGroupQuotasStmt {
229    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
230        write!(f, "ALTER WORKLOAD GROUP {}", self.name)?;
231
232        if !self.quotas.is_empty() {
233            write!(f, " SET")?;
234
235            for (idx, (key, value)) in self.quotas.iter().enumerate() {
236                if idx != 0 {
237                    write!(f, ",")?;
238                }
239
240                write!(f, " {} = '{}'", key, value)?;
241            }
242        }
243
244        Ok(())
245    }
246}
247
248#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
249pub struct UnsetWorkloadGroupQuotasStmt {
250    pub name: Identifier,
251    #[drive(skip)]
252    pub quotas: Vec<Identifier>,
253}
254
255impl Display for UnsetWorkloadGroupQuotasStmt {
256    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
257        write!(f, "ALTER WORKLOAD GROUP {}", self.name)?;
258
259        if !self.quotas.is_empty() {
260            write!(f, " UNSET")?;
261
262            for (idx, name) in self.quotas.iter().enumerate() {
263                if idx != 0 {
264                    write!(f, ",")?;
265                }
266
267                write!(f, " {}", name)?;
268            }
269        }
270
271        Ok(())
272    }
273}
274
275#[cfg(test)]
276mod tests {
277    use std::time::Duration;
278
279    use crate::ast::QuotaValueStmt;
280
281    #[test]
282    fn test_parse_percentage() {
283        // Valid cases
284        // Valid cases
285        assert_eq!(
286            QuotaValueStmt::parse_percentage("50%"),
287            Some(QuotaValueStmt::Percentage(50))
288        );
289        assert_eq!(
290            QuotaValueStmt::parse_percentage("100%"),
291            Some(QuotaValueStmt::Percentage(100))
292        );
293        assert_eq!(
294            QuotaValueStmt::parse_percentage(" 0% "),
295            Some(QuotaValueStmt::Percentage(0))
296        );
297
298        // Invalid cases
299        assert_eq!(QuotaValueStmt::parse_percentage("101%"), None);
300        assert_eq!(QuotaValueStmt::parse_percentage("50"), None);
301        assert_eq!(QuotaValueStmt::parse_percentage("1/2"), None);
302        assert_eq!(QuotaValueStmt::parse_percentage(""), None);
303        assert_eq!(QuotaValueStmt::parse_percentage(" % "), None);
304        assert_eq!(QuotaValueStmt::parse_percentage("-10%"), None);
305        assert_eq!(QuotaValueStmt::parse_percentage("abc%"), None);
306    }
307
308    #[test]
309    fn test_parse_human_size() {
310        // Valid cases
311        assert_eq!(
312            QuotaValueStmt::parse_human_size("0"),
313            Some(QuotaValueStmt::Bytes(0))
314        );
315        assert_eq!(
316            QuotaValueStmt::parse_human_size("1024"),
317            Some(QuotaValueStmt::Bytes(1024))
318        );
319        assert_eq!(
320            QuotaValueStmt::parse_human_size("1.5KB"),
321            Some(QuotaValueStmt::Bytes(1536))
322        );
323        assert_eq!(
324            QuotaValueStmt::parse_human_size(" 2 MB "),
325            Some(QuotaValueStmt::Bytes(2 * 1024 * 1024))
326        );
327        assert_eq!(
328            QuotaValueStmt::parse_human_size("1G"),
329            Some(QuotaValueStmt::Bytes(1024 * 1024 * 1024))
330        );
331        assert_eq!(
332            QuotaValueStmt::parse_human_size("0.5gb"),
333            Some(QuotaValueStmt::Bytes(512 * 1024 * 1024))
334        );
335
336        // Invalid cases
337        assert_eq!(QuotaValueStmt::parse_human_size(""), None);
338        assert_eq!(QuotaValueStmt::parse_human_size("-1"), None);
339        assert_eq!(QuotaValueStmt::parse_human_size("1.2.3"), None);
340        assert_eq!(QuotaValueStmt::parse_human_size("1TB"), None);
341        assert_eq!(QuotaValueStmt::parse_human_size("abc"), None);
342    }
343
344    #[test]
345    fn test_parse_human_timeout() {
346        // Valid cases
347        assert_eq!(
348            QuotaValueStmt::parse_human_timeout("0"),
349            Some(QuotaValueStmt::Duration(Duration::from_secs(0)))
350        );
351        assert_eq!(
352            QuotaValueStmt::parse_human_timeout("30"),
353            Some(QuotaValueStmt::Duration(Duration::from_secs(30)))
354        );
355        assert_eq!(
356            QuotaValueStmt::parse_human_timeout("1.5s"),
357            Some(QuotaValueStmt::Duration(Duration::from_secs_f64(1.5)))
358        );
359        assert_eq!(
360            QuotaValueStmt::parse_human_timeout(" 30 MIN "),
361            Some(QuotaValueStmt::Duration(Duration::from_secs(30 * 60)))
362        );
363        assert_eq!(
364            QuotaValueStmt::parse_human_timeout("2h"),
365            Some(QuotaValueStmt::Duration(Duration::from_secs(2 * 3600)))
366        );
367        assert_eq!(
368            QuotaValueStmt::parse_human_timeout("100ms"),
369            Some(QuotaValueStmt::Duration(Duration::from_millis(100)))
370        );
371        assert_eq!(
372            QuotaValueStmt::parse_human_timeout("0.5d"),
373            Some(QuotaValueStmt::Duration(Duration::from_secs(12 * 3600)))
374        );
375
376        // Invalid cases
377        assert_eq!(QuotaValueStmt::parse_human_timeout(""), None);
378        assert_eq!(QuotaValueStmt::parse_human_timeout("-1"), None);
379        assert_eq!(QuotaValueStmt::parse_human_timeout("1.2.3"), None);
380        assert_eq!(QuotaValueStmt::parse_human_timeout("1y"), None);
381        assert_eq!(QuotaValueStmt::parse_human_timeout("abc"), None);
382    }
383
384    #[test]
385    fn test_parse_number() {
386        // Valid cases
387        assert_eq!(
388            QuotaValueStmt::parse_number("0"),
389            Some(QuotaValueStmt::Number(0))
390        );
391        assert_eq!(
392            QuotaValueStmt::parse_number("123"),
393            Some(QuotaValueStmt::Number(123))
394        );
395        assert_eq!(
396            QuotaValueStmt::parse_number(" 456 "),
397            Some(QuotaValueStmt::Number(456))
398        );
399
400        // Invalid cases
401        assert_eq!(QuotaValueStmt::parse_number(""), None);
402        assert_eq!(QuotaValueStmt::parse_number("-1"), None);
403        assert_eq!(QuotaValueStmt::parse_number("1.2"), None);
404        assert_eq!(QuotaValueStmt::parse_number("abc"), None);
405    }
406}
databend_common_ast/ast/statements/workload.rs

databend_common_ast/ast/statements/
workload.rs
