locations: collect locations of expressions and SelectItems

Author: lustefaniak

src/ast/mod.rs

@@ -346,6 +346,10 @@ where
     pub fn unwrap(self) -> T {
         self.inner
     }
+
+    pub fn span_location(&self) -> Span {
+        self.span
+    }
 }
 
 pub trait SpanWrapped: Clone + Eq + Ord + std::hash::Hash + PartialOrd + PartialEq {
@@ -1272,7 +1276,7 @@ pub enum Statement {
         table: bool,
         on: Option<OnInsert>,
         /// RETURNING
-        returning: Option<Vec<SelectItem>>,
+        returning: Option<Vec<WithSpan<SelectItem>>>,
     },
     // TODO: Support ROW FORMAT
     Directory {
@@ -1332,7 +1336,7 @@ pub enum Statement {
         /// WHERE
         selection: Option<Expr>,
         /// RETURNING
-        returning: Option<Vec<SelectItem>>,
+        returning: Option<Vec<WithSpan<SelectItem>>>,
     },
     /// DELETE
     Delete {
@@ -1345,7 +1349,7 @@ pub enum Statement {
         /// WHERE
         selection: Option<Expr>,
         /// RETURNING
-        returning: Option<Vec<SelectItem>>,
+        returning: Option<Vec<WithSpan<SelectItem>>>,
     },
     /// CREATE VIEW
     CreateView {

src/ast/query.rs

@@ -204,7 +204,7 @@ pub struct Select {
     /// MSSQL syntax: `TOP (<N>) [ PERCENT ] [ WITH TIES ]`
     pub top: Option<Top>,
     /// projection expressions
-    pub projection: Vec<SelectItem>,
+    pub projection: Vec<WithSpan<SelectItem>>,
     /// INTO
     pub into: Option<SelectInto>,
     /// FROM
@@ -298,7 +298,7 @@ pub struct LateralView {
     /// LATERAL VIEW table name
     pub lateral_view_name: ObjectName,
     /// LATERAL VIEW optional column aliases
-    pub lateral_col_alias: Vec<Ident>,
+    pub lateral_col_alias: Vec<WithSpan<Ident>>,
     /// LATERAL VIEW OUTER
     pub outer: bool,
 }
@@ -382,9 +382,9 @@ impl fmt::Display for Cte {
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub enum SelectItem {
     /// Any expression, not followed by `[ AS ] alias`
-    UnnamedExpr(Expr),
+    UnnamedExpr(WithSpan<Expr>),
     /// An expression, followed by `[ AS ] alias`
-    ExprWithAlias { expr: Expr, alias: Ident },
+    ExprWithAlias { expr: WithSpan<Expr>, alias: WithSpan<Ident> },
     /// `alias.*` or even `schema.table.*`
     QualifiedWildcard(ObjectName, WildcardAdditionalOptions),
     /// An unqualified `*`
@@ -697,7 +697,7 @@ pub enum TableFactor {
         alias: Option<TableAlias>,
         array_expr: Box<Expr>,
         with_offset: bool,
-        with_offset_alias: Option<Ident>,
+        with_offset_alias: Option<WithSpan<Ident>>,
     },
     /// Represents a parenthesized table factor. The SQL spec only allows a
     /// join expression (`(foo <JOIN> bar [ <JOIN> baz ... ])`) to be nested,
@@ -857,7 +857,7 @@ impl fmt::Display for TableFactor {
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub struct TableAlias {
-    pub name: Ident,
+    pub name: WithSpan<Ident>,
     pub columns: Vec<Ident>,
 }
 

src/parser.rs

@@ -192,7 +192,7 @@ impl fmt::Display for ParserError {
 impl std::error::Error for ParserError {}
 
 // By default, allow expressions up to this deep before erroring
-const DEFAULT_REMAINING_DEPTH: usize = 50;
+const DEFAULT_REMAINING_DEPTH: usize = 48;
 
 #[derive(Debug, Default, Clone, PartialEq, Eq)]
 pub struct ParserOptions {
@@ -2291,7 +2291,7 @@ impl<'a> Parser<'a> {
     }
 
     /// Parse a comma-separated list of 1+ SelectItem
-    pub fn parse_projection(&mut self) -> Result<Vec<SelectItem>, ParserError> {
+    pub fn parse_projection(&mut self) -> Result<Vec<WithSpan<SelectItem>>, ParserError> {
         // BigQuery allows trailing commas, but only in project lists
         // e.g. `SELECT 1, 2, FROM t`
         // https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#trailing_commas
@@ -4728,7 +4728,7 @@ impl<'a> Parser<'a> {
     pub fn parse_optional_alias(
         &mut self,
         reserved_kwds: &[Keyword],
-    ) -> Result<Option<Ident>, ParserError> {
+    ) -> Result<Option<WithSpan<Ident>>, ParserError> {
         let after_as = self.parse_keyword(Keyword::AS);
         let next_token = self.next_token();
         match next_token.token {
@@ -4738,7 +4738,7 @@ impl<'a> Parser<'a> {
             // (For example, in `FROM t1 JOIN` the `JOIN` will always be parsed as a keyword,
             // not an alias.)
             Token::Word(w) if after_as || !reserved_kwds.contains(&w.keyword) => {
-                Ok(Some(w.to_ident()))
+                Ok(Some(w.to_ident().spanning(next_token.span)))
             }
             // MSSQL supports single-quoted strings as aliases for columns
             // We accept them as table aliases too, although MSSQL does not.
@@ -4752,9 +4752,9 @@ impl<'a> Parser<'a> {
             //    character. When it sees such a <literal>, your DBMS will
             //    ignore the <separator> and treat the multiple strings as
             //    a single <literal>."
-            Token::SingleQuotedString(s) => Ok(Some(Ident::with_quote('\'', s))),
+            Token::SingleQuotedString(s) => Ok(Some(Ident::with_quote('\'', s).spanning(next_token.span))),
             // Support for MySql dialect double quoted string, `AS "HOUR"` for example
-            Token::DoubleQuotedString(s) => Ok(Some(Ident::with_quote('\"', s))),
+            Token::DoubleQuotedString(s) => Ok(Some(Ident::with_quote('\"', s).spanning(next_token.span))),
             _ => {
                 if after_as {
                     return self.expected("an identifier after AS", next_token);
@@ -5273,7 +5273,7 @@ impl<'a> Parser<'a> {
 
     /// Parse a CTE (`alias [( col1, col2, ... )] AS (subquery)`)
     pub fn parse_cte(&mut self) -> Result<Cte, ParserError> {
-        let name = self.parse_identifier()?.unwrap();
+        let name = self.parse_identifier()?;
 
         let mut cte = if self.parse_keyword(Keyword::AS) {
             self.expect_token(&Token::LParen)?;
@@ -5444,7 +5444,7 @@ impl<'a> Parser<'a> {
                 let outer = self.parse_keyword(Keyword::OUTER);
                 let lateral_view = self.parse_expr()?;
                 let lateral_view_name = self.parse_object_name()?;
-                let lateral_col_alias = self
+                let lateral_col_alias: Vec<WithSpan<Ident>> = self
                     .parse_comma_separated(|parser| {
                         parser.parse_optional_alias(&[
                             Keyword::WHERE,
@@ -6582,7 +6582,8 @@ impl<'a> Parser<'a> {
     }
 
     /// Parse a comma-delimited list of projections after SELECT
-    pub fn parse_select_item(&mut self) -> Result<SelectItem, ParserError> {
+    pub fn parse_select_item(&mut self) -> Result<WithSpan<SelectItem>, ParserError> {
+        let start_span = self.index;
         match self.parse_wildcard_expr()? {
             WildcardExpr::Expr(expr) => {
                 let expr: Expr = if self.dialect.supports_filter_during_aggregation()
@@ -6603,19 +6604,29 @@ impl<'a> Parser<'a> {
                 } else {
                     expr
                 };
+                let expr_with_location = expr.spanning(self.span_from_index(start_span));
                 self.parse_optional_alias(keywords::RESERVED_FOR_COLUMN_ALIAS)
                     .map(|alias| match alias {
-                        Some(alias) => SelectItem::ExprWithAlias { expr, alias },
-                        None => SelectItem::UnnamedExpr(expr),
+                        Some(alias) => SelectItem::ExprWithAlias {
+                            expr: expr_with_location,
+                            alias,
+                        }
+                        .spanning(self.span_from_index(start_span)),
+                        None => {
+                            SelectItem::UnnamedExpr(expr_with_location)
+                                .spanning(self.span_from_index(start_span))
+                        }
                     })
             }
             WildcardExpr::QualifiedWildcard(prefix) => Ok(SelectItem::QualifiedWildcard(
                 prefix,
                 self.parse_wildcard_additional_options()?,
-            )),
+            )
+            .spanning(self.span_from_index(start_span))),
             WildcardExpr::Wildcard => Ok(SelectItem::Wildcard(
                 self.parse_wildcard_additional_options()?,
-            )),
+            )
+            .spanning(self.span_from_index(start_span))),
         }
     }
 
@@ -7293,6 +7304,50 @@ impl<'a> Parser<'a> {
             representation: UserDefinedTypeRepresentation::Composite { attributes },
         })
     }
+
+    fn span_from_index(&mut self, mut start_index: usize) -> Span {
+        let mut start_token = &self.tokens[start_index];
+        loop {
+            match start_token {
+                TokenWithLocation {
+                    token: Token::Whitespace(_),
+                    span: _,
+                } => {
+                    start_index += 1;
+                    start_token = &self.tokens[start_index];
+                    continue;
+                }
+                _ => break,
+            }
+        }
+        let start_span = start_token.span;
+
+        let mut idx = self.index.max(start_index).min(self.tokens.len());
+        loop {
+            if idx <= start_index || idx >= self.tokens.len() {
+                break;
+            }
+            let curr_token = &self.tokens[idx];
+            match curr_token {
+                TokenWithLocation {
+                    token: Token::Whitespace(_),
+                    span: _,
+                } => {
+                    idx -= 1;
+                    continue;
+                }
+                TokenWithLocation {
+                    token: Token::Word(word),
+                    span: _,
+                } if word.keyword != Keyword::NoKeyword => {
+                    idx -= 1;
+                    continue;
+                }
+                non_whitespace => return non_whitespace.span.union(start_span),
+            }
+        }
+        return start_span;
+    }
 }
 
 impl Word {

src/test_utils.rs

@@ -207,7 +207,7 @@ pub fn number(n: &'static str) -> Value {
 
 pub fn table_alias(name: impl Into<String>) -> Option<TableAlias> {
     Some(TableAlias {
-        name: Ident::new(name),
+        name: Ident::new(name).empty_span(),
         columns: vec![],
     })
 }

src/tokenizer.rs

@@ -358,7 +358,7 @@ pub struct Location {
 }
 
 impl Location {
-    fn is_valid(&self) -> bool {
+    pub fn is_valid(&self) -> bool {
         self.line > 0 && self.column > 0
     }
 }

tests/sqlparser_bigquery.rs

@@ -337,7 +337,7 @@ fn test_select_wildcard_with_except() {
             additional_elements: vec![],
         }),
         ..Default::default()
-    });
+    }).empty_span();
     assert_eq!(expected, select.projection[0]);
 
     let select = bigquery_and_generic()
@@ -348,7 +348,7 @@ fn test_select_wildcard_with_except() {
             additional_elements: vec![Ident::new("employee_id")],
         }),
         ..Default::default()
-    });
+    }).empty_span();
     assert_eq!(expected, select.projection[0]);
 
     assert_eq!(
@@ -382,7 +382,7 @@ fn test_select_wildcard_with_replace() {
             })],
         }),
         ..Default::default()
-    });
+    }).empty_span();
     assert_eq!(expected, select.projection[0]);
 
     let select = bigquery_and_generic().verified_only_select(
@@ -417,7 +417,7 @@ fn test_select_wildcard_with_replace() {
             ],
         }),
         ..Default::default()
-    });
+    }).empty_span();
     assert_eq!(expected, select.projection[0]);
 }
 
@@ -461,7 +461,7 @@ fn parse_map_access_offset() {
                 order_by: vec![],
                 null_treatment: None,
             })],
