SQLite, DuckDB, PostgreSQL: Accept ._name after whitespace in a qualified name

src/tokenizer.rs

@@ -950,8 +950,7 @@ impl<'a> Tokenizer<'a> {
         self.tokenize_with_location_into_buf_with_mapper(buf, |token| token)
     }
 
-    /// Tokenize the statement and produce a vector of tokens, mapping each token
-    /// with provided `mapper`
+    /// Tokenize into the provided buffer, using existing and mapped tokens as context for subsequent tokens.
     pub fn tokenize_with_location_into_buf_with_mapper(
         &mut self,
         buf: &mut Vec<TokenWithSpan>,
@@ -963,8 +962,15 @@ impl<'a> Tokenizer<'a> {
             col: 1,
         };
 
+        let mut prev_non_ws = buf
+            .iter()
+            .rposition(|t| !matches!(t.token, Token::Whitespace(_)));
         let mut location = state.location();
-        while let Some(token) = self.next_token(&mut state, buf.last().map(|t| &t.token))? {
+        while let Some(token) = self.next_token(
+            &mut state,
+            buf.last().map(|t| &t.token),
+            prev_non_ws.map(|index| &buf[index].token),
+        )? {
             let span = location.span_to(state.location());
 
             // Check if this is a multiline comment hint that should be expanded
@@ -974,10 +980,14 @@ impl<'a> Tokenizer<'a> {
                         && comment.starts_with('!') =>
                 {
                     // Re-tokenize the hints and add them to the buffer
-                    self.tokenize_comment_hints(comment, span, buf, &mut mapper)?;
+                    self.tokenize_comment_hints(comment, span, buf, &mut prev_non_ws, &mut mapper)?;
                 }
                 _ => {
-                    buf.push(mapper(TokenWithSpan { token, span }));
+                    let token = mapper(TokenWithSpan { token, span });
+                    if !matches!(token.token, Token::Whitespace(_)) {
+                        prev_non_ws = Some(buf.len());
+                    }
+                    buf.push(token);
                 }
             }
 
@@ -993,6 +1003,7 @@ impl<'a> Tokenizer<'a> {
         comment: &str,
         span: Span,
         buf: &mut Vec<TokenWithSpan>,
+        prev_non_ws: &mut Option<usize>,
         mut mapper: impl FnMut(TokenWithSpan) -> TokenWithSpan,
     ) -> Result<(), TokenizerError> {
         // Strip the leading '!' and any version digits (e.g., "50110")
@@ -1018,12 +1029,20 @@ impl<'a> Tokenizer<'a> {
 
         // Tokenize the hint content and add tokens to the buffer
         let mut location = state.location();
-        while let Some(token) = inner.next_token(&mut state, buf.last().map(|t| &t.token))? {
+        while let Some(token) = inner.next_token(
+            &mut state,
+            buf.last().map(|t| &t.token),
+            prev_non_ws.map(|index| &buf[index].token),
+        )? {
             let token_span = location.span_to(state.location());
-            buf.push(mapper(TokenWithSpan {
+            let token = mapper(TokenWithSpan {
                 token,
                 span: token_span,
-            }));
+            });
+            if !matches!(token.token, Token::Whitespace(_)) {
+                *prev_non_ws = Some(buf.len());
+            }
+            buf.push(token);
             location = state.location();
         }
 
@@ -1061,6 +1080,7 @@ impl<'a> Tokenizer<'a> {
         &self,
         chars: &mut State,
         prev_token: Option<&Token>,
+        prev_non_ws_token: Option<&Token>,
     ) -> Result<Option<Token>, TokenizerError> {
         match chars.peek() {
             Some(&ch) => match ch {
@@ -1338,12 +1358,9 @@ impl<'a> Tokenizer<'a> {
                 }
                 // numbers and period
                 '0'..='9' | '.' => {
-                    // special case where if ._ is encountered after a word then that word
-                    // is a table and the _ is the start of the col name.
-                    // if the prev token is not a word, then this is not a valid sql
-                    // word or number.
+                    // ._ after a word is a field access, not a decimal point
                     if ch == '.' && chars.peekable.clone().nth(1) == Some('_') {
-                        if let Some(Token::Word(_)) = prev_token {
+                        if let Some(Token::Word(_)) = prev_non_ws_token {
                             chars.next();
                             return Ok(Some(Token::Period));
                         }
@@ -4385,15 +4402,30 @@ mod tests {
 
         compare(expected, tokens);
 
-        let sql = String::from("SELECT ._123");
-        if let Ok(tokens) = Tokenizer::new(&dialect, &sql).tokenize() {
-            panic!("Tokenizer should have failed on {sql}, but it succeeded with {tokens:?}");
-        }
+        // Whitespace between the word and ._ is accepted
+        let sql = String::from("SELECT table ._col");
+        let tokens = Tokenizer::new(&dialect, &sql).tokenize().unwrap();
+        let expected = vec![
+            Token::make_keyword("SELECT"),
+            Token::Whitespace(Whitespace::Space),
+            Token::Word(Word {
+                value: "table".to_string(),
+                quote_style: None,
+                keyword: Keyword::TABLE,
+            }),
+            Token::Whitespace(Whitespace::Space),
+            Token::Period,
+            Token::Word(Word {
+                value: "_col".to_string(),
+                quote_style: None,
+                keyword: Keyword::NoKeyword,
+            }),
+        ];
+        compare(expected, tokens);
 
-        let sql = String::from("SELECT ._abc");
-        if let Ok(tokens) = Tokenizer::new(&dialect, &sql).tokenize() {
-            panic!("Tokenizer should have failed on {sql}, but it succeeded with {tokens:?}");
-        }
+        // Without a preceding identifier ._ is still rejected
+        let err = Tokenizer::new(&dialect, "._abc").tokenize();
+        assert!(err.is_err(), "expected tokenizer error for bare ._abc");
     }
 
     #[test]
@@ -4497,6 +4529,47 @@ mod tests {
         );
     }
 
+    #[test]
+    fn tokenize_comment_hint_underscore_after_non_word_is_error() {
+        let sql = "SELECT 1 /*!50110 ._a */ FROM t";
+        let dialect = MySqlDialect {};
+        let err = Tokenizer::new(&dialect, sql).tokenize().unwrap_err();
+        assert_eq!(
+            err,
+            TokenizerError {
+                message: "Unexpected character '_'".to_string(),
+                location: Location::new(1, 11),
+            }
+        );
+    }
+
+    #[test]
+    fn tokenize_into_existing_buffer_keeps_qualified_name_context() {
+        let dialect = MySqlDialect {};
+        let mut tokens = Tokenizer::new(&dialect, "SELECT t ")
+            .tokenize_with_location()
+            .unwrap();
+        Tokenizer::new(&dialect, "._col FROM t")
+            .tokenize_with_location_into_buf(&mut tokens)
+            .unwrap();
+        let tokens = tokens
+            .into_iter()
+            .map(|token| token.token)
+            .filter(|token| !matches!(token, Token::Whitespace(_)))
+            .collect();
+        compare(
+            vec![
+                Token::make_keyword("SELECT"),
+                Token::make_word("t", None),
+                Token::Period,
+                Token::make_word("_col", None),
+                Token::make_keyword("FROM"),
+                Token::make_word("t", None),
+            ],
+            tokens,
+        );
+    }
+
     #[test]
     fn tokenize_lt() {
         all_dialects().tokenizes_to(

tests/sqlparser_common.rs

@@ -20803,3 +20803,84 @@ fn display_nested_unary_minus_in_linear_time() {
         .expect("display should not take exponential time");
     assert_eq!(printed, sql);
 }
+
+#[test]
+fn parse_period_underscore_after_whitespace() {
+    let dialects = all_dialects_where(|d| d.is_identifier_start('_'));
+    dialects.verified_stmt("SELECT t._a FROM t");
+    dialects.one_statement_parses_to("SELECT t ._a FROM t", "SELECT t._a FROM t");
+    dialects.one_statement_parses_to("SELECT t . _a FROM t", "SELECT t._a FROM t");
+
+    let no_expression =
+        || ParserError::ParserError("Expected: an expression, found: .".to_string());
+    for (sql, error) in [
+        ("SELECT ._abc", no_expression()),
+        ("SELECT ._123", no_expression()),
+        (
+            "SELECT 1 ._a",
+            ParserError::TokenizerError(
+                "Unexpected character '_' at Line: 1, Column: 10".to_string(),
+            ),
+        ),
+    ] {
+        assert_eq!(dialects.parse_sql_statements(sql), Err(error), "{sql}");
+    }
+}
+
+#[test]
+fn tokenize_long_whitespace_run_completes() {
+    use std::sync::mpsc;
+    use std::thread;
+    use std::time::Duration;
+
+    let request_id = "tokenize-long-whitespace-run";
+    let timeout = Duration::from_secs(5);
+    let sql = format!("SELECT{}t._a FROM t", " ".repeat(1024 * 1024));
+    let (tx, rx) = mpsc::channel();
+    thread::spawn(move || {
+        let result = Tokenizer::new(&GenericDialect {}, &sql)
+            .tokenize_with_location()
+            .map(drop);
+        let _ = tx.send((request_id, result));
+    });
+
+    let (response_id, result) = rx
+        .recv_timeout(timeout)
+        .expect("tokenizer exceeded the whitespace-run deadline");
+    assert_eq!(response_id, request_id);
+    result.expect("whitespace-separated qualified name should tokenize");
+}
+
+#[test]
+fn tokenize_long_whitespace_run_in_comment_hint_completes() {
+    use core::time::Duration;
+    use std::sync::mpsc;
+    use std::thread;
+
+    let request_id = "tokenize-long-whitespace-run-in-comment-hint";
+    let timeout = Duration::from_secs(5);
+    let sql = format!("SELECT /*!{}t ._a */ FROM t", " ".repeat(1024 * 1024));
+    let (tx, rx) = mpsc::channel();
+    thread::spawn(move || {
+        let result = [&GenericDialect {} as &dyn Dialect, &MySqlDialect {}]
+            .into_iter()
+            .try_for_each(|dialect| {
+                Tokenizer::new(dialect, &sql)
+                    .tokenize_with_location()
+                    .map(drop)
+            });
+        let _ = tx.send((request_id, result));
+    });
+
+    let (response_id, result) = rx
+        .recv_timeout(timeout)
+        .expect("hint tokenizer exceeded the whitespace-run deadline");
+    assert_eq!(response_id, request_id);
+    result.expect("whitespace-separated qualified name in hint should tokenize");
+}
+
+#[test]
+fn parse_qualified_underscore_column_from_comment_hint() {
+    let dialects = all_dialects_where(|d| d.supports_multiline_comment_hints());
+    dialects.one_statement_parses_to("SELECT t /*!50110 ._col */ FROM t", "SELECT t._col FROM t");
+}