Skip to content

Improve support for cursors for SQL Server #1831

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 12 commits into from
May 2, 2025
Merged

Improve support for cursors for SQL Server #1831

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
ac298a6
Add additional cursor parsing support for SQL Server
aharpervc Apr 25, 2025
f1e8ac7
Add `statements_parse_to` helper
aharpervc Apr 28, 2025
5ec1463
Add support for parsing `WHILE` statements
aharpervc Apr 25, 2025
a72e1c1
Make `OPEN` reserved for table aliases
aharpervc Apr 29, 2025
01d85a0
Introduce `FetchPosition` to simplify FROM vs IN
aharpervc Apr 29, 2025
9965777
Remove unnecessary comment
aharpervc Apr 29, 2025
648c024
Introduce a `OpenStatement` struct for the `OPEN` statement
aharpervc Apr 29, 2025
bf0036a
Expand test example to assert AST
aharpervc Apr 29, 2025
2941291
Add test for `OPEN` statement
aharpervc Apr 29, 2025
3d2001f
Merge conditions into single `match`
aharpervc Apr 29, 2025
dbf7606
Fix incorrect trailing comma
aharpervc Apr 29, 2025
3608d8c
Remove statement count assertion from `statements_parse_to`
aharpervc Apr 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 72 additions & 2 deletions src/ast/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2226,7 +2226,33 @@ impl fmt::Display for IfStatement {
}
}

/// A block within a [Statement::Case] or [Statement::If]-like statement
/// A `WHILE` statement.
///
/// Example:
/// ```sql
/// WHILE @@FETCH_STATUS = 0
/// BEGIN
/// FETCH NEXT FROM c1 INTO @var1, @var2;
/// END
/// ```
///
/// [MsSql](https://learn.microsoft.com/en-us/sql/t-sql/language-elements/while-transact-sql)
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct WhileStatement {
pub while_block: ConditionalStatementBlock,
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't absolutely need a WhileStatement struct; we could be doing Statement::While(ConditionalStatementBlock) instead. I'm following the example of CASE & IF, which also do it this way.

}

impl fmt::Display for WhileStatement {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let WhileStatement { while_block } = self;
write!(f, "{while_block}")?;
Ok(())
}
}

