operator/statement/
set.rs

1// Copyright 2023 Greptime Team
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::str::FromStr;
16use std::time::Duration;
17
18use common_time::Timezone;
19use lazy_static::lazy_static;
20use regex::Regex;
21use session::ReadPreference;
22use session::context::Channel::Postgres;
23use session::context::QueryContextRef;
24use session::session_config::{PGByteaOutputValue, PGDateOrder, PGDateTimeStyle, PGIntervalStyle};
25use snafu::{OptionExt, ResultExt, ensure};
26use sql::ast::{Expr, Ident, Value};
27use sql::statements::set_variables::SetVariables;
28use sqlparser::ast::ValueWithSpan;
29
30use crate::error::{InvalidConfigValueSnafu, InvalidSqlSnafu, NotSupportedSnafu, Result};
31
32lazy_static! {
33    // Regex rules:
34    // The string must start with a number (one or more digits).
35    // The number must be followed by one of the valid time units (ms, s, min, h, d).
36    // The string must end immediately after the unit, meaning there can be no extra
37    // characters or spaces after the valid time specification.
38    static ref PG_TIME_INPUT_REGEX: Regex = Regex::new(r"^(\d+)(ms|s|min|h|d)$").unwrap();
39}
40
41pub fn set_read_preference(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
42    let read_preference_expr = exprs.first().context(NotSupportedSnafu {
43        feat: "No read preference find in set variable statement",
44    })?;
45
46    match read_preference_expr {
47        Expr::Value(ValueWithSpan {
48            value: Value::SingleQuotedString(expr),
49            ..
50        })
51        | Expr::Value(ValueWithSpan {
52            value: Value::DoubleQuotedString(expr),
53            ..
54        }) => {
55            match ReadPreference::from_str(expr.as_str().to_lowercase().as_str()) {
56                Ok(read_preference) => ctx.set_read_preference(read_preference),
57                Err(_) => {
58                    return NotSupportedSnafu {
59                        feat: format!(
60                            "Invalid read preference expr {} in set variable statement",
61                            expr,
62                        ),
63                    }
64                    .fail();
65                }
66            }
67            Ok(())
68        }
69        expr => NotSupportedSnafu {
70            feat: format!(
71                "Unsupported read preference expr {} in set variable statement",
72                expr
73            ),
74        }
75        .fail(),
76    }
77}
78
79pub fn set_timezone(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
80    let tz_expr = exprs.first().context(NotSupportedSnafu {
81        feat: "No timezone find in set variable statement",
82    })?;
83    match tz_expr {
84        Expr::Value(ValueWithSpan {
85            value: Value::SingleQuotedString(tz),
86            ..
87        })
88        | Expr::Value(ValueWithSpan {
89            value: Value::DoubleQuotedString(tz),
90            ..
91        }) => {
92            match Timezone::from_tz_string(tz.as_str()) {
93                Ok(timezone) => ctx.set_timezone(timezone),
94                Err(_) => {
95                    return NotSupportedSnafu {
96                        feat: format!("Invalid timezone expr {} in set variable statement", tz),
97                    }
98                    .fail();
99                }
100            }
101            Ok(())
102        }
103        expr => NotSupportedSnafu {
104            feat: format!(
105                "Unsupported timezone expr {} in set variable statement",
106                expr
107            ),
108        }
109        .fail(),
110    }
111}
112
113pub fn set_bytea_output(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
114    let Some((var_value, [])) = exprs.split_first() else {
115        return (NotSupportedSnafu {
116            feat: "Set variable value must have one and only one value for bytea_output",
117        })
118        .fail();
119    };
120    let Expr::Value(value) = var_value else {
121        return (NotSupportedSnafu {
122            feat: "Set variable value must be a value",
123        })
124        .fail();
125    };
126    ctx.configuration_parameter().set_postgres_bytea_output(
127        PGByteaOutputValue::try_from(value.value.clone()).context(InvalidConfigValueSnafu)?,
128    );
129    Ok(())
130}
131
132pub fn validate_client_encoding(set: SetVariables) -> Result<()> {
133    let Some((encoding, [])) = set.value.split_first() else {
134        return InvalidSqlSnafu {
135            err_msg: "must provide one and only one client encoding value",
136        }
137        .fail();
138    };
139    let encoding = match encoding {
140        Expr::Value(ValueWithSpan {
141            value: Value::SingleQuotedString(x),
142            ..
143        })
144        | Expr::Identifier(Ident {
145            value: x,
146            quote_style: _,
147            span: _,
148        }) => x.to_uppercase(),
149        _ => {
150            return InvalidSqlSnafu {
151                err_msg: format!("client encoding must be a string, actual: {:?}", encoding),
152            }
153            .fail();
154        }
155    };
156    // For the sake of simplicity, we only support "UTF8" ("UNICODE" is the alias for it,
157    // see https://www.postgresql.org/docs/current/multibyte.html#MULTIBYTE-CHARSET-SUPPORTED).
158    // "UTF8" is universal and sufficient for almost all cases.
159    // GreptimeDB itself is always using "UTF8" as the internal encoding.
160    ensure!(
161        encoding == "UTF8" || encoding == "UNICODE",
162        NotSupportedSnafu {
163            feat: format!("client encoding of '{}'", encoding)
164        }
165    );
166    Ok(())
167}
168
169// if one of original value and new value is none, return the other one
170// returns new values only when it equals to original one else return error.
171// This is only used for handling datestyle
172fn merge_datestyle_value<T>(value: Option<T>, new_value: Option<T>) -> Result<Option<T>>
173where
174    T: PartialEq,
175{
176    match (&value, &new_value) {
177        (None, _) => Ok(new_value),
178        (_, None) => Ok(value),
179        (Some(v1), Some(v2)) if v1 == v2 => Ok(new_value),
180        _ => InvalidSqlSnafu {
181            err_msg: "Conflicting \"datestyle\" specifications.",
182        }
183        .fail(),
184    }
185}
186
187fn try_parse_datestyle(expr: &Expr) -> Result<(Option<PGDateTimeStyle>, Option<PGDateOrder>)> {
188    enum ParsedDateStyle {
189        Order(PGDateOrder),
190        Style(PGDateTimeStyle),
191    }
192    fn try_parse_str(s: &str) -> Result<ParsedDateStyle> {
193        PGDateTimeStyle::try_from(s)
194            .map_or_else(
195                |_| PGDateOrder::try_from(s).map(ParsedDateStyle::Order),
196                |style| Ok(ParsedDateStyle::Style(style)),
197            )
198            .context(InvalidConfigValueSnafu)
199    }
200    match expr {
201        Expr::Identifier(Ident {
202            value: s,
203            quote_style: _,
204            span: _,
205        })
206        | Expr::Value(ValueWithSpan {
207            value: Value::SingleQuotedString(s),
208            ..
209        })
210        | Expr::Value(ValueWithSpan {
211            value: Value::DoubleQuotedString(s),
212            ..
213        }) => s
214            .split(',')
215            .map(|s| s.trim())
216            .try_fold((None, None), |(style, order), s| match try_parse_str(s)? {
217                ParsedDateStyle::Order(o) => Ok((style, merge_datestyle_value(order, Some(o))?)),
218                ParsedDateStyle::Style(s) => Ok((merge_datestyle_value(style, Some(s))?, order)),
219            }),
220        _ => NotSupportedSnafu {
221            feat: "Not supported expression for datestyle",
222        }
223        .fail(),
224    }
225}
226
227/// Set the allow query fallback configuration parameter to true or false based on the provided expressions.
228///
229pub fn set_allow_query_fallback(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
230    let allow_fallback_expr = exprs.first().context(NotSupportedSnafu {
231        feat: "No allow query fallback value find in set variable statement",
232    })?;
233    match allow_fallback_expr {
234        Expr::Value(ValueWithSpan {
235            value: Value::Boolean(allow),
236            span: _,
237        }) => {
238            ctx.configuration_parameter()
239                .set_allow_query_fallback(*allow);
240            Ok(())
241        }
242        expr => NotSupportedSnafu {
243            feat: format!(
244                "Unsupported allow query fallback expr {} in set variable statement",
245                expr
246            ),
247        }
248        .fail(),
249    }
250}
251
252pub fn set_intervalstyle(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
253    let Some((var_value, [])) = exprs.split_first() else {
254        return NotSupportedSnafu {
255            feat: "Set variable value must have one and only one value for intervalstyle",
256        }
257        .fail();
258    };
259    let Expr::Value(value) = var_value else {
260        return NotSupportedSnafu {
261            feat: "Set variable value must be a value",
262        }
263        .fail();
264    };
265    ctx.configuration_parameter().set_pg_intervalstyle_format(
266        PGIntervalStyle::try_from(&value.value).context(InvalidConfigValueSnafu)?,
267    );
268    Ok(())
269}
270
271pub fn set_datestyle(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
272    // ORDER,
273    // STYLE,
274    // ORDER,ORDER
275    // ORDER,STYLE
276    // STYLE,ORDER
277    let (style, order) = exprs
278        .iter()
279        .try_fold((None, None), |(style, order), expr| {
280            let (new_style, new_order) = try_parse_datestyle(expr)?;
281            Ok((
282                merge_datestyle_value(style, new_style)?,
283                merge_datestyle_value(order, new_order)?,
284            ))
285        })?;
286
287    let (old_style, older_order) = *ctx.configuration_parameter().pg_datetime_style();
288    ctx.configuration_parameter()
289        .set_pg_datetime_style(style.unwrap_or(old_style), order.unwrap_or(older_order));
290    Ok(())
291}
292
293pub fn set_query_timeout(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
294    let timeout_expr = exprs.first().context(NotSupportedSnafu {
295        feat: "No timeout value find in set query timeout statement",
296    })?;
297    match timeout_expr {
298        Expr::Value(ValueWithSpan {
299            value: Value::Number(timeout, _),
300            ..
301        }) => {
302            match timeout.parse::<u64>() {
303                Ok(timeout) => ctx.set_query_timeout(Duration::from_millis(timeout)),
304                Err(_) => {
305                    return NotSupportedSnafu {
306                        feat: format!("Invalid timeout expr {} in set variable statement", timeout),
307                    }
308                    .fail();
309                }
310            }
311            Ok(())
312        }
313        // postgres support time units i.e. SET STATEMENT_TIMEOUT = '50ms';
314        Expr::Value(ValueWithSpan {
315            value: Value::SingleQuotedString(timeout),
316            ..
317        })
318        | Expr::Value(ValueWithSpan {
319            value: Value::DoubleQuotedString(timeout),
320            ..
321        }) => {
322            if ctx.channel() != Postgres {
323                return NotSupportedSnafu {
324                    feat: format!("Invalid timeout expr {} in set variable statement", timeout),
325                }
326                .fail();
327            }
328            let timeout = parse_pg_query_timeout_input(timeout)?;
329            ctx.set_query_timeout(Duration::from_millis(timeout));
330            Ok(())
331        }
332        expr => NotSupportedSnafu {
333            feat: format!(
334                "Unsupported timeout expr {} in set variable statement",
335                expr
336            ),
337        }
338        .fail(),
339    }
340}
341
342// support time units in ms, s, min, h, d for postgres protocol.
343// https://www.postgresql.org/docs/8.4/config-setting.html#:~:text=Valid%20memory%20units%20are%20kB,%2C%20and%20d%20(days).
344fn parse_pg_query_timeout_input(input: &str) -> Result<u64> {
345    match input.parse::<u64>() {
346        Ok(timeout) => Ok(timeout),
347        Err(_) => {
348            if let Some(captures) = PG_TIME_INPUT_REGEX.captures(input) {
349                let value = captures[1].parse::<u64>().expect("regex failed");
350                let unit = &captures[2];
351
352                match unit {
353                    "ms" => Ok(value),
354                    "s" => Ok(value * 1000),
355                    "min" => Ok(value * 60 * 1000),
356                    "h" => Ok(value * 60 * 60 * 1000),
357                    "d" => Ok(value * 24 * 60 * 60 * 1000),
358                    _ => unreachable!("regex failed"),
359                }
360            } else {
361                NotSupportedSnafu {
362                    feat: format!(
363                        "Unsupported timeout expr {} in set variable statement",
364                        input
365                    ),
366                }
367                .fail()
368            }
369        }
370    }
371}
372
373#[cfg(test)]
374mod test {
375    use crate::statement::set::parse_pg_query_timeout_input;
376
377    #[test]
378    fn test_parse_pg_query_timeout_input() {
379        assert!(parse_pg_query_timeout_input("").is_err());
380        assert!(parse_pg_query_timeout_input(" 50 ms").is_err());
381        assert!(parse_pg_query_timeout_input("5s 1ms").is_err());
382        assert!(parse_pg_query_timeout_input("3a").is_err());
383        assert!(parse_pg_query_timeout_input("1.5min").is_err());
384        assert!(parse_pg_query_timeout_input("ms").is_err());
385        assert!(parse_pg_query_timeout_input("a").is_err());
386        assert!(parse_pg_query_timeout_input("-1").is_err());
387
388        assert_eq!(50, parse_pg_query_timeout_input("50").unwrap());
389        assert_eq!(12, parse_pg_query_timeout_input("12ms").unwrap());
390        assert_eq!(2000, parse_pg_query_timeout_input("2s").unwrap());
391        assert_eq!(60000, parse_pg_query_timeout_input("1min").unwrap());
392    }
393}
operator/statement/set.rs

operator/statement/
set.rs
