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::{IntoError, 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
41/// Sets the session WAL policy for ordinary inserts without changing table options.
42pub fn set_skip_wal(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
43    let [Expr::Value(value)] = exprs.as_slice() else {
44        return NotSupportedSnafu {
45            feat: "SET skip_wal requires exactly one boolean value",
46        }
47        .fail();
48    };
49    let skip_wal = match &value.value {
50        Value::Boolean(value) => *value,
51        Value::SingleQuotedString(value) | Value::DoubleQuotedString(value) => {
52            value.parse::<bool>().map_err(|_| {
53                NotSupportedSnafu {
54                    feat: format!("Invalid skip_wal value {value:?}: expected true or false"),
55                }
56                .build()
57            })?
58        }
59        _ => {
60            return NotSupportedSnafu {
61                feat: "SET skip_wal requires true or false",
62            }
63            .fail();
64        }
65    };
66    ctx.set_skip_wal(skip_wal);
67    Ok(())
68}
69
70pub fn set_read_preference(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
71    let read_preference_expr = exprs.first().context(NotSupportedSnafu {
72        feat: "No read preference find in set variable statement",
73    })?;
74
75    match read_preference_expr {
76        Expr::Value(ValueWithSpan {
77            value: Value::SingleQuotedString(expr),
78            ..
79        })
80        | Expr::Value(ValueWithSpan {
81            value: Value::DoubleQuotedString(expr),
82            ..
83        }) => {
84            match ReadPreference::from_str(expr.as_str().to_lowercase().as_str()) {
85                Ok(read_preference) => ctx.set_read_preference(read_preference),
86                Err(_) => {
87                    return NotSupportedSnafu {
88                        feat: format!(
89                            "Invalid read preference expr {} in set variable statement",
90                            expr,
91                        ),
92                    }
93                    .fail();
94                }
95            }
96            Ok(())
97        }
98        expr => NotSupportedSnafu {
99            feat: format!(
100                "Unsupported read preference expr {} in set variable statement",
101                expr
102            ),
103        }
104        .fail(),
105    }
106}
107
108pub fn set_timezone(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
109    let tz_expr = exprs.first().context(NotSupportedSnafu {
110        feat: "No timezone find in set variable statement",
111    })?;
112    match tz_expr {
113        Expr::Value(ValueWithSpan {
114            value: Value::SingleQuotedString(tz),
115            ..
116        })
117        | Expr::Value(ValueWithSpan {
118            value: Value::DoubleQuotedString(tz),
119            ..
120        }) => {
121            match Timezone::from_tz_string(tz.as_str()) {
122                Ok(timezone) => ctx.set_timezone(timezone),
123                Err(_) => {
124                    return NotSupportedSnafu {
125                        feat: format!("Invalid timezone expr {} in set variable statement", tz),
126                    }
127                    .fail();
128                }
129            }
130            Ok(())
131        }
132        expr => NotSupportedSnafu {
133            feat: format!(
134                "Unsupported timezone expr {} in set variable statement",
135                expr
136            ),
137        }
138        .fail(),
139    }
140}
141
142pub fn set_bytea_output(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
143    let Some((var_value, [])) = exprs.split_first() else {
144        return (NotSupportedSnafu {
145            feat: "Set variable value must have one and only one value for bytea_output",
146        })
147        .fail();
148    };
149    let Expr::Value(value) = var_value else {
150        return (NotSupportedSnafu {
151            feat: "Set variable value must be a value",
152        })
153        .fail();
154    };
155    ctx.configuration_parameter().set_postgres_bytea_output(
156        PGByteaOutputValue::try_from(value.value.clone()).context(InvalidConfigValueSnafu)?,
157    );
158    Ok(())
159}
160
161pub fn validate_client_encoding(set: SetVariables) -> Result<()> {
162    let Some((encoding, [])) = set.value.split_first() else {
163        return InvalidSqlSnafu {
164            err_msg: "must provide one and only one client encoding value",
165        }
166        .fail();
167    };
168    let encoding = match encoding {
169        Expr::Value(ValueWithSpan {
170            value: Value::SingleQuotedString(x),
171            ..
172        })
173        | Expr::Identifier(Ident {
174            value: x,
175            quote_style: _,
176            span: _,
177        }) => x.to_uppercase(),
178        _ => {
179            return InvalidSqlSnafu {
180                err_msg: format!("client encoding must be a string, actual: {:?}", encoding),
181            }
182            .fail();
183        }
184    };
185    // For the sake of simplicity, we only support "UTF8" ("UNICODE" is the alias for it,
186    // see https://www.postgresql.org/docs/current/multibyte.html#MULTIBYTE-CHARSET-SUPPORTED).
187    // "UTF8" is universal and sufficient for almost all cases.
188    // GreptimeDB itself is always using "UTF8" as the internal encoding.
189    ensure!(
190        encoding == "UTF8" || encoding == "UNICODE",
191        NotSupportedSnafu {
192            feat: format!("client encoding of '{}'", encoding)
193        }
194    );
195    Ok(())
196}
197
198// if one of original value and new value is none, return the other one
199// returns new values only when it equals to original one else return error.
200// This is only used for handling datestyle
201fn merge_datestyle_value<T>(value: Option<T>, new_value: Option<T>) -> Result<Option<T>>
202where
203    T: PartialEq,
204{
205    match (&value, &new_value) {
206        (None, _) => Ok(new_value),
207        (_, None) => Ok(value),
208        (Some(v1), Some(v2)) if v1 == v2 => Ok(new_value),
209        _ => InvalidSqlSnafu {
210            err_msg: "Conflicting \"datestyle\" specifications.",
211        }
212        .fail(),
213    }
214}
215
216fn try_parse_datestyle(expr: &Expr) -> Result<(Option<PGDateTimeStyle>, Option<PGDateOrder>)> {
217    enum ParsedDateStyle {
218        Order(PGDateOrder),
219        Style(PGDateTimeStyle),
220    }
221    fn try_parse_str(s: &str) -> Result<ParsedDateStyle> {
222        PGDateTimeStyle::try_from(s)
223            .map_or_else(
224                |_| PGDateOrder::try_from(s).map(ParsedDateStyle::Order),
225                |style| Ok(ParsedDateStyle::Style(style)),
226            )
227            .context(InvalidConfigValueSnafu)
228    }
229    match expr {
230        Expr::Identifier(Ident {
231            value: s,
232            quote_style: _,
233            span: _,
234        })
235        | Expr::Value(ValueWithSpan {
236            value: Value::SingleQuotedString(s),
237            ..
238        })
239        | Expr::Value(ValueWithSpan {
240            value: Value::DoubleQuotedString(s),
241            ..
242        }) => s
243            .split(',')
244            .map(|s| s.trim())
245            .try_fold((None, None), |(style, order), s| match try_parse_str(s)? {
246                ParsedDateStyle::Order(o) => Ok((style, merge_datestyle_value(order, Some(o))?)),
247                ParsedDateStyle::Style(s) => Ok((merge_datestyle_value(style, Some(s))?, order)),
248            }),
249        _ => NotSupportedSnafu {
250            feat: "Not supported expression for datestyle",
251        }
252        .fail(),
253    }
254}
255
256/// Applies supported query options and rejects unsupported owned namespaces.
257pub fn set_query_option(variable: &str, exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<bool> {
258    let lower = variable.to_ascii_lowercase();
259    if !lower.starts_with("query.")
260        && !lower.starts_with("datafusion.")
261        && session::query_options::canonical_query_option_name(variable).is_none()
262    {
263        return Ok(false);
264    }
265    let [Expr::Value(value)] = exprs.as_slice() else {
266        return Err(InvalidConfigValueSnafu.into_error(
267            session::session_config::Error::InvalidConfigValue {
268                name: variable.to_string(),
269                value: exprs
270                    .iter()
271                    .map(ToString::to_string)
272                    .collect::<Vec<_>>()
273                    .join(", "),
274                hint: "Expected exactly one literal value".to_string(),
275                location: snafu::Location::new(file!(), line!(), column!()),
276            },
277        ));
278    };
279    let value = match &value.value {
280        Value::Boolean(value) => value.to_string(),
281        Value::Number(value, _)
282        | Value::SingleQuotedString(value)
283        | Value::DoubleQuotedString(value) => value.clone(),
284        _ => {
285            return Err(InvalidConfigValueSnafu.into_error(
286                session::session_config::Error::InvalidConfigValue {
287                    name: variable.to_string(),
288                    value: value.value.to_string(),
289                    hint: "Expected a boolean, number, or quoted string literal".to_string(),
290                    location: snafu::Location::new(file!(), line!(), column!()),
291                },
292            ));
293        }
294    };
295    ctx.configuration_parameter()
296        .set_query_option(variable, &value)
297        .context(InvalidConfigValueSnafu)?;
298    Ok(true)
299}
300
301pub fn set_intervalstyle(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
302    let Some((var_value, [])) = exprs.split_first() else {
303        return NotSupportedSnafu {
304            feat: "Set variable value must have one and only one value for intervalstyle",
305        }
306        .fail();
307    };
308    let intervalstyle = match var_value {
309        Expr::Identifier(Ident {
310            value,
311            quote_style: _,
312            span: _,
313        }) => PGIntervalStyle::try_from(value.as_str()).context(InvalidConfigValueSnafu)?,
314        Expr::Value(value) => {
315            PGIntervalStyle::try_from(&value.value).context(InvalidConfigValueSnafu)?
316        }
317        _ => {
318            return NotSupportedSnafu {
319                feat: "Set variable value must be a value or identifier",
320            }
321            .fail();
322        }
323    };
324    ctx.configuration_parameter()
325        .set_pg_intervalstyle_format(intervalstyle);
326    Ok(())
327}
328
329pub fn set_datestyle(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
330    // ORDER,
331    // STYLE,
332    // ORDER,ORDER
333    // ORDER,STYLE
334    // STYLE,ORDER
335    let (style, order) = exprs
336        .iter()
337        .try_fold((None, None), |(style, order), expr| {
338            let (new_style, new_order) = try_parse_datestyle(expr)?;
339            Ok((
340                merge_datestyle_value(style, new_style)?,
341                merge_datestyle_value(order, new_order)?,
342            ))
343        })?;
344
345    let (old_style, older_order) = *ctx.configuration_parameter().pg_datetime_style();
346    ctx.configuration_parameter()
347        .set_pg_datetime_style(style.unwrap_or(old_style), order.unwrap_or(older_order));
348    Ok(())
349}
350
351pub fn set_query_timeout(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
352    let timeout_expr = exprs.first().context(NotSupportedSnafu {
353        feat: "No timeout value find in set query timeout statement",
354    })?;
355    match timeout_expr {
356        Expr::Value(ValueWithSpan {
357            value: Value::Number(timeout, _),
358            ..
359        }) => {
360            match timeout.parse::<u64>() {
361                Ok(timeout) => ctx.set_query_timeout(Duration::from_millis(timeout)),
362                Err(_) => {
363                    return NotSupportedSnafu {
364                        feat: format!("Invalid timeout expr {} in set variable statement", timeout),
365                    }
366                    .fail();
367                }
368            }
369            Ok(())
370        }
371        // postgres support time units i.e. SET STATEMENT_TIMEOUT = '50ms';
372        Expr::Value(ValueWithSpan {
373            value: Value::SingleQuotedString(timeout),
374            ..
375        })
376        | Expr::Value(ValueWithSpan {
377            value: Value::DoubleQuotedString(timeout),
378            ..
379        }) => {
380            if ctx.channel() != Postgres {
381                return NotSupportedSnafu {
382                    feat: format!("Invalid timeout expr {} in set variable statement", timeout),
383                }
384                .fail();
385            }
386            let timeout = parse_pg_query_timeout_input(timeout)?;
387            ctx.set_query_timeout(Duration::from_millis(timeout));
388            Ok(())
389        }
390        expr => NotSupportedSnafu {
391            feat: format!(
392                "Unsupported timeout expr {} in set variable statement",
393                expr
394            ),
395        }
396        .fail(),
397    }
398}
399
400// support time units in ms, s, min, h, d for postgres protocol.
401// https://www.postgresql.org/docs/8.4/config-setting.html#:~:text=Valid%20memory%20units%20are%20kB,%2C%20and%20d%20(days).
402fn parse_pg_query_timeout_input(input: &str) -> Result<u64> {
403    match input.parse::<u64>() {
404        Ok(timeout) => Ok(timeout),
405        Err(_) => {
406            if let Some(captures) = PG_TIME_INPUT_REGEX.captures(input) {
407                let value = captures[1].parse::<u64>().expect("regex failed");
408                let unit = &captures[2];
409
410                match unit {
411                    "ms" => Ok(value),
412                    "s" => Ok(value * 1000),
413                    "min" => Ok(value * 60 * 1000),
414                    "h" => Ok(value * 60 * 60 * 1000),
415                    "d" => Ok(value * 24 * 60 * 60 * 1000),
416                    _ => unreachable!("regex failed"),
417                }
418            } else {
419                NotSupportedSnafu {
420                    feat: format!(
421                        "Unsupported timeout expr {} in set variable statement",
422                        input
423                    ),
424                }
425                .fail()
426            }
427        }
428    }
429}
430
431#[cfg(test)]
432mod test {
433    use session::Session;
434    use session::context::{Channel, QueryContext};
435    use sql::ast::{Expr, Value};
436
437    use super::set_skip_wal;
438    use crate::statement::set::parse_pg_query_timeout_input;
439
440    #[test]
441    fn test_set_skip_wal() {
442        let ctx = QueryContext::arc();
443        for value in [true, false] {
444            set_skip_wal(vec![Expr::Value(Value::Boolean(value).into())], ctx.clone()).unwrap();
445            assert_eq!(ctx.skip_wal(), value);
446            set_skip_wal(
447                vec![Expr::Value(
448                    Value::SingleQuotedString(value.to_string()).into(),
449                )],
450                ctx.clone(),
451            )
452            .unwrap();
453            assert_eq!(ctx.skip_wal(), value);
454        }
455        ctx.set_skip_wal(true);
456        for values in [
457            vec![],
458            vec![Expr::Value(Value::Number("1".to_string(), false).into())],
459            vec![Expr::Value(
460                Value::SingleQuotedString("invalid".to_string()).into(),
461            )],
462            vec![Expr::Value(Value::Boolean(false).into()); 2],
463        ] {
464            assert!(set_skip_wal(values, ctx.clone()).is_err());
465            assert!(ctx.skip_wal());
466        }
467    }
468
469    #[test]
470    fn test_set_skip_wal_session_isolation() {
471        for channel in [Channel::Mysql, Channel::Postgres] {
472            let session = Session::new(None, channel, Default::default(), 0);
473            let other = Session::new(None, channel, Default::default(), 1);
474            assert!(!session.new_query_context().skip_wal());
475            set_skip_wal(
476                vec![Expr::Value(Value::Boolean(true).into())],
477                session.new_query_context(),
478            )
479            .unwrap();
480            assert!(session.new_query_context().skip_wal());
481            assert!(!other.new_query_context().skip_wal());
482        }
483    }
484
485    #[test]
486    fn test_parse_pg_query_timeout_input() {
487        assert!(parse_pg_query_timeout_input("").is_err());
488        assert!(parse_pg_query_timeout_input(" 50 ms").is_err());
489        assert!(parse_pg_query_timeout_input("5s 1ms").is_err());
490        assert!(parse_pg_query_timeout_input("3a").is_err());
491        assert!(parse_pg_query_timeout_input("1.5min").is_err());
492        assert!(parse_pg_query_timeout_input("ms").is_err());
493        assert!(parse_pg_query_timeout_input("a").is_err());
494        assert!(parse_pg_query_timeout_input("-1").is_err());
495
496        assert_eq!(50, parse_pg_query_timeout_input("50").unwrap());
497        assert_eq!(12, parse_pg_query_timeout_input("12ms").unwrap());
498        assert_eq!(2000, parse_pg_query_timeout_input("2s").unwrap());
499        assert_eq!(60000, parse_pg_query_timeout_input("1min").unwrap());
500    }
501}
operator/statement/set.rs

operator/statement/
set.rs
