risingwave_frontend/expr/
session_timezone.rs

1// Copyright 2025 RisingWave 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 risingwave_common::types::DataType;
16pub use risingwave_pb::expr::expr_node::Type as ExprType;
17
18pub use crate::expr::expr_rewriter::ExprRewriter;
19pub use crate::expr::function_call::FunctionCall;
20use crate::expr::{Expr, ExprImpl, ExprVisitor};
21use crate::session::current;
22
23/// `SessionTimezone` will be used to resolve session
24/// timezone-dependent casts, comparisons or arithmetic.
25pub struct SessionTimezone {
26    timezone: String,
27    /// Whether or not the session timezone was used
28    used: bool,
29}
30
31impl ExprRewriter for SessionTimezone {
32    fn rewrite_function_call(&mut self, func_call: FunctionCall) -> ExprImpl {
33        let (func_type, inputs, ret) = func_call.decompose();
34        let inputs: Vec<ExprImpl> = inputs
35            .into_iter()
36            .map(|expr| self.rewrite_expr(expr))
37            .collect();
38        if let Some(expr) = self.with_timezone(func_type, &inputs, &ret) {
39            self.mark_used();
40            expr
41        } else {
42            FunctionCall::new_unchecked(func_type, inputs, ret).into()
43        }
44    }
45}
46
47impl SessionTimezone {
48    pub fn new(timezone: String) -> Self {
49        Self {
50            timezone,
51            used: false,
52        }
53    }
54
55    pub fn timezone(&self) -> String {
56        self.timezone.clone()
57    }
58
59    pub fn used(&self) -> bool {
60        self.used
61    }
62
63    fn mark_used(&mut self) {
64        if !self.used {
65            self.used = true;
66            current::notice_to_user(format!(
67                "Your session timezone is {}. It was used in the interpretation of timestamps and dates in your query. If this is unintended, \
68                change your timezone to match that of your data's with `set timezone = [timezone]` or \
69                rewrite your query with an explicit timezone conversion, e.g. with `AT TIME ZONE`.\n",
70                self.timezone
71            ));
72        }
73    }
74
75    // Inlines conversions based on session timezone if required by the function
76    fn with_timezone(
77        &self,
78        func_type: ExprType,
79        inputs: &[ExprImpl],
80        return_type: &DataType,
81    ) -> Option<ExprImpl> {
82        match func_type {
83            // `input_timestamptz::varchar`
84            // => `cast_with_time_zone(input_timestamptz, zone_string)`
85            // `input_varchar::timestamptz`
86            // => `cast_with_time_zone(input_varchar, zone_string)`
87            // `input_date::timestamptz`
88            // => `input_date::timestamp AT TIME ZONE zone_string`
89            // `input_timestamp::timestamptz`
90            // => `input_timestamp AT TIME ZONE zone_string`
91            // `input_timestamptz::date`
92            // => `(input_timestamptz AT TIME ZONE zone_string)::date`
93            // `input_timestamptz::time`
94            // => `(input_timestamptz AT TIME ZONE zone_string)::time`
95            // `input_timestamptz::timestamp`
96            // => `input_timestamptz AT TIME ZONE zone_string`
97            ExprType::Cast => {
98                assert_eq!(inputs.len(), 1);
99                let mut input = inputs[0].clone();
100                let input_type = input.return_type();
101                match (input_type, return_type) {
102                    (DataType::Timestamptz, DataType::Varchar)
103                    | (DataType::Varchar, DataType::Timestamptz) => {
104                        Some(self.cast_with_timezone(input, return_type.clone()))
105                    }
106                    (DataType::Date, DataType::Timestamptz)
107                    | (DataType::Timestamp, DataType::Timestamptz) => {
108                        input = input.cast_explicit(&DataType::Timestamp).unwrap();
109                        Some(self.at_timezone(input))
110                    }
111                    (DataType::Timestamptz, DataType::Date)
112                    | (DataType::Timestamptz, DataType::Time)
113                    | (DataType::Timestamptz, DataType::Timestamp) => {
114                        input = self.at_timezone(input);
115                        input = input.cast_explicit(return_type).unwrap();
116                        Some(input)
117                    }
118                    _ => None,
119                }
120            }
121            // `lhs_date CMP rhs_timestamptz`
122            // => `(lhs_date::timestamp AT TIME ZONE zone_string) CMP rhs_timestamptz`
123            // `lhs_timestamp CMP rhs_timestamptz`
124            // => `(lhs_timestamp AT TIME ZONE zone_string) CMP rhs_timestamptz`
125            // `lhs_timestamptz CMP rhs_date`
126            // => `lhs_timestamptz CMP (rhs_date::timestamp AT TIME ZONE zone_string)`
127            // `lhs_timestamptz CMP rhs_timestamp`
128            // => `lhs_timestamptz CMP (rhs_timestamp AT TIME ZONE zone_string)`
129            ExprType::Equal
130            | ExprType::NotEqual
131            | ExprType::LessThan
132            | ExprType::LessThanOrEqual
133            | ExprType::GreaterThan
134            | ExprType::GreaterThanOrEqual
135            | ExprType::IsDistinctFrom
136            | ExprType::IsNotDistinctFrom => {
137                assert_eq!(inputs.len(), 2);
138                let mut inputs = inputs.to_vec();
139                for idx in 0..2 {
140                    if matches!(inputs[(idx + 1) % 2].return_type(), DataType::Timestamptz)
141                        && matches!(
142                            inputs[idx % 2].return_type(),
143                            DataType::Date | DataType::Timestamp
144                        )
145                    {
146                        let mut to_cast = inputs[idx % 2].clone();
147                        // Cast to `Timestamp` first, then use `AT TIME ZONE` to convert to
148                        // `Timestamptz`
149                        to_cast = to_cast.cast_explicit(&DataType::Timestamp).unwrap();
150                        inputs[idx % 2] = self.at_timezone(to_cast);
151                        return Some(
152                            FunctionCall::new_unchecked(func_type, inputs, return_type.clone())
153                                .into(),
154                        );
155                    }
156                }
157                None
158            }
159            // `add(lhs_interval, rhs_timestamptz)`
160            // => `add_with_time_zone(rhs_timestamptz, lhs_interval, zone_string)`
161            // `add(lhs_timestamptz, rhs_interval)`
162            // => `add_with_time_zone(lhs_timestamptz, rhs_interval, zone_string)`
163            // `subtract(lhs_timestamptz, rhs_interval)`
164            // => `subtract_with_time_zone(lhs_timestamptz, rhs_interval, zone_string)`
165            ExprType::Subtract | ExprType::Add => {
166                assert_eq!(inputs.len(), 2);
167                let canonical_match = matches!(inputs[0].return_type(), DataType::Timestamptz)
168                    && matches!(inputs[1].return_type(), DataType::Interval);
169                let inverse_match = matches!(inputs[1].return_type(), DataType::Timestamptz)
170                    && matches!(inputs[0].return_type(), DataType::Interval);
171                assert!(!(inverse_match && func_type == ExprType::Subtract)); // This should never have been parsed.
172                if canonical_match || inverse_match {
173                    let (orig_timestamptz, interval) =
174                        if func_type == ExprType::Add && inverse_match {
175                            (inputs[1].clone(), inputs[0].clone())
176                        } else {
177                            (inputs[0].clone(), inputs[1].clone())
178                        };
179                    let new_type = match func_type {
180                        ExprType::Add => ExprType::AddWithTimeZone,
181                        ExprType::Subtract => ExprType::SubtractWithTimeZone,
182                        _ => unreachable!(),
183                    };
184                    let rewritten_expr = FunctionCall::new(
185                        new_type,
186                        vec![
187                            orig_timestamptz,
188                            interval,
189                            ExprImpl::literal_varchar(self.timezone()),
190                        ],
191                    )
192                    .unwrap()
193                    .into();
194                    return Some(rewritten_expr);
195                }
196                None
197            }
198            // `date_trunc(field_string, input_timestamptz)`
199            // => `date_trunc(field_string, input_timestamptz, zone_string)`
200            ExprType::DateTrunc | ExprType::Extract | ExprType::DatePart => {
201                if !(inputs.len() == 2 && inputs[1].return_type() == DataType::Timestamptz) {
202                    return None;
203                }
204                assert_eq!(inputs[0].return_type(), DataType::Varchar);
205                if let ExprImpl::Literal(lit) = &inputs[0]
206                    && matches!(func_type, ExprType::Extract | ExprType::DatePart)
207                    && lit
208                        .get_data()
209                        .as_ref()
210                        .is_none_or(|v| v.as_utf8().eq_ignore_ascii_case("epoch"))
211                {
212                    // No need to rewrite when field is `null` or `epoch`.
213                    // This is optional but avoids false warning in common case.
214                    return None;
215                }
216                let mut new_inputs = inputs.to_vec();
217                new_inputs.push(ExprImpl::literal_varchar(self.timezone()));
218                Some(FunctionCall::new(func_type, new_inputs).unwrap().into())
219            }
220            // `char_to_timestamptz(input_string, format_string)`
221            // => `char_to_timestamptz(input_string, format_string, zone_string)`
222            ExprType::CharToTimestamptz => {
223                if !(inputs.len() == 2
224                    && inputs[0].return_type() == DataType::Varchar
225                    && inputs[1].return_type() == DataType::Varchar)
226                {
227                    return None;
228                }
229                let mut new_inputs = inputs.to_vec();
230                new_inputs.push(ExprImpl::literal_varchar(self.timezone()));
231                Some(FunctionCall::new(func_type, new_inputs).unwrap().into())
232            }
233            // `to_char(input_timestamptz, format_string)`
234            // => `to_char(input_timestamptz, format_string, zone_string)`
235            ExprType::ToChar => {
236                if !(inputs.len() == 2
237                    && inputs[0].return_type() == DataType::Timestamptz
238                    && inputs[1].return_type() == DataType::Varchar)
239                {
240                    return None;
241                }
242                let mut new_inputs = inputs.to_vec();
243                new_inputs.push(ExprImpl::literal_varchar(self.timezone()));
244                Some(FunctionCall::new(func_type, new_inputs).unwrap().into())
245            }
246            _ => None,
247        }
248    }
249
250    fn at_timezone(&self, input: ExprImpl) -> ExprImpl {
251        FunctionCall::new(
252            ExprType::AtTimeZone,
253            vec![input, ExprImpl::literal_varchar(self.timezone.clone())],
254        )
255        .unwrap()
256        .into()
257    }
258
259    fn cast_with_timezone(&self, input: ExprImpl, return_type: DataType) -> ExprImpl {
260        FunctionCall::new_unchecked(
261            ExprType::CastWithTimeZone,
262            vec![input, ExprImpl::literal_varchar(self.timezone.clone())],
263            return_type,
264        )
265        .into()
266    }
267}
268
269#[derive(Default)]
270pub struct TimestamptzExprFinder {
271    has: bool,
272}
273
274impl TimestamptzExprFinder {
275    pub fn has(&self) -> bool {
276        self.has
277    }
278}
279
280impl ExprVisitor for TimestamptzExprFinder {
281    fn visit_function_call(&mut self, func_call: &FunctionCall) {
282        if func_call.return_type() == DataType::Timestamptz {
283            self.has = true;
284            return;
285        }
286
287        for input in &func_call.inputs {
288            if input.return_type() == DataType::Timestamptz {
289                self.has = true;
290                return;
291            }
292        }
293
294        func_call
295            .inputs()
296            .iter()
297            .for_each(|expr| self.visit_expr(expr));
298    }
299}
risingwave_frontend/expr/session_timezone.rs

risingwave_frontend/expr/
session_timezone.rs