-        })
+        }.empty_span()).empty_span()
     );
 
     // test other operators

tests/sqlparser_clickhouse.rs

@@ -59,7 +59,7 @@ fn parse_array_access_expr() {
                     order_by: vec![],
                     null_treatment: None,
                 })],
-            })],
+            }.empty_span()).empty_span()],
             into: None,
             from: vec![TableWithJoins {
                 relation: Table {
@@ -182,7 +182,7 @@ fn parse_delimited_identifiers() {
             with_hints,
         } => {
             assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);
-            assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name);
+            assert_eq!(Ident::with_quote('"', "alias").empty_span(), alias.unwrap().name);
             assert!(args.is_none());
             assert!(with_hints.is_empty());
         }
@@ -209,13 +209,13 @@ fn parse_delimited_identifiers() {
         }),
         expr_from_projection(&select.projection[1]),
     );
-    match &select.projection[2] {
+    match select.projection[2].clone().unwrap() {
         SelectItem::ExprWithAlias { expr, alias } => {
             assert_eq!(
-                &Expr::Identifier(Ident::with_quote('"', "simple id").empty_span()),
+                Expr::Identifier(Ident::with_quote('"', "simple id").empty_span()).empty_span(),
                 expr
             );
-            assert_eq!(&Ident::with_quote('"', "column alias"), alias);
+            assert_eq!(Ident::with_quote('"', "column alias").empty_span(), alias);
         }
         _ => panic!("Expected ExprWithAlias"),
     }