/// A block within a [Statement::Case] or [Statement::If] or [Statement::While]-like statement
///
/// Example 1:
/// ```sql
Expand All @@ -2242,6 +2268,14 @@ impl fmt::Display for IfStatement {
/// ```sql
/// ELSE SELECT 1; SELECT 2;
/// ```
///
/// Example 4:
/// ```sql
/// WHILE @@FETCH_STATUS = 0
/// BEGIN
/// FETCH NEXT FROM c1 INTO @var1, @var2;
/// END
/// ```
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
Expand Down Expand Up @@ -2981,6 +3015,8 @@ pub enum Statement {
Case(CaseStatement),
/// An `IF` statement.
If(IfStatement),
/// A `WHILE` statement.
While(WhileStatement),
/// A `RAISE` statement.
Raise(RaiseStatement),
/// ```sql
Expand Down Expand Up @@ -3032,6 +3068,14 @@ pub enum Statement {
partition: Option<Box<Expr>>,
},
/// ```sql
/// OPEN cursor_name
/// ```
/// Opens a cursor.
Open {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I placed this next to CLOSE because they're semantically paired, rather than alphabetical. Not sure what's preferred on this project

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we wrap this new statement in a named struct?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done 👍. I didn't do that originally so as to more closely mimic the existing code

/// Cursor name
cursor_name: Ident,
},
/// ```sql
/// CLOSE
/// ```
/// Closes the portal underlying an open cursor.
Expand Down Expand Up @@ -3403,6 +3447,10 @@ pub enum Statement {
/// Cursor name
name: Ident,
direction: FetchDirection,
/// Differentiate between dialects that fetch `FROM` vs fetch `IN`
///
/// [MsSql](https://learn.microsoft.com/en-us/sql/t-sql/language-elements/fetch-transact-sql)
from_or_in: AttachedToken,
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure what's best here, it could also be two separate Optional fields

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we represent it with an explicit enum?
e.g

enum FetchPosition {
    From
    In
}

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done 👍

/// Optional, It's possible to fetch rows form cursor to the table
into: Option<ObjectName>,
},
Expand Down Expand Up @@ -4225,11 +4273,25 @@ impl fmt::Display for Statement {
Statement::Fetch {
name,
direction,
from_or_in,
into,
} => {
write!(f, "FETCH {direction} ")?;

write!(f, "IN {name}")?;
match &from_or_in.0.token {
Token::Word(w) => match w.keyword {
Keyword::FROM => {
write!(f, "FROM {name}")?;
}
Keyword::IN => {
write!(f, "IN {name}")?;
}
_ => unreachable!(),
},
_ => {
unreachable!()
}
}

if let Some(into) = into {
write!(f, " INTO {into}")?;
Expand Down Expand Up @@ -4319,6 +4381,9 @@ impl fmt::Display for Statement {
Statement::If(stmt) => {
write!(f, "{stmt}")
}
Statement::While(stmt) => {
write!(f, "{stmt}")
}
Statement::Raise(stmt) => {
write!(f, "{stmt}")
}
Expand Down Expand Up @@ -4488,6 +4553,11 @@ impl fmt::Display for Statement {
Ok(())
}
Statement::Delete(delete) => write!(f, "{delete}"),
Statement::Open { cursor_name } => {
write!(f, "OPEN {cursor_name}")?;

Ok(())
}
Statement::Close { cursor } => {
write!(f, "CLOSE {cursor}")?;

Expand Down
12 changes: 11 additions & 1 deletion src/ast/spans.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ use super::{
ReferentialAction, RenameSelectItem, ReplaceSelectElement, ReplaceSelectItem, Select,
SelectInto, SelectItem, SetExpr, SqlOption, Statement, Subscript, SymbolDefinition, TableAlias,
TableAliasColumnDef, TableConstraint, TableFactor, TableObject, TableOptionsClustered,
TableWithJoins, UpdateTableFromKind, Use, Value, Values, ViewColumnDef,
TableWithJoins, UpdateTableFromKind, Use, Value, Values, ViewColumnDef, WhileStatement,
WildcardAdditionalOptions, With, WithFill,
};

Expand Down Expand Up @@ -338,6 +338,7 @@ impl Spanned for Statement {
} => source.span(),
Statement::Case(stmt) => stmt.span(),
Statement::If(stmt) => stmt.span(),
Statement::While(stmt) => stmt.span(),
Statement::Raise(stmt) => stmt.span(),
Statement::Call(function) => function.span(),
Statement::Copy {
Expand All @@ -364,6 +365,7 @@ impl Spanned for Statement {
from_query: _,
partition: _,
} => Span::empty(),
Statement::Open { cursor_name } => cursor_name.span,
Statement::Close { cursor } => match cursor {
CloseCursor::All => Span::empty(),
CloseCursor::Specific { name } => name.span,
Expand Down Expand Up @@ -774,6 +776,14 @@ impl Spanned for IfStatement {
}
}

impl Spanned for WhileStatement {
fn span(&self) -> Span {
let WhileStatement { while_block } = self;

while_block.span()
}
}

impl Spanned for ConditionalStatements {
fn span(&self) -> Span {
match self {
Expand Down
3 changes: 3 additions & 0 deletions src/keywords.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -981,6 +981,7 @@ define_keywords!(
WHEN,
WHENEVER,
WHERE,
WHILE,
WIDTH_BUCKET,
WINDOW,
WITH,
Expand Down Expand Up @@ -1064,6 +1065,8 @@ pub const RESERVED_FOR_TABLE_ALIAS: &[Keyword] = &[
Keyword::SAMPLE,
Keyword::TABLESAMPLE,
Keyword::FROM,
// Reserved for SQL Server cursors
Keyword::OPEN,
];

/// Can't be used as a column alias, so that `SELECT <expr> alias`
Expand Down
60 changes: 56 additions & 4 deletions src/parser/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -536,6 +536,10 @@ impl<'a> Parser<'a> {
self.prev_token();
self.parse_if_stmt()
}
Keyword::WHILE => {
self.prev_token();
self.parse_while()
}
Keyword::RAISE => {
self.prev_token();
self.parse_raise_stmt()
Expand Down Expand Up @@ -570,6 +574,10 @@ impl<'a> Parser<'a> {
Keyword::ALTER => self.parse_alter(),
Keyword::CALL => self.parse_call(),
Keyword::COPY => self.parse_copy(),
Keyword::OPEN => {
self.prev_token();
self.parse_open()
}
Keyword::CLOSE => self.parse_close(),
Keyword::SET => self.parse_set(),
Keyword::SHOW => self.parse_show(),
Expand Down Expand Up @@ -700,8 +708,18 @@ impl<'a> Parser<'a> {
}))
}

/// Parse a `WHILE` statement.
///
/// See [Statement::While]
fn parse_while(&mut self) -> Result<Statement, ParserError> {
self.expect_keyword_is(Keyword::WHILE)?;
let while_block = self.parse_conditional_statement_block(&[Keyword::END])?;

Ok(Statement::While(WhileStatement { while_block }))
}

/// Parses an expression and associated list of statements
/// belonging to a conditional statement like `IF` or `WHEN`.
/// belonging to a conditional statement like `IF` or `WHEN` or `WHILE`.
///
/// Example:
/// ```sql
Expand All @@ -716,20 +734,36 @@ impl<'a> Parser<'a> {

let condition = match &start_token.token {
Token::Word(w) if w.keyword == Keyword::ELSE => None,
Token::Word(w) if w.keyword == Keyword::WHILE => {
let expr = self.parse_expr()?;
Some(expr)
}
_ => {
let expr = self.parse_expr()?;
then_token = Some(AttachedToken(self.expect_keyword(Keyword::THEN)?));
Some(expr)
}
};

let statements = self.parse_statement_list(terminal_keywords)?;
let conditional_statements = if self.peek_keyword(Keyword::BEGIN) {
let begin_token = self.expect_keyword(Keyword::BEGIN)?;
let statements = self.parse_statement_list(terminal_keywords)?;
let end_token = self.expect_keyword(Keyword::END)?;
Comment on lines +749 to +751
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We seem to have this pattern upcoming in a few places, like #1810 maybe it would be good to pull it out into a method and reuse it both here and the preexisting usage here? We can probably do so in the former PR instead

ConditionalStatements::BeginEnd(BeginEndStatements {
begin_token: AttachedToken(begin_token),
statements,
end_token: AttachedToken(end_token),
})
} else {
let statements = self.parse_statement_list(terminal_keywords)?;
ConditionalStatements::Sequence { statements }
};

Ok(ConditionalStatementBlock {
start_token: AttachedToken(start_token),
condition,
then_token,
conditional_statements: ConditionalStatements::Sequence { statements },
conditional_statements,
})
}

Expand Down Expand Up @@ -4453,6 +4487,9 @@ impl<'a> Parser<'a> {
break;
}
}
if let Token::EOF = self.peek_nth_token_ref(0).token {
break;
}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we collapse this into above to use a match statement?

match &self.peek_nth_token_ref(0).token {
    Token::Word(n) if ...
    Token::Eof 
}

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done 👍. I was probably trying to minimize the diff for review here

values.push(self.parse_statement()?);
self.expect_token(&Token::SemiColon)?;
}
Expand Down Expand Up @@ -6609,7 +6646,13 @@ impl<'a> Parser<'a> {
}
};

self.expect_one_of_keywords(&[Keyword::FROM, Keyword::IN])?;
let from_or_in_token = if self.peek_keyword(Keyword::FROM) {
self.expect_keyword(Keyword::FROM)?
} else if self.peek_keyword(Keyword::IN) {
self.expect_keyword(Keyword::IN)?
} else {
return parser_err!("Expected FROM or IN", self.peek_token().span.start);
};

let name = self.parse_identifier()?;

Expand All @@ -6622,6 +6665,7 @@ impl<'a> Parser<'a> {
Ok(Statement::Fetch {
name,
direction,
from_or_in: AttachedToken(from_or_in_token),
into,
})
}
Expand Down Expand Up @@ -8735,6 +8779,14 @@ impl<'a> Parser<'a> {
})
}

/// Parse [Statement::Open]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe I missed it, we seem to be lacking test cases for the open statement feature?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's part of test_mssql_cursor, but I'll make a separate test function just for OPEN for clarity

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done 👍

fn parse_open(&mut self) -> Result<Statement, ParserError> {
self.expect_keyword(Keyword::OPEN)?;
Ok(Statement::Open {
cursor_name: self.parse_identifier()?,
})
}

pub fn parse_close(&mut self) -> Result<Statement, ParserError> {
let cursor = if self.parse_keyword(Keyword::ALL) {
CloseCursor::All
Expand Down
26 changes: 26 additions & 0 deletions src/test_utils.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,8 @@ impl TestedDialects {
///
/// 2. re-serializing the result of parsing `sql` produces the same
/// `canonical` sql string
///
/// For multiple statements, use [`statements_parse_to`].
pub fn one_statement_parses_to(&self, sql: &str, canonical: &str) -> Statement {
let mut statements = self.parse_sql_statements(sql).expect(sql);
assert_eq!(statements.len(), 1);
Expand All @@ -166,6 +168,30 @@ impl TestedDialects {
only_statement
}

/// The same as [`one_statement_parses_to`] but it works for a multiple statements
pub fn statements_parse_to(
&self,
sql: &str,
statement_count: usize,
canonical: &str,
) -> Vec<Statement> {
let statements = self.parse_sql_statements(sql).expect(sql);
assert_eq!(statements.len(), statement_count);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this assertion seems to already be covered by the if/else below? so that we can skip the statement_count argument requirement?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hm, I don't fully understand. Without this line you can't guarantee that the string you feed in has exactly the number of statements you intend it to parse. Also, one_statement_parses_to has this same assertion before the if/else.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah so I meant that in both cases when asserting :

in this case we're already explicitly check that both statements lists are identical

assert_eq!(self.parse_sql_statements(canonical).unwrap(), statements);

Then in this case, we're doing so implicitly, reconstructing the input sql based off the returned statement

assert_eq!(
                sql,
                statements
                    .iter()
                    .map(|s| s.to_string())
                    .collect::<Vec<_>>()
                    .join("; ")
            );

So that i imagine it shouldn't be possible for the count assertion to fail and either of the subsequent assertion to pass?

one_statement_parse_to uses the count internally and there it makes sense to sanity check that since the expected count is always one, in this case we're exposing the count as a function argument which makes for suboptimal API that the user has to manually increment/decrement a counter when the sql changes. So that if the count argument isn't strictly necessary we would want to skip it

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Then in this case, we're doing so implicitly, reconstructing the input sql based off the returned statement ... So that i imagine it shouldn't be possible for the count assertion to fail and either of the subsequent assertion to pass?

I will remove the assertion here to get this branch merged. However, I think removing it removes a level of safety that is beneficial. Part of my thinking here is motivated by my upcoming branch on making semi colon statement delimiters optional. So any code that is making assumptions about "number of statements" becomes even more useful.

But perhaps that branch can re-introduce that assertion if necessary.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

BTW, this helper was introduced over on the GO branch, does your opinion change at all seeing the usage over there?

if !canonical.is_empty() && sql != canonical {
assert_eq!(self.parse_sql_statements(canonical).unwrap(), statements);
} else {
assert_eq!(
sql,
statements
.iter()
.map(|s| s.to_string())
.collect::<Vec<_>>()
.join("; ")
);
}
statements
}

/// Ensures that `sql` parses as an [`Expr`], and that
/// re-serializing the parse result produces canonical
pub fn expr_parses_to(&self, sql: &str, canonical: &str) -> Expr {
Expand Down
Loading