Add support for MySQL auto_increment offset (apache#950)

serprex · Nov 6, 2023 · 1ff2e70 · 1ff2e70
1 parent aa9f991
commit 1ff2e70
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 0 deletions.
diff --git a/src/ast/helpers/stmt_create_table.rs b/src/ast/helpers/stmt_create_table.rs
@@ -66,6 +66,7 @@ pub struct CreateTableBuilder {
     pub clone: Option<ObjectName>,
     pub engine: Option<String>,
     pub comment: Option<String>,
+    pub auto_increment_offset: Option<u32>,
     pub default_charset: Option<String>,
     pub collation: Option<String>,
     pub on_commit: Option<OnCommit>,
@@ -98,6 +99,7 @@ impl CreateTableBuilder {
             clone: None,
             engine: None,
             comment: None,
+            auto_increment_offset: None,
             default_charset: None,
             collation: None,
             on_commit: None,
@@ -204,6 +206,11 @@ impl CreateTableBuilder {
         self
     }
 
+    pub fn auto_increment_offset(mut self, offset: Option<u32>) -> Self {
+        self.auto_increment_offset = offset;
+        self
+    }
+
     pub fn default_charset(mut self, default_charset: Option<String>) -> Self {
         self.default_charset = default_charset;
         self
@@ -257,6 +264,7 @@ impl CreateTableBuilder {
             clone: self.clone,
             engine: self.engine,
             comment: self.comment,
+            auto_increment_offset: self.auto_increment_offset,
             default_charset: self.default_charset,
             collation: self.collation,
             on_commit: self.on_commit,
@@ -296,6 +304,7 @@ impl TryFrom<Statement> for CreateTableBuilder {
                 clone,
                 engine,
                 comment,
+                auto_increment_offset,
                 default_charset,
                 collation,
                 on_commit,
@@ -324,6 +333,7 @@ impl TryFrom<Statement> for CreateTableBuilder {
                 clone,
                 engine,
                 comment,
+                auto_increment_offset,
                 default_charset,
                 collation,
                 on_commit,

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -1418,6 +1418,7 @@ pub enum Statement {
         clone: Option<ObjectName>,
         engine: Option<String>,
         comment: Option<String>,
+        auto_increment_offset: Option<u32>,
         default_charset: Option<String>,
         collation: Option<String>,
         on_commit: Option<OnCommit>,
@@ -2451,6 +2452,7 @@ impl fmt::Display for Statement {
                 default_charset,
                 engine,
                 comment,
+                auto_increment_offset,
                 collation,
                 on_commit,
                 on_cluster,
@@ -2605,6 +2607,9 @@ impl fmt::Display for Statement {
                 if let Some(comment) = comment {
                     write!(f, " COMMENT '{comment}'")?;
                 }
+                if let Some(auto_increment_offset) = auto_increment_offset {
+                    write!(f, " AUTO_INCREMENT {auto_increment_offset}")?;
+                }
                 if let Some(order_by) = order_by {
                     write!(f, " ORDER BY ({})", display_comma_separated(order_by))?;
                 }

diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -3586,6 +3586,17 @@ impl<'a> Parser<'a> {
             None
         };
 
+        let auto_increment_offset = if self.parse_keyword(Keyword::AUTO_INCREMENT) {
+            let _ = self.consume_token(&Token::Eq);
+            let next_token = self.next_token();
+            match next_token.token {
+                Token::Number(s, _) => Some(s.parse::<u32>().expect("literal int")),
+                _ => self.expected("literal int", next_token)?,
+            }
+        } else {
+            None
+        };
+
         let order_by = if self.parse_keywords(&[Keyword::ORDER, Keyword::BY]) {
             if self.consume_token(&Token::LParen) {
                 let columns = if self.peek_token() != Token::RParen {
@@ -3667,6 +3678,7 @@ impl<'a> Parser<'a> {
             .clone_clause(clone)
             .engine(engine)
             .comment(comment)
+            .auto_increment_offset(auto_increment_offset)
             .order_by(order_by)
             .default_charset(default_charset)
             .collation(collation)

diff --git a/tests/sqlparser_mysql.rs b/tests/sqlparser_mysql.rs
@@ -312,6 +312,31 @@ fn parse_create_table_comment() {
     }
 }
 
+#[test]
+fn parse_create_table_auto_increment_offset() {
+    let canonical =
+        "CREATE TABLE foo (bar INT NOT NULL AUTO_INCREMENT) ENGINE=InnoDB AUTO_INCREMENT 123";
+    let with_equal =
+        "CREATE TABLE foo (bar INT NOT NULL AUTO_INCREMENT) ENGINE=InnoDB AUTO_INCREMENT=123";
+
+    for sql in [canonical, with_equal] {
+        match mysql().one_statement_parses_to(sql, canonical) {
+            Statement::CreateTable {
+                name,
+                auto_increment_offset,
+                ..
+            } => {
+                assert_eq!(name.to_string(), "foo");
+                assert_eq!(
+                    auto_increment_offset.expect("Should exist").to_string(),
+                    "123"
+                );
+            }
+            _ => unreachable!(),
+        }
+    }
+}
+
 #[test]
 fn parse_create_table_set_enum() {
     let sql = "CREATE TABLE foo (bar SET('a', 'b'), baz ENUM('a', 'b'))";