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::{ResultExt, ensure};
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
72            .parser
73            .parse_one_of_keywords(&[Keyword::FROM, Keyword::IN]);
74
75        let cursor_name = self
76            .parser
77            .parse_object_name(false)
78            .context(error::SyntaxSnafu)?;
79
80        Ok(Statement::FetchCursor(FetchCursor {
81            cursor_name: ParserContext::canonicalize_object_name(cursor_name)?,
82            fetch_size,
83        }))
84    }
85
86    pub(crate) fn parse_close_cursor(&mut self) -> Result<Statement> {
87        let _ = self.parser.expect_keyword(Keyword::CLOSE);
88        let cursor_name = self
89            .parser
90            .parse_object_name(false)
91            .context(error::SyntaxSnafu)?;
92
93        Ok(Statement::CloseCursor(CloseCursor {
94            cursor_name: ParserContext::canonicalize_object_name(cursor_name)?,
95        }))
96    }
97}
98
99#[cfg(test)]
100mod tests {
101
102    use super::*;
103    use crate::dialect::GreptimeDbDialect;
104    use crate::parser::ParseOptions;
105
106    #[test]
107    fn test_parse_declare_cursor() {
108        let sql = "DECLARE c1 CURSOR FOR\nSELECT * FROM numbers";
109        let result =
110            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
111                .unwrap();
112
113        if let Statement::DeclareCursor(dc) = &result[0] {
114            assert_eq!("c1", dc.cursor_name.to_string());
115            assert_eq!(
116                "DECLARE c1 CURSOR FOR SELECT * FROM numbers",
117                dc.to_string()
118            );
119        } else {
120            panic!("Unexpected statement");
121        }
122
123        let sql = "DECLARE c1 CURSOR FOR\nINSERT INTO numbers VALUES (1);";
124        let result =
125            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
126        assert!(result.is_err());
127    }
128
129    #[test]
130    fn test_parese_fetch_cursor() {
131        let sql = "FETCH 1000 FROM c1";
132        let result =
133            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
134                .unwrap();
135
136        if let Statement::FetchCursor(fc) = &result[0] {
137            assert_eq!("c1", fc.cursor_name.to_string());
138            assert_eq!("1000", fc.fetch_size.to_string());
139            assert_eq!(sql, fc.to_string());
140        } else {
141            panic!("Unexpected statement")
142        }
143    }
144
145    #[test]
146    fn test_close_fetch_cursor() {
147        let sql = "CLOSE c1";
148        let result =
149            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
150                .unwrap();
151
152        if let Statement::CloseCursor(cc) = &result[0] {
153            assert_eq!("c1", cc.cursor_name.to_string());
154            assert_eq!(sql, cc.to_string());
155        } else {
156            panic!("Unexpected statement")
157        }
158    }
159}
sql/parsers/cursor_parser.rs

sql/parsers/
cursor_parser.rs
