sql/parsers/
cursor_parser.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 snafu::{ensure, ResultExt};
16use sqlparser::keywords::Keyword;
17use sqlparser::tokenizer::Token;
18
19use crate::error::{self, Result};
20use crate::parser::ParserContext;
21use crate::statements::cursor::{CloseCursor, DeclareCursor, FetchCursor};
22use crate::statements::statement::Statement;
23
24impl ParserContext<'_> {
25    pub(crate) fn parse_declare_cursor(&mut self) -> Result<Statement> {
26        let _ = self.parser.expect_keyword(Keyword::DECLARE);
27        let cursor_name = self
28            .parser
29            .parse_object_name(false)
30            .context(error::SyntaxSnafu)?;
31        let _ = self
32            .parser
33            .expect_keywords(&[Keyword::CURSOR, Keyword::FOR]);
34
35        let mut is_select = false;
36        if let Token::Word(w) = self.parser.peek_token().token {
37            match w.keyword {
38                Keyword::SELECT | Keyword::WITH => {
39                    is_select = true;
40                }
41                _ => {}
42            }
43        };
44        ensure!(
45            is_select,
46            error::InvalidSqlSnafu {
47                msg: "Expect select query in cursor statement".to_string(),
48            }
49        );
50
51        let query_stmt = self.parse_query()?;
52        match query_stmt {
53            Statement::Query(query) => Ok(Statement::DeclareCursor(DeclareCursor {
54                cursor_name: ParserContext::canonicalize_object_name(cursor_name),
55                query,
56            })),
57            _ => error::InvalidSqlSnafu {
58                msg: format!("Expect query, found {}", query_stmt),
59            }
60            .fail(),
61        }
62    }
63
64    pub(crate) fn parse_fetch_cursor(&mut self) -> Result<Statement> {
65        let _ = self.parser.expect_keyword(Keyword::FETCH);
66
67        let fetch_size = self
68            .parser
69            .parse_literal_uint()
70            .context(error::SyntaxSnafu)?;
71        let _ = self.parser.parse_keyword(Keyword::FROM);
72
73        let cursor_name = self
74            .parser
75            .parse_object_name(false)
76            .context(error::SyntaxSnafu)?;
77
78        Ok(Statement::FetchCursor(FetchCursor {
79            cursor_name: ParserContext::canonicalize_object_name(cursor_name),
80            fetch_size,
81        }))
82    }
83
84    pub(crate) fn parse_close_cursor(&mut self) -> Result<Statement> {
85        let _ = self.parser.expect_keyword(Keyword::CLOSE);
86        let cursor_name = self
87            .parser
88            .parse_object_name(false)
89            .context(error::SyntaxSnafu)?;
90
91        Ok(Statement::CloseCursor(CloseCursor {
92            cursor_name: ParserContext::canonicalize_object_name(cursor_name),
93        }))
94    }
95}
96
97#[cfg(test)]
98mod tests {
99
100    use super::*;
101    use crate::dialect::GreptimeDbDialect;
102    use crate::parser::ParseOptions;
103
104    #[test]
105    fn test_parse_declare_cursor() {
106        let sql = "DECLARE c1 CURSOR FOR\nSELECT * FROM numbers";
107        let result =
108            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
109                .unwrap();
110
111        if let Statement::DeclareCursor(dc) = &result[0] {
112            assert_eq!("c1", dc.cursor_name.to_string());
113            assert_eq!(
114                "DECLARE c1 CURSOR FOR SELECT * FROM numbers",
115                dc.to_string()
116            );
117        } else {
118            panic!("Unexpected statement");
119        }
120
121        let sql = "DECLARE c1 CURSOR FOR\nINSERT INTO numbers VALUES (1);";
122        let result =
123            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
124        assert!(result.is_err());
125    }
126
127    #[test]
128    fn test_parese_fetch_cursor() {
129        let sql = "FETCH 1000 FROM c1";
130        let result =
131            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
132                .unwrap();
133
134        if let Statement::FetchCursor(fc) = &result[0] {
135            assert_eq!("c1", fc.cursor_name.to_string());
136            assert_eq!("1000", fc.fetch_size.to_string());
137            assert_eq!(sql, fc.to_string());
138        } else {
139            panic!("Unexpected statement")
140        }
141    }
142
143    #[test]
144    fn test_close_fetch_cursor() {
145        let sql = "CLOSE c1";
146        let result =
147            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
148                .unwrap();
149
150        if let Statement::CloseCursor(cc) = &result[0] {
151            assert_eq!("c1", cc.cursor_name.to_string());
152            assert_eq!(sql, cc.to_string());
153        } else {
154            panic!("Unexpected statement")
155        }
156    }
157}
sql/parsers/cursor_parser.rs

sql/parsers/
cursor_parser.rs
