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 intervalstyle = match var_value {
260        Expr::Identifier(Ident {
261            value,
262            quote_style: _,
263            span: _,
264        }) => PGIntervalStyle::try_from(value.as_str()).context(InvalidConfigValueSnafu)?,
265        Expr::Value(value) => {
266            PGIntervalStyle::try_from(&value.value).context(InvalidConfigValueSnafu)?
267        }
268        _ => {
269            return NotSupportedSnafu {
270                feat: "Set variable value must be a value or identifier",
271            }
272            .fail();
273        }
274    };
275    ctx.configuration_parameter()
276        .set_pg_intervalstyle_format(intervalstyle);
277    Ok(())
278}
279
280pub fn set_datestyle(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
281    // ORDER,
282    // STYLE,
283    // ORDER,ORDER
284    // ORDER,STYLE
285    // STYLE,ORDER
286    let (style, order) = exprs
287        .iter()
288        .try_fold((None, None), |(style, order), expr| {
289            let (new_style, new_order) = try_parse_datestyle(expr)?;
290            Ok((
291                merge_datestyle_value(style, new_style)?,
292                merge_datestyle_value(order, new_order)?,
293            ))
294        })?;
295
296    let (old_style, older_order) = *ctx.configuration_parameter().pg_datetime_style();
297    ctx.configuration_parameter()
298        .set_pg_datetime_style(style.unwrap_or(old_style), order.unwrap_or(older_order));
299    Ok(())
300}
301
302pub fn set_query_timeout(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
303    let timeout_expr = exprs.first().context(NotSupportedSnafu {
304        feat: "No timeout value find in set query timeout statement",
305    })?;
306    match timeout_expr {
307        Expr::Value(ValueWithSpan {
308            value: Value::Number(timeout, _),
309            ..
310        }) => {
311            match timeout.parse::<u64>() {
312                Ok(timeout) => ctx.set_query_timeout(Duration::from_millis(timeout)),
313                Err(_) => {
314                    return NotSupportedSnafu {
315                        feat: format!("Invalid timeout expr {} in set variable statement", timeout),
316                    }
317                    .fail();
318                }
319            }
320            Ok(())
321        }
322        // postgres support time units i.e. SET STATEMENT_TIMEOUT = '50ms';
323        Expr::Value(ValueWithSpan {
324            value: Value::SingleQuotedString(timeout),
325            ..
326        })
327        | Expr::Value(ValueWithSpan {
328            value: Value::DoubleQuotedString(timeout),
329            ..
330        }) => {
331            if ctx.channel() != Postgres {
332                return NotSupportedSnafu {
333                    feat: format!("Invalid timeout expr {} in set variable statement", timeout),
334                }
335                .fail();
336            }
337            let timeout = parse_pg_query_timeout_input(timeout)?;
338            ctx.set_query_timeout(Duration::from_millis(timeout));
339            Ok(())
340        }
341        expr => NotSupportedSnafu {
342            feat: format!(
343                "Unsupported timeout expr {} in set variable statement",
344                expr
345            ),
346        }
347        .fail(),
348    }
349}
350
351// support time units in ms, s, min, h, d for postgres protocol.
352// https://www.postgresql.org/docs/8.4/config-setting.html#:~:text=Valid%20memory%20units%20are%20kB,%2C%20and%20d%20(days).
353fn parse_pg_query_timeout_input(input: &str) -> Result<u64> {
354    match input.parse::<u64>() {
355        Ok(timeout) => Ok(timeout),
356        Err(_) => {
357            if let Some(captures) = PG_TIME_INPUT_REGEX.captures(input) {
358                let value = captures[1].parse::<u64>().expect("regex failed");
359                let unit = &captures[2];
360
361                match unit {
362                    "ms" => Ok(value),
363                    "s" => Ok(value * 1000),
364                    "min" => Ok(value * 60 * 1000),
365                    "h" => Ok(value * 60 * 60 * 1000),
366                    "d" => Ok(value * 24 * 60 * 60 * 1000),
367                    _ => unreachable!("regex failed"),
368                }
369            } else {
370                NotSupportedSnafu {
371                    feat: format!(
372                        "Unsupported timeout expr {} in set variable statement",
373                        input
374                    ),
375                }
376                .fail()
377            }
378        }
379    }
380}
381
382#[cfg(test)]
383mod test {
384    use crate::statement::set::parse_pg_query_timeout_input;
385
386    #[test]
387    fn test_parse_pg_query_timeout_input() {
388        assert!(parse_pg_query_timeout_input("").is_err());
389        assert!(parse_pg_query_timeout_input(" 50 ms").is_err());
390        assert!(parse_pg_query_timeout_input("5s 1ms").is_err());
391        assert!(parse_pg_query_timeout_input("3a").is_err());
392        assert!(parse_pg_query_timeout_input("1.5min").is_err());
393        assert!(parse_pg_query_timeout_input("ms").is_err());
394        assert!(parse_pg_query_timeout_input("a").is_err());
395        assert!(parse_pg_query_timeout_input("-1").is_err());
396
397        assert_eq!(50, parse_pg_query_timeout_input("50").unwrap());
398        assert_eq!(12, parse_pg_query_timeout_input("12ms").unwrap());
399        assert_eq!(2000, parse_pg_query_timeout_input("2s").unwrap());
400        assert_eq!(60000, parse_pg_query_timeout_input("1min").unwrap());
401    }
402}
operator/statement/set.rs

operator/statement/
set.rs